The Neural Moving Average Model for Scalable Variational Inference of State Space Models

A SSM is based on a latent Markov chain $x=x_{1{:}T}$. We focus on the case of continuous states $x_{i}\in\mathbb{R}^{d}$. States evolve through a transition density $p(x_{i}|x_{i-1},\theta),$ with parameters $\theta\in\mathbb{R}^{p}$. We assume the initial state is $x_{0}(\theta)$, a deterministic function of $\theta$. (This allows examples with initial state known – $x_{0}$ is a constant – or unknown – $x_{0}$ depends on unknown parameters.) Observations $y_{i}\in\mathbb{R}^{d_{y}}$ are available for $i\in\mathcal{S}\subseteq 0{:}T$ following an observation density $p(y_{i}|x_{i},\theta).$

In the Bayesian framework, after specifying a prior density $p(\theta)$, interest lies in the posterior density

One application of SSMs, which we use in our examples, is as discrete approximations to stochastic differential equations (SDEs), as follows:

where $\epsilon_{i}\sim N(0,I_{d})$ are independent random vectors. Here $\alpha$ is a $d$-dimensional drift vector, $\beta$ is a $d\times d$ positive-definite diffusion matrix and $\sqrt{\beta}$ denotes its Cholesky factor. The state $x_{i}$ approximates the state of the SDE process at time $i\Delta t$. Taking the limit $\Delta t\to 0$ in an appropriate way recovers the exact SDE [Øksendal, 2003, Särkkä and Solin, 2019].

A normalising flow represents a random object $x$ as $g_{m}\circ\ldots g_{2}\circ g_{1}(z)$: a composition of learnable bijections of a base random object $z$. Here we suppose $x=x_{1:T}$ and $x_{i}\in\mathbb{R}$. (Later we consider $x_{i}$ as a vector.) We take $z=z_{1:T}$ as independent $N(0,1)$ variables. By the standard change of variable result, the log-density of $x$ is

where $\varphi$ is the $N(0,1)$ log-density function and $J_{j}$ is the Jacobian matrix of transformation $g_{j}$ given input $g_{j-1}\circ\ldots g_{2}\circ g_{1}(z)$.

The bijections in an IAF are mainly affine layers, which transform input $z^{\text{in}}$ to output $z^{\text{out}}$ by

with $\sigma_{i}>0$. This transformation scales and shifts each $z^{\text{in}}_{i}$. The shift and scale shift values, $\mu_{i}$ and $\sigma_{i}$, are typically neural network outputs. An efficient approach is to use a single neural network to output all the $\mu_{i},\sigma_{i}$ values for a particular affine layer. This network uses masked dense layers so that $(\mu_{i},\sigma_{i})$ depends only on $z^{\text{in}}_{1{:}i-1}$ as required [Germain et al., 2015, Kingma et al., 2016, Papamakarios et al., 2017]. In the resulting IAF each affine layer is based on a different neural network of this form. We’ll refer to this as a masked IAF.

The shift and scale functions for $z^{\text{out}}_{i}$ in (4) have an autoregressive property: they depend on $z^{\text{in}}$ only through $z^{\text{in}}_{j}$ with $j<i$. Hence the Jacobian matrix of the transformation is diagonal with non-zero entries $\sigma_{1:T}$. The log-density of an IAF made of $m$ affine layers is

where $\sigma^{j}_{i}$ is the shift value for the $i$th input to the $j$th affine layer.

IAFs typically alternate between affine layers and permutation layers, using order reversing or random permutations. Such layers have Jacobians with absolute determinant 1. Thus the log-density calculation is unchanged (interpreting $j$ in (5) to index the $j$th affine layer not the $j$th layer of any type). The supplement (Section G.1) details methods to restrict the output of a IAF e.g. to ensure all $x_{i}$s are positive.

IAFs are flexible and, for small $T$, allow fast sampling and calculation of a sample’s log-density. However they are expensive for large $T$ as large neural networks are needed to map between length $T$ sequences.

Our neural moving average (nMA) model reduces the number of weights that IAFs require by using a CNN to calculate the $\mu_{i}$ and $\sigma_{i}$ values in an affine layer. Thus it can be thought of as a kind of local IAF. Here we explain the main idea by presenting a version for scalar $x_{i}$. Section 3.3 extends this to the vector $x_{i}$ case.

