marginparsep has been altered.
topmargin has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Markovian Gaussian Process Variational Autoencoders

Anonymous Authors1

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preliminary work. Under review by the International Conference on Machine
Learning (ICML). Do not distribute.

###### Abstract

Sequential VAEs have been successfully considered for many high-dimensional
time series modelling problems, with many variant models relying on discrete-
time mechanisms such as recurrent neural networks (RNNs). On the other hand,
continuous-time methods have recently gained attraction, especially in the
context of irregularly-sampled time series, where they can better handle the
data than discrete-time methods. One such class are Gaussian process
variational autoencoders (GPVAEs), where the VAE prior is set as a Gaussian
process (GP). However, a major limitation of GPVAEs is that it inherits the
cubic computational cost as GPs, making it unattractive to practioners. In
this work, we leverage the equivalent discrete state space representation of
Markovian GPs to enable linear time GPVAE training via Kalman filtering and
smoothing. We show on a variety of high-dimensional temporal and
spatiotemporal tasks that our method performs favourably compared to existing
approaches whilst being computationally highly scalable.

##  1 Introduction

Figure 1: Illustration of the MGPVAE model. (Top) The state posterior
$q_{\phi}(\mathbf{s}_{1:T})$ is parameterised by encoder outputs and computed
using filtering and smoothing. (Bottom) At prediction time, posterior
predictive distributions can be calculated at any $t$.

Modelling multivariate time series data has extensive applications in e.g.,
video and audio generation (Li & Mandt, 2018), climate data analysis (Ravuri
et al., 2021) and finance (Sims, 1980). Among existing deep generative models
for time series, a popular class of model is sequential variational auto-
encoders (VAEs) (Chung et al., 2015; Fraccaro et al., 2017; Fortuin et al.,
2020), which extend VAEs (Kingma & Welling, 2014) to sequential data.
Originally proposed for image generation, a VAE is a latent variable model
which encodes data via an encoder to a low-dimensional latent space and then
decodes via a decoder to reconstruct the original data. To extend VAEs to
sequential data, the latent space must also include temporal information.
Sequential VAEs accomplish this by modelling the latent variables as a
multivariate time series, where many existing approaches define a state-space
model which governs the latent dynamics. These state-space model-based
sequential VAE approaches can be classified into two subgroups:

  * •

Discrete-time: The first approach relies on building a discrete-time state-
space model for the latent variables. The transition distribution is often
parameterised by a recurrent neural network (RNN) such as Long Short-Term
Memory (LSTM; Hochreiter & Schmidhuber (1997)) or Gated Recurrent Units (GRU;
Cho et al. (2014)). Notable methods include the Variational Recurrent Neural
Network (VRNN; Chung et al. (2015)) and Kalman VAE (KVAE; Fraccaro et al.
(2017)). However, these approaches may suffer from training issues, such as
vanishing gradients (Pascanu et al., 2013), and may struggle with irregularly-
sampled time series data (Rubanova et al., 2019).

  * •

Continuous-time: The second approach involves continuous-time representations,
where the latent space is modelled using a continuous-time dynamic model. A
notable class of such methods are neural differential equations (Chen et al.,
2018; Rubanova et al., 2019; Li et al., 2020; Kidger, 2022), which model the
latent variables using a system of differential equations, described by its
initial conditions, drift and diffusion. As remarked in Li et al. (2020), one
can construct a neural stochastic differential equation (neuralSDE) that can
be interpreted as an infinite-dimensional noise VAE. However, although these
models can flexibly handle irregularly-sampled data, they also require
numerical solvers to solve the underlying latent processes, which may cause
training difficulties (Park et al., 2021), memory issues (Chen et al., 2018;
Li et al., 2020) and slow computation times.

Another line of continuous-time approaches, related to neuralSDEs, that treat
the latent multivariate time series as a random function of time, and model
the random function as a tractable stochastic process, are Gaussian Process
Variational Autoencoders (GPVAEs; Casale et al. (2018); Pearce (2020); Fortuin
et al. (2020); Ashman et al. (2020); Jazbec et al. (2021)) which model the
latent variables using Gaussian processes (GPs) (Rasmussen, 2003). As compared
with dynamic model-based approaches which focus on modelling the latent
variable transitions (reflecting local properties mainly), the GP model for
the latent variables describes, in a better way, the global properties of the
time series, such as smoothness and periodicity. Therefore GPVAEs may be
better suited for e.g., climate time series data which clearly exhibits
periodic behaviour. Unfortunately GPVAEs are not directly applicable to long
sequences as they suffer from $\mathcal{O}(T^{3})$ computational cost,
therefore approximations need to be made. Indeed, Ashman et al. (2020);
Fortuin et al. (2020); Jazbec et al. (2021) proposed approximations based on
sparse Gaussian processes (Titsias, 2009; Hensman et al., 2013; 2015) or
variational inference (Fortuin et al., 2020) to improve the scalability of
GPVAEs.

In this work we propose Markovian GPVAEs (MGPVAEs) to bridge state-space
model-based and stochastic process based approaches of sequential VAEs, aiming
to achieve the best in both worlds. Our approach is inspired by the key fact
that, when the GP is over time, a large class of GPs can be written as a
linear SDE (Särkkä & Solin, 2019), for which there exists exact and unique
solutions (Øksendal, 2003). As a result, there exists an equivalent discrete
linear state space representation of GPs. Therefore the dynamic model for the
latent variables has both discrete and continuous-time representations. This
brings the following key advantages to the latent dynamic model of an MGPVAE:

  * •

The continuous-time representation allows the incorporation of inductive
biases via the GP kernel design (e.g., smoothness, periodic and monotonic
trends), to achieve better prediction results and training efficiency. It also
enables modelling irregularly sampled time series data.

  * •

The equivalent discrete-time representation, which is linear, enables Kalman
filtering and smoothing (Särkkä & Solin, 2019; Adam et al., 2020; Chang et
al., 2020; Wilkinson et al., 2020; 2021; Hamelijnck et al., 2021) that
computes the posterior distributions in $\mathcal{O}(T)$ time. As the observed
data is assumed to come from non-linear transformations of the latent
variables, we further apply site-based approximations (Chang et al., 2020) for
the non-linear likelihood terms to enable analytic solutions for the filtering
and smoothing procedures.

In our experiments, We study much longer datasets ($T\approx 100)$ compared to
many previous GPVAE and discrete-time works, which are only are of the
magnitude of $T\approx 10$. We include a range of datasets that describe
different properties of MGPVAE compared to existing approaches:

  * •

We deliver competitive performance compared to many existing methods on
corrupt and irregularly-sampled video and robot action data at a fraction of
the cost of many existing models.

  * •

We extend our work to spatiotemporal climate data, where none of the discrete-
time sequential VAEs are suited for modelling. We show that it outperforms
traditional GP and existing sparse GPVAE models in terms of both predictive
performance and speed.

##  2 Background

Consider building generative models for high-dimensional time series (e.g.,
video data). Here an observed sequence of length $T$ is denoted as
$\mathbf{Y}_{t_{1}}\ldots,\mathbf{Y}_{t_{T}}\in\mathbb{R}^{D_{y}}$, where
$t_{i}$ represents the timestamp of the $i$th observation in the sequence.
Note that in general $t_{i}\neq i$ for irregularly sampled time series. As the
proposed MGPVAE has both discrete state-space model based and stochastic
process based formulations, below we introduce these two types of sequential
VAEs and the key relevant techniques.

#### Sequential VAEs with state-space models:

Consider $t_{i}=i$ w.l.o.g., and assume each of the latent states in
$\mathbf{Z}_{1:T}=(\mathbf{z}_{1:T}^{1},\ldots,\mathbf{z}_{1:T}^{L})\in\mathbb{R}^{T\times
L}$ has $L$ latent dimensions. Then the generative model is defined as

$\displaystyle
p(\mathbf{Y}_{1:T},\mathbf{Z}_{1:T})=p(\mathbf{Z}_{1:T})\prod_{t=1}^{T}p(\mathbf{Y}_{t}|\mathbf{Z}_{t}),$
(1)

where we choose
$p(\mathbf{Y}_{t}|\mathbf{Z}_{t})=\mathcal{N}(\mathbf{Y}_{t};\varphi(\mathbf{Z}_{t}),\sigma^{2}\mathbf{I})$
to be a multivariate Gaussian distribution, and
$\varphi:\mathbb{R}^{L}\rightarrow\mathbb{R}^{D_{y}}$ is decoder network that
transforms the latent state to the Gaussian mean. The prior
$p(\mathbf{Z}_{1:T})$ is defined by the transition probabilities, e.g.,
$p(\mathbf{Z}_{1:T})=\prod_{t=1}^{T}p(\mathbf{Z}_{t}|\mathbf{Z}_{<t})$.
Training is done by maximising the variational lower bound $\mathcal{L}$
(Ranganath et al., 2014):

$\textstyle\log p(\mathbf{Y}_{1:T})\geq$
$\textstyle\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{Z}_{1:T})}[\log
p(\mathbf{Y}_{t}|\mathbf{Z}_{t})]$
$\textstyle-\text{KL}(q(\mathbf{Z}_{1:T}|\mathbf{Y}_{1:T})||p(\mathbf{Z}_{1:T})):=\mathcal{L},$
(2)

where $q(\mathbf{Z}_{1:T}|\mathbf{Y}_{1:T})$ is the approximate posterior
parameterised by an encoder network. Often this $q$ distribution is defined by
mean-field approximation over time, i.e.,
$q(\mathbf{Z}_{1:T}|\mathbf{Y}_{1:T})=\prod_{t=1}^{T}q(\mathbf{Z}_{t}|\mathbf{Y}_{t})$,
or by using transition probabilities as well, e.g.,
$q(\mathbf{Z}_{1:T}|\mathbf{Y}_{1:T})=\prod_{t=1}^{T}q(\mathbf{Z}_{t}|\mathbf{Z}_{t-1},\mathbf{Y}_{\leq
t})$. Below we also write
$q(\mathbf{Z}_{1:T})=q(\mathbf{Z}_{1:T}|\mathbf{Y}_{1:T})$ to simplify
notation.

#### Gaussian Process Variational Autoencoders:

A Gaussian process is a stochastic process denoted by
$\mathbf{Z}\sim\mathcal{GP}(0,k)$, where
$k:\mathbb{R}\times\mathbb{R}\rightarrow\mathbb{R}$ is the kernel or
covariance function that specifies the similarity between two time stamps.
This allows us to explicitly enforce inductive biases or global behaviour. For
instance, if $\mathbf{Y}_{t}$ is a periodic system, then we may use a periodic
kernel $k$ or a longer initial kernel lengthscale to incorporate this
knowledge in the model; if the underlying process is smooth, then a kernel can
also be chosen so that it induces smooth functions (Kanagawa et al., 2018).

GPVAEs (Casale et al., 2018; Pearce, 2020; Fortuin et al., 2020; Ashman et
al., 2020; Jazbec et al., 2021) define the decoder network $\varphi$ and the
conditional distribution $p(\mathbf{Y}_{t}|\mathbf{Z}_{t})$ in the same way as
presented above. However, instead of using transition probabilities, a GPVAE
places a multi-output GP prior on the latent variables $\\{\mathbf{Z}_{t}\\}$:
$\mathbf{Z}_{t}\sim\mathcal{GP}(0,\mathbf{k})$, where one can choose the
kernel of the multi-output GP to be $\mathbf{k}=\mathbf{I}\otimes k$, i.e.,
the output GPs across dimensions share the same kernel $k$. However, each
dimension may also be induced with separate kernels $k^{1},\ldots,k^{L}$,
which we adopt in this work, giving block diagonal kernel matrices.

Again we use the variational lower-bound (Eq. (2) when $t_{i}=i$) as the
training objective, but with a different approximate posterior $q$. Some
examples include GPVAE (Pearce, 2020; Fortuin et al., 2020)

$\textstyle
q(\mathbf{Z}_{1:T})=\prod_{l=1}^{L}N(\mathbf{Z}^{l}_{1:T};\tilde{\mathbf{Y}}_{1:T}^{l},\tilde{\mathbf{V}}_{1:T}^{l}),$

where
$(\tilde{\mathbf{Y}}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})_{l=1}^{L}=(\mu_{\phi}^{l}(\mathbf{Y}_{t}),\mathbf{\Sigma}^{l}_{\phi}(\mathbf{Y}_{t}))_{l=1}^{L}=\phi(\mathbf{Y}_{t})$
are outputs of the encoder network $\phi$. This corresponds to a mean-field
approximation over latent dimensions instead of time. To avoid direct
parameterisation of the full covariance matrix
$\mathbf{\Sigma}^{l}_{\phi}(\mathbf{Y}_{1:T})\in\mathbb{R}^{T\times T}$ which
can be expensive for long sequences, Fortuin et al. (2020) proposed a banded
parameterisation of the precision matrix (Blei & Lafferty, 2006; Bamler &
Mandt, 2017), reducing both the time and memory complexity to
$\mathcal{O}(T)$. However, this choice makes it more difficult to work with
irregularly-sampled data and previous works only focused on corrupt video
frames.

Another more flexible option is to use sparse GP approximations with inducing
points (Jazbec et al., 2021):

$\textstyle q(\mathbf{Z}_{1:T})$
$\textstyle=\prod_{l=1}^{L}p(\mathbf{Z}^{l}_{1:T}|\mathbf{F}^{l}_{m})q(\mathbf{F}^{l}_{m}),$
$\textstyle q(\mathbf{F}^{l}_{m})$
$\textstyle=N(\mathbf{F}^{l}_{m}|\mathbf{m}^{l}_{m},\mathbf{A}_{m}^{l}),$
$\textstyle\mathbf{S}^{l}$
$\textstyle=\mathbf{K}_{mm}^{l}+\mathbf{K}_{mT}^{l}\text{diag}(\tilde{\mathbf{V}}^{l}_{1:T})^{-1}\mathbf{K}_{Tm}^{l},$
$\textstyle\mathbf{m}^{l}_{m}$
$\textstyle=\mathbf{K}_{mm}^{l}(\mathbf{S}^{l})^{-1}\mathbf{K}_{mT}^{l}\text{diag}(\tilde{\mathbf{V}}^{l}_{1:T})^{-1}\tilde{\mathbf{Y}}^{l}_{1:T},$
$\textstyle\mathbf{A}^{l}_{m}$
$\textstyle=\mathbf{K}_{mm}^{l}(\mathbf{S}^{l})^{-1}\mathbf{K}_{mm}^{l},$

where $p(\mathbf{Z}^{l}_{1:T}|\mathbf{F}_{m}^{l})$ is the standard
multivariate Gaussian conditional distribution and
$[\mathbf{K}_{mT}^{l}]_{ij}:=k^{l}(\mathbf{U}_{i},j)$ with $i=1,\ldots,m$ and
$j=1,\ldots,T$, and
$[\mathbf{K}_{mm}^{l}]_{ij}:=k^{l}(\mathbf{U}_{i},\mathbf{U}_{j})$ with
$i,j=1,\ldots,m$, for pre-determined inducing time locations
$\mathbf{U}=[\mathbf{U}_{1},\ldots,\mathbf{U}_{m}]^{\intercal}$. However, the
time complexity scales with $\mathcal{O}(m^{3}+m^{2}T)$, where in practice to
attain good performance $m=\mathcal{O}(\log T)$ (Burt et al., 2019), and
therefore the complexity increases massively when dealing with longer time
series.

#### Markovian Gaussian Processes:

Interestingly, the banded parameterisation coincides with the structure of the
Markovian GP state space $\mathbf{s}_{t}$. w.l.o.g. suppose $\mathbf{Z}_{t}$
is one-dimensional in this subsection, then with a conjugate likelihood (e.g.
linear Gaussian) $p(\mathbf{Y}_{t}|\mathbf{Z}_{t})$ and a Markovian kernel $k$
(Särkkä & Solin, 2019), we can write the GP regression problem as an Itô SDE
of latent dimension $d$

$\textstyle\text{d}\mathbf{s}_{t}=\mathbf{F}\mathbf{s}_{t}\text{d}t+\mathbf{L}\text{d}B_{t},\quad\mathbf{Z}_{t}=\mathbf{H}\mathbf{s}_{t},$
$\textstyle\mathbf{Y}_{t}|\mathbf{Z}_{t}\sim
p(\mathbf{Y}_{t}|\mathbf{Z}_{t}),$ (3)

where $\mathbf{F}\in\mathbb{R}^{d\times d},\mathbf{L}\in\mathbb{R}^{d\times
e},\mathbf{H}\in\mathbb{R}^{1\times d}$ are the feedback, noise effect and
emission matrices, respectively, and $B_{t}$ is an $e$-dimensional
(correlated) Brownian motion with diffusion $\mathbf{Q}_{c}$.
$\mathbf{s}_{0}\sim\mathcal{N}(0,\mathbf{P}_{\infty})$, where
$\mathbf{P}_{\infty}$ is the stationary state covariance, which satisfies the
Lyapunov equation (Solin, 2016). The state is typically the $d$ derivatives
$\mathbf{s}_{t}=(\mathbf{Z}_{t},\mathbf{Z}^{(1)}_{t},\ldots,\mathbf{Z}_{t}^{(d-1)})^{\intercal}$
with the subsequent emission matrix $\mathbf{H}=(1,0,\ldots,0)$.

The linear SDE in Eq. (2) admits a unique closed form solution, allowing the
recursive updates of $\mathbf{s}_{t_{i+1}}$ given $\mathbf{s}_{t_{i}}$:

$\textstyle\mathbf{s}_{t_{i+1}}=\mathbf{A}_{i,i+1}\mathbf{s}_{t_{i}}+\mathbf{q}_{i},\quad\mathbf{q}_{i}\sim\mathcal{N}(0,\mathbf{Q}_{i,i+1}),$
$\textstyle\mathbf{Z}_{t_{i}}=\mathbf{H}\mathbf{s}_{t_{i}},\quad\mathbf{Y}_{t_{i}}|\mathbf{Z}_{t_{i}}\sim
p(\mathbf{Y}_{t_{i}}|\mathbf{Z}_{t_{i}})$ (4) $\textstyle\text{with
}\mathbf{A}_{i,i+1}=e^{\Delta_{i}\mathbf{F}},\quad\text{where
}\Delta_{i}=t_{i+1}-t_{i},$
$\textstyle\mathbf{Q}_{i,i+1}=\int_{t_{0}}^{\Delta_{i}+t_{0}}e^{(\Delta_{i}+t_{0}-\tau)\mathbf{F}}\mathbf{L}\mathbf{Q}_{c}\mathbf{L}^{\intercal}[e^{(\Delta_{i}+t_{0}-\tau)\mathbf{F}}]^{\intercal}\text{d}\tau.$

Note that the $\mathbf{Q}_{i,i+1}$ can be easily obtained in closed form. See
Appendix A for a detailed explanation. For conjugate likelihood
$p(\mathbf{Y}_{t}|\mathbf{Z}_{t})\equiv p(\mathbf{Y}_{t}|\mathbf{s}_{t})$, we
can use the recursive Kalman filtering and smoothing equations (also known as
the forward-backward equations) to obtain the posterior distribution (Särkkä &
Solin, 2019) $p(\mathbf{s}_{t}|\mathbf{Y}_{1:T})$ with $\mathcal{O}(Td^{3})$
complexity. The corresponding filter prediction, filtering and smoothing
equations are:

$\textstyle p(\mathbf{s}_{t+1}|\mathbf{Y}_{1:t})$ $\textstyle=\int
p(\mathbf{s}_{t+1}|\mathbf{s}_{t})p(\mathbf{s}_{t}|\mathbf{Y}_{1:t})\text{d}\mathbf{s}_{t},$
(5) $\textstyle p(\mathbf{s}_{t}|\mathbf{Y}_{1:t})$
$\textstyle=\frac{1}{\ell_{t}}p(\mathbf{Y}_{t}|\mathbf{s}_{t})p(\mathbf{s}_{t}|\mathbf{Y}_{1:t-1}),$
$\textstyle p(\mathbf{s}_{t}|\mathbf{Y}_{1:T})$
$\textstyle=p(\mathbf{s}_{t}|\mathbf{Y}_{1:t})\int\frac{p(\mathbf{s}_{t+1}|\mathbf{s}_{t})p(\mathbf{s}_{t+1}|\mathbf{Y}_{1:T})}{p(\mathbf{s}_{t+1}|\mathbf{Y}_{1:t})}\text{d}\mathbf{s}_{t+1},$

where $\ell_{t}=\int
p(\mathbf{Y}_{t}|\mathbf{s}_{t})p(\mathbf{s}_{t}|\mathbf{Y}_{1:t-1})\text{d}\mathbf{s}_{t}$
is tractable as the integrand is a product of Gaussians. If
$p(\mathbf{Y}_{t}|\mathbf{s}_{t})$ is non-conjugate (e.g. a Poisson
distribution), then the posterior cannot be obtained analytically.

The size of $d$ depends on the kernel. For example, the
Matern-$\nicefrac{{3}}{{2}}$ kernel yields $d=2$, since the GP sample paths
lie in an RKHS that is norm equivalent to a space of functions with 1
derivative (Kanagawa et al., 2018). The periodic or quasi-periodic kernels
(Solin & Särkkä, 2014) may yield larger $d$’s. However, in comparison to
sparse Gaussian process approximations in Ashman et al. (2020); Jazbec et al.
(2021) that have complexity $\mathcal{O}(m^{3}+m^{2}T)$, where $m$ is the
number of inducing points, $d$ does not depend on $T$ (unlike sparse GPs (Burt
et al., 2019) that depend on $\mathcal{O}(\log^{D}T)$, where $D$ is the time
variable dimension) and thus does not need to grow as $T$ increases.

##  3 Markovian Gaussian Process Variational Autoencoders

In this section, we propose Markovian GPVAEs (MGPVAEs) and a corresponding
variational inference scheme for model learning.

###  3.1 Model

Let $L$ be the dimensionality of the latent variables and $k^{l}$ be the
kernel for the $l$th channel with state dimension $d_{l}$. Then we have a
total state space dimension of $\sum_{l=1}^{L}d_{l}$. Let
$\varphi:\mathbb{R}^{L}\rightarrow\mathbb{R}^{D_{y}}$ be the decoder network.
Then the generative model, under the linear SDE form of Markovian GP in the
latent space, is (with
$\mathbf{Z}_{t}=(\mathbf{Z}_{t}^{1},\ldots,\mathbf{Z}_{t}^{L})^{\intercal}$)

$\textstyle\text{d}\mathbf{s}^{l}_{t}=\mathbf{F}^{l}\mathbf{s}^{l}_{t}\text{d}t+\mathbf{L}^{l}\text{d}B_{t}^{l},\quad\mathbf{Z}_{t}^{l}=\mathbf{H}^{l}\mathbf{s}^{l}_{t},\quad
l=1,\ldots,L$ $\textstyle\mathbf{Y}_{t}|\mathbf{Z}_{t}\sim
p(\mathbf{Y}_{t}|\mathbf{Z}_{t})\equiv
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t})),$ (6)

which equivalently becomes the linear discrete state space model with
nonlinear likelihood

$\textstyle\mathbf{s}_{t_{i+1}}^{l}=\mathbf{A}^{l}_{i,i+1}\mathbf{s}^{l}_{t_{i}}+\mathbf{q}^{l}_{i},\quad\mathbf{q}^{l}_{i}\sim\mathcal{N}(0,\mathbf{Q}^{l}_{i,i+1}),$
$\textstyle\mathbf{Z}_{t}^{l}=\mathbf{H}^{l}\mathbf{s}^{l}_{t},\quad\mathbf{Y}_{t}|\mathbf{Z}_{t}\sim
p(\mathbf{Y}_{t}|\mathbf{Z}_{t})\equiv
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t})).$ (7)

Note that the transformation
$\mathbf{Z}_{t}^{l}=\mathbf{H}^{l}\mathbf{s}_{t}^{l}$ is deterministic, and
the stochasticity arises from the $\mathbf{s}_{t}$ variables.

###  3.2 Variational inference

Suppose $q(\mathbf{s})$ is the approximate posterior over
$\mathbf{s}:=\mathbf{s}_{1:T}\in\mathbb{R}^{T\times Ld}$. Then minimising
$\text{KL}(q(\mathbf{s})||p(\mathbf{s}|\mathbf{Y}))$ is equivalent to
maximising the lower bound
$\mathcal{L}:=\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{s}_{t})}\log
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t}))-\text{KL}(q(\mathbf{s})||p(\mathbf{s}))$.
Note that since $\mathbf{Z}_{t}$ is a deterministic transformation of
$\mathbf{s}_{t}$, the KL-divergence is between the posterior and prior
distributions of $\mathbf{s}_{t}$. We wish to compute $q(\mathbf{s})$ in
linear time using Kalman filtering and smoothing. However, due to the presence
of a nonlinear decoder network $\varphi$ in the likelihood, it is no longer
possible to obtain the exact posterior $p(\mathbf{s}|\mathbf{Y})$ due to non-
conjugacy. However, if we approximate the likelihood
$p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t}))$ with Gaussian sites, as is done in
Pearce (2020); Ashman et al. (2020); Jazbec et al. (2021); Chang et al.
(2020), Kalman filtering and smoothing can be performed as conjugacy is
reintroduced in the filtering and smoothing equations.

We propose the Gaussian-site approximation

$\textstyle q(\mathbf{s})\propto
p(\mathbf{s})\prod_{l=1}^{L}\prod_{t=1}^{T}N(\tilde{\mathbf{Y}}_{t}^{l}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t}),$

where $\tilde{\mathbf{Y}}^{l}_{t}\in\mathbb{R}$ and
$\tilde{\mathbf{V}}^{l}_{t}\in\mathbb{R}$ for $l=1,\ldots,L$. Instead of
optimising $\tilde{\mathbf{Y}}_{t}$ and $\tilde{\mathbf{V}}_{t}$ as free-form
parameters using conjugate-computation variational inference (Khan & Lin,
2017), we encode them using outputs of an encoder network $\phi$ i.e.
$(\tilde{\mathbf{Y}}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})_{l=1}^{L}=\phi(\mathbf{Y}_{t})$.

Then with straightforward computations,

$\textstyle\text{KL}(q(\mathbf{s})||p(\mathbf{s}))=\mathbb{E}_{q(\mathbf{s})}\log\frac{q(\mathbf{s})}{p(\mathbf{s})}$
$\textstyle=\mathbb{E}_{q(\mathbf{s})}\log\frac{p(\mathbf{s})\prod_{l=1}^{L}\prod_{t=1}^{T}N(\tilde{\mathbf{Y}}_{t}^{l}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})}{p(\mathbf{s})\int
p(\mathbf{s})\prod_{l=1}^{L}\prod_{t=1}^{T}N(\tilde{\mathbf{Y}}_{t}^{l}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})\text{d}\mathbf{s}}$
$\textstyle=\mathbb{E}_{q(\mathbf{s})}\sum_{l=1}^{l}\sum_{t=1}^{T}\log
N(\tilde{\mathbf{Y}}_{t}^{l}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})$
$\textstyle-\log\mathbb{E}_{p(\mathbf{s})}\prod_{t=1}^{T}\prod_{l=1}^{L}N(\tilde{\mathbf{Y}}_{t}^{l}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t}).$

The ELBO (2) thus becomes

$\displaystyle\mathcal{L}=\underbrace{\log\mathbb{E}_{p(\mathbf{s})}\left[\prod_{t=1}^{T}\prod_{l=1}^{L}N(\tilde{\mathbf{Y}}^{l}_{t}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})\right]}_{\text{E3}}$
$\displaystyle+\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{s}_{t})}\bigg{[}\underbrace{\log
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t}))}_{\text{E2}}-\sum_{l=1}^{L}\underbrace{\log
N(\tilde{\mathbf{Y}}^{l}_{t}|\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t})}_{\text{E1}}\bigg{]}.$

(E1) $q(\mathbf{s}_{t})$ can be obtained using the Kalman smoothing
distributions and can be computed in linear time (see the Kalman filtering and
smoothing equations in Appendix 1). The reason is because $q(\mathbf{s})$ is
constructed by replacing the likelihood in Eq. (2) with Gaussian sites, making
it possible to analytically evaluate the filtering and smoothing equations
(5). Therefore using the same calculations as Chang et al. (2020); Hamelijnck
et al. (2021), the first term in the ELBO can be computed analytically:

$\textstyle\mathbb{E}_{q(\mathbf{s}_{t})}[\text{E1}]=-\frac{1}{2}\log|2\pi\tilde{\mathbf{V}}^{l}_{t}|-\frac{1}{2}(\tilde{\mathbf{Y}}^{l}_{t})^{\intercal}(\tilde{\mathbf{V}}^{l}_{t})^{-1}\tilde{\mathbf{Y}}_{t}^{l}$
$\textstyle+(\tilde{\mathbf{Y}}^{l}_{t})^{\intercal}\mathbf{H}^{l}\mathbf{m}_{t}^{l,s}-\frac{1}{2}[\text{Tr}((\tilde{\mathbf{V}}^{l})^{-1}\mathbf{H}^{l}\mathbf{P}_{t}^{l,s}(\mathbf{H}^{l})^{\intercal})]$
$\textstyle+(\mathbf{m}_{t}^{l,s})^{\intercal}(\mathbf{H}^{l})^{\intercal}\mathbf{H}^{l}\mathbf{m}_{t}^{l,s},$

where $\mathbf{m}_{t}^{l,s}$ and $\mathbf{P}_{t}^{l,s}$ are the smoothing mean
and covariances respectively at time $t$.

(E2) is intractable but we can estimate it using Monte Carlo with $K$ samples
$\mathbf{s}_{t,j}\sim q(\mathbf{s}_{t})$, and thus samples
$\mathbf{Z}_{t,j}=(\mathbf{H}^{1}\mathbf{s}_{t,j}^{1},\ldots,\mathbf{H}^{L}\mathbf{s}_{t,j}^{L})^{\intercal}$:

$\textstyle\mathbb{E}_{q(\mathbf{s}_{t})}\log
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t}))\approx\frac{1}{K}\sum_{j=1}^{K}\log
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t,j})).$

(E3) is the log partition function of $q(\mathbf{s})$, which is also the log
marginal likelihood of the approximate model $\sum_{l=1}^{L}\log
p(\tilde{\mathbf{Y}}^{l})$. Note that the latent channels are independent of
each other, allowing us to sum over the log marginal likelihood over each
channel. We can further decompose each term of the sum into

$\textstyle\log p(\tilde{\mathbf{Y}}^{l})$ $\textstyle=\log
p(\tilde{\mathbf{Y}}_{1}^{l})\prod_{t=2}^{T}p(\tilde{\mathbf{Y}}^{l}_{t}|\tilde{\mathbf{Y}}^{l}_{1:t-1})$
$\textstyle=\sum_{t=1}^{T}\log\mathbb{E}_{p(\mathbf{s}_{t}^{l}|\tilde{\mathbf{Y}}^{l}_{1:t-1})}N(\tilde{\mathbf{Y}}^{l}_{t};\mathbf{H}^{l}\mathbf{s}^{l}_{t},\tilde{\mathbf{V}}^{l}_{t}),$

where $p(\mathbf{s}_{t}^{l}|\tilde{\mathbf{Y}}^{l}_{1:t-1})$ is the predictive
filter distribution obtained with Kalman filtering. Fortunately, $\log
p(\tilde{\mathbf{Y}}^{l})$ can be computed during the filtering stage (see
Algorithm 1 for a full breakdown of Kalman filtering and smoothing). In
addition, a graphical representation of the Markovian GPVAE is in Figure 1.

###  3.3 Spatiotemporal Modelling

Spatiotemporal modelling is an important task with many real world
applications (Cressie, 2015). Traditional methods such as kriging, or Gaussian
process regression, incurs cubic computational costs and are even more costly
and difficult when multiple variables need to be modelled jointly. GPVAEs may
ameliorate this issue by effectively simplifying the task via an encoder-
decoder model and has been proven to be effective in Ashman et al. (2020).
Following section 4.2 of Hamelijnck et al. (2021), given a separable
spatiotemporal kernel
$k(r,t,r^{\prime},t^{\prime})=k_{r}(r,r^{\prime})k_{t}(t,t^{\prime})$, it is
straightforward to extend MGPVAE to model spatiotemporal data, which will be
make it a highly scalable spatiotemporal model. We consider the model:

$\textstyle\mathbf{Z}(r,t)\sim\mathcal{GP}(0,k),\mathbf{Y}(r,t)|\mathbf{Z}(r,t)\sim
p(\mathbf{Y}(r,t)|\varphi(\mathbf{Z}(r,t))),$

where for classical kriging (Cressie, 2015), $\varphi$ is the identity map and
$\mathbf{Z}(r,t)$ is of the same dimensionality as $\mathbf{Y}(r,t)$.

For convenience of notation, we avoid introducing subscripts $l$ and only
write down 1 latent dimension with $k$. Suppose we have $N_{s}$ spatial
coordinates observed over time, denoted by the spatial matrix
$\mathbf{R}\in\mathbb{R}^{N_{r}\times D_{x}}$, it is possible to rewrite the
GP regression model by stacking the states for each spatial location on top of
each other to get:

$\textstyle\mathbf{s}_{t_{i+1}}=\mathbf{A}_{i,i+1}\mathbf{s}_{t_{i}}+\mathbf{q}_{i},\quad\mathbf{q}_{i}\sim\mathcal{N}(0,\mathbf{Q}_{i,i+1}),$
$\textstyle\mathbf{Y}_{t}|\mathbf{Z}_{t}\sim
p(\mathbf{Y}_{t}|\varphi(\mathbf{Z}_{t})),$ (8)

where
$\mathbf{s}_{t}=[\mathbf{s}_{t}(r_{1}),\ldots,\mathbf{s}_{t}(r_{N_{s}})]^{\intercal}$,
$\mathbf{Z}_{t}=[\mathbf{L}_{\mathbf{R}\mathbf{R}}^{r}\otimes\mathbf{H}^{t}]\mathbf{s}_{t_{i}}$
and $\mathbf{A}_{i,i+1}=\mathbf{I}_{N_{r}}\otimes\mathbf{A}_{i,i+1}^{t}$,
$\mathbf{Q}_{i,i+1}=\mathbf{I}_{N_{r}}\otimes\mathbf{Q}_{i,i+1}^{t}$,
$\mathbf{K}^{r}_{\mathbf{R}\mathbf{R}}=\mathbf{L}_{\mathbf{R}\mathbf{R}}^{r}(\mathbf{L}_{\mathbf{R}\mathbf{R}}^{r})^{\intercal}$,
where superscripts $t$ and $r$ indicate the temporal state space and spatial
kernel matrices respectively. The graphical model for MGPVAE is shown in
Figure 2, demonstrating how the states for each spatial location are
independently filtered and smoothed over time, and then spatially mixed by the
emission matrix. Lastly, we approximate the likelihood with a mean-field
amortised approximation
$\prod_{t=1}^{T}N(\tilde{\mathbf{Y}}_{t}|[\mathbf{L}_{\mathbf{R}\mathbf{R}}^{r}\otimes\mathbf{H}^{t}]\mathbf{s}_{t},\tilde{\mathbf{V}}_{t})$,
where $\tilde{\mathbf{Y}}_{t},\tilde{\mathbf{V}}_{t}\in\mathbb{R}^{N_{x}}$.
See Appendix A.3 for further details.

Figure 2: Graphical model of the separable spatiotemporal MGPVAE model. The
temporal dynamics of the states $\mathbf{s}_{t}(r)$ at each location $r$ are
independently handled. At each time step $t$, these states are spatially mixed
to produce $\mathbf{Z}_{t}$, which is then transformed by a non-linear mapping
to $\mathbf{Y}_{t}$.

###  3.4 Computational Complexity

For simplicity let us assume that each channel has the same kernel but is
modelled independently. Then the computational complexity of MGPVAE is
$\mathcal{O}(Ld^{3}T)$. For GPVAE, it is also $\mathcal{O}(Ld^{3}T)$; for
SVGPVAE, $\mathcal{O}(L(Tm^{2}+m^{3}))$ with $m$ inducing points. For KVAE and
VRNN, the complexity is $\mathcal{O}(LT)$, but there may be large big-O
constants due to the RNN network sizes. For neuralODEs and neuralSDEs, the
complexity is linear with respect to the number of discretisation steps, which
can potentially be much larger than $T$.

For spatiotemporal modelling, the computational complexity for MGPVAE will be
$\mathcal{O}(Ld^{3}TN_{t}^{3})$ in this case, but can be further lowered to
$\mathcal{O}(Ld^{3}T(N_{t}m^{2}+m^{3}))$ if we sparsify the spatial domain by
using $M$ spatial inducing points. For SVGPVAE, it is
$\mathcal{O}(L(TN_{t}m^{2}+m^{3}))$ with $m$ inducing points over space and
time.

##  4 Related Work

GPVAEs have been explored for time series modelling in Fortuin et al. (2020);
Ashman et al. (2020). In our work, we experiment with the state-of-the-art
GPVAEs of Fortuin et al. (2020) and SVGPVAE (Jazbec et al., 2021), and explore
a wider variety of tasks. Unlike Fortuin et al. (2020), MGPVAE is capable of
tackling both corrupt and missing frames imputation tasks, as well as
spatiotemporal modelling tasks, with great scalability. Compared to sparse
GPVAEs (Ashman et al., 2020; Jazbec et al., 2021), MGPVAE does not require
inducing points.

NeuralODEs (Chen et al., 2018; Rubanova et al., 2019) and neuralSDEs (Li et
al., 2020) are also continuous-time models that can tackle the same tasks as
MGPVAE. However, they depend heavily on the time discretisation, how well the
initial conditions are learned and the expressiveness of the drift and
diffusion functions. In our work, we experiment with neuralODEs, which have
previously been used for similar missing frames imputation tasks, and find
that it is the slowest model whithout achieving good predictive performance.

Many discrete-time and continuous-time models, such as VRNN (Chung et al.,
2015), KVAE (Fraccaro et al., 2017) and latentODE (Rubanova et al., 2019), are
only designed to model temporal, but not spatiotemporal, data. On the other
hand, classical multioutput GPs can only handle lower-dimensional
spatiotemporal datasets as there cannot be any dimensionality reduction to a
latent space, whereas GPVAE enables the encoder-decoder networks to learn
meaningful low-dimensional representations for high-dimensional data. SVGPVAE
is able to handle spatiotemporal data, but has the disadvantage of being less
efficient than MGPVAE due to the use of inducing points over space and time
jointly.

Chang et al. (2020); Hamelijnck et al. (2021) considered modelling non-
conjugate likelihoods with Gaussian approximations, which would allow for
Kalman filtering and smoothing operations. We adopt this strategy to allow
flexible decoder choices with non-linear mapping, where a key difference is
that the encoder helps us construct a low-dimensional approximation to the
likelihood function, which allows us to work with Kalman filtering and
smoothing in the lower-dimensional latent space.

##  5 Experiments

We present 3 sets of experiments: rotating MNIST, Mujoco action data and
spatiotemporal data modelling. We benchmark MGPVAE against a variety of
continuous and discrete time models, such as GPVAE (Fortuin et al., 2020),
SVGPVAE (Jazbec et al., 2021), KVAE (Fraccaro et al., 2017), VRNN (Chung et
al., 2015), LatentODE (Rubanova et al., 2019) and sparse variational
multioutput GP (MOGP). We evaluate the performances using both test negative
log-likelihood (NLL) and root mean squared error (RMSE). See further
experimental setup details and results in Appendix B.

###  5.1 Rotating MNIST

In this experiment, we produce sequences of MNIST frames in which the digits
are rotated with a periodic length of 50, over $T=100$ frames. We tackle 2
imputation tasks for: corrupted frames where the frame pixels are randomly set
to 0, and missing frames where frames are randomly dropped out of each
sequence. Each task has 4000 /1000 train/test sequences respectively. The
underlying dynamics are simple (rotation) which may favour models with strong
inductive biases, such as GPVAE and MGPVAE. To test this, for both models, we
use Matern-$\nicefrac{{3}}{{2}}$ kernels with lengthscale 40.

Figure 3: (Left) Corrupt frames imputation results for an unseen sequence of
5’s. (Right) Missing frames imputation results for an unseen sequence of 9’s.
Missing frames are red frames.

Table 1: Test NLL and RMSE for both the corrupt (Cor) and missing frames (Mis)
imputation tasks.

Model NLL-Cor $(\downarrow)$ RMSE-Cor $(\downarrow)$ Time-Cor (s/epoch
$\downarrow$) NLL-Mis $(\downarrow)$ RMSE-Mis $(\downarrow)$ Time-Mis (s/epoch
$\downarrow$) VRNN $9898\pm 162.0$ $0.1768\pm 0.001563$ 63.51 $16240\pm 2090$
$0.1796\pm 0.008002$ 103.6 KVAE $12500\pm 83.13$ $0.2025\pm 0.0006077$ 139.2
$10730\pm 1232$ $0.1582\pm 0.008688$ 149.0 GPVAE $9026\pm 48.70$
$\mathbf{0.1340\pm 0.0004529}$ 48.93 NA NA NA MGPVAE $\mathbf{8556\pm 69.66}$
$\mathbf{0.1468\pm 0.0006738}$ 50.45 $\mathbf{8925\pm 53.40}$
$\mathbf{0.1508\pm 0.0005190}$ 59.43

#### Corrupt frames imputation:

This is a task that highly suits standard RNN-based models such as VRNN and
KVAE, since the frames are observed at regular time steps. Even so, we see
from Table 1 that both GPVAE and MGPVAE perform significantly better than VRNN
and KVAE in both NLL and RMSE, validating our hypothesis of inductive biases
helping the learning dynamics. Furthermore, we observe that the RMSE for
MGPVAE is worse than GPVAE, although it has a slightly better NLL. However,
the left panel of Figure 3 shows that the images generated do not visually
differ significantly, which implies that the model performance is comparable.

#### Missing frames imputation:

Rubanova et al. (2019) showed that RNN-based models fail in imputing
irregularly sampled time-series, as they struggle to correctly update the
hidden states at time steps of unobserved frames. This is confirmed by our
results in Table 1: MGPVAE outperforms both VRNN and KVAE in terms of NLL and
RMSE.111We omit GPVAE here as the implementation by Fortuin et al. (2020)
cannot efficiently handle missing frames in batches. In the right panel of
Figure 3, we illustrate the posterior mean imputations, and again VRNN fails.
These results are expected since VRNN implements filtering (only includes past
observations), while MGPVAE and KVAE include a smoothing step (includes both
past and future observations).

###  5.2 Mujoco Action Data

The Mujoco dataset is a physical simulation dataset generated using the
Deepmind Control Suite (Tunyasuvunakool et al., 2020). We obtained the
modified data from Rubanova et al. (2019), which has 14 dimensions, and for
all models we use 15-dimensional latent dimensions. We modify the task so that
we only train on the observed time steps, whereas in Rubanova et al. (2019)
the models have access to data at all the time steps. This makes the task
harder as the model has less information to work with during training. Similar
to the previous experiment, we have 4000 training sequences and 1000 test
sequences. Compared to rotating MNIST, the underlying nonlinear dynamics are
more complex, and each dimension can behave differently.

#### Results:

We see from Table 2 that the performance of VRNN, KVAE and latentODE are
significantly worse than GP-based models, and overall MGPVAE achieves the best
NLL and RMSE. This is because missing data imputation is a difficult task for
the discrete-time RNN-based models. On the other hand, as ODE sample paths are
deterministic once the initial conditions are fixed, latentODE significantly
underestimates the predictive uncertainty caused by heterogeneity between
sequences. Figure 4 confirms the results; indeed the non-GP models struggle to
simultaneously fit the data and estimate uncertainty well.

In terms of time complexity, VRNN, KVAE and SVGPVAE are comparable to each
other. LatentODE is significantly slower than all the other models due to the
need for time discretisation. The time complexity for SVGPVAE is dependent on
the number of inducing points, which we choose to be $m=20$, which is already
a modest number and yet is still slower than MGPVAE. For longer sequences,
this is even more undesirable as we would have to choose more inducing points.
In comparison, MGPVAE does not require any inducing points and is the most
time-efficient model.

Figure 4: 95$\%$ posterior credible intervals for unseen mujoco sequences in
its 6,7 and 8th dimensions. The red dots show observed data. Note that some
LatentODE predictions are not showing as they fall outside the limits.

Table 2: Imputation results for the Mujoco task.

Model NLL-Mis $(\downarrow)$ RMSE-Mis $(\downarrow)$ Time (s/epoch
$\downarrow$) VRNN $-1230\pm 120.3$ $0.09347\pm 0.01449$ 133.8 KVAE $-148.3\pm
62.22$ $0.1595\pm 0.01142$ 110.0 LatentODE $101.9\pm 3.776$ $0.06519\pm
0.001403$ 440.1 SVGPVAE-20 $-1766\pm 81.55$ $0.05582\pm 0.004030$ 156.1 MGPVAE
$\mathbf{-2129\pm 40.48}$ $\mathbf{0.03662\pm 0.001939}$ $\mathbf{68.19}$

###  5.3 Spatiotemporal Climate Data

We obtained climate data, including temperature and precipitation, from ERA5
using Google Earth Engine API (Gorelick et al., 2017). The task is to
condition on observed data $\\{\mathbf{Y}(r,t)\\}_{r,t}$ (temperature and air
pressure), and predict on unknown spatial locations
$\\{\mathbf{Y}(r_{*},t)\\}_{r_{*},t}$ for all time steps. Here we focus on GP-
based models since they, unlike RNN-based models, can flexibly handle
spatiotemporal data by combining spatial and temporal kernels to efficiently
model correlated structures. In particular, for GPVAE models we use a
separable kernel $k(r,t,r^{\prime},t)=k_{r}(r,r^{\prime})k_{t}(t,t^{\prime})$
for each latent channel and the spatial kernel $k_{r}$ is shared across
channels. We also consider Gaussian process prediction which is also known as
kriging in spatial statistics (Cressie, 2015). Sparse approximation is needed
for scalability for the MOGP baseline due to computational feasibility. We
emphasise that the dimensionality of this problem, which is 8, is high
relative to traditional spatiotemporal modelling problems (Cressie, 2015).

Figure 5: 95$\%$ posterior credible intervals for climate variables at an
unseen spatial location.  Figure 6: Posterior means over space for climate
variables over 5 time steps.

#### Results:

We report the corresponding quantitative results in Table 3 and visualise the
prediction results at an unseen spatial location in Figure 5. We observe that
MOGP underfits and overestimates the uncertainty, which is expected as
traditional GP regression models struggle to fit a complicated multi-
dimensional time series with a single GP function. In comparison, the encoder-
decoder networks allow GPVAEs to learn simpler dynamics, resulting in better
performance. The conclusions are similar when we fix the time steps and plot
the corresponding posterior mean over space in Figure 6, where GPVAE models
better predict the spatial patterns.

SVGPVAE underperforms MGPVAE even when pushing the number of inducing points
to the memory limit of our hardware (500). Interestingly, SVGPVAE
underestimates the uncertainty more as we increase the number of latent
channels to 8, as can also be seen from the large negative log-likelihood per
dim (NLPD) values. SVGPVAE results can be improved if having well-placed and
sufficiently-many inducing points, but this also means SVGPVAE is memory
inefficient when compared with MGPVAE. Similarly for computational time,
MGPVAE is faster than the other models with more inducing points, but slower
than the ones with less of inducing points (which also perform worse).

Table 3: ERA5 prediction results.

Model NLPD $(\downarrow)$ RMSE $(\downarrow)$ Time $(\text{s/epoch
}\downarrow)$ MOGP-100 $10.38\pm 0.1783$ $1119\pm 21.29$ 0.1220 MOGP-1000
$10.38\pm 0.1784$ $1119\pm 21.13$ 0.6001 SVGPVAE-4-100 $14.29\pm 0.6817$
$433.1\pm 15.58$ 0.07786 SVGPVAE-4-500 $6.543\pm 0.5982$ $399.0\pm 18.23$
1.007 SVGPVAE-8-100 $1624\pm 670.7$ $388.8\pm 18.68$ 0.1508 SVGPVAE-8-500
$2349\pm 1819$ $\mathbf{306.2\pm 22.53}$ 1.913 MGPVAE-4 $2.454\pm 0.1149$
$402.1\pm 15.24$ 0.45876 MGPVAE-8 $\mathbf{-0.3070\pm 0.1021}$
$\mathbf{352.7\pm 24.86}$ 0.5128

##  6 Conclusion and Discussion

We propose MGPVAE, a GPVAE model using Markovian Gaussian processes that uses
Kalman filtering and smoothing with linear time complexity. This is achieved
by approximating the non-Gaussian likelihoods using Gaussian sites. Compared
to Fortuin et al. (2020), which also achieves linear time complexity by using
an approximate covariance structure, our model leaves the original Gaussian
process covariance structure intact, and additionally work with both
irregularly-sampled time series and spatiotemporal data. Experiments on video,
Mujoco action and climate modelling tasks show that our method is both
competitive and scalable, compared to modern discrete and continuous-time
models. Future work can explore the use of parallel filtering (Särkkä &
García-Fernández, 2020), other nonlinear filtering approaches and forecasting
applications.

#### Limitations:

MGPVAE requires the use of kernels that admit a Markovian decomposition, and
therefore kernels such that the squared exponential kernel will not be
permissible (Särkkä & Solin, 2019) without additional approximations. However,
we argue that most commonly used kernels are indeed Markovian, which should be
sufficient for most applications. In addition, we need to use Gaussian site
approximations for the likelihood, in order to allow for Kalman filtering and
smoothing. This may lower the approximation accuracy, though in our
experiments we see that MGPVAE still performs strongly.

## References

  * Adam et al. (2020) Adam, V., Eleftheriadis, S., Artemev, A., Durrande, N., and Hensman, J.  Doubly sparse variational Gaussian processes.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 2874–2884. PMLR, 2020.
  * Ashman et al. (2020) Ashman, M., So, J., Tebbutt, W., Fortuin, V., Pearce, M., and Turner, R. E.  Sparse Gaussian process variational autoencoders.  _arXiv preprint arXiv:2010.10177_ , 2020.
  * Bamler & Mandt (2017) Bamler, R. and Mandt, S.  Dynamic word embeddings.  In _International conference on Machine learning_ , pp. 380–389. PMLR, 2017.
  * Blei & Lafferty (2006) Blei, D. M. and Lafferty, J. D.  Dynamic topic models.  In _Proceedings of the 23rd international conference on Machine learning_ , pp. 113–120, 2006.
  * Burt et al. (2019) Burt, D., Rasmussen, C. E., and Van Der Wilk, M.  Rates of convergence for sparse variational gaussian process regression.  In _International Conference on Machine Learning_ , pp. 862–871. PMLR, 2019.
  * Casale et al. (2018) Casale, F. P., Dalca, A., Saglietti, L., Listgarten, J., and Fusi, N.  Gaussian process prior variational autoencoders.  _Advances in neural information processing systems_ , 31, 2018.
  * Chang et al. (2020) Chang, P. E., Wilkinson, W. J., Khan, M. E., and Solin, A.  Fast variational learning in state-space Gaussian process models.  In _2020 IEEE 30th International Workshop on Machine Learning for Signal Processing (MLSP)_ , pp. 1–6. IEEE, 2020.
  * Chen et al. (2018) Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K.  Neural ordinary differential equations.  _Advances in neural information processing systems_ , 31, 2018.
  * Cho et al. (2014) Cho, K., van Merriënboer, B., Bahdanau, D., and Bengio, Y.  On the properties of neural machine translation: Encoder–decoder approaches.  In _Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation_ , pp. 103–111, Doha, Qatar, October 2014. Association for Computational Linguistics.  doi: 10.3115/v1/W14-4012.  URL https://aclanthology.org/W14-4012.
  * Chung et al. (2015) Chung, J., Kastner, K., Dinh, L., Goel, K., Courville, A. C., and Bengio, Y.  A recurrent latent variable model for sequential data.  _Advances in neural information processing systems_ , 28, 2015.
  * Cressie (2015) Cressie, N.  _Statistics for spatial data_.  John Wiley & Sons, 2015.
  * Evans (2006) Evans, L. C.  An introduction to stochastic differential equations version 1.2.  _Lecture Notes, UC Berkeley_ , 2006.
  * Fortuin et al. (2020) Fortuin, V., Baranchuk, D., Rätsch, G., and Mandt, S.  Gp-vae: Deep probabilistic time series imputation.  In _International conference on artificial intelligence and statistics_ , pp. 1651–1661. PMLR, 2020.
  * Fraccaro et al. (2017) Fraccaro, M., Kamronn, S., Paquet, U., and Winther, O.  A disentangled recognition and nonlinear dynamics model for unsupervised learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Gorelick et al. (2017) Gorelick, N., Hancher, M., Dixon, M., Ilyushchenko, S., Thau, D., and Moore, R.  Google earth engine: Planetary-scale geospatial analysis for everyone.  _Remote Sensing of Environment_ , 2017.  doi: 10.1016/j.rse.2017.06.031.  URL https://doi.org/10.1016/j.rse.2017.06.031.
  * Hamelijnck et al. (2021) Hamelijnck, O., Wilkinson, W., Loppi, N., Solin, A., and Damoulas, T.  Spatio-temporal variational Gaussian processes.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Hensman et al. (2013) Hensman, J., Fusi, N., and Lawrence, N. D.  Gaussian processes for big data.  _UAI_ , 2013.
  * Hensman et al. (2015) Hensman, J., Matthews, A., and Ghahramani, Z.  Scalable variational gaussian process classification.  In _Artificial Intelligence and Statistics_ , pp. 351–360. PMLR, 2015.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Jazbec et al. (2021) Jazbec, M., Ashman, M., Fortuin, V., Pearce, M., Mandt, S., and Rätsch, G.  Scalable gaussian process variational autoencoders.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3511–3519. PMLR, 2021.
  * Kanagawa et al. (2018) Kanagawa, M., Hennig, P., Sejdinovic, D., and Sriperumbudur, B. K.  Gaussian Processes and Kernel Methods: A Review on Connections and Equivalences.  _arXiv:1807.02582 [cs, stat]_ , July 2018.  URL http://arxiv.org/abs/1807.02582.  arXiv: 1807.02582.
  * Khan & Lin (2017) Khan, M. and Lin, W.  Conjugate-computation variational inference: Converting variational inference in non-conjugate models to inferences in conjugate models.  In _Artificial Intelligence and Statistics_ , pp. 878–887. PMLR, 2017.
  * Kidger (2022) Kidger, P.  On neural differential equations.  _PhD Thesis_ , 2022.
  * Kingma & Welling (2014) Kingma, D. P. and Welling, M.  Auto-Encoding Variational Bayes.  _arXiv:1312.6114 [cs, stat]_ , May 2014.  URL http://arxiv.org/abs/1312.6114.  arXiv: 1312.6114.
  * Li et al. (2020) Li, X., Wong, T.-K. L., Chen, R. T., and Duvenaud, D.  Scalable gradients for stochastic differential equations.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3870–3882. PMLR, 2020.
  * Li & Mandt (2018) Li, Y. and Mandt, S.  Disentangled sequential autoencoder.  _ICML_ , 2018.
  * Park et al. (2021) Park, S., Kim, K., Lee, J., Choo, J., Lee, J., Kim, S., and Choi, E.  Vid-ode: Continuous-time video generation with neural ordinary differential equation.  _arXiv preprint arXiv:2010.08188_ , pp. online, 2021.
  * Pascanu et al. (2013) Pascanu, R., Mikolov, T., and Bengio, Y.  On the difficulty of training recurrent neural networks.  In _International conference on machine learning_ , pp. 1310–1318. PMLR, 2013.
  * Pearce (2020) Pearce, M.  The gaussian process prior vae for interpretable latent dynamics from pixels.  In _Symposium on Advances in Approximate Bayesian Inference_ , pp. 1–12. PMLR, 2020.
  * Ranganath et al. (2014) Ranganath, R., Gerrish, S., and Blei, D.  Black box variational inference.  In _Artificial intelligence and statistics_ , pp. 814–822. PMLR, 2014.
  * Rasmussen (2003) Rasmussen, C. E.  Gaussian processes in machine learning.  In _Summer school on machine learning_ , pp. 63–71. Springer, 2003\.
  * Ravuri et al. (2021) Ravuri, S., Lenc, K., Willson, M., Kangin, D., Lam, R., Mirowski, P., Fitzsimons, M., Athanassiadou, M., Kashem, S., Madge, S., et al.  Skilful precipitation nowcasting using deep generative models of radar.  _Nature_ , 597(7878):672–677, 2021.
  * Rubanova et al. (2019) Rubanova, Y., Chen, R. T., and Duvenaud, D. K.  Latent ordinary differential equations for irregularly-sampled time series.  _Advances in neural information processing systems_ , 32, 2019.
  * Särkkä & García-Fernández (2020) Särkkä, S. and García-Fernández, Á. F.  Temporal parallelization of bayesian smoothers.  _IEEE Transactions on Automatic Control_ , 66(1):299–306, 2020.
  * Särkkä & Solin (2019) Särkkä, S. and Solin, A.  _Applied stochastic differential equations_ , volume 10.  Cambridge University Press, 2019.
  * Sims (1980) Sims, C. A.  Macroeconomics and reality.  _Econometrica: journal of the Econometric Society_ , pp. 1–48, 1980\.
  * Solin (2016) Solin, A.  Stochastic differential equation methods for spatio-temporal Gaussian process regression.  2016\.  PhD Thesis: Aalto University.
  * Solin & Särkkä (2014) Solin, A. and Särkkä, S.  Explicit link between periodic covariance functions and state space models.  In _AISTATS_ , 2014.
  * Tebbutt et al. (2021) Tebbutt, W., Solin, A., and Turner, R. E.  Combining pseudo-point and state space approximations for sum-separable Gaussian processes.  In _Uncertainty in Artificial Intelligence_ , pp. 1607–1617. PMLR, 2021.
  * Titsias (2009) Titsias, M.  Variational learning of inducing variables in sparse gaussian processes.  In _Artificial intelligence and statistics_ , pp. 567–574. PMLR, 2009.
  * Tunyasuvunakool et al. (2020) Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y.  dmcontrol: Software and tasks for continuous control.  _Software Impacts_ , 6:100022, 2020.  ISSN 2665-9638.  doi: https://doi.org/10.1016/j.simpa.2020.100022.  URL https://www.sciencedirect.com/science/article/pii/S2665963820300099.
  * Wilkinson et al. (2020) Wilkinson, W., Chang, P., Andersen, M., and Solin, A.  State Space Expectation Propagation: Efficient Inference Schemes for Temporal Gaussian Processes.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 10270–10281. PMLR, July 2020.  URL https://proceedings.mlr.press/v119/wilkinson20a.html.
  * Wilkinson et al. (2021) Wilkinson, W. J., Solin, A., and Adam, V.  Sparse Algorithms for Markovian Gaussian Processes.  _arXiv:2103.10710 [cs, stat]_ , June 2021.  URL http://arxiv.org/abs/2103.10710.  arXiv: 2103.10710.
  * Øksendal (2003) Øksendal, B.  Stochastic differential equations.  In _Stochastic differential equations_ , pp. 65–84. Springer, 2003\.

##  Appendix A Markovian Gaussian Process Background

In this section, we provide a rigorous treatment of the derivation of the
state space Markovian Gaussian process equations as previously presented in
Särkkä & Solin (2019) with stochastic differential equation technicalities
from Øksendal (2003); Evans (2006).

###  A.1 Stochastic Differential Equations

######  Definition A.1 (Univariate Brownian Motion; (Øksendal, 2003; Evans,
2006)).

Let $(\Omega,\mathcal{F},\mathbb{P})$ be a probability space.
$B:[0,T]\times\Omega\rightarrow\mathbb{R}$ is a univariate Brownian motion if:

  1. 1.

$B_{0}=0$, almost surely.

  2. 2.

$t\mapsto B(t,\cdot)$ is continuous almost surely.

  3. 3.

For $t_{4}>t_{3}\geq t_{2}>t_{1}$, $B_{t_{4}}-B_{t_{3}}\perp\\!\\!\\!\perp
B_{t_{2}}-B_{t_{1}}$.

  4. 4.

For any $\delta>0$, $B_{t+\delta}-B_{t}\sim\mathcal{N}(0,\delta q)$, where $q$
the correlation factor.

In a similar manner, we may extend univariate Brownian motions into correlated
multi-dimensional forms.

######  Definition A.2 (Multi-dimensional Brownian Motion).

$B_{t}$ is an $e$-dimensional (correlated) Brownian motion: each $B^{k}$, for
$k=1,\ldots,e$, is a univariate Brownian motion. We assume that each $B^{k}$
is correlated, resulting in
$B_{t+\delta}-B_{t}\sim\mathcal{N}(0,\delta\mathbf{Q}_{c})$, where
$\mathbf{Q}_{c}\in\mathbb{R}^{e\times e}$ is the spectral density matrix.

Define $\mathcal{F}_{t}$ to be the $\sigma$-algebra generated by
$B_{s}(\cdot)$ for $s\leq t$. Intuitively, this is the ‘history of
information’ up to time $t$ created by $B_{s}$ for $s\leq t$. Therefore for a
stochastic process driven by $B_{s}$ up to time $t$, we would like it to be
measurable at all times $t$, which motivates the next definition of
measurability for stochastic processes:

######  Definition A.3.

Let {$\mathcal{F}_{t}\\}_{t\geq 0}$ be an increasing family of
$\sigma$-algebras generated by $B_{t}$ i.e. for $s<t$,
$\mathcal{F}_{s}\subset\mathcal{F}_{t}$. Then
$v(t,\omega):[0,\infty)\times\Omega\rightarrow\mathbb{R}^{e}$ is
$\mathcal{F}_{t}$-adapted if $v(t,\omega)$ is measurable in $\mathcal{F}_{t}$
for all $t\geq 0$.

With the correlated formulation, we can rederive Itô isometry using the same
proof as in the uncorrelated case (Øksendal, 2003; Evans, 2006). We use the
definition of the multi-dimensional Itô integral in Definition 3.3.1 Øksendal
(2003). Let $\mathbb{E}$ be the expectation with respect to $\mathbb{P}$.

######  Definition A.4 (Multi-dimensional Itô Integral (Øksendal, 2003)).

Let $\mathcal{V}^{d\times e}(S,T)$, for $d\in\mathbb{N}$ and $S<T$, be a set
of matrix-valued stochastic processes $v(t,\omega)\in\mathbb{R}^{d\times e}$
where each entry $v_{ij}(t,\omega)$ satisfies:

  * •

$(t,\omega)\rightarrow v_{ij}(t,\omega)$ is
$\mathcal{B}(\mathbb{R})\times\mathcal{F}$-measurable, where
$\mathcal{B}(\mathbb{R})$ is the Borel $\sigma$-algebra on $[0,\infty)$.

  * •

$v_{ij}(\tau,\omega)\in L^{2}([S,T])$ in expectation i.e.
$\mathbb{E}[\int_{S}^{T}v_{ij}^{2}(\tau,\cdot)\text{d}\tau]$.

  * •

Let $W_{t}$ be a univariate Brownian motion. There exists an increasing family
of $\sigma$-algebras {$\mathcal{H}_{t}\\}_{t\geq 0}$ such that (1) $W_{t}$ is
a martingale with respect to $\mathcal{H}_{t}$ and (2) $v_{ij}(t,\omega)$ is
$\mathcal{H}_{t}$-adapted.

Then, we can define the multi-dimensional Itô integral as follows: For all
$v\in\mathcal{V}^{d\times e}(S,T)$,

$\displaystyle\int_{S}^{T}v(\tau,\omega)\text{d}B_{\tau}=\int_{S}^{T}\begin{pmatrix}v_{11}(\tau,\omega)&\cdots&v_{1e}(\tau,\omega)\\\
\vdots&\ddots&\vdots\\\ v_{d1}(\tau,\omega)&\cdots&v_{de}(\tau,\omega)\\\
\end{pmatrix}\begin{pmatrix}\text{d}B^{1}_{\tau}\\\ \vdots\\\
\text{d}B^{e}_{\tau}\\\ \end{pmatrix},$

where
$\bigg{[}\int_{S}^{T}v(\tau,\omega)\text{d}B_{\tau}\bigg{]}_{i}=\sum_{j=1}^{e}\int_{S}^{T}v_{ij}(\tau,\omega)dB^{j}_{\tau}$.

######  Proposition A.5 (Multi-dimensional Itô Isometry).

For $\mathbf{F},\mathbf{G}\in\mathcal{V}^{d\times e}$ (as defined in Øksendal
(2003)) and $S<T$, then

$\textstyle\mathbb{E}[\int_{S}^{T}\mathbf{F}(\tau,\cdot)\text{d}B_{\tau}][\int_{S}^{T}\mathbf{G}(\tau,\cdot)\text{d}B_{\tau}]^{\intercal}=\mathbb{E}\int_{S}^{T}\mathbf{F}(\tau,\cdot)\mathbf{Q}_{c}\mathbf{G}(\tau,\cdot)^{\intercal}\text{d}\tau.$

Furthermore, if $\mathbf{F}$ and $\mathbf{G}$ are deterministic, then for
$t_{1},t_{2}\geq t_{0}$,

$\textstyle\mathbb{E}[\int_{t_{0}}^{t_{1}}\mathbf{F}(\tau)\text{d}B_{\tau}][\int_{t_{0}}^{t_{2}}\mathbf{G}(\tau)\text{d}B_{\tau}]^{\intercal}=\mathbb{E}\int_{t_{0}}^{\min(t_{1},t_{2})}\mathbf{F}(\tau)\mathbf{Q}_{c}\mathbf{G}(\tau)^{\intercal}\text{d}\tau.$

###### Proof.

Let $\\{e_{i}\\}_{i=1}^{e}$ be the canonical basis in $\mathbb{R}^{e}$. We
first show that $ik$-th component of
$\bigg{[}\mathbb{E}[\int_{S}^{T}\mathbf{F}(\tau,\cdot)\text{d}B_{\tau}][\int_{S}^{T}\mathbf{G}(\tau,\cdot)\text{d}B_{\tau}]^{\intercal}\bigg{]}\in\mathbb{R}^{d\times
d}$, for $S<T$, is equal to

$\textstyle\int_{S}^{T}\mathbf{F}_{i}^{\intercal}(\tau,\cdot)\mathbf{Q}_{c}[\mathbf{G}(\tau,\cdot)]_{k}^{\intercal}\text{d}\tau$
$\textstyle=\mathbb{E}\int_{S}^{T}e_{i}^{\intercal}\mathbf{F}^{\intercal}(\tau,\cdot)\mathbf{Q}_{c}[\mathbf{G}(\tau,\cdot)]^{\intercal}_{k}\text{d}\tau$
$\textstyle=e_{i}^{\intercal}\bigg{[}\mathbb{E}\int_{S}^{T}\mathbf{F}^{\intercal}(\tau,\cdot)\mathbf{Q}_{c}\mathbf{G}(\tau,\cdot)^{\intercal}\text{d}\tau\bigg{]}e_{k}$
$\textstyle=\bigg{[}\mathbb{E}\int_{S}^{T}\mathbf{F}^{\intercal}(\tau,\cdot)\mathbf{Q}_{c}\mathbf{G}(\tau,\cdot)^{\intercal}\text{d}\tau\bigg{]}_{ik},$

which would complete the proof. Indeed, the $ik$-th component is equal to

$\textstyle\mathbb{E}\bigg{[}\sum_{j=1}^{e}\int_{S}^{T}\mathbf{F}_{ij}(\tau,\cdot)\text{d}B_{\tau}^{j}\bigg{]}\bigg{[}\sum_{l=1}^{e}\int_{S}^{T}\mathbf{G}_{kl}(\tau,\cdot)\text{d}B_{\tau}^{l}\bigg{]}$
$\textstyle=\bigg{[}\sum_{j,l=1}^{e}\mathbb{E}\int_{S}^{T}\mathbf{F}_{ij}(\tau,\cdot)\text{d}B_{\tau}^{j}\int_{S}^{T}\mathbf{G}_{kl}(\tau,\cdot)\text{d}B_{\tau}^{l}\bigg{]}$
$\textstyle\stackrel{{\scriptstyle\mathclap{\text{1D It\^{o}
Isometry}}}}{{=}}\qquad\bigg{[}\sum_{j,l=1}^{e}\mathbb{E}\int_{S}^{T}\mathbf{F}_{ij}(\tau,\cdot)[\mathbf{Q}_{c}]_{jl}\mathbf{G}_{kl}(\tau,\cdot)\text{d}\tau\bigg{]}$
$\textstyle=\sum_{j,l=1}^{e}\mathbb{E}\int_{S}^{T}\mathbf{F}_{ij}(\tau,\cdot)[\mathbf{Q}_{c}]_{jl}[\mathbf{G}^{\intercal}(\tau,\cdot)]_{lk}\text{d}\tau$
$\textstyle=\mathbb{E}\int_{S}^{T}e_{i}^{\intercal}\mathbf{F}(\tau,\cdot)\mathbf{Q}_{c}[\mathbf{G}^{\intercal}(\tau,\cdot)]e_{k}\text{d}\tau$
$\textstyle=\bigg{[}\mathbb{E}\int_{S}^{T}\mathbf{F}(\tau,\cdot)\mathbf{Q}_{c}[\mathbf{G}^{\intercal}(\tau,\cdot)]\text{d}\tau\bigg{]}_{ik},$

as required. Next, if $\mathbf{F}$ and $\mathbf{G}$ are deterministic and
suppose that $t_{1}<t_{2}$ without loss of generality, then

$\textstyle\mathbb{E}[\int_{t_{0}}^{t_{1}}\mathbf{F}(\tau)\text{d}B_{\tau}][\int_{t_{0}}^{t_{2}}\mathbf{G}(\tau)\text{d}B_{\tau}]^{\intercal}$
$\textstyle=\mathbb{E}[\int_{t_{0}}^{t_{1}}\mathbf{F}(\tau)\text{d}B_{\tau}][\int_{t_{0}}^{t_{1}}\mathbf{G}(\tau)\text{d}B_{\tau}+\cancel{\int_{t_{1}}^{t_{2}}\mathbf{G}(\tau)\text{d}B_{\tau}}]^{\intercal}$
$\textstyle=\mathbb{E}[\int_{t_{0}}^{\min(t_{1},t_{2})}\mathbf{F}(\tau)\text{d}B_{\tau}][\int_{t_{0}}^{\min(t_{1},t_{2})}\mathbf{G}(\tau)\text{d}B_{\tau}]^{\intercal}$
$\textstyle=\mathbb{E}\int_{t_{0}}^{\min(t_{1},t_{2})}\mathbf{F}(\tau)\mathbf{Q}_{c}\mathbf{G}^{\intercal}(\tau)\text{d}B_{\tau},$

where for the first line we used property (iii) of Definition A.1 of Brownian
motions in multi-dimensions. ∎

###  A.2 Markovian Gaussian Processes

A Markovian Gaussian process $f\sim\mathcal{GP}(0,k)$ can be written with an
SDE of latent dimension $d$

$\textstyle\text{d}\mathbf{s}(t)=\mathbf{F}\mathbf{s}(t)\text{d}t+\mathbf{L}\text{d}B_{t},\quad
f(x)=\mathbf{H}\mathbf{s}(t),$ (9)

where $\mathbf{F}\in\mathbb{R}^{d\times d},\mathbf{L}\in\mathbb{R}^{d\times
e},\mathbf{H}\in\mathbb{R}^{1\times d}$ are the feedback, noise effect and
emission matrices, and $B_{t}$ is an $e$-dimensional (correlated) Brownian
motion with spectral density matrix $\mathbf{Q}_{c}$. Suppose that
$\mathbf{s}(t_{0})\sim\mathcal{N}(\mathbf{m}_{0},\mathbf{P}_{0})$ with the
stationary state mean and covariance. Note that we assume that $s(t_{0})$
independent of $\mathcal{F}_{t_{0}}^{+}$, the $\sigma$-algebra generated by
$B_{t}-B_{s}$ for $t\geq s\geq t_{0}$, in order to invoke existence and
uniquess of SDE solutions (Existence and Uniqueness Theorem, page 90, Evans
(2006)). Thus the linear SDE equation 2 admits a unique closed form solution:

$\textstyle\mathbf{s}(t)=e^{(t-t_{0})\mathbf{F}}\mathbf{s}(t_{0})+\int_{t_{0}}^{t}e^{(t-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau}.$

We have

$\textstyle
m(t)=\mathbb{E}[\mathbf{s}(t)]=e^{(t-t_{0})\mathbf{F}}\mathbf{m}_{0},$

since for any $f\in\mathcal{V}^{d\times e}$,
$\mathbb{E}[\int_{S}^{T}f(\tau)\text{d}B_{\tau}]=0$ (Øksendal, 2003; Evans,
2006). To calculate the covariance, we have

$\textstyle\kappa(t,t^{\prime})$
$\textstyle=\mathbb{E}[\mathbf{s}(t)-m(t))(\mathbf{s}(t^{\prime})-m(t^{\prime}))^{\intercal}]$
$\textstyle=e^{(t-t_{0})\mathbf{F}}\mathbb{E}[(\mathbf{s}(t_{0})-\mathbf{m}_{0})(\mathbf{s}(t_{0})-\mathbf{m}_{0})^{\intercal}][e^{(t^{\prime}-t_{0})\mathbf{F}}]^{\intercal}$
$\textstyle+\mathbb{E}[\int_{t_{0}}^{t}e^{(t-\tau)\mathbf{F}}\text{d}B_{\tau}][\int_{t_{0}}^{t^{\prime}}e^{(t^{\prime}-\tau)\mathbf{F}}\text{d}B_{\tau^{\prime}}]^{\intercal}$
$\textstyle+\cancel{\mathbb{E}\bigg{[}e^{(t-t_{0})\mathbf{F}}\mathbf{L}(\mathbf{s}(t_{0})-\mathbf{m}_{0})(\int_{t_{0}}^{t^{\prime}}e^{(t^{\prime}-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau})^{\intercal}\bigg{]}}\quad\mbox{(Independence)}$
$\textstyle+\cancel{\mathbb{E}\bigg{[}(\int_{t_{0}}^{t}e^{(t-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau})(\mathbf{s}(t_{0})-\mathbf{m}_{0})^{\intercal}\mathbf{L}^{\intercal}[e^{(t-t_{0})\mathbf{F}}]^{\intercal}\bigg{]}}\quad\mbox{(Independence)}$
$\textstyle=e^{(t-t_{0})\mathbf{F}}\mathbf{P}_{0}[e^{(t^{\prime}-t_{0})\mathbf{F}}]^{\intercal}+\int_{t_{0}}^{\min(t,t^{\prime})}e^{(t-\tau)\mathbf{F}}\mathbf{L}\mathbf{Q}_{c}\mathbf{L}^{\intercal}[e^{(t^{\prime}-\tau)\mathbf{F}}]^{\intercal}\text{d}\tau\quad\mbox{(It\^{o}
Isometry; Proposition\leavevmode\nobreak\ \ref{prop:Ito_isometry})},$

where we note that
$\int_{t_{0}}^{t}e^{(t-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau}$ and
$\int_{t_{0}}^{t^{\prime}}e^{(t^{\prime}-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau}$
are $\mathcal{F}_{t}$-measurable for
$\mathcal{F}_{t}\subset\mathcal{F}_{t_{0}}^{+}$, and thus independent to
$\mathbf{s}(t_{0})$.

Therefore for all $t\in\\{t_{0},\ldots,t_{T}\\}$, if we solve for each $t_{i}$
with initial time $t_{i-1}$, then $s(t)\sim\mathcal{N}(m(t),k(t,t))$ with

$\textstyle m(t_{i+1})$
$\textstyle=e^{(t_{i+1}-t_{i})\mathbf{F}}\mathbf{m}_{0}$ $\textstyle
k(t_{i+1},t_{i+1})$
$\textstyle=e^{(t_{i+1}-t_{i})\mathbf{F}}\mathbf{P}_{0}[e^{(t_{i+1}-t_{i})\mathbf{F}}]^{\intercal}+\int_{t_{i}}^{t_{i+1}}e^{(t_{i+1}-\tau)\mathbf{F}}\mathbf{L}\mathbf{Q}_{c}\mathbf{L}^{\intercal}[e^{(t_{i+1}-\tau)\mathbf{F}}]^{\intercal}\text{d}\tau.$

Denote $\Delta_{t_{i+1}}=t_{i+1}-t_{i}$,
$\mathbf{A}_{i,i+1}=e^{(t_{i+1}-t_{i})\mathbf{F}}=e^{\Delta_{t_{i+1}}\mathbf{F}}$
and
$\mathbf{Q}_{i,i+1}=\int_{t_{0}}^{\Delta_{t_{i+1}}+t_{0}}e^{(\Delta_{t_{i+1}}+t_{0}-\tau)\mathbf{F}}\mathbf{L}\mathbf{Q}_{c}\mathbf{L}^{\intercal}[e^{(\Delta_{t_{i+1}}+t_{0}-\tau)\mathbf{F}}]^{\intercal}\text{d}\tau$.
Then we can derive the recursive equations

$\textstyle\mathbf{s}(t_{i+1})$
$\textstyle=\mathbf{A}_{i,i+1}\mathbf{s}(t_{i})+\int_{t_{i}}^{t_{i+1}}e^{(t_{i+1}-\tau)\mathbf{F}}\mathbf{L}\text{d}B_{\tau}$
$\textstyle=\mathbf{A}_{i,i+1}\mathbf{s}(t_{i})+\int_{t_{0}}^{\Delta_{t_{i+1}}+t_{0}}e^{\Delta_{t_{i+1}}+t_{0}-\tau}\mathbf{L}\text{d}B_{\tau}$
$\textstyle=\mathbf{A}_{i,i+1}\mathbf{s}(t_{i})+\mathbf{q}_{i},$

where $\mathbf{q}_{i}\sim\mathcal{N}(0,\mathbf{Q}_{i,i+1})$.

#### Types of Kernels:

A large number of kernels do allow for the Markovian property to be satisfied.
A full list of them could be found in Solin (2016); Särkkä & Solin (2019).
Here, we list 2 common kernels:

  * •

Matern-$\nicefrac{{3}}{{2}}$: The kernel is
$k(t,t^{\prime})=\sigma^{2}(1+\frac{\sqrt{3}d}{\rho})\exp(-\frac{\sqrt{3}d}{\rho})$,
where $d=|t-t^{\prime}|$. With $\lambda=\frac{\sqrt{3}}{\ell}$:

$\textstyle\mathbf{F}=\begin{pmatrix}0&1\\\
-\lambda^{2}&-2\lambda\end{pmatrix},\quad\mathbf{m}_{0}=\mathbf{0},\quad\mathbf{P}_{0}=\begin{pmatrix}\sigma^{2}&0\\\
0&\sigma^{2}\lambda^{2}\end{pmatrix},\quad\mathbf{Q}_{c}=\frac{12\sqrt{3}\sigma^{2}}{\ell^{2}},\quad\mathbf{L}=\begin{pmatrix}0\\\
1\end{pmatrix},\quad\mathbf{H}=\begin{pmatrix}1\\\ 0\end{pmatrix}.$

  * •

Matern-$\nicefrac{{5}}{{2}}$: The kernel is
$k(t,t^{\prime})=\sigma^{2}(1+\frac{\sqrt{5}d}{\rho}+\frac{5d^{2}}{3\rho^{2}})\exp(-\frac{\sqrt{5}d}{\rho})$,
where $d=|t-t^{\prime}|$. With $\lambda=\frac{\sqrt{5}}{\ell}$ and
$\kappa=\frac{5\sigma^{2}}{3\ell^{2}}$:

$\textstyle\mathbf{F}=\begin{pmatrix}0&1&0\\\ 0&0&1\\\
-\lambda^{3}&-3\lambda^{2}&-3\lambda\end{pmatrix},\quad\mathbf{m}_{0}=\mathbf{0},\quad\mathbf{P}_{0}=\begin{pmatrix}\sigma^{2}&0&\kappa\\\
0&\kappa&0\\\
-\kappa&0&\frac{25\sigma^{2}}{\ell^{4}}\end{pmatrix},\quad\mathbf{Q}_{c}=\frac{400\sqrt{5}\sigma^{2}}{3\ell^{5}},\quad\mathbf{L}=\begin{pmatrix}0\\\
0\\\ 1\end{pmatrix},\quad\mathbf{H}=\begin{pmatrix}1\\\ 0\\\ 0\end{pmatrix}.$

#### Kalman Filtering and Smoothing:

The full Kalman filtering and smoothing algorithm is delineated in Algorithm
1.

Algorithm 1 Kalman Filtering and Smoothing

1:Inputs: Variational parameters
$\tilde{\mathbf{Y}}_{1:T},\tilde{\mathbf{V}}_{1:T}$. Initial conditions
$\mathbf{m}_{0}^{f}=\mathbf{m}_{0}$, $\mathbf{P}_{0}^{f}=\mathbf{P}_{0}$,
transition matrices $\\{\mathbf{A}_{i-1,i},\mathbf{Q}_{i-1,i}\\}_{i=1}^{T}$
and emission matrix $\mathbf{H}$.

2:Filtering:

3:for $i=1,\ldots,T$ do

4: Compute predictive filter distribution
$p(\mathbf{s}_{t}|\tilde{\mathbf{Y}}_{1:i-1})=N(\mathbf{s}_{t}|\mathbf{m}^{p}_{t},\mathbf{P}^{p}_{t})$:
$\textstyle\mathbf{m}^{p}_{t_{i}}$
$\textstyle=\mathbf{A}_{i-1,i}\mathbf{m}_{t_{i-1}}^{f}$
$\textstyle\mathbf{P}^{p}_{t_{i}}$
$\textstyle=\mathbf{A}_{i-1,i}\mathbf{P}_{t_{i-1}}^{f}\mathbf{A}_{i-1,i}^{\intercal}+\mathbf{Q}_{i-1,i}.$

5: Let
$\Lambda_{t_{i}}=\mathbf{H}\mathbf{P}^{p}_{t_{i}}\mathbf{H}^{\intercal}+\tilde{\mathbf{V}}_{t_{i}}$.
Compute log marginal likelihood:
$\textstyle\ell_{t_{i}}=\log\mathbb{E}_{p(\mathbf{s}_{t_{i}}|\tilde{\mathbf{Y}}_{1:i-1})}N(\tilde{\mathbf{Y}}_{t_{i}};\mathbf{H}\mathbf{s}_{t_{i}},\tilde{\mathbf{V}}_{t_{i}})=\log
N(\tilde{\mathbf{Y}}_{t_{i}};\mathbf{H}\mathbf{m}^{p}_{t_{i}},\Lambda_{t_{i}})$

6: Compute updated filter distribution
$p(\mathbf{s}_{t_{i}}|\tilde{\mathbf{Y}}_{1:i})=N(\mathbf{s}_{t_{i}}|\mathbf{m}^{f}_{t_{i}},\mathbf{P}^{f}_{t_{i}})$:
$\textstyle\mathbf{W}_{t_{i}}=\mathbf{P}_{t_{i}}^{p}\mathbf{H}^{\intercal}\Lambda_{t_{i}}^{-1}\Rightarrow\mathbf{m}_{t_{i}}^{f}=\mathbf{m}_{t_{i}}^{p}+\mathbf{W}_{t_{i}}(\tilde{\mathbf{Y}}_{t_{i}}-\mathbf{H}\mathbf{m}_{t_{i}}^{p}),\quad\mathbf{P}_{t_{i}}^{f}=\mathbf{P}_{t_{i}}^{p}-\mathbf{W}_{t_{i}}\Lambda_{t_{i}}\mathbf{W}_{t_{i}}^{\intercal}.$

7:end for

8:Smoothing, with initial conditions $\mathbf{m}_{T}^{s}=\mathbf{m}_{T}^{f}$,
$\mathbf{P}_{T}^{s}=\mathbf{m}_{T}^{f}$:

9:for $i=T-1,\ldots,1$ do

10: Compute the smoothing distribution
$q(\mathbf{s}_{t_{i}})=p(\mathbf{s}_{t_{i}}|\tilde{\mathbf{Y}}_{1:T})=N(\mathbf{s}_{t_{i}}|\mathbf{m}^{s}_{t_{i}},\mathbf{P}^{s}_{t_{i}})$
with the RTS smoother:
$\textstyle\mathbf{G}_{t_{i}}=\mathbf{P}_{t_{i}}^{f}\mathbf{A}_{i,i+1}[\mathbf{P}_{t_{i+1}}^{p}]^{-1}\Rightarrow\mathbf{m}_{t_{i}}^{s}=\mathbf{m}_{t_{i}}^{f}+\mathbf{G}_{t_{i}}(\mathbf{m}_{t_{i+1}}^{s}-\mathbf{m}_{t_{i+1}}^{p}),\quad\mathbf{P}_{t_{i}}^{s}=\mathbf{P}_{t_{i}}^{f}+\mathbf{G}_{t_{i}}(\mathbf{P}_{t_{i+1}}^{s}-\mathbf{P}_{t_{i+1}}^{p})\mathbf{G}_{t_{i}}^{\intercal}$

11:end for

12:Return: log marginal likelihood $\sum_{t=1}^{T}\ell_{t_{i}}$, marginal
posteriors $\\{q(\mathbf{s}_{t_{i}})\\}_{i=1}^{T}$.

###  A.3 Further details on spatiotemporal MGPVAE

To perform prediction at arbitrary spatiotemporal locations $(r_{*},t_{*})$,
we can use a conditional independence property proven in Tebbutt et al. (2021)
for separable spatiotemporal Markovian GPs that yields the predictive
distribution

$\textstyle p(\mathbf{Z}(r_{*},t_{*})|\mathbf{Y})\approx
q(\mathbf{Z}(r_{*},t_{*}))$ (10) $\textstyle:=$ $\textstyle\int
p(\mathbf{Z}(r_{*},t_{*})|\mathbf{s}(\mathbf{R},t_{*}))q(\mathbf{s}(\mathbf{R},t_{*})|\mathbf{s}(\mathbf{R},t_{1:T}))\text{d}\mathbf{s}(\mathbf{R},t_{*}),$

where
$q(\mathbf{s}(\mathbf{R},t_{1:T}))=N(\mathbf{s}(\mathbf{R},t_{1:T})|m(\mathbf{R},t_{1:T}),c(\mathbf{R},t_{1:T}))$
is the approximate state posterior $s(\mathbf{R},t_{*})$ conditioned on
$\mathbf{s}(\mathbf{R},t_{1:T})$. Given in Tebbutt et al. (2021) and Wilkinson
et al. (2021),

$\textstyle p(\mathbf{Z}(r_{*},t_{*})|\mathbf{s}(\mathbf{R},t_{*}))=$
$\textstyle
N(\mathbf{Z}(r_{*},t_{*})|\mathbf{B}_{r_{*}}\mathbf{s}(\mathbf{R},t_{*}),\mathbf{C}_{r_{*}}),$

where

$\textstyle\mathbf{B}_{r_{*}}$
$\textstyle=[\mathbf{K}^{r}_{r_{*}\mathbf{R}}(\mathbf{K}^{r}_{\mathbf{R}\mathbf{R}})^{-1}]\otimes[\mathbf{L}^{r}_{\mathbf{R}\mathbf{R}}\otimes\mathbf{H}^{t}]$
$\textstyle\mathbf{C}_{r_{*}}$
$\textstyle=k_{t}(0,0)[\mathbf{K}_{r_{*}r_{*}}^{r}-\mathbf{K}^{r}_{r_{*}\mathbf{R}}(\mathbf{K}^{r}_{\mathbf{R}\mathbf{R}})^{-1}\mathbf{K}^{r}_{\mathbf{R}r_{*}}].$

Therefore the integral in Equation (10) yields
$q(\mathbf{Z}(r_{*},t_{*}))=N(\mathbf{Z}(r_{*},t_{*})|\mathbf{B}_{r_{*}}m(t_{1:T},\mathbb{R}),\mathbf{B}_{r_{*}}c(t_{1:T,\mathbb{R}})\mathbf{B}_{r_{*}}^{\intercal}+\mathbf{C}_{r_{*}})$.

##  Appendix B Additional Experimental and Implementation Details

For the generating modelling tasks of corrupt image (re)-generation, the NLL
on an test sequence $\mathbf{Y}=(\mathbf{Y}_{1},\ldots,\mathbf{Y}_{100})$ is

$\textstyle\log p(\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Y}|\mathbf{Z})p(\mathbf{Z})\text{d}\mathbf{Z}$ $\textstyle=\log\int
p(\mathbf{Y}|\mathbf{Z})\frac{p(\mathbf{Z})}{q(\mathbf{Z}|\mathbf{Y}_{\text{corrupt}})}q(\mathbf{Z}|\mathbf{Y}_{\text{corrupt}})\text{d}\mathbf{Z}$
$\textstyle=\log\frac{1}{K}\sum_{k=1}^{K}p(\mathbf{Y}|\mathbf{Z}_{k})\frac{p(\mathbf{Z}_{k})}{q(\mathbf{Z}_{k}|\mathbf{Y}_{\text{corrupt}})},\quad\mathbf{Z}_{k}\sim
q(\mathbf{Z}|\mathbf{Y}_{\text{corrupt}}),\quad k=1\ldots,K$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\bigg{[}\log
p(\mathbf{Y}|\mathbf{Z}_{k})-\log\frac{q(\mathbf{Z}_{k}|\mathbf{Y})}{p(\mathbf{Z}_{k})}\bigg{]}.$
(11)

For the tasks of missing frame imputation, given test sequence
$\mathbf{Y}=(\mathbf{Y}_{1},\ldots,\mathbf{Y}_{100})$, write
$\mathbf{Y}\equiv(\mathbf{Y},\mathbf{Y}^{c})$, where $\mathbf{Y}^{c}$ are the
unseen frames and $\mathbf{Y}^{c}$ are seen frames.

$\textstyle\log p(\mathbf{Y})$ $\textstyle\equiv\log
p(\mathbf{Y}^{c}|\mathbf{Y})+\log p(\mathbf{Y}).$

$\log p(\mathbf{Y})$ can be computed by Equation B and

$\textstyle\log p(\mathbf{Y}^{c}|\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Y}^{c}|\mathbf{Z}^{c})p(\mathbf{Z}^{c}|\mathbf{Y})\text{d}\mathbf{Z}^{c}$
$\textstyle=\log\int p(\mathbf{Y}^{c}|\mathbf{Z}^{c})\bigg{[}\int
p(\mathbf{Z}^{c}|\mathbf{Z})p(\mathbf{Z}|\mathbf{Y})\text{d}\mathbf{Z}\bigg{]}\text{d}\mathbf{Z}^{c}$
$\textstyle=\log\int p(\mathbf{Y}^{c}|\mathbf{Z}^{c})\bigg{[}\underbrace{\int
p(\mathbf{Z}^{c}|\mathbf{Z})\underbrace{q(\mathbf{Z}|\mathbf{Y})}_{\text{encoder}}\text{d}\mathbf{Z}}_{=:q(\mathbf{Z}^{c}|\mathbf{Y})\approx
p(\mathbf{Z}^{c}|\mathbf{Y})}\bigg{]}\text{d}\mathbf{Z}^{c}$
$\textstyle\approx\log\int
p(\mathbf{Y}^{c}|\mathbf{Z}^{c})q(\mathbf{Z}^{c}|\mathbf{Y})\text{d}\mathbf{Z}^{c}$
$\textstyle=\log\frac{1}{K}\sum_{k=1}^{K}p(\mathbf{Y}^{c}|\mathbf{Z}^{c}_{k}),\quad\mathbf{Z}_{k}^{c}\sim
q(\mathbf{Z}^{c}|\mathbf{Y}),\quad k=1\ldots,K,$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}^{c}|\mathbf{Z}^{c}_{k}),$

$q(\mathbf{Z}^{c}|\mathbf{Y})$ can be analytically obtained (e.g. GP posterior
distribution). For VRNN/KVAE, $q(\mathbf{Z}^{c}|\mathbf{Y})$ would just be
determined by the outputs of the RNN at each time step $t=1,\ldots,100$.

For the spatiotemporal modelling task, we compute the negative log predictive
distribution (NLPD) at new spatial locations $s^{*}\notin\mathcal{S}$ and
temporal locations $t\in\mathcal{T}$ observed during training, conditioned on
observed data
$\mathcal{D}=\\{\mathbf{Y}(s,t)\\}_{s\in\mathcal{S},t\in\mathcal{T}}$, which
is

$\textstyle\log p(\mathbf{Y}(t,s_{*})|\mathcal{D})$ $\textstyle=\log\int
p(\mathbf{Y}(t,s_{*})|\mathbf{Z}(t,s_{*}))p(\mathbf{Z}(t,s_{*})|\mathcal{D})\text{d}\mathbf{Z}(t,s_{*})$
$\textstyle\approx\log\int
p(\mathbf{Y}(t,s_{*})|\mathbf{Z}(t,s_{*}))q(\mathbf{Z}(t,s_{*})|\mathcal{D})\text{d}\mathbf{Z}(t,s_{*})$
$\textstyle=\log\frac{1}{K}\sum_{k=1}^{K}p(\mathbf{Y}(t,s_{*})|\mathbf{Z}_{k}(t,s_{*})),\quad\mathbf{Z}_{k}(t,s_{*})\sim
q(\mathbf{Z}(t,s_{*})|\mathcal{D}),\quad k=1,\ldots,K,$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}(t,s_{*})|\mathbf{Z}_{k}(t,s_{*})).$

We hereby provide a derivation of Equation B for each model (using the
original notation as much as possible).

#### MGPVAE:

$\textstyle\log p(\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Y}|\mathbf{s})\frac{p(\mathbf{s})}{q(\mathbf{s})}q(\mathbf{s})\text{d}\mathbf{s},$
$\textstyle=\log\int p(\mathbf{Y}|\mathbf{s})\frac{p(\mathbf{s})\int
p(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s},\tilde{\mathbf{V}})p(\mathbf{s})\text{d}\mathbf{s}}{p(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s},\tilde{\mathbf{V}})p(\mathbf{s})}q(\mathbf{s})\text{d}\mathbf{s},$
$\textstyle\approx\log\frac{1}{K}+\log\sum_{k=1}^{K}p(\mathbf{Y}|\mathbf{s})\frac{\int
p(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s}_{k},\tilde{\mathbf{V}})p(\mathbf{s}_{k})\text{d}\mathbf{s}}{p(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s}_{k},\tilde{\mathbf{V}})},\quad\mathbf{s}_{k}\sim
q(\mathbf{s}_{k})$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}|\mathbf{s}_{k})-\log
p(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s}_{k},\tilde{\mathbf{V}})+\log\mathbb{E}_{p(\mathbf{s})}N(\tilde{\mathbf{Y}}|\mathbf{H}\mathbf{s},\tilde{\mathbf{V}})$

#### KVAE:

When encountering a missing frame, the Kalman gain is zero in the Kalman
filtering and smoothing operations. The filter prediction distribution is then
used to recompute the $A$ and $C$ matrices, which are finally fed back into
the usual filtering algorithm.

$\textstyle\log p(\mathbf{Y})$
$\textstyle=\log\int\frac{p(\mathbf{Y}|\mathbf{a})p(\mathbf{a}|\mathbf{Z})p(\mathbf{Z})q(\mathbf{a},\mathbf{Z}|\mathbf{Y})}{q(\mathbf{a},\mathbf{Z}|\mathbf{Y})}\text{d}\mathbf{a}\text{d}\mathbf{Z},\quad\bigg{[}q(\mathbf{a},\mathbf{Z}|\mathbf{Y})=q(\mathbf{a}|\mathbf{Y})p(\mathbf{Z}|\mathbf{a})\bigg{]}$
$\textstyle\approx\log\sum_{k=1,\ldots
K}\frac{p(\mathbf{Y}|\mathbf{a}_{k})p(\mathbf{a}_{k}|\mathbf{Z}_{k})p(\mathbf{Z}_{k})}{q(\mathbf{a}_{k}|\mathbf{Y}_{k})p(\mathbf{Z}_{k}|\mathbf{Y}_{k})}+\log\frac{1}{K},\quad(\mathbf{a}_{k},\mathbf{Z}_{k})\sim
q(\mathbf{a}_{k},\mathbf{Z}_{k}|\mathbf{Y})$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}|\mathbf{a}_{k})-\log q(\mathbf{a}_{k}|\mathbf{Y})+\log
p(\mathbf{a}_{k}|\mathbf{Z}_{k})+\log p(\mathbf{Z}_{k})-\log
p(\mathbf{Z}_{k}|\mathbf{a}_{k}).$

#### VRNN:

When encountering a missing frame, the previous hidden state is fed into prior
network and $\mathbf{Z}_{t}$ is sampled from the prior. It is then decoded and
the decoded output is then used as a pseudo-datapoint for autoencoding.

$\textstyle\log p(\mathbf{Y})$
$\textstyle=\log\prod_{t=1}^{T}p(\mathbf{Y}_{t}|\mathbf{Y}_{<t})$
$\textstyle=\log\int\prod_{t=1}^{T}p(\mathbf{Y}_{t}|\mathbf{Y}_{<t},\mathbf{Z}_{\leq
t})\frac{p(\mathbf{Z}_{t}|\mathbf{Y}_{<t},\mathbf{Z}_{<t})}{q(\mathbf{Z}_{t}|\mathbf{Y}_{\leq
t},\mathbf{Z}_{<t})}q(\mathbf{Z}_{t}|\mathbf{Y}_{\leq
t},\mathbf{Z}_{<t})\text{d}\mathbf{Z},$
$\textstyle\approx\log\frac{1}{K}+\log\sum_{k=1}^{K}\prod_{t=1}^{T}p(\mathbf{Y}_{t}|\mathbf{Y}_{<t},\mathbf{Z}^{k}_{\leq
t})\frac{p(\mathbf{Z}_{t}^{k}|\mathbf{Y}_{<t}\mathbf{Z}^{k}_{<t})}{q(\mathbf{Z}^{k}_{t}|\mathbf{Y}_{\leq
t},\mathbf{Z}^{k}_{<t})},\quad\mathbf{Z}_{1:T}^{k}\sim q(\mathbf{Z}_{\leq
T}|\mathbf{Y}_{\leq T})=\prod_{t=1}^{T}q(\mathbf{Z}_{t}|\mathbf{Y}_{\leq
t},\mathbf{Z}_{<t})$
$\textstyle=\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\sum_{t=1}^{T}\log
p(\mathbf{Y}_{t}|\mathbf{Y}_{<t},\mathbf{Z}^{k}_{\leq
t})-\log\frac{q(\mathbf{Z}_{t}^{k}|\mathbf{Y}_{\leq
t}\mathbf{Z}^{k}_{<t})}{p(\mathbf{Z}^{k}_{t}|\mathbf{Y}_{<t},\mathbf{Z}^{k}_{<t})}$

#### LatentODE:

The latentODE places a prior over the initial conditions $\mathbf{Z}_{0}\sim
p(\mathbf{Z}_{0})$ and a posterior via ODERNN
$q(\mathbf{Z}_{0}|\mathbf{Y}_{1:T})\equiv q(\mathbf{Z}_{0})$. Thus the NLL is

$\textstyle\log p(\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Y}|\mathbf{Z}_{0})\frac{p(\mathbf{Z}_{0})}{q(\mathbf{Z}_{0})}q(\mathbf{Z}_{0})d\mathbf{Z}$
$\textstyle\approx\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}|\mathbf{Z}_{0}^{k})-\log\frac{q(\mathbf{Z}_{0}^{k})}{p(\mathbf{Z}_{0}^{k})},\quad\mathbf{Z}_{0}^{k}\sim
p(\mathbf{Z}_{0}).$

#### GPVAE:

$\textstyle\log p(\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Y}|\mathbf{Z})\frac{p(\mathbf{Z})}{q(\mathbf{Z})}q(\mathbf{Z})d\mathbf{Z}$
$\textstyle\approx\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}|\mathbf{Z}_{k})-\log\frac{q(\mathbf{Z}_{k}|\mathbf{Y})}{p(\mathbf{Z}_{k})},\quad\mathbf{Z}_{k}\sim
q(\mathbf{Z}_{k}|\mathbf{Y}).$

#### SVGPVAE:

We are given that
$q(\mathbf{Z},\mathbf{F}_{m})=p(\mathbf{Z}|\mathbf{F}_{m})q(\mathbf{F}_{m}|\mu,\mathbf{A})$,
where $\mathbf{F}_{m}$ is the inducing variable and $(\mu,\mathbf{A})$ are the
inducing variable mean and covariance.

$\textstyle\log p(\mathbf{Y})$ $\textstyle=\log\int
p(\mathbf{Z},\mathbf{F}_{m},\mathbf{Y})\frac{q(\mathbf{Z},\mathbf{F}_{m})}{q(\mathbf{Z},\mathbf{F}_{m})}\text{d}\mathbf{Z}\text{d}\mathbf{F}_{m}$
$\textstyle=\log\int\frac{p(\mathbf{Y}|\mathbf{Z})p(\mathbf{Z}|\mathbf{F}_{m})p(\mathbf{F}_{m})}{p(\mathbf{Z}|\mathbf{F}_{m})q(\mathbf{F}_{m})}p(\mathbf{Z}|\mathbf{F}_{m})q(\mathbf{F}_{m})\text{d}\mathbf{Z}\text{d}\mathbf{F}_{m}$
$\textstyle=\log\int\frac{p(\mathbf{Y}|\mathbf{Z})p(\mathbf{F}_{m})}{q(\mathbf{F}_{m})}p(\mathbf{Z}|\mathbf{F}_{m})q(\mathbf{F}_{m})\text{d}\mathbf{Z}\text{d}\mathbf{F}_{m}$
$\textstyle\approx\frac{1}{P}\sum_{p=1}^{P}\log\frac{1}{K}+\text{logsumexp}_{k=1,\ldots,K}\log
p(\mathbf{Y}|\mathbf{Z}_{k}^{p})-\log\frac{q(\mathbf{F}_{m}^{k})}{p(\mathbf{F}_{m}^{k})},\quad\mathbf{F}_{m}^{k}\sim
q(\mathbf{F}_{m})\quad\mathbf{Z}_{k}^{p}\sim
p(\mathbf{Z}_{k}|\mathbf{F}_{m}^{k}),$

where $p(\mathbf{Z}_{k}|\mathbf{F}_{m})$ is a standard GP conditional
distribution.

###  B.1 Video Data

For each of the 4000 train and 1000 test MNIST images, we create a $T=100$
length sequence of rotating MNIST images by applying a clockwise rotation with
period $t=50$, meaning that the image gets rotated twice. For corrupt video
sequences, we randomly remove 40$\%$ of the pixels and replace them with the
value 0. For the missing frames imputation task, we randomly remove 40$\%$ of
the frames and keep track of the time steps for which the frames are removed.

For each model during training, we used the following configurations (note
that we use the standard PyTorch-TensorFlow-Jax notation for network layers).

  * •

Batch size: 40

  * •

Training epochs: 300

  * •

Number of latent channels: $L=16$

  * •

Adam optimizer learning rate: 1e-3

  * •

clipGradNorm(model parameters, 100)

  * •

Encoder structure: Conv(out=32, k=3, strides=2), ReLU(), Conv2D(out=32, k=3,
strides=2), ReLU(), Flatten(), hiddenToVariationalParams(), where
hiddenToVariationalParams() depends on the model.

  * •

Decoder structure: Linear($L$, 8*8*32), Reshape((8,32,32)),
Conv2DTranspose(out=64, k=3, strides=2, padding=same), ReLU(),
Conv2DTranspose(out=32, k=3, strides=2, padding=same), ReLU(),
Conv2DTranspose(out=1, k=3, strides=1, padding=same), Reshape((32, 32, 1)).

  * •

If KVAE or VRNN, we had to applying KL and Adam optimizer step size schedulers
in order to make them work, meaning that they are more intricately optimised
than the other models.

During test time, we used $K=20$ latent samples to compute the NLL and RMSE
(via the posterior mean). The NLL and RMSE of the predictive results on the
test set are calculated over the entire sequence.

Figure 7: Additional corrupt frames imputation results.

Figure 8: Additional missing frames imputation results.

###  B.2 Action Data

We have 4000 train and 1000 test Mujoco sequences obtained from Rubanova et
al. (2019). We created missing time steps by randomly dropping out $40\%$ of
them.

For each model during training, we used the following configurations (note
that we use the standard PyTorch-TensorFlow-Jax notation for network layers).
Note that unlike Rubanova et al. (2019), we do not use a masked autoencoder
approach for training and only use the non-missing time steps to compute the
ELBO.

  * •

Batch size: 40

  * •

Training epochs: 300

  * •

Number of latent channels: $L=15$

  * •

Adam optimizer learning rate: 5e-3 (1e-3 for MGPVAE since this yielded better
results).

  * •

clipGradNorm(model parameters, 100)

  * •

Encoder structure: Linear($L$, 32)-ReLU-Linear(32, L). If latentODE then
slightly different due to the ODE-RNN.

  * •

Decoder structure: Linear($L$, 16)-ReLU-Linear(16, 14).

  * •

If KVAE or VRNN, we had to applying KL and Adam optimizer step size schedulers
in order to make them work, meaning that they are more intricately optimised
than the other models. LatentODE also used a KL scheduler.

During test time, we used $K=20$ latent samples to compute the NLL and RMSE
(via the posterior mean). The NLL and RMSE of the predictive results on the
test set are calculated over the entire sequence.

Figure 9: Additional results figures for unseen Mujoco sequences. Note that
some LatentODE predictions are not showing as they fall outside the graph
limits.

###  B.3 Spatiotemporal Data

We download the ECMWF ERA5 data (ECMWF/ERA5/DAILY) from Google Earth Engine
API (Gorelick et al., 2017). The data vectors contain:

  * •

Mean 2m air temperature

  * •

Minimum 2m air temperature

  * •

Maximum 2m air temperature

  * •

Dewpoint 2m air temperature

  * •

Surface pressure

  * •

Mean sea level pressure

  * •

$u$ component of wind 10m

  * •

$v$ component of wind 10m

The spatial locations are represented by [longitude, latitude, elevation]. We
index time by $0,1,\ldots,100$.

The model configurations for SVGPVAE and MGPVAE are:

  * •

clipGradNorm(model parameters, 100)

  * •

Encoder structure: Linear($L$, 16)-ReLU-Linear(16, L).

  * •

Decoder structure: Linear($L$, 16)-ReLU-Linear(16, 14).

#### MOGP and SVGPVAE:

We stack time and space together in the dataset, and have inducing points over
space-time jointly. During training, we use only 1 latent sample to compute
the ELBO for SVGPVAE, a minibatch size of 40 and 20000 epochs over the entire
dataset. We use an Adam optimizer learning rate of 1e-3.

#### MGPVAE:

We don’t do any minibatching and directly use the entire training dataset,
which is of dimension roughly 60-100-8. During training, we compute the ELBO
with 20 latent samples and train over 20000 epochs. We use an Adam optimizer
learning rate of 1e-3.

#### Prediction:

We compute the NLL with 20 latent samples over a test set of unseen spatial
locations at all time steps.

Figure 10: Additional figures of the results for the ERA5 experiment given
fixed spatial locations.