Overcoming Mean-Field Approximations in Recurrent Gaussian Process Models

TL;DR

This paper introduces VCDT, a novel variational inference method for recurrent Gaussian process models that explicitly models dependencies between states and transition functions, improving accuracy and calibration without increasing computational complexity.

Contribution

It eliminates the mean-field factorisation in variational inference for Gaussian process dynamical systems, enabling better posterior calibration and predictive performance.

Findings

VCDT outperforms mean-field methods in predictive accuracy.

VCDT provides more calibrated estimates of transition functions.

The method maintains computational efficiency comparable to existing approaches.

Abstract

We identify a new variational inference scheme for dynamical systems whose transition function is modelled by a Gaussian process. Inference in this setting has either employed computationally intensive MCMC methods, or relied on factorisations of the variational posterior. As we demonstrate in our experiments, the factorisation between latent system states and transition function can lead to a miscalibrated posterior and to learning unnecessarily large noise terms. We eliminate this factorisation by explicitly modelling the dependence between state trajectories and the Gaussian process posterior. Samples of the latent states can then be tractably generated by conditioning on this representation. The method we obtain (VCDT: variationally coupled dynamics and trajectories) gives better predictive performance and more calibrated estimates of the transition function, yet maintains the same time and space complexities as mean-field methods. Code is available at: github.com/ialong/GPt.

Figures

Tables

Table 1. Table 1: Variations of approximate posteriors. 𝐀 t , 𝐛 t , 𝐒 t subscript 𝐀 𝑡 subscript 𝐛 𝑡 subscript 𝐒 𝑡 \mathbf{A}_{t},\mathbf{b}_{t},\mathbf{S}_{t} are free parameters in all cases. 𝐂 f ​ ( 𝐱 t ) subscript 𝐂 𝑓 subscript 𝐱 𝑡 \mathbf{C}_{f}(\mathbf{x}_{t}) is the sparse GP’s marginal posterior variance 𝐂 f ​ ( 𝐱 t ) ≜ K 𝐱 t , 𝐱 t + K 𝐱 t , Z ​ K Z , Z − 1 ​ ( 𝚺 𝐮 − K Z , Z ) ​ K Z , Z − 1 ​ K Z , 𝐱 t ≜ subscript 𝐂 𝑓 subscript 𝐱 𝑡 subscript 𝐾 subscript 𝐱 𝑡 subscript 𝐱 𝑡 subscript 𝐾 subscript 𝐱 𝑡 𝑍 superscript subscript 𝐾 𝑍 𝑍 1 subscript 𝚺 𝐮 subscript 𝐾 𝑍 𝑍 superscript subscript 𝐾 𝑍 𝑍 1 subscript 𝐾 𝑍 subscript 𝐱 𝑡 \mathbf{C}_{f}(\mathbf{x}_{t})\triangleq K_{\mathbf{x}_{t},\mathbf{x}_{t}}+K_{\mathbf{x}_{t},Z}K_{Z,Z}^{-1}\mathopen{}\mathclose{{}\left(\mathbf{\Sigma}_{\mathbf{u}}-K_{Z,Z}}\right)K_{Z,Z}^{-1}K_{Z,\mathbf{x}_{t}} whereas 𝐂 f | 𝐮 ​ ( 𝐱 t ) subscript 𝐂 conditional 𝑓 𝐮 subscript 𝐱 𝑡 \mathbf{C}_{f|\mathbf{u}}(\mathbf{x}_{t}) is the conditional variance of f ​ ( 𝐱 t ) | 𝐮 conditional 𝑓 subscript 𝐱 𝑡 𝐮 f(\mathbf{x}_{t})\,|\,\mathbf{u} : 𝐂 f | 𝐮 ​ ( 𝐱 t ) ≜ K 𝐱 t , 𝐱 t − K 𝐱 t , Z ​ K Z , Z − 1 ​ K Z , 𝐱 t ≜ subscript 𝐂 conditional 𝑓 𝐮 subscript 𝐱 𝑡 subscript 𝐾 subscript 𝐱 𝑡 subscript 𝐱 𝑡 subscript 𝐾 subscript 𝐱 𝑡 𝑍 superscript subscript 𝐾 𝑍 𝑍 1 subscript 𝐾 𝑍 subscript 𝐱 𝑡 \mathbf{C}_{f|\mathbf{u}}(\mathbf{x}_{t})\triangleq K_{\mathbf{x}_{t},\mathbf{x}_{t}}-K_{\mathbf{x}_{t},Z}K_{Z,Z}^{-1}K_{Z,\mathbf{x}_{t}} .

	$q\left(X|f\right)$	${𝐟}_t$	${𝐒}_t^{\ast }$	sampling
1) Factorised - linear	$q\left(X\right)$	${𝐱}_t$	${𝐒}_t$	$𝒪\left(T\right)$
2) Factorised - non-linear	$q\left(X\right)$	$K_{{𝐱}_t,Z}{K}_{Z,Z}^{-1}{𝝁}_{𝐮}$	${𝐒}_t+{𝐀}_t{𝐂}_f\left({𝐱}_t\right){𝐀}_t^{\top }$	$𝒪\left(T\right)$
3) Non-Factorised - non-linear	$q\left(X|f\right)$	$f\left({𝐱}_t\right)$	${𝐒}_t$	$𝒪\left(T^3\right)$
4) VCDT	$q\left(X|𝐮\right)$	$K_{{𝐱}_t,Z}{K}_{Z,Z}^{-1}𝐮$	${𝐒}_t+{𝐀}_t{𝐂}_{f|𝐮}\left({𝐱}_t\right){𝐀}_t^{\top }$	$𝒪\left(T\right)$

Table 2. Table 2: Test set performance (NLPP; RMSE) for the system identification datasets. Lower is better.

Model	Actuator	Ballbeam	Drive	Dryer	Gas Furnace
Factorised - linear	-0.364; 0.154	-0.486; 0.075	0.770; 0.439	-0.709; 0.098	-0.296; 0.170
Factorised - non-linear	-0.641; 0.142	-1.379; 0.073	0.283; 0.246	-1.310; 0.049	0.264; 0.168
VCDT	-0.644; 0.141	-1.395; 0.072	0.238; 0.285	-1.282; 0.050	0.377; 0.169
PR-SSM (initialised)	3.653; 1.976	24.765; 1.118	-0.649; 0.139	-1.265; 0.053	0.144; 0.162

Table 3. Table 3: Test set performance for the Cart and Pole dataset. Results averaged over 30 different test sequences (prediction 20 steps into the future). Lower is better.

Model	NLPP	RMSE
Factorised - linear	2.268	3.548
Factorised - non-linear	1.847	2.974
VCDT	0.694	2.139
PR-SSM	-0.049	1.613
PR-SSM (corrected sampling)	-0.053	1.628

Equations

$$f$$

$${\mathbf{x}}_1$$

$${\mathbf{x}}_{t+1}|f,{\mathbf{x}}_t$$

$${\mathbf{y}}_t|{\mathbf{x}}_t$$

$$\displaystyle p(f(Z))\triangleq p(f(\mathbf{z}_{1}),\dots,f(\mathbf{z}_{N}))=\mathcal{N}\mathopen{}\mathclose{{}\left({\boldsymbol{m}_{Z},K_{Z,Z}}}\right)$$

$$\begin{array}{rl}p(X,Y,f(X))& =\\ p(x_1)p(f(& X))\prod _{t=1}^{T-1}p({\mathbf{x}}_{t+1}|f,{\mathbf{x}}_t)\prod _{t=1}^{T}p({\mathbf{y}}_t|{\mathbf{x}}_t).\end{array}$$

$$p(f(\mathbf{x}_{t+1})\,|\,f(\mathbf{x}_{1:t}))=\mathcal{N}\mathopen{}\mathclose{{}\left({f(\mathbf{x}_{t+1})\,|\,\boldsymbol{\mu}_{f_{t+1}},\mathbf{\Sigma}_{f_{t+1}}}}\right)$$

$$f({\mathbf{x}}_{1:t})$$

$${\mathbf{\mu }}_{f_{t+1}}$$

$${\mathbf{\Sigma }}_{f_{t+1}}$$

$$f(\mathbf{x}_{t})\triangleq\mathopen{}\mathclose{{}\left\{f_{d}(\mathbf{x}_{t})}\right\}_{d=1}^{D}\,,\hskip 5.69046ptp(f(X))=\prod_{d=1}^{D}p(f_{d}(X))$$

$$\displaystyle\log p(Y)\geq\int q(X,f)\Bigg{[}\log\prod_{t=1}^{T}p(\mathbf{y}_{t}|\mathbf{x}_{t})+$$

$$\displaystyle\qquad\log\frac{p(\mathbf{x}_{1})p(f)\prod_{t=1}^{T-1}p(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})}{q(X,f)}\Bigg{]}\mathrm{d}f\mathrm{d}X$$

$$\displaystyle=\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{x}_{t})}\mathopen{}\mathclose{{}\left[\log p(\mathbf{y}_{t}\,|\,\mathbf{x}_{t})}\right]-\mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(f)\,\middle\|\,p(f)\,}\right]+$$

$$\displaystyle\qquad\mathbb{E}_{q(X,f)}\mathopen{}\mathclose{{}\left[\log\frac{p(\mathbf{x}_{1})\prod_{t=1}^{T-1}p(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})}{q(X\,|\,f)}}\right]\,.$$

$$q(f^{\ne \mathbf{u}},\mathbf{u},X)=q(X|f^{\ne \mathbf{u}},\mathbf{u})p(f^{\ne \mathbf{u}}|\mathbf{u})q(\mathbf{u}).$$

$$p(X|f,Y)=p({\mathbf{x}}_1|f,Y)\prod _{t=2}^{T}p({\mathbf{x}}_t|f({\mathbf{x}}_{t-1}),{\mathbf{y}}_{t:T}).$$

$$q(X|f)=q({\mathbf{x}}_1)\prod _{t=1}^{T-1}q({\mathbf{x}}_{t+1}|f,{\mathbf{x}}_t)$$

$$\displaystyle\mathcal{L}=\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{x}_{t})}\mathopen{}\mathclose{{}\left[\log p(\mathbf{y}_{t}\,|\,\mathbf{x}_{t})}\right]-\mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(\mathbf{u})\,\middle\|\,p(\mathbf{u})\,}\right]+$$

$$\displaystyle\scalebox{0.75}[1.0]{$-$}\mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(\mathbf{x}_{1})\,\middle\|\,p(\mathbf{x}_{1})\,}\right]$$

$$\displaystyle\scalebox{0.75}[1.0]{$-$}\sum_{t=1}^{T-1}\mathbb{E}_{q(f,\mathbf{x}_{t})}\mathopen{}\mathclose{{}\left[\mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(\mathbf{x}_{t+1}|f,\mathbf{x}_{t})\,\middle\|\,p(\mathbf{x}_{t+1}|f,\mathbf{x}_{t})\,}\right]}\right]\,.$$

$${\mathbf{S}}_t^{\ast }$$

$${\mathbf{A}}_t$$

$${\mathbf{b}}_t$$

$$q(X|f)$$

$$\displaystyle\mathopen{}\mathclose{{}\left[q(\mathbf{x}_{\tau_{1}+1})\prod_{t=\tau_{1}+1}^{\tau_{1}+\tau_{2}-1}q(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})}\right]\dots\,.$$

$$p(f({\mathbf{x}}_1),\dots ,f({\mathbf{x}}_{T-1})|\mathbf{u})=\prod _{t=1}^{T-1}p(f({\mathbf{x}}_t)|\mathbf{u}).$$

$$f(x)=0.8+(x+0.2)(1-\frac{5}{1+\exp (-2x)}).$$

Full text

Gaussian processes, time series, dynamical systems, variational inference, Gaussian process state-space model, recurrent Gaussian process model

1 Introduction

Many time series are well explained by assuming their progression is determined by an underlying dynamical system. A model of this dynamical system can be used to make strong predictions of the future behaviour of the time series. We often require predictions of systems with unknown dynamics, be it for economic models of financial markets, or models of physical systems to be controlled with model predictive control or model-based reinforcement learning. Particularly in low-data regimes, estimates of the uncertainty in the predictions are crucial for making robust decisions, as the expected utility of an action can depend strongly on the distribution of the outcome (von Neumann et al., 1944). A striking example of this is model based policy search (Deisenroth & Rasmussen, 2011).

There are many approaches to time series prediction. Within the machine learning community, autoregressive (AR) (Billings, 2013) and state-space (SSMs) models are popular, in part due to the relative ease with which predictions can be made. AR predictions are obtained by a learned mapping from the $H$ last observations to the next one, whereas SSMs model the underlying dynamical system by learning a transition function which maps a system state forward in time. At each point in time, the state contains sufficient information for predicting both future states and observations. While AR models can be easier to train, SSMs have the potential to be more data efficient, due to their minimal state representation.

We aim to learn complex, possibly stochastic, non-linear dynamical systems from noisy data, following the Bayesian paradigm in order to capture model uncertainty. We use a Gaussian process (GP) prior (Rasmussen & Williams, 2005) for the transition function, giving the Gaussian Process State Space Model (GPSSM) (Frigola et al., 2013). Aside from capturing uncertainty, GPs are non-parametric, guaranteeing that our model complexity will not saturate as we observe more data.

Despite the challenge of performing accurate approximate inference in the GPSSM, an impressive amount of progress has been made (Frigola et al., 2014; McHutchon, 2014; Eleftheriadis et al., 2017; Ialongo et al., 2017; Doerr et al., 2018), helped along by the development of elegant variational inference techniques (Titsias, 2009; Hensman et al., 2013) which retain the key non-parametric property of GPs. In this work, we improve variational approximations by critically examining existing methods and their failure modes. We propose a family of non-factorised variational posteriors which alleviate crucial problems with earlier approaches while maintaining the same efficiency. We refer to the proposed approach as VCDT: variationally coupled dynamics and trajectories.

2 Background

State-space models are common in machine learning, and appear in many forms. At the most basic level, linear-Gaussian state-space models can be learned by maximum likelihood, combining Kalman smoothing and EM for instance (Roweis & Ghahramani, 1999). Extensions have been developed for deterministic non-linear transitions (Ghahramani & Roweis, 1999). Recurrent neural networks, like the popular LSTM (Hochreiter & Schmidhuber, 1996) learn deterministic mappings with state space structure for sequence prediction, which have seen successful recent use in a wide range of tasks including translation (Sutskever et al., 2014).

2.1 Bayesian non-parametric approaches

We are particularly interested in prediction tasks which require good estimates of uncertainty, for example, for use in model-based reinforcement learning systems (Deisenroth & Rasmussen, 2011). These applications distinguish themselves by requiring 1) continual learning, as datasets are incrementally gathered, and 2) uncertainty estimates to ensure that policies learned are robust to the many dynamics that are consistent with a small dataset. Bayesian non-parametric models provide a unified and elegant way of solving both these problems: model complexity scales with the size of the dataset and is controlled by the Bayesian Occam’s razor (Rasmussen & Ghahramani, 2000). In this work, we focus on approximations (Titsias, 2009; Hensman et al., 2013; Matthews et al., 2016) which offer improved computational scaling with the ability to fully recover the original non-parametric model.

2.2 Gaussian process State Space Models

Recurrent models with GP transitions (GPSSM) have been proposed in a variety of ways, each with its own inference method (Wang et al., 2005; Ko & Fox, 2009; Turner et al., 2010; Frigola et al., 2013). Early methods relied on maximum a posteriori inference for either the latent states (Wang et al., 2005; Ko & Fox, 2009) or the transition function (Turner et al., 2010). Frigola et al. (2013) presented the first fully Bayesian treatment with a particle MCMC method that sampled over the latent states and the GP transition function. Due to issues with computational complexity and sampling in higher dimensional latent spaces, later attention turned to variational methods. All variational inference schemes have so far relied on independence assumptions between latent states and transition function (Frigola et al., 2014; McHutchon, 2014; Ialongo et al., 2017), sometimes even factorising the state distribution over time (Mattos et al., 2016). Eleftheriadis et al. (2017) introduced a recognition model to help with optimisation of the variational distribution, while keeping the mean-field assumption. Recently, Doerr et al. (2018) introduced the first method to account for the dependence between transition function and latent states by using a doubly stochastic inference scheme similar to Salimbeni & Deisenroth (2017). However, their approach has severe limitations which we will discuss and show experimentally. Bui et al. (2016) investigated Power Expectation Propagation as an alternative approach for fitting factorised approximate posteriors.

2.3 Variational inference approaches

Variational methods often make simplifying assumptions about the form of the approximate posterior to improve computational tractability. This may result in significantly biased solutions. The bias is particularly severe if independence assumptions are made where strong correlations are actually present (Turner & Sahani, 2011). Clearly, in dynamical systems, the trajectory of the latent states depends strongly on the dynamics. Hence, we focus on improving existing variational inference schemes by removing the independence assumption between latent states and transition function. We identify a general method that performs well across varying noise scales.

3 The Model

3.1 State-space model structure

We model discrete-time sequences of observations Y=\mathopen{}\mathclose{{}\left\{\mathbf{y}_{t}}\right\}_{t=1}^{T}, where $\mathbf{y}_{t}\in\mathbb{R}^{E}$, by positing that each data point is generated by a corresponding latent variable $\mathbf{x}_{t}\in\mathbb{R}^{D}$ (the system’s “state”). These latent variables form a Markov chain, implying that, for any time-step $t$, we can generate $\mathbf{x}_{t+1}$ by conditioning only on $\mathbf{x}_{t}$ and the transition function $f$. We take the system to evolve according to a single, time-invariant transition function, which we would like to infer. While our inference scheme provides the freedom to choose arbitrary transition $p(\mathbf{x}_{t+1}|f,\mathbf{x}_{t})$ and observation $p(\mathbf{y}_{t}|\mathbf{x}_{t})$ density functions, in keeping with previous literature, we use Gaussians. We also assume a linear mapping between $\mathbf{x}_{t}$ and the mean of $\mathbf{y}_{t}\,|\,\mathbf{x}_{t}$. This does not limit the range of systems that can be modelled, though it may require us to choose a higher dimensional latent state (Frigola, 2015). We make this choice to reduce the non-identifiabilities between transitions and emissions. To reduce them further we may also impose constraints on the scale of $f$ or of the linear mapping $\mathbf{C}$. This allows the model to be expressed by the following equations:

[TABLE]

where we take the “process noise” covariance matrix $\mathbf{Q}$ to be diagonal, encouraging the transition function to account for all correlations between latent dimensions. The crucial difference from a standard state-space model is that we have placed a GP prior over the transition function. Thus, for a given mean function $m(\cdot)$ and positive-definite covariance function $k(\cdot,\cdot)$, any finite collection of function evaluations, at arbitrary inputs Z=\mathopen{}\mathclose{{}\left\{\mathbf{z}_{n}}\right\}_{n=1}^{N},\mathbf{z}_{n}\in\mathbb{R}^{D}, will have a Gaussian joint distribution:

[TABLE]

where \mathopen{}\mathclose{{}\left[\boldsymbol{m}_{Z}}\right]_{n,1}=m(\mathbf{z}_{n}) and \mathopen{}\mathclose{{}\left[K_{Z,Z}}\right]_{m,n}=k(\mathbf{z}_{m},\mathbf{z}_{n}). This gives the joint density of our model:

[TABLE]

In order to draw samples, $f(\mathbf{x}_{t+1})$ has to be drawn by conditioning on all the previously sampled function values:

[TABLE]

[TABLE]

By multiplying the Gaussian conditionals, we can return to our GP prior as in eq. 5. Although the function values are jointly Gaussian, the states $X$ are not necessarily so, since the GP’s mean and kernel functions specify non-linear relations between inputs $\mathbf{x}_{t}$ and outputs $f(\mathbf{x}_{t})$.

3.2 Multivariate latent states

If our state has more than one dimension (i.e. $D>1$), we model the system’s transitions $f:\mathbb{R}^{D}\to\mathbb{R}^{D}$ by placing a GP prior on each of the $f_{d}:\mathbb{R}^{D}\to\mathbb{R}$ dimension-specific functions, and we write:

[TABLE]

where each independent GP has its own mean and kernel functions: $m_{d}(\cdot),k_{d}(\cdot,\cdot)$. For ease of notation, we drop subscripts indexing state dimensions $d$. Since the GP’s are independent from each other, each process only has to condition on its own function evaluations. Thus, the multivariate conditional $p(f(\mathbf{x}_{t+1})\,|\,f(\mathbf{x}_{1:t}))$ is a diagonal Gaussian where the mean and variance of dimension $d$ are as in eq. 9 and eq. 10, using as mean function $m_{d}(\cdot)$ and as kernel function $k_{d}(\cdot,\cdot)$.

3.3 Control inputs

A control input $\mathbf{c}_{t}\in\mathbb{R}^{{D_{\mathbf{c}}}}$ influences the transition in a Markovian manner. We can therefore view it as being part of an “augmented” latent state. This makes our transition function $f:\mathbb{R}^{{D_{\mathbf{x}}}+{D_{\mathbf{c}}}}\to\mathbb{R}^{D_{\mathbf{x}}}$, and requires only a change to $f$’s kernel and mean functions to be defined over ${D_{\mathbf{x}}}+{D_{\mathbf{c}}}$ input dimensions. Our transition probability is now also conditioned on $\mathbf{c}_{t}$: $p(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t},\mathbf{c}_{t})$, though, for brevity, we drop $\mathbf{c}_{t}$ from the notation, as it is not a random variable and we always condition on it.

3.4 Computational costs of sampling

Despite the conceptual similarity between the non-parametric GPSSM and parametric non-linear state space models, it is considerably more expensive to sample from the GPSSM prior. The main reason for this can be seen in eqs. 7, 8, 9 and 10: to sample $\mathbf{x}_{T}$ we have to condition on the $T-1$ previous points. This requires incrementally building up a $T\times T$ matrix inverse, which costs \mathcal{O}\mathopen{}\mathclose{{}\left(T^{3}}\right) in time. Since our non-parametric transition function is implicitly defined by the function values $f(\mathbf{x}_{1:t})$ sampled so far, it is necessary to condition on all previous samples to ensure our overall function is self-consistent over the whole time series.

This poses a problem for MCMC based inference methods, which rely on drawing samples from distributions closely related to the prior. Prediction from a GP posterior also has cubic cost. One way to avoid this cost is to sample from an adapted GP prior (e.g. FITC (Snelson & Ghahramani, 2005)), which was suggested by Frigola et al. (2013).

Modifying the model in this way can have unintended consequences (Bauer et al., 2016), so we focus on performing inference without assumptions on the structure of the GP or its kernel. Our inference scheme should approximate the correct model, while avoiding the \mathcal{O}\mathopen{}\mathclose{{}\left(T^{3}}\right) cost at training time.

4 Variational inference

4.1 General variational bound

We are interested in approximating the posterior $p(X,f\,|\,Y)$111We abuse notation by directly denoting a density over the infinite-dimensional $f$, even though such a density does not exist. We only ever concern ourselves with finite subsets of $f$. using variational inference. The general procedure is to construct a variational lower bound $\mathcal{L}$ to the log marginal likelihood that has the KL divergence from an approximation to the true posterior as its slack: \log p(Y)-\mathcal{L}=\mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(X,f)\,\middle\|\,p(X,f\,|\,Y)\,}\right]. By maximising $\mathcal{L}$ w.r.t. $q$, we improve the quality of the approximation.

Practical stochastic variational inference places two requirements on approximating distributions. First, we need to be able to generate samples from them and, second, we need to be able to evaluate their density. We start by deriving the lower bound for a general joint approximate posterior $q(X,f)$, to which we will later add constraints in order to produce different tractability-accuracy trade-offs. Following the usual derivation using Jensen’s inequality (Blei et al., 2017), we obtain

[TABLE]

For $q(f)$ we will use a sparse Gaussian process (Titsias, 2009; Hensman et al., 2013; Matthews, 2016). This constraint is ubiquitous in models where the inputs to a GP are random variables (Damianou et al., 2016), such as the GPLVM (Titsias & Lawrence, 2010), or deep GP (Damianou & Lawrence, 2013), as it provides analytical as well as computational tractability. A sparse GP approximation specifies a free Gaussian density on the function values $\mathbf{u}$ at $M$ locations q(\mathbf{u})=\mathcal{N}\mathopen{}\mathclose{{}\left({\mathbf{u};\boldsymbol{\mu}_{\mathbf{u}},\mathbf{\Sigma}_{\mathbf{u}}}}\right), and uses the prior conditional for the rest of the Gaussian process, which we denote $f^{\neq\mathbf{u}}$. At this time, we let $X$ depend on the entire process $f$ for generality:

[TABLE]

Choosing the posterior in this way results in a finite KL between the approximate posterior and the prior, which becomes \mathrm{KL}\mathopen{}\mathclose{{}\left[\,q(\mathbf{u})\,\middle\|\,p(\mathbf{u})\,}\right] (Matthews et al., 2016).

4.2 Approximate posterior design choices

So far, the only simplifying assumptions have been reducing $q(f^{\neq\mathbf{u}}|\mathbf{u})$ to the prior conditional $p(f^{\neq\mathbf{u}}|\mathbf{u})$ and $q(\mathbf{u})$ to a Gaussian. Now we have to specify the form of $q(X|f^{\neq\mathbf{u}},\mathbf{u})$. Most existing approximate inference schemes focused on factorised posteriors $q(X|f^{\neq\mathbf{u}},\mathbf{u})=q(X)$ to obtain an efficient sampling scheme or even a closed form for $\mathcal{L}$. In particular, Frigola et al. (2014) find the optimal variational $q(X)$ by calculus of variations and sample from it, whereas McHutchon (2014) and Ialongo et al. (2017) constrain it to be Markov-Gaussian. A factorised, Gaussian posterior together with Gaussian emissions leads to a Gaussian optimal $q(\mathbf{u})$, which, for certain kernels, gives us a closed-form bound. Eleftheriadis et al. (2017) reduced the $\mathcal{O}(TD^{2})$ memory cost associated with storing the Markov-Gaussian $q(X)$ and improved optimisation behaviour by employing a recurrent neural network as a recognition model. This also parallels the work of (Krishnan et al., 2017), where recurrent neural networks perform smoothing in order to learn a deterministic model of the dynamics.

In our experiments, we show that imposing independence between $X$ and $f$ does in fact encourage over-confident beliefs about the dynamics, which is a well-known phenomenon in variational inference (Turner & Sahani, 2011). To avoid this we focus on approximate posteriors which preserve correlations between $X$ and $f$ in various forms. We impose one more general constraint: Markovian structure in $X$. This structure is required to eliminate the impractical amount of correlations that can be expressed between the large number of random variables in $X$. Furthermore, the true posterior is also Markovian in $X$:

[TABLE]

Hence the general form of our approximate posterior:

[TABLE]

where we are free to specify $q(\mathbf{x}_{1})$ and $q(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})$. We can thus simplify the ELBO from eq. 13 to

[TABLE]

Now we add two restrictions: 1) we use a Gaussian q(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})=\mathcal{N}\mathopen{}\mathclose{{}\left({\mathbf{x}_{t+1};g_{t}(f,\mathbf{x}_{t}),h_{t}(f,\mathbf{x}_{t})}}\right) to obtain a tractable KL, and 2) we assume $g_{t}$ and $h_{t}$ are linear, and depend only on a finite number of points on $f(\cdot)$, q(\mathbf{x}_{t+1}\,|\,f,\mathbf{x}_{t})=\mathcal{N}\mathopen{}\mathclose{{}\left({\mathbf{x}_{t+1};\mathbf{A}_{t}\mathbf{f}_{t}+\mathbf{b}_{t},\mathbf{S}_{t}^{*}}}\right), to make the expectation over $q(f)$ tractable. We choose $\mathbf{A}_{t}$, $\mathbf{b}_{t}$, and $\mathbf{S}_{t}$ to be free variational parameters in all cases (see table 1 for how $\mathbf{S}_{t}$ and $\mathbf{S}_{t}^{*}$ are related). They are crucial to allow information from the observations to make its way into our latent state posterior. We now discuss various choices for $\mathbf{f}_{t}$ and $\mathbf{S}_{t}^{*}$, and the influence they have on the flexibility of the approximate posterior, which we summarise in table 1. The densities of all these variational distributions can still be evaluated, and exact samples can be drawn for unbiased evaluation of the bound.

4.3 A natural choice

Another way to justify our choice of posterior is that, for Gaussian emissions, the exact posterior “filtering” factors: $p(\mathbf{x}_{t+1}|f,\mathbf{x}_{t},\mathbf{y}_{1:t+1})$ are Gaussians of precisely the form we propose. In fact, we can compute the corresponding $\mathbf{A}_{t}$, $\mathbf{b}_{t}$, and $\mathbf{S}_{t}^{*}$ in closed form:

[TABLE]

Notice that here $\mathbf{A}_{t}$ and $\mathbf{S}_{t}^{*}$ are constant through time, reducing the time and memory footprint of the algorithm. While the exact posterior “smoothing” factors $p(\mathbf{x}_{t+1}|f,\mathbf{x}_{t},\mathbf{y}_{1:T})$ are not generally Gaussian or closed-form (hence leaving $\mathbf{A}_{t},\mathbf{b}_{t},\mathbf{S}_{t}$ free), in settings where the observation noise is small, filtering and smoothing will give very similar results, making the filtering factors very close to optimal. Likewise, if future observations hold little information about the present due to process noise corruption, the filtering factors will provide good approximations to the conditional state posterior.

4.4 Factorised approximations

We can recover the existing Gaussian factorised approximation of McHutchon (2014) by taking $\mathbf{f}_{t}=\mathbf{x}_{t}$ (see 1 in table 1). The linear relationship between consecutive states implies a linearisation of the transition function. Yet, the true $p(\mathbf{x}_{t+1}|f(\mathbf{x}_{t}),\mathbf{y}_{t+1:T})$ will contain influence from $\mathbf{x}_{t}$ only through the non-linear transition function.

We can break the linearisation constraint by correlating $\mathbf{x}_{t+1}$ with $\mathbf{x}_{t}$ “pushed through” the approximate GP mean (2 in table 1). This results in a non-Gaussian $q(X)$ which is still factorised from $q(f)$. Since this corresponds to marginalising out $f$ from $q(X\,|\,f)$, we let $\mathbf{S}_{t}^{*}=\mathbf{S}_{t}+\mathbf{A}_{t}\mathbf{C}_{f}(\mathbf{x}_{t})\mathbf{A}_{t}^{\top}$ to allow our state posterior to be more “uncertain” precisely where our GP posterior tells us to. If the observations provide sufficient evidence to discount the GP’s uncertainty, optimisation will drive $\mathbf{A}_{t}$ and $\mathbf{S}_{t}$ down, giving us a concentrated state posterior.

The main reason we introduce this modified factorised method, is so we can directly isolate the effect of introducing dependence between $X$ and $f$, compared to simply making $q(X)$ non-Gaussian.

Both approximate posteriors can be sampled in \mathcal{O}\mathopen{}\mathclose{{}\left(T}\right) time due to their Markovian structure. The factorised - linear posterior is particularly convenient as it allows us to exploit GPU-optimised banded-matrix operations as in Durrande et al. (2019).

4.5 Direct dependence on $f$

The factorised, non-linear approximation from the previous section is attempting to indirectly incorporate the learned transition function into the approximation over $X$. Instead of summarising the information using the mean of $f$, we want to break the factorisation between $X$ and $f$ by using the actual, stochastic $f$ in the approximation. This is summarised in table 1, line 3. The actual function value $f(\mathbf{x}_{t})$ of a function in the posterior is weighed by $\mathbf{A}_{t}$ against $\mathbf{b}_{t}$.

Sampling from this approximate posterior poses a challenge. From eq. 16, we note that every time a $f(\mathbf{x}_{t})$ is sampled, it has to be consistent with a single function throughout the entire time series. This means that $f(\mathbf{x}_{t})$ has to be sampled by conditioning on the function values sampled for earlier time-steps, resulting in an \mathcal{O}\mathopen{}\mathclose{{}\left(T^{3}}\right) cost (as per section 3.4). This high cost prohibits learning in long time series. We can circumvent this issue by cutting the approximate posterior into subsequences with lengths $\tau_{1},\dots,\tau_{n}$. The state immediately after a subsequence, e.g. $\mathbf{x}_{\tau_{1}+1}$ is given a non-conditional posterior:

[TABLE]

This reduces the cost of sampling to \mathcal{O}\mathopen{}\mathclose{{}\left(\frac{T}{\tau}\tau^{3}}\right) (for $\tau_{1},\dots,\tau_{n}=\tau$), and allows for mini-batching over subsequences. However, this factorisation catastrophically “breaks” long-range dependences, again introducing a potential mismatch between our state and transition function posteriors. In the extreme case of $\tau=1$, we would obtain an approximate posterior which factorises completely across time, as in Mattos et al. (2016).

4.6 VCDT and “U-dependence”

The previous section introduces an approximate posterior exhibiting dependence between $X$ and $f$. This will generally come at a cubic cost since our transition function has potentially as many degrees of freedom as time-steps. Parametric models avoid this cubic cost since they can sample a transition function with fixed degrees of freedom before sampling the entire time series. Conveniently, our sparse variational GP posterior separates the finite degrees of freedom that we can learn from the data, from the rest of the potentially infinite degrees of freedom of our prior.

In order to gain an approximate posterior with a tractable sampling scheme, we choose $q(X|f)=q(X|\mathbf{u})$ ((4) in table 1). Crucially, we can sample from this posterior in linear-time, as only a single $\mathbf{u}$ needs to be sampled from $q(\mathbf{u})$ for an entire sample trajectory. We are effectively being parametric about our sample trajectories while minimising a KL distance from the full non-parametric model. By the properties of sparse variational GPs (Matthews et al., 2016), by adding more inducing points we can reduce this distance.

Intuitively, by finely tiling with inducing points the parts of the state-space we are likely to traverse with our sample paths we can fully specify the behaviour of our transition function in those regions. In this case, conditioning on $\mathbf{u}$ is sufficient to “tie” down the transitions arbitrarily tightly leaving very little to be gained by additionally conditioning on sampled $f$ values. The elegance of this approach lies in the fact that we can assess how many inducing points are required by finding the saturation point where the converged ELBO values stop improving as more points are added.

4.7 Comparison to PR-SSM

Doerr et al. (2018) were, to the best of our knowledge, the first to consider a non-factorised variational posterior for the GPSSM. Their work, however, has two significant shortcomings. Firstly, $q(\mathbf{x}_{t+1}|f,\mathbf{x}_{t})$ is taken to be the same as the prior transition. This leaves no free parameters, except for those in $q(\mathbf{x}_{1})$ and $q(\mathbf{u})$, to incorporate information from the observations into our posterior: no filtering or smoothing is possible. This gives a good approximation only when the process noise is low and/or the observed sequence is short, as low noise levels can compound in a long sequence. Even in the absence of process noise, using the observations to update our beliefs about the states can help optimisation greatly by effectively setting up easier “targets” for the GP to go through. Notice that by setting $\mathbf{A}_{t}=\bf{I}$, $\mathbf{b}_{t}=\bf{0}$, and $\mathbf{S}_{t}=\mathbf{Q}$ in our posterior transitions we can also force them to match the prior, making PR-SSM a special case of our more flexible parameterisation.

Secondly, Doerr et al. (2018) employ a sampling scheme which gives incorrect marginal samples of $q(\mathbf{x}_{t})$, even assuming, as they did, that:

[TABLE]

Despite this conditional factorisation assumption, once $\mathbf{u}$ is integrated out, as in PR-SSM, the $f(\mathbf{x}_{t})$’s become correlated. Thus, in order to correctly sample $f(\mathbf{x}_{t})$, we need to condition on all previous samples $f(\mathbf{x}_{1:t-1})$, introducing a cubic time cost. This is ignored in PR-SSM and samples from the posterior over $f$ are effectively drawn independently for every time-step $t$. The resulting samples from $q(\mathbf{x}_{t})$ are biased. A mismatch is thus introduced between the functional form of the approximate posterior (which has correlations between function values), and the samples we use to compute its expectations. Hence, PR-SSM’s objective does not correspond to a valid variational lower bound.

We solve this by explicitly sampling $\mathbf{u}$ and conditioning on it. Conditioning on the same $\mathbf{u}$ enforces consistency of the dynamics along a sample trajectory. However, we also wish to avoid the factorisation assumption of eq. 22 as that does not generally correspond to a valid GP prior. Explicitly assuming $q(X\,|\,f)=q(X\,|\,\mathbf{u})$ we obtain a valid variational bound where we only need to condition on $\mathbf{u}$ to sample $X$, giving us an \mathcal{O}\mathopen{}\mathclose{{}\left(T}\right) cost while maintaining a full GP prior.

5 Experiments

In the experiments we set out to test how the proposed posterior (VCDT) compares with the Factorised - linear, Factorised - non-linear, and PR-SSM (Doerr et al., 2018) approximations. Each model serves as a “control” for a specific experimental question:

Factorised - linear: what is gained by making our approximation non-linear? I.e. if our posterior $q(X)$ were not jointly Gaussian, but only its Markovian conditional factors. 2. 2.

Factorised - non-linear: what is gained by introducing dependence between our states and our transition function? I.e. if we do not marginalise both $f(\mathbf{x}_{t})$ and $\mathbf{u}$ from our posterior over $X$. 3. 3.

PR-SSM: what is gained, or lost, by forcing our model to find transition functions that explain the data well, without accounting for process noise?

To answer these questions, we test model calibration and predictive performance in a number of environments: the “kink” function, a set of five system identification benchmark datasets (also used in (Doerr et al., 2018)), and a cart and pole system.

5.1 “Kink” Function

We generate data according to the transition function displayed in red in fig. 1:

[TABLE]

This function generates cyclic trajectories through the state-space, as it combines two linear segments with slope greater than $1$ (left) and less than $-1$ (right). It also contains a non-linearity at the joining of these segments, giving it a variety of dynamics and serving as an interesting testing environment. We train our models on the data described in fig. 1, fixing their emissions ($\mathbf{C},\mathbf{d},\mathbf{R}$) to the true, generative ones so as to enable direct latent space comparison222This is to side-step any scale invariances built into our model..

In order to make inference challenging, we add a small amount of process noise (s.d. $0.05$) and a significant amount of observation noise (s.d. $\sqrt{0.8}$). As we can see from fig. 2, in this regime PR-SSM is unable to learn, given that its posterior encourages it to find trajectories with no process noise. Indeed, regardless of which data we are training on, during optimisation, PR-SSM always drives the process noise to zero. This leads it to converge to poor solutions when process noise exists and/or to drive up the observation noise (across all experiments, PR-SSM learned the highest values for the observation noise). Even initialising PR-SSM at the solution found by other methods leads to the same sub-optimal fit of fig. 2. On the other hand, the GPSSM models can handle process noise and recover inflated estimates of it in order to ascribe some transitions to noise and justify smoother dynamics. The Factorised - linear model’s estimates of the process noise are consistently the highest (across many settings of noise parameters and sequence lengths), followed by the Factorised - non-linear model. As can be seen in fig. 3, VCDT finds the most well-calibrated posterior. The factorised approaches (the linear one is shown, though they gave a very similar fit) result in a posterior that is not only less accurate, but also more confidently wrong. In general, the posteriors found by the factorised approaches favour higher process noise to achieve smoother, over-confident dynamics. This is consistent with the behaviour discussed in (Turner & Sahani, 2011): mean-field variational inference favours more concentrated posteriors. By matching our true posterior’s structure more closely and modelling the dependence between $X$ and $f$, we can overcome this issue.

5.2 System identification benchmarks

We also train our models on five benchmark datasets taken from (De Moor et al., 1997). Test set results are reported in table 2. As in (Doerr et al., 2018), we train on the first half of the sequence, we normalise the data, and we use a 4-dimensional latent state (plus 1 non-stochastic dimension for the control inputs). No mini-batching was used, the bound is evaluated using 100 samples from the posterior, and we use 100 variational inducing points. The test results are for predictions 30 steps into the future from the end of the training sequence.

We initialise all models at the solution found by the Factorised - non-linear method to assess differences in optima of the variational objectives. This turned out to be necessary for PR-SSM as learning on such long sequences, without mini-batching, proved very challenging without any filtering from the observations. As can be seen from table 2, PR-SSM can often find good solutions despite its constraints. This is particularly true for low noise regimes, as in the “Drive” dataset. The GPSSM models are more robust in general though, with VCDT in the lead among GPSSMs for NLPP in all but the last two datasets. No error-bars are displayed since the initialisation was deterministic, no mini-batching was used, and 100 samples from the posterior gave low variance estimates of the training objective. To compute the test statistics accurately, we used many more samples ($10^{5}$).

5.3 Cart and pole

Finally, we consider a cart and pole system (or inverted pendulum). We add Gaussian noise to the states in order to obtain our observations. Standard deviations are: 0.03; 0.1047; 0.4; 1.3963; for the respective dimensions of our system: cart position, pendulum angle, cart velocity and angular velocity of the pendulum. The noise s.d. levels correspond to 3cm; 6 degrees; 3cm/0.075sec; 6 degrees/0.075sec. Results for models fit using 100 samples from the posterior (per optimisation step) and 300 inducing points are shown in table 3. Because there is no process noise and the training sequences are short (20 time-steps), we see that PR-SSM performs the best in terms of predictive performance. Of course PR-SSM can be viewed as a special case of our approximation, albeit with a mismatched sampling scheme. Running the experiment with the correct sampling scheme and fixing the variational parameters to: $\mathbf{A}_{t}=\bf{I}$, $\mathbf{b}_{t}=\bf{0}$, and $\mathbf{S}_{t}=\mathbf{Q}$, we learn a solution with the same performance as PR-SSM. In the GPSSM models, VCDT performs the best, and, as can be seen in fig. 4, it also recovers the lowest noise levels, using the dynamics to explain more of the structure in the data.

6 Conclusion

The GPSSM is a powerful formalism. The main challenge is to perform fast, accurate inference. Variational inference maintains many of the benefits of using non-parametric models, but faces particular difficulties in time series models due to their sensitivity to factorisation assumptions. Naive non-factorising approximations regain the poor computational scaling that variational methods were introduced to avoid. By exploiting the low-rank structure of the approximate GP posterior, we were able to construct a non-factorised posterior with the desired computational scaling. This often leads to better predictive performance, calibration, and an improved estimation of model parameters.

Acknowledgements

ADI would like to acknowledge the generous support of the Cambridge-Tübingen and Qualcomm Innovation fellowships.