To define the nMA model we describe how a single affine layer produces its shift and scale values. The affine layer uses a CNN with input $z^{\text{in}}$, a vector of length $T$. Let $h^{k}$ represent the $k$th hidden layer of the CNN, a matrix of dimension $(T,n_{k})$ where $n_{k}$ is a tuning choice. The first layer applies a convolution with receptive field length $\ell$. This is an off-centre convolution so that row $i$ of $h^{1}$ is a transformation of $z^{\text{in}}_{i-\ell{:}i-1}$. We use zero-padding by taking $z^{\text{in}}_{i}=0$ for $i<0$. The following hidden layers are length-$1$ convolutions, so row $i$ of $h^{k+1}$ is a transformation of row $i$ of $h^{k}$. The output, $h^{n}$, is a matrix of dimension $(T,2)$ whose $i$th row contains $\mu_{i}$ and $\sigma_{i}$. The final layer applies a softplus activation to produce the $\sigma_{i}$ values, ensuring they are positive. An identity activation is used to produce the $\mu_{i}$ values. The $\mu_{i}$ and $\sigma_{i}$ values are used in (4) to produce the output of the affine layer.

A nMA model composes several affine layers of the form just described. Some properties of the distribution for the output sequence $x$ are:

1. 1.

   No long-range dependence: $x_{i}$ and $x_{j}$ are independent if $|i-j|>m\ell$, where $m$ is the number of affine layers.

2. 2.

   Stationary local dependence: the distributions of $x_{i{:}j}$ and $x_{i+a{:}j+a}$ are the same for most choices of $a$. (Subsequences near to the start of $x$ can differ due to zero-padding.)

To improve the flexibility of the nMA model, affine layers can be alternated with order-reversing permutations. (Random permutations would not be suitable, as they would disrupt our ability to sample subsequences quickly, as described in Section 3.4.) Throughout the paper we consider nMA models without order reversing permutation layers, as we found these models already sufficiently flexible for our examples. (The supplement, Section C, details how to incorporate these layers.)

We relax stationary local dependence by injecting local side information to the CNN i.e. giving an extra feature vector $s_{i}$ as input for each position $i$ in the first CNN layer. We also use global side information to allow $x$ to depend on the parameter values $\theta$ i.e. giving $\theta$ as extra input for every position $i$. See the supplement (Sections D, E) for details of the side information we use in practice.

Here we generalise the nMA model to the case where $x_{i}\in\mathbb{R}^{d}$. We now let $z$ be a sequence $z_{1},z_{2},\ldots,z_{T}$ of independent random $N(0,I_{d})$ vectors. A nMA affine layer makes the transformation

scaling the vector $z^{\text{in}}_{i}$ (elementwise multiplication by vector $\sigma_{i}$) then shifting it (adding vector $\mu_{i}$).

In the scalar case it was important to allow complex dependencies between $z^{\text{out}}_{i}$ values. Now we must also allow dependencies within each $z^{\text{out}}_{i}$ vector. To do so we use coupling layers as in Dinh et al. [2016].

We use an extra $k$ subscript to denote the $k$th component of a vector e.g. $z^{\text{in}}_{ik}$. We select some $a\approx d/2$. For $k\leq a$, we take $\mu_{ik}=0$ and $\sigma_{ik}=1$, so that $z^{\text{out}}_{ik}=z^{\text{in}}_{ik}$. For $k>a$, we compute $\mu_{ik}$ and $\sigma_{ik}$ using a CNN, modifying the scalar case as follows. Now row $i$ of $h^{1}$ is a transformation of $z^{\text{in}}_{i-\ell{:}i-1}$ (the $\ell$ vectors preceding $z^{\text{in}}_{i}$), and also $z^{\text{in}}_{ik}$ for $k\leq a$ (the part of $z^{\text{in}}_{i}$ not being modified). The output $h^{n}$ is now a tensor of dimension $(T,d-a,2)$ containing $\mu_{ik}$ and $\sigma_{ik}$ values for $k>a$.

This affine layer does not transform the first $a$ components of $z^{\text{in}}_{i}$. To allow different components to be transformed in each layer, we permute components between affine layers. For example, a $d=2$ permutation layer transforms $z^{\text{in}}$ to $z^{\text{out}}$ by $z^{\text{out}}_{i1}=z^{\text{in}}_{i2}$, $z^{\text{out}}_{i2}=z^{\text{in}}_{i1}$. The log-density is now

where $\varphi$ is the $N(0,I_{d})$ log-density function and $\sigma^{j}_{ik}$ is the $k$th entry of the shift vector for position $i$ output by the $j$th affine layer. Decomposing $\log q(x)$ into $\lambda_{i}$ contributions will be useful in Section 4.

Sampling from a nMA model is straightforward. First sample the base random object $z$. This is a sequence of length $T$ (of scalars or vectors – the sampling process is similar in either case). Now apply the IAF’s layers to this in turn. To apply an affine layer, pass the input (and any side information) through the layer’s CNN to calculate shift and scale values, then apply the affine transformation. The final output is the sampled sequence $x$. The cost of sampling in this way is $O(T)$.

In the next section, we will often wish to sample a short subsequence $x_{u{:}v}$. It is possible to do this at $O(1)$ cost with respect to $T$. Algorithm A in the supplement gives the details. In brief, the key insight is that $x_{u{:}v}$ only depends on $z$ through $z_{u-m\ell{:}v}$. Therefore we sample $z_{u-m\ell{:}v}$ and apply the layers to this subsequence. The output will contain the correct values of $x_{u{:}v}$.

We wish to infer the joint posterior density $p(\theta,x|y)$. We introduce a family of approximations indexed by $\phi$, $q(\theta,x;\phi)$. Optimisation is used to find $\phi$ minimising the Kullback-Leibler divergence $KL[q(\theta,x;\phi)||p(\theta,x|y)]$. This is equivalent to maximising the ELBO (evidence lower bound) [Jordan et al., 1999],

Here $r$ is a log-density ratio. The optimal $q(\theta,x;\phi)$ approximates the posterior density. It is typically overconcentrated, unless the approximating family is expressive enough to allow particularly close matches to the posterior.

Optimisation for VI can be performed efficiently using the reparameterisation trick [Kingma and Welling, 2014, Rezende et al., 2014, Titsias and Lázaro-Gredilla, 2014]. That is, letting $(\theta,x)$ be the output of an invertible deterministic function $g(\varepsilon,\phi)$ for some random variable $\varepsilon$ with a fixed distribution. Then the ELBO gradient and unbiased Monte Carlo estimate are

where $(\theta^{(j)},x^{(j)})=g(\varepsilon^{(j)},\phi)$ and $\varepsilon^{(1)},\ldots,\varepsilon^{(n)}$ are independent $\varepsilon$ samples. This gradient estimate can be used in stochastic gradient optimisation algorithms.

Our variational family for the SSM posterior (1) is

where $\phi=(\phi_{\theta},\phi_{x})$. We use a masked IAF for $q(\theta;\phi_{\theta})$ and a nMA model for $q(x|\theta;\phi_{x})$. For the latter we inject $\theta$ as side information. See the supplement (Sections D and E) for more details.

The masked IAF maps a base random vector $z_{\theta}$ to $\theta$ using parameters $\phi_{\theta}$, as described in Section 3.1. The nMA model maps $\theta$ and a sequence of vectors $z_{x}$ to $x$ using parameters $\phi_{x}$ as described in Section 3.3. Hence we have a mapping from $\varepsilon=(z_{\theta},z_{x})$ to $g(\varepsilon,\phi)=(\theta,x)$, allowing us to use the reparameterisation trick below.

This section derives a mini-batch optimisation algorithm to train $\phi$ based on sampling short $x$ subsequences, so that the cost-per-training-iteration is $O(1)$. The algorithm is applicable for scalar or multivariate $x_{i}$. In this presentation we assume that $\mathcal{S}=0{:}T$ i.e. there are observations for all $i$ values. To relax this assumption remove any terms involving $y_{i}$ for $i\not\in\mathcal{S}$.

For our variational family (12), the ELBO is (8) with

Substituting (1) and (7) into (13) gives

Now introduce batches $B_{1},B_{2},\ldots,B_{b}$: length $M$ sequences of consecutive integers partitioning $1{:}T$. Draw $\kappa$ uniformly from $1{:}b$. Then an unbiased estimate of $r$ is

Hence an unbiased estimate of the ELBO gradient is

where $(\theta^{(j)},x^{(j)})=g(\varepsilon^{(j)},\phi)$ and $\varepsilon^{(1)},\ldots,\varepsilon^{(n)}$ are independent $\varepsilon$ samples.

Algorithm 1 presents our mini-batch training procedure. Each iteration of Algorithm 1 involves sampling a subsequence of $x$ values of length $M+1$. The cost is $O(1)$ with respect to the total length of the sequence $T$. This compares favourably to the $O(T)$ cost of sampling the entire $x$ sequence. Discussion of implementation details is given in the supplement (Section E).

First we consider the AR(1) model $x_{i+1}=\theta_{1}+\theta_{2}x_{i}+\theta_{3}\epsilon$, with $\epsilon\sim N(0,1)$ and $x_{0}=10$. We assume observations $y_{i}\sim N(x_{i},1)$ for $i\in 0{:}T$, and independent $N(0,10^{2})$ priors on $\theta_{1},\theta_{2},\log\theta_{3}$. We use this model to investigate how our method scales with larger $T$, and the effect of receptive field length $\ell$. To judge the accuracy of our results we compare to near-exact posterior inference using MCMC, as described in the supplement (Section A).

Next we test our method on short time series with complex dynamics. We use a version of the Lotka-Volterra model for predator-prey population dynamics under three events: prey reproduction, predation (in which prey are consumed and predators get resources to reproduce) and predator death. A SDE Lotka-Volterra model (for derivation see e.g. Golightly and Wilkinson [2011]) is defined by drift and diffusion

where $x=(u,v)$ represents population sizes of prey and predators. The parameters $\theta=\left(\theta_{1},\theta_{2},\theta_{3}\right)$ control the rates of the three events described above.

We consider a discretised version of this SDE, as described in Section 2.2, with $\Delta t=0.1$ and $u_{0}=v_{0}=100$. We simulated realisations under parameters $\theta=(0.5,0.0025,0.3)$ of $x_{i}$ for $i\in 1{:}500$, and construct two datasets with observations at: (a) $i=0,10,20,\ldots,500$ (dense); (b) $i=0,100,200,\ldots,500$ (sparse). We assume noisy observations $y_{i}\sim N(x_{i},I_{2})$ and independent $N(0,10^{2})$ priors for $\log\theta_{1},\log\theta_{2},\log\theta_{3}$.

Unlike previous examples, we needed to restrict $u_{i}$ and $v_{i}$ to be positive. Also, we found multiple posterior modes, and needed to pretrain carefully to control which we converged to. See Section G of the supplement for more details of both these issues.

For observation setting (a), we compared our results to near-exact posterior samples from MCMC [Golightly and Wilkinson, 2008, Fuchs, 2013]. These papers use a Metropolis-within-Gibbs MCMC scheme with carefully chosen proposal constructs. Designing suitable proposals can be challenging, particularly in sparse observation regimes [Whitaker et al., 2017]. Consequently we were unable to use MCMC in setting (b).

Figure 2 (left plot) displays the visual similarity between marginal densities estimates from variational and MCMC output in setting (a). The VI output is taken from the 30,000th iteration, after $\approx 10$ minutes of computation. Figure 2 (right plots) shows variational output for $\theta$ and $x$ in setting (b) after $\approx 20$ minutes of computation. These results are consistent with the ground truth parameter values and $x$ path. VI using an autoregressive distribution for $x$ has also performed well in a similar scenario [Ryder et al., 2018], but required more training time (roughly 2 hours).

Here we test our method on a long time series with an unobserved component, using the FitzHugh-Nagumo model. A SDE version, following Jensen et al. [2012], van der Meulen and Schauer [2017], is defined by drift and diffusion

where $x=(v,w)$ represents the current membrane potential and latent recovery variables.

We consider a discretised version of this SDE, as in Section 2.2, with $\Delta t=0.1,v_{0}=2,w_{0}=3$. We simulate synthetic data under parameter values $\theta=(2.0,1.0,1.5,0.5,0.3)$ up to $T=1,000,000$, recording observations at every $i$ to mimic a high frequency observation scenario. We assume independent observations $y_{i}\sim N(v_{i},0.1^{2})$ and independent $N(0,10^{2})$ priors for $\log\theta_{1},\theta_{2},\theta_{3},\log\theta_{4},\log\theta_{5}$.

Figure 3 displays estimates of the unobserved component $w$, and marginal density estimates for $\theta$. The results are consistent with the ground truth parameter values and $w$ path. The approximate posterior is sampled after roughly $180$ minutes of training.

We analyse a real data under a log-Gaussian stochastic volatility model presented as a discretised SDE with drift and diffusion

where $x=\left(r,z\right)$ is the returns process and latent volatility factor, respectively. Similar discrete-time models have been analysed by Andersen and Lund [1997], Eraker [2001], Durham and Gallant [2002], but we use the form presented in Golightly and Wilkinson [2006].

Similarly to Golightly and Wilkinson [2006], we use 1508 weekly observations on the three-month U.S. Treasury bill rate for August 13, 1967 – August 30, 1996, and perform inference under independent $N(0,10^{2})$ priors for $\theta_{1},\theta_{2},\log\theta_{3},\log\theta_{4}$ and $z_{0}$. We assume the returns process is fully observed without error and set $\Delta t=1.0$. Our analysis took  20 minutes on a single GPU. Figure 4 shows the results, which are consistent with those obtained from the MCMC analysis in Golightly and Wilkinson [2006].