GraphGrad: Efficient Estimation of Sparse Polynomial Representations for General State-Space Models

The general state-space model can be described by

where $t=1,\dots,T$ denotes discrete time, ${\mathbf{x}}_{t}\in\mathbb{R}^{N_{x}}$ is the state of the system at time $t$, ${\mathbf{y}}_{t}\in\mathbb{R}^{N_{y}}$ is the observation at time $t$, $p({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1};\bm{\theta})$ is the density of the hidden state ${\mathbf{x}}_{t}$, given the previous state ${\mathbf{x}}_{t-1}$, $p({\mathbf{y}}_{t}|{\mathbf{x}}_{t};\bm{\theta})$ is the density of the observation ${\mathbf{y}}_{t}$ given the hidden state ${\mathbf{x}}_{t}$, and $\bm{\theta}$ is a set of model parameters. The initial value of the hidden state is distributed ${\mathbf{x}}_{0}\sim p({\mathbf{x}}_{0}|\bm{\theta})$. The state sequence ${\mathbf{x}}_{0:T}$ is typically hidden, whilst the related sequence of observations ${\mathbf{y}}_{1:T}$ is known.

Filtering methods aim at estimating the hidden state at time $t$, denoted ${\mathbf{x}}_{t}$, typically utilising the posterior density function of the state conditional on the observations up to time $t$, denoted ${\mathbf{y}}_{1:t}$. Particle filters approximate this density, $p({\mathbf{x}}_{t}|{\mathbf{y}}_{1:t};\bm{\theta})$, using a set of $K$ Monte Carlo samples (particles) and their associated weights, $\{{\mathbf{x}}_{1:T}^{(k)},\widetilde{w}_{1:T}^{(k)}\}_{k=1}^{K}$. The posterior density can then be approximated by

A commonly used particle filtering method is the sequential importance resampling algorithm, given in Alg. 1. At every time-step $t$, the $K$ particles and normalised weights, $\{{\mathbf{x}}_{1:T}^{(k)},\overline{w}_{1:T}^{(k)}\}_{k=1}^{K}$, are calculated. First, we perform the resampling step (line 6), which generates $K$ samples, sampling ${\mathbf{x}}_{t-1}^{(k)},\ k=1,\dots,K$, with probability $\overline{w}_{t-1}^{(k)}$. The resampling step is vital to avoid the degeneracy of the filter, i.e., to ensure diversity in the particle set and obtain more accurate approximations of the posterior distribution, $p({\mathbf{x}}_{t}|{\mathbf{y}}_{1:t};\bm{\theta})$. Next, $K$ particles ${\mathbf{x}}_{t}^{(k)}$, $k=1,\dots,K$, are drawn from the proposal distribution $\pi({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1},{\mathbf{y}}_{t};\bm{\theta})$ (line 8). Finally, we incorporate the observation ${\mathbf{y}}_{t}$, which is done via the particle weights, given by ${w}_{t}^{(k)}$, $k=1,\dots,K$, in line 9, and the normalised weights (line 10).

In many problems of interest, the parameter $\bm{\theta}$ is not known, and must be estimated. The posterior distribution of $\bm{\theta}$ in the state-space model can be factorised as $p(\bm{\theta}|\mathbf{y}_{1:T})\propto p(\bm{\theta})p(\mathbf{y}_{1:T}|\bm{\theta})$, where $\bm{\theta}$ is the parameter of interest, $p(\bm{\theta})$ is the prior distribution of the parameter $\bm{\theta}$, and $p(\mathbf{y}_{1:T}|\bm{\theta})$ is given by

where $p(\mathbf{y}_{1}|\mathbf{y}_{1:0};\bm{\theta}):=p(\mathbf{y}_{1}|\bm{\theta})$ [15]. The prior distribution $p(\bm{\theta})$ encodes our pre-existing beliefs as to the value and structure of the parameter $\bm{\theta}$, and can be used for regularisation, e.g., by the Lasso [16].

There are many methods to estimate parameters $\bm{\theta}$ given its posterior density function $p(\bm{\theta}|\mathbf{y}_{1:T})$. We can broadly classify these parameter estimation methods as point estimation methods and distributional methods. Point estimation methods provide a single estimate that is, in some defined way, the optimal value. An example of a point estimation method is the maximum-a-posteriori (MAP) estimator, that defines the optimal value of $\bm{\theta}$ as the one that maximises the posterior density $p(\bm{\theta}|\mathbf{y}_{1:T})$. The method proposed in this work yields a MAP estimator.

In the case of a linear-Gaussian state-space model, we have closed form methods for the MAP estimator assuming a diffuse prior [15], and efficient methods for obtaining the MAP estimator for sparsity promoting priors [17]. Distributional methods estimate the posterior distribution of the parameter, with common methods being importance sampling [18], Markov chain Monte Carlo [19], and variational inference [20]. For the linear-Gaussian state-space model, we can utilise reversible jump Markov chain Monte Carlo to obtain an estimate of the distribution of sparsity [21], however no similar method exists for the general state-space model. Our proposed method in this work is a point estimation method.

For general SSMs, the likelihood of the parameter $\bm{\theta}$, necessary to compute the posterior and hence to design a procedure to maximise it, cannot be obtained in closed form. We can estimate it using the particle filter, by the Monte Carlo estimate

where $w_{t}^{(k)}$ and $\widetilde{w}_{t-1}^{(k)}$ are the weights of the particle filter in Alg. 1 [15, Chapter 12]. Note that the weights are dependent on the parameter $\bm{\theta}$ through their computation in Alg. 1. Using these weights, we construct an estimate of the log-likelihood by

where $c$ results from the proportionality in Eq. (4), and we note that if resampling occurred at time $t-1$, which is always the case if following Alg. 1, we have $\widetilde{w}_{t-1}=1/K$. Note that, in practice, log-weights are used for numerical stability, and we compute $\sum_{k=1}^{K}w_{t}^{(k)}\widetilde{w}_{t-1}^{(k)}$ using a logsumexp reparameterisation.

The outputs of the particle filter, as given in Alg. 1, namely $\{{\mathbf{x}}^{(k)}_{1:T},w^{(k)}_{1:T}\}_{k=1}^{K}$, are not differentiable with respect to $\bm{\theta}$ [12], because of the resampling step on line 6.  The resampling step requires sampling a multinomial distribution. Sampling a multinomial distribution is not differentiable, as an infinitesimal change in the input probabilities can lead to a discrete change in the output sample value [12].

Differentiable particle filters (DPFs) are recently introduced tools that modify the resampling step of the SIR particle filter to be differentiable, and therefore lead to a particle filtering algorithm that is differentiable with respect to $\bm{\theta}$ [11, 22, 12, 23]. There exist a number of DPF methods, based on various techniques for making the resampling step differentiable. These range from weight retention in soft resampling [10], to optimal transport [11]. An overview of the DPF and its interplay with deep learning methods can be found in [22], which motivates our choice of differentiable particle filter, and our usage of stochastic gradient descent methods.

In this work, we built upon the DPF approach from [12], which utilises a stop gradient operator to make the resampling step differentiable. We summarise this method in Alg. 2. This algorithm yields unbiased gradient estimates with minimal computational overhead, and does not modify the behaviour of the forward pass of the particle filter.

In Alg. 2, at each time-step $t$, we first sample the previous particle set, sampling $a_{t}^{(k)}$, $k=1,\dots,K$, with probability $\overline{w}_{t-1}$ (line 6). The value of $a_{t}^{(k)}$ determines the ancestry of the $k$-th particle at time $t$. Note that we apply a stop gradient operator to the weights in the sampling method, and therefore do not attempt to propagate gradients through sampling a discrete distribution. We then set the weights of all $K$ particles to $1/K$ whilst preserving gradient information in the weights by dividing the weights by themselves, applying a stop-gradient operation to the divisor, and then multiplying by the constant $1/K$ (line 7). The rest of the DPF proceeds as Alg. 1, described in Sec. II-A, with the particle sampling, weighting, and normalisation steps unmodified.

The DPF is particularly useful when estimating parameters, as we can compute the gradients of functions of the likelihood and of the particle trajectories, and therefore compute parameters optimising a chosen loss function. In particular, we typically use likelihood-based losses when performing inference where the hidden state is unknown in the training data, and trajectory-based losses when the sequence of hidden states is known for the training data. Examples of likelihood-based losses are the negative log-likelihood and the ELBO. The mean-square error of the inferred state is a typical trajectory-based loss. Note that the trajectory of the hidden states depends on the weights, and therefore the weights must be differentiable even if the loss function is based only on the trajectory.

Given a loss function based on the weights and/or particles of a differentiable particle filter, we can compute the minimiser using a gradient-based optimisation scheme. This is efficient, especially compared to the gradient-free methods previously required, and is more robust to the stochasticity of the likelihood and state estimates, as many gradient schemes are designed with noisy gradients in mind [24, 25].

We here present our notation, used throughout the paper. We denote by ${\mathbf{A}}_{\cdot,i}$ the $i$-th column vector of matrix ${\mathbf{A}}$, and by ${\mathbf{A}}_{i,\cdot}$ the $i$-th row vector of matrix ${\mathbf{A}}$. We denote by $\mathbf{Id}_{n}$ the $n\times n$ identity matrix.

We denote by $\mathds{1}_{\mathrm{cond}}(x)$ the binary indicator function, which returns $1$ when $x$ satisfies $\mathrm{cond}$, and $0$ otherwise. We denote by $\mathrm{sgn}(x)=\mathds{1}_{\geq 0}(x)-\mathds{1}_{\leq 0}(x)$ the sign function. We denote by $\mathrm{abs}(x)$ the absolute value function. All three functions can be applied element-wise to a matrix or vector.

We denote by $\mathbb{N}_{0}$ the natural numbers including $0$, and by $\mathbb{N}$ the natural numbers excluding $0$.

We denote by $\mathrm{count}(a,b)$ the number of times the item $a$ occurs in the list $b$, and by $\mathrm{cwr}(S,r)$ the set of all length $r$ combinations of the elements of the set $S$ with replacement, up to reordering. For example, we have $\mathrm{count}(1,[1,2,2,3,1])=2$ and $\mathrm{cwr}(\{a,b,c\},2)=\{[a,a],[a,b],[a,c],[b,b],[b,c],[c,c]\}$.

We denote by $\bot(x)$ the stop-gradient operator applied to $x$, with properties as defined in [12].

In order to learn a polynomial approximation to the state transition distribution, let us first define the learnt and static parameters of our polynomials. A polynomial is the result of summing a number of monomials, and hence can be constructed as a sum of estimated monomial terms. The degree of a monomial is given by the sum of the powers of its constituent terms, with a polynomial having degree equal to the maximum degree of its constituent monomials. In our model, we assume a fixed maximum degree $d$ for our polynomial approximation.

Once $d$ is set, we construct all length $N_{x}$ sequences of positive integers that sum to $n\leq d,\,n\in\mathbb{N}_{0}$, resulting in

unique sequences. This simple procedure allows us to generate the powers of all monomial terms in a polynomial of degree $d$, that we store in an $N_{x}\times M$ matrix, denoted ${\mathbf{D}}$, with the term $D_{i,j}$ corresponding to the power of state dimension $i$ in the $j$-th monomial term. The polynomial expression is then defined by matrix ${\mathbf{C}}$, following Eq. (7). The ordering of the monomials is arbitrary but must be consistent, as it implies the order of the columns of ${\mathbf{C}}$ and ${\mathbf{D}}$.

Within time series modelling, we can often interpret the driving parameters of a sparse model as a graph encoding the connectivity of the system [17, 21, 26, 27, 28, 29]. This is the case here as well, where we observe that, for $(a,b)\in\{1,\dots,N_{x}\}^{2}$, state dimension $b$ affects state dimension $a$ in our estimated dynamics if $(\mathrm{abs}({\mathbf{C}}){\mathbf{D}}^{\top})_{a,b}\neq 0$, where we note that $\mathrm{abs}({\mathbf{C}}){\mathbf{D}}^{\top}$ is an $N_{x}\times N_{x}$ matrix. We can interpret this in terms of Granger causality, where we see that including information from state $b$ improves the knowledge on state $a$, and therefore we say that state $b$ Granger-causes state $a$ if $(\mathrm{abs}({\mathbf{C}}){\mathbf{D}}^{\top})_{a,b}\neq 0$.

We are therefore able to construct a directed graph of our estimated system from the matrices ${\mathbf{C}}$ and ${\mathbf{D}}$. This graph has adjacency matrix ${\mathbf{A}}$, defined by ${\mathbf{A}}=\mathds{1}_{\neq 0}(\mathrm{abs}({\mathbf{C}}){\mathbf{D}}^{\top})\in\{0,1\}^{N_{x}\times N_{x}}$, where we have an edge from node $b$ to node $a$ if $A_{a,b}=1$, and no edge if $A_{a,b}=0$. If this graph is not fully connected, or equivalently if there exist $(a,b)\in\{1,\dots,N_{x}\}^{2}$ such that $A_{a,b}=A_{b,a}=0$, then some state dimensions do not directly interact in our estimated system.

We can also interpret our system as a collection of $M$ graphs, $\{\mathcal{G}_{j}\}_{j=1}^{M}$, where the $j$-th graph $\mathcal{G}_{j}$ has adjacency matrix $\mathrm{abs}({\mathbf{C}}_{\cdot,j}){\mathbf{D}}_{\cdot,j}^{\top}$. Each $\mathcal{G}_{j}$ can be interpreted as encoding the connectivity resulting from the $j$-th monomial. A sparsity promoting prior on ${\mathbf{C}}$ also promotes sparsity in these graphs. We can therefore interpret our method of estimating ${\mathbf{C}}$ as estimating multiple interacting sparse graphical models, one for each of the underlying monomials.

We present here a worked example utilising the Lorenz 63 model. The Lorenz 63 model [13] is a popular chaotic oscillator model, with a discrete time variant given by,

where $\Delta t$ is the time elapsed between observations, ${\mathbf{v}}_{t}\sim\mathcal{N}(\bm{0},\Sigma_{v})$ the state noise term, and ${\mathbf{r}}_{t}\sim\mathcal{N}(\bm{0},\Sigma_{r})$ the observation noise term, and $\beta,\,\rho,\,\sigma$ are real scalar parameters, often taken as $\beta=8/3,\ \sigma=10,\ \rho=28$. The initial condition ${\mathbf{x}}_{0}$ is arbitrary, so long as it is non-zero, and is often taken to be $[1,0,0]$.