Spatio-Temporal Variational Gaussian Processes
Year: 2021

Abstract:
We introduce a scalable approach to Gaussian process inference that combines spatio-temporal filtering with natural gradient variational inference, resulting in a non-conjugate GP method for multivariate data that scales linearly with respect to time. Our natural gradient approach enables application of parallel filtering and smoothing, further reducing the temporal span complexity to be logarithmic in the number of time steps. We derive a sparse approximation that constructs a state-space model over a reduced set of spatial inducing points, and show that for separable Markov kernels the full and sparse cases exactly recover the standard variational GP, whilst exhibiting favourable computational properties. To further improve the spatial scaling we propose a mean-field assumption of independence between spatial locations which, when coupled with sparsity and parallelisation, leads to an efficient and accurate method for large spatio-temporal problems.

heading: 1 Introduction
text: Most real-world processes occur across space and time, exhibit complex dependencies, and are observed through noisy irregular samples. Take, for example, the task of modelling air pollution across a city. Such a task involves large amounts of noisy, partially-observed data with strong seasonal effects governed by weather, traffic, human movement, etc. This setting motivates a probabilistic perspective, allowing for the incorporation of prior knowledge and the quantification of uncertainty.
Gaussian processes (GPS, [38]) provide such a probabilistic modelling paradigm, but their inherent cubic computational scaling in the number of data, N , limits their applicability to spatio-temporal tasks. Arguably the most successful methods to address this issue are sparse GPS [37], which summarise the true GP posterior through a reduced set of M inducing points and have dominant computational scaling O(NM2), and spatio-temporal GPS [43], which rewrite the GP prior as a state-space model and use filtering to perform inference in O(Nd3), where d is the dimensionality of the state-space. Sparse GPS and spatio-temporal GPS have been combined by constructing a Markovian system in which a set of spatial inducing points are tracked over time [24, 51].
However, existing methods for spatio-temporal GPS make approximations to the prior conditional model [24] or do not exploit natural gradients [45], meaning they do not provide the same inference and learning results as state-of-the-art variational GPS [26] in the presence of non-conjugate likelihoods or sparsity, which has hindered their widespread adoption. We introduce spatio-temporal
∗equal contribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
variational GPS (ST-VGP), which provide the exact same results as standard variational GPS, whilst reducing the computational scaling in the temporal dimension from cubic to linear. ST-VGP is derived using a natural gradient variational inference approach based on filtering and smoothing. We also derive this method’s sparse variant, and demonstrate how it enables the use of significantly more inducing points than the standard approach, leading to improved predictive performance.
We then show how the spatio-temporal structure can be exploited even further to improve both the temporal and spatial scaling. We demonstrate for the first time how to apply parallel filtering and smoothing [41] to non-conjugate GPS to reduce the temporal (span) complexity to be logarithmic. We then reformulate the model to enable an efficient mean-field approximation across space, improving the complexity with respect to the number of spatial points. We analyse the practical performance and scalability of our proposed methods, demonstrating how they make it possible to apply GPS to large-scale spatio-temporal scenarios without sacrificing inference quality.

heading: 1.1 Related Work
text: GPS are commonly used for spatio-temporal modelling in both machine learning and spatial statistics [38, 22, 14, 6]. Many approaches to overcome their computational burden have been proposed, from nearest neighbours [17] to parallel algorithms on GPUs [48]. Within machine learning, the sparse GP approach is perhaps the most popular [37, 46], and is typically combined with mini-batching to allow training on massive datasets [26]. However, it fails in practical cases where the number of inducing points must grow with the size of the data, such as for time series [49].
When the data lie on a grid, separable kernels exhibit Kronecker structure which can be exploited for efficient inference [39]. This approach has been generalised to the partial grid setting [53], and to structured kernel interpolation (SKI, [52]) which requires only that inducing points be on a grid. Generally, these approaches are limited to the conjugate case, although Laplace-based extensions exist [19]. Bruinsma et al. [10] present an approach to spatio-temporal modelling that performs an orthogonal projection of the data to enforce independence between the latent processes.
It has been shown that variational GPS can be computed in linear time either by exploiting sparse precision structure [18] or via filtering and smoothing [11]. Other inference schemes such as Laplace and expectation propagation have also been proposed [35, 51]. In the spatio-temporal case, sparsity has been used in the spatial dimension [24, 43]. These methods historically suffered from the fact that i) filtering was not amenable to fast automatic differentiation due to its recursive nature, and ii) state-of-the-art inference schemes had not been developed to make them directly comparable to other methods. The first is no longer an issue since many machine learning frameworks are now capable of efficiently differentiating recursive models [11]. We address the second point with this paper. A similar algorithm to ours that is also sparse in the temporal dimension has been developed [50, 2], and relevant properties of the spatio-temporal model presented here are also analysed in [45]. Fourier features [28] are an alternative approach to scalable GPS, but are not suited to very long time series with high variability due to the need for impractically many inducing features.

heading: 2 Background
text: We consider data lying on a spatio-temporal grid comprising input–output pairs, {X(st) ∈ R
Nt×Ns×D,Y(st) ∈ RNt×Ns}, where Nt is the number of temporal points, Ns the number of spatial points, and D = 1 + Ds the input dimensionality (with Ds being the number of spatial dimensions). We use t and s to represent time and space respectively. The assumption of the grid structure is relaxed via the introduction of sparse methods in Sec. 3, and by the natural handling of missing data.
For consistency with the GP literature we let X = vec(X(st)) ∈ RN×D, Y = vec(Y(st)) ∈ RN×1, where N = NtNs is the total number of data points. We use the operator vec(·) to simply convert data from a spatio-temporal grid into vector form, whilst keeping observations ordered by time and then space. For notational convenience we define Xn,k = X (st) n,k , Yn,k = Y (st) n,k , which indexes data at time index n and spatial index k. We use tn to denote the n’th time point, S ∈ RNs×Ds to denote all spatial grid points and Sk the k’th one. Let f : RD → R to be a random function with a zero-mean GP prior, then for a given likelihood p(Y | f(X)) the generative model is,
f(x) ∼ GP(0, κ(x, x′)), Y | f ∼ ∏Nt
n=1 ∏Ns k=1 p(Yn,k | fn,k), (1)
where fn,k = f(Xn,k), and we let fn be the function values of all spatial points at time tn. When the kernel κ is evaluated at given inputs we write the corresponding gram matrix as KXX′ = κ(X,X′). To make it explicit that f takes spatio-temporal inputs we also abuse the notation slightly to write f(x) = f(t, s) and κ(x, x′) = κ(t, s, t′, s′). A summary of all notation used is provided in App. A. For Gaussian likelihoods the posterior, p(f |Y), is available in closed form, otherwise approximations must be used. In either case, inference typically comes at a cubic cost of O(N3t N 3 s ).

heading: 2.1 State Space Spatio-Temporal Gaussian Processes
text: One method for handling the cubic scaling of GPS is to reformulate the prior in Eq. (1) as a state space model, reducing the computational scaling to linear in the number of time points [43]. The enabling assumption is that the kernel is both Markovian and separable between time and space: κ(t, s, t′, s′) = κt(t, t ′)κs(s, s ′). We use the term Markovian kernel to refer to a kernel which can be re-written in state-space form (see [44] for an overview). First, we write down the GP prior as a stochastic partial differential equation (SPDE, see [15]) ∂tf̄(t, s) = As f̄(t, s) + Ls w(t, s), where w(t, s) is a (spatio-temporal) white noise process and As a suitable (pseudo-)differential operator [see 42]. By appropriately defining the model matrices and the white noise spectral density function, SPDEs of this form can represent a large class of separable and non-separable GP models.
When the kernel is separable, this SPDE can be simplified to a finite-dimensional SDE [24] by marginalising to a finite set of spatial locations, S ∈ RNs×Ds , giving, df̄(t) = F f̄(t) dt+ L dβ(t), where f̄(t) is the Gaussian distributed state at the spatial points S at time t, with dimensionality d = Nsdt, where dt is the dimensionality of the state-space model induced by κt(·, ·). dβ(t) has spectral density Qc, and the matrix H extracts the function value from the state: fn = Hf̄(tn). F and L are the feedback and noise effect matrices [42]. This simplification to an SDE is possible due to the independence between spatial points at time t and all other time steps, given the current state [45]. This follows from the fact that for any separable kernel, f(t, s) and f(t′, s′) are independent given f(t′, s) [36]. For a step size ∆n = tn+1 − tn, the discrete-time model matrices are,
An = Φ(F∆n), Qn = ∫∆n 0 Φ(∆n − τ)LQc L ⊤Φ(∆n − τ) ⊤ dτ, (2)
where Φ(·) is the matrix exponential. The resulting discrete model is,
f̄(tn+1) = An f̄(tn) + qn, Yn | f̄(tn) ∼ p(Yn |H f̄(tn)), (3)
where qn ∼ N(0,Qn). If p(Yn |H f̄(tn)) is Gaussian then Kalman smoothing algorithms can be employed to perform inference in Eq. (3) in O(Ntd3) = O(NtN3s d 3 t ).
Markovian GPs with Spatial Sparsity Sparse GPS re-define the GP prior over a smaller set of M inducing points: let u = f(Z) ∈ RM×1 be the inducing variables at inducing locations Z ∈ RM×D, then the augmented prior is p(f ,u) = p(f |u)p(u), where p(u) = N(u |0,KZZ), and with Gaussian conditional p(f |u). If the inducing points are placed on a spatio-temporal grid, with Zs ∈ RMs×Ds being the spatial inducing locations, the conditional p(f |u) can be simplified to (see App. D):
p(f |u) = N ( f | [ I⊗ (K
(s) SS ⊗K −(s) ZsZs
) ] u,K
(t) tt ⊗ Q̃s
) , (4)
where Q̃s = K (s) SS − K (s) SZs K −(s) ZsZs K (s) ZsS
(see App. A for notational details). The fully independent training conditional (FITC) method [37] approximates the full conditional covariance with its diagonal, leading to the following convenient property: qFITC(f |u) =
∏Nt n=1 qFITC(fn |u) =∏Nt
n=1 qFITC(fn |un), where the last equality holds because I⊗ (K (s) SS ⊗K −(s) ZsZs
) is block diagonal. This factorisation across time allows the model to be cast into the state-space form of Eq. (3), but where the state f̄(t) is defined over the reduced set of spatial inducing points [24]. Inference can be performed in O(NtM3s d 3 t ).

heading: 2.2 Sparse Variational GPs
text: To perform approximate inference in the presence of sparsity or non-Gaussian likelihoods, variational methods cast inference as optimisation through minimisation of the Kullback–Leibler divergence (KLD) from the true posterior to the approximate posterior [8]. Although direct computation of the KLD is intractable, it can be rewritten as the maximisation of the evidence lower bound (ELBO).
Unlike FITC, the sparse variational GP (SVGP, [46]) does not approximate the conditional p(f |u), but instead approximates the posterior as q(f ,u) = p(f |u) q(u), where q(u) = N(u |m,P) is a Gaussian whose parameters are to be optimised. The SVGP ELBO is:
LSVGP = Eq(u) [ Eq(f |u) [log p(Y | f)] ] −KL [q(u) ‖ p(u)] , (5)
which can be computed in O(NM2 + M3). SVGP has many benefits over methods such as FITC, including: non-Gaussian likelihoods can be handled through quadrature or Monte-Carlo approximations [27, 33], it is applicable to big data through stochastic VI and mini-batching [26], and the inducing locations are ‘variationally protected’ and hence prevent overfitting [7].
Natural Gradients Natural gradient descent calculates gradients in distribution space rather than parameter space, and has been shown to improve inference time and quality for variational GPS [26, 40]. A natural descent direction is obtained by scaling the standard gradient by the inverse of the Fisher information matrix, Eq(·) [ ∇2 log q(·) ] [5]. For a Gaussian approximate posterior, the natural gradient of target L with respect to the natural parameters λ can be calculated without directly forming the Hessian, since it can be shown to be equivalent to the gradient with respect to the mean parameters µ = [m,mm⊤ +P] [25]. The natural parameter update, with learning rate β, becomes,
λ← λ+ β ∂L
∂µ . (6)
To update the approximation posterior, λ can be simply transformed to the moment parameterisation [m,P]. A table of mappings between the various parametrisations is given in App. G.
CVI and the Approximate Likelihood Khan and Lin [31] show that when the prior and approximate posterior are conjugate (as in the GP case), further elegant properties of exponential family distributions mean that Eq. (6) is equivalent to a two step Bayesian update:
λ̃← (1− β) λ̃+ β ∂ Eq(f)[log p(Y | f)]
∂µ , λ← η + λ̃ , (7)
where η are the natural parameters of the prior and λ̃ are the natural parameters of the likelihood contribution. Letting g(·) = ∂Eq [log p(Y | f)]
∂· , the gradients can be computed in terms of the mean and covariance via the chain rule: g(µ) = [ g(m)− 2g(P)m, g(P) ] . Eq. (7) shows that, since the prior parameters η are known, natural gradient variational inference is completely characterised by updates to an approximate likelihood, which we denote N(Ỹ | f , Ṽ), parameterised by covariance Ṽ = (−2λ̃(2))−1 and mean Ỹ = Ṽλ̃(1) (see App. A). The approximate posterior then has the form,
q(f) = N(Ỹ | f , Ṽ) p(f)
∫ N(Ỹ | f , Ṽ) p(f) df . (8)
Computing q(f) amounts to performing conjugate GP regression with the model prior and the approximation likelihood. This approach is called conjugate-computation variational inference (CVI, [31]). To re-emphasise that the CVI updates compute the exact same quantity as Eq. (6), we provide an alternative derivation in App. H by directly applying the chain rule to Eq. (6).

heading: 3 Spatio-Temporal Variational Gaussian Processes
text: In this section we introduce a spatio-temporal VGP that has linear complexity with respect to time whilst obtaining the identical variational posterior as the standard VGP. We will then go on to derive this method’s sparse variant, which gives the same posterior as SVGP when the inducing points are set similarly (i.e., on a spatio-temporal grid), but is capable of scaling to much larger values of M .

heading: 3.1 The Spatio-Temporal VGP ELBO
text: We first derive our proposed spatio-temporal VGP ELBO. We do this by exploiting the form of the approximate posterior after a natural gradient step in order to write the ELBO as a sum of three terms, each of which can be efficiently computed through filtering and smoothing. As shown in Sec. 2.2, after a natural gradient step, the approximate posterior q(f) ∝ N(Ỹ | f , Ṽ) p(f) decomposes as a Bayesian update applied to the model prior with an approximate likelihood. Following Chang et al. [11] we substitute Eq. (8) into the VGP ELBO:
LVGP = Eq(f)
[ log p(Y | f) p(f)
q(f)
] = Eq(f) [ log p(Y | f)✟✟p(f) ∫ N(Ỹ | f , Ṽ) p(f) df
N(Ỹ | f , Ṽ)✟✟p(f)
] (9)
=
Nt∑
n=1
Ns∑
k=1
Eq(fn,k)
[ log p(Yn,k | fn,k) ] − Eq(f) [ log N(Ỹ | f , Ṽ) ] + logEp(f) [ N(Ỹ | f , Ṽ) ] .
The first term is the expected log likelihood, the second is the expected log approximate likelihood, and the final term is the log marginal likelihood of the approximation posterior, log p(Ỹ) = logEp(f) [ N(Ỹ | f , Ṽ) ] . Naïvely evaluating LVGP requires O(N3) computation for both the update to q(f) and the marginal likelihood. We now show how to compute this with linear scaling in Nt.
We observe that after a natural gradient update, Ṽ, the approximate likelihood covariance, has the same form as the gradient g(P) because, as seen in Eq. (7), λ̃ is only additively updated by g(µ). Since the expected likelihood, Eq(f)[log p(Y | f)], factorises across observations, g(P) is diagonal
and hence so is Ṽ. The approximate likelihood therefore factorises in the same way as the true one:
log N(Ỹ | f , Ṽ) =
Nt∑
n=1
Ns∑
k=1
log N(Ỹn,k | fn,k, Ṽn,k). (10)
We now turn our attention to computing the posterior and the marginal likelihood. Recall that if the kernel is separable between time and space, κ(t, s, t′, s′) = κt(t, t′)κs(s, s′), then the GP prior can be re-written as Eq. (3). This separability property further results in the state-space model matrices having a convenient Kronecker structure,
f̄(tn+1) = [ INs ⊗A (t) n ] f̄(tn) + qn , Ỹn | f̄(tn) ∼ p(Ỹn |H f̄(tn)), (11)
where qn ∼ N(0,K (s) SS ⊗Q (t) n ) and H = INs ⊗H (t). Here A(t)n ∈ Rdt×dt , Q (t) n ∈ Rdt×dt , and H(t) ∈ R1×dt are the transition matrix, process noise covariance, and measurement model of the SDE (see Sec. 2.1) induced by the kernel κt(·, ·), respectively.
Because the GP prior is Markov and the approximate likelihood factorises across time, the approximate GP posterior is also Markov [45]. Hence marginals q(fn) can be computed through linear filtering and smoothing applied to Eq. (11). Furthermore, the marginal likelihood of a linear Gaussian state-space model, p(Ỹ) = p(Ỹ1) ∏Nt n=2 p(Ỹn | Ỹ1:n−1), can be computed sequentially by run-
ning the forward filter, since p(Ỹn | Ỹ1:n−1) = ∫ p(Ỹn |Hf̄(tn)) p(f̄(tn) | Ỹ1:n−1) df̄(tn), where
p(f̄(tn) | Ỹ1:n−1) is the predictive filtering distribution. By combining all of the above properties we can now write the ELBO as,
LST-VGP =
Nt∑
n=1
Ns∑
k=1
Eq(fn,k)
[ log p(Yn,k | fn,k) ] − Nt∑
n=1
Ns∑
k=1
Eq(fn,k)
[ log N(Ỹn,k | fn,k, Ṽn,k) ]
+
Nt∑
n=1
logE p(f̄(tn) | Ỹ1:n−1)
[ N(Ỹn |Hf̄(tn), Ṽn) ] . (12)
Algorithm 1 Spatio-temporal sparse VGP
Input: Data:{X,Y}, Initial params.:{Ỹ, Ṽ}, Learning rates:{β, ρ} while ELBO (L) not converged do
⊲ CVI natural gradient step:
q(u), ℓ = Alg. 2(Ỹ, Ṽ) E = Eq(u)[Ep(f |u)[log[p(Y | f)]] λ̃ = (1− β)λ̃+ β ∂E ∂µ Ṽ = (−2λ̃(2))−1, Ỹ = Ṽλ̃(1)
⊲ Hyperparameter gradient step:
q(u), ℓ = Alg. 2(Ỹ, Ṽ) E = Eq(u)[Ep(f |u)[log[p(Y | f)]]
L = E − Eq(u)[log N(Ỹ |u, Ṽ)] + ℓ ⊲ ELBO
θ = θ + ρ ∂L ∂θ
end while
Algorithm 2 Sparse spatio-temporal smoothing
Input: Likelihood:{Ỹ, Ṽ}, Space prior:{K(s) ZsZs }, Time prior:{A(t),Q(t),H(t)}
⊲ Construct model matrices: An = IMs ⊗A (t) n , Qn = K (s) ZsZs ⊗Q (t) n , H = IMs ⊗H (t)
⊲ Filtering and smoothing: if using parallel filter / smoother then
q(u), ℓ = Alg. 4(Ỹ, Ṽ,A,Q,H) else
q(u), ℓ = Alg. 3(Ỹ, Ṽ,A,Q,H) end if
⊲ Return posterior marginals and log likelihood: return q(u), ℓ
This ELBO can be computed with linear scaling in Nt: O(NtN3s d 3 t ). We now show that the natural gradient step for updating the parameters of N(Ỹ | f , Ṽ) can be computed with the same complexity.

heading: 3.2 Efficient Natural Gradient Updates
text: As discussed in Sec. 2.2, a natural gradient update to the posterior, q(f) ∝ p(f)N(Ỹ | f , Ṽ), has superior convergence properties to gradient descent, and is completely characterised by an update to the approximate likelihood, N(Ỹ | f , Ṽ), whose mean and covariance are the free parameters of the model, and implicitly define the same variational parameters as VGP. Since the likelihood factorises across the data points, these updates only require computation of the marginal distribution q(fn,k) to obtain Eq(fn,k)[log p(Yn,k | fn,k)] and its gradients.
As we have shown, computation of the marginal posterior amounts to smoothing over the state, f̄ ∼ N(f̄ | m̄, P̄), with the model in Eq. (11). The time marginals are given by applying the measurement model to the state: q(fn) = N(fn |mn = Hm̄n,Pn = HP̄nH⊤) after which q(fn,k) = ∫ q(fn) dfn, 6=k can then be obtained by integrating out the other spatial points. Given the
marginal, Eq. (7) can be used to give the new likelihood parameters Ỹ and Ṽ. The full learning algorithm iterates this process alternately with a hyperparameter update via gradient descent applied to the ELBO, Eq. (12), and has computational complexity O(NtN3s d 3 t ). We call this method the spatio-temporal variational GP (ST-VGP).
3.3 Spatial Sparsity: from O(NtN 3 s d 3 t ) to O(NtM 3 s d 3 t )
We now introduce spatial inducing points, Zs, in order to reduce the effective dimensionality of the state-space model. Whilst we maintain the same notation for consistency, it should be noted that the sparse model no longer requires the data to be on a spatio-temporal grid, only that the inducing points are. In this case, letting q(u) = N(u |m(u),P(u)) be the sparse variational posterior, the marginal q(fn) = N(f |mn,Pn) only depends on m (u) n ,P (u) n due to the conditional independence property for separable kernels discussed in Sec. 2.1. We compute the posterior q(u) via filtering and smoothing over the state f̄(t) in a similar way to ST-VGP by setting,
An = IMs ⊗A (t) n , Qn = K (s) ZsZs ⊗Q(t)n , H = IMs ⊗H (t). (13)
Alg. 2 gives the smoothing algorithm. However, the natural gradient update, Eq. (7), now becomes,
λ̃← (1− β) λ̃+ β ∂ Eq(u)
[ Ep(f |u) [log p(Y | f)] ]
∂µ(u) , (14)
which results in λ̃(2)n , and hence also Ṽn, being a dense matrix (i.e., Ṽ is block-diagonal) due to the conditional mapping, p(fn |un). Therefore the approximate likelihood for the sparse model factorises across time, but not space (see App. J for details): log N(Ỹ |u, Ṽ) = ∑Nt n=1 log N(Ỹn |un, Ṽn).
The Spatio-Temporal Sparse VGP ELBO Adding inducing points in space is equivalent to placing the inducing points on a spatio-temporal grid (i.e., inducing points exist at all time steps), and hence the variational objective directly follows from LSVGP using a similar argument to Sec. 3.1:
LST-SVGP = Eq(f ,u)
[ log p(Y | f)✘✘ ✘p(f |u) ✟ ✟p(u) ∫ N(Ỹ |u, Ṽ) p(u) du
N(Ỹ |u, Ṽ)✘✘ ✘p(f |u) ✟ ✟p(u)
]
=
Nt∑
n=1
Ns∑
k=1
Eq(un) [ Ep(fn,k |un) [ log p(Yn,k | fn,k) ]] − Nt∑
n=1
Eq(un)
[ log N(Ỹn |un, Ṽn) ]
+
Nt∑
n=1
logE p(f̄(tn) | Ỹ1:n−1)
[ N(Ỹn |Hf̄(tn), Ṽn) ] , (15)
where the final term is given by the forward filter.
Efficient Natural Gradient Updates The marginal required to compute the ELBO and natural gradient, q(fn,k) = ∫∫ p(f |u) q(u) du df 6=n,k = ∫∫ p(f |un) q(un) dun df 6=n,k , is the predictive distribution at input Xn,k from the posterior q(u). Because the inducing points have only been placed in space, this can be simplified through the Kronecker structure given by the state-space model. As shown in App. I, the marginal mean and covariance are,
mn,k = K (s) SkZs K −(s) ZsZs m(u)n , Pn,k = K (s) SkZs K −(s) ZsZs P(u)n K −(s) ZsZs K (s) ZsSk +K (t) Xn,kXn,k ( K (s) SkSk −K (s) SkZs K −(s) ZsZs K (s) ZsSk ) , (16)
where m(u)n = Hm̄n, P (u) n = HP̄nH ⊤ are given by filtering and smoothing. By combining the above properties we see that all the terms required for the natural gradient updates and hyperparameter learning can be computed efficiently in O(NtM3s d 3 t ). We call this approach the spatio-temporal sparse variational GP (ST-SVGP). The full algorithm is given in Alg. 1.

heading: 4 Further Improving the Temporal and Spatial Scaling
text: We now propose two approaches to further improve the computational properties of ST-VGP and ST-SVGP. First, we show how parallel filtering and smoothing can be used for non-conjugate GP inference, which results in a theoretical span complexity of O(logNtd3). We then present a spatial mean-field approximation, which can be used independently, or in combination with sparsity.

heading: 4.1 Parallel Bayesian Filtering and Smoothing
text: The associative scan algorithm [9] uses a divide-and-conquer approach combined with parallelisation to convert N sequential associative operations into logN sequential steps (for an operator ∗, associativity implies (a ∗ b) ∗ c = a ∗ (b ∗ c)). This algorithm has been made applicable to conjugate Markov
GPS by deriving a new form of the Kalman filtering and smoothing operations that are associative [41, 13]. We give the form of these associative operators in App. E and show, for the first time, how these methods can be adapted to the non-conjugate setting. This follows directly from the use of the CVI approach to natural gradient VI, which requires only conjugate computations, i.e., the linear filter and smoother. Consider the ST-SVGP ELBO:
LST-SVGP = Eq(u) [ Ep(f |u) [ log p(Y | f) ]] ︸ ︷︷ ︸
factorises across time and space, compute in parallel
−Eq(u) [ log N(Ỹ |u, Ṽ) ] ︸ ︷︷ ︸
factorises across time, compute in parallel
+ logEp(u) [ N(Ỹ |u, Ṽ) ] ︸ ︷︷ ︸
compute with parallel filter
.
The first two terms can be computed in parallel, since they decompose across time given the marginals q(un). The final term can be computed via the parallel filter, and the required marginals q(un) via the parallel smoother, which makes ST-SVGP a highly parallelisable algorithm. Alg. 4 gives the filtering and smoothing algorithm and Alg. 2 shows how this method can be used in place of the sequential filter when performing inference. One drawback of the parallel filter is that when the state dimension is large, many of the available computational resources may be consumed by the arithmetic operations involved in a single filtering step, and the logarithmic scaling may be lost. Fortunately, the spatial mean-field approximation presented in Sec. 4.2 helps to alleviate this issue. In App. E we provide more details on the method as well as a detailed examination of its practical properties.

heading: 4.2 Spatial Mean-Field Approximation
text: We reconsider the state space model for the spatio-temporal GP derived in Sec. 2.1, which has process noise qn ∼ N(0,K (s) SS ⊗ Q (t) n ). This Kronecker structure implies Ns independent processes are linearly mixed using spatial covariance K(s) SS
. The linearity of this operation makes it possible to reformulate the model to include the mixing as part of the measurement, rather than the prior:
f̄(tn+1) = An f̄(tn) + qn, Yn | f̄(tn) ∼ p(Yn | [C (s) SS ⊗H(t)] f̄(tn)), (17)
where qn ∼ N(0, INs ⊗Q (t) n ) and C (s) SS is the Cholesky factor of K(s) SS (see App. F for the derivation). This has the benefit that now both An and Qn = INs ⊗Q (t) n , are block diagonal, such that under the prior the latent processes are fully independent. This enables a mean-field assumption between the Ns latent posterior processes: q(f̄(t)) ≈ ∏Ns k=1 q(f̄k(t)), where f̄k(t) is the dt-dimensional state corresponding to spatial point Sk. This approximation enforces block-diagonal structure in the state covariance, such that matrix operations acting on the full state can be avoided. Dependence between the latent processes is still captured via the measurement model (likelihood), and our experiments show that this approach still approximates the true posterior well (see Sec. 5 and App. K.6) whilst providing significant computational gains when Ns is large.

heading: 5 Experiments
text: We examine the scalability and performance of ST-VGP and its variants. Throughout, we use a Matérn-3/2 kernel and optimise the hyperparameters by maximising the ELBO using Adam [32].
We use learning rates of ρ = 0.01, β = 1 in the conjugate case, and ρ = 0.01, β = 0.1 in the non-conjugate case. We use 5-fold cross-validation (i.e., 80–20 train-test split), train for 500 iterations (except for AIR-QUALITY where we train for 300) and report RMSE, negative log predictive density (NLPD, see App. K.1) and average per-iteration training times on CPU and GPU. When using a GPU, the parallel filter and smoother are used.
Synthetic Experiment We construct 7 toy datasets with rich temporal structure but smooth spatial structure (see App. K.2) and varying size: Nt = 10, 82, 155, 227, 300, 500, 1000, and construct a 500×100 grid that serves as a test set for all cases. As the dataset size increases we expect the predictive performance of all methods to improve at the expense of run time. We compare against SKI and SVGP (see Sec. 1.1). Fig. 2a shows that whilst SVGP becomes infeasible for more than 300 time steps, the ST-SVGP variants scale linearly with time and are faster than SKI (except for the very small datasets, in which the model compile time in JAX dominates). In App. K we show the test performance of all models.
NYC-CRIME – Count Dataset We model crime numbers across New York City, USA (NYC), using daily complaint data from [1]. Crime data has seasonal trends and is spatially dependent. Accurate modelling can lead to more efficient allocation of police resources [20, 3]. We first consider a small subset of the data to highlight when our methods exactly recover SVGP. We bin the data from 1st to 10th of January 2014 (Nt = 10) into a spatial grid of size 30×30 and drop cells that do not intersect with land (Ns = 447, N = 4470). We run ST-SVGP and SVGP with inducing points initialised to the same locations. We plot the training ELBO in Fig. 2b and performance in Table 1. For fixed inducing points, both models have the same training curve and provide the same predictions (up to numerical differences). Optimising the inducing points improves both methods. A comparable inference method for non-conjugate likelihoods has not yet been developed for SKI.
We next consider observations from January to July 2014, with daily binning (Nt = 182) and the same spatial grid (Ns = 447, N = 81,354). We run ST-SVGP and its mean-field variant (MF-STSVGP) with 30 spatial inducing points (equivalent to SVGP with M = 30×182 = 5460). Table 2 shows that our methods outperform SVGP (with M = 1500, M = 3000 and batch sizes 1500, 3000 respectively) because they can afford more inducing points for the same computational budget.
Regression: AIR-QUALITY Finally, we model PM10 (µg/m3) air quality across London, UK. The measurements exhibit periodic fluctuations and highly irregular behaviour due to events like weather and traffic jams. Using hourly data from the London air quality network [29] between January 2019 and April 2019 (Nt = 2159), we drop sensors that are not within the London boundaries or have more than 40% of missing data (Ns = 72, N = 155,448). We run ST-SVGP and MF-ST-SVGP with 30 inducing points in space (equivalent to SVGP with M = 30×2159 = 64,770 inducing points). To ensure the run times are comparable on both CPU and GPU, we run SVGP with 2000, 2500, 5000, and 8000 inducing points with mini-batch sizes of 600, 800, 2000, and 3000 respectively. We run SKI with Nt temporal inducing points and 6 inducing points in each spatial dimension.
Our methods significantly outperform the SVGP baselines because they can afford considerably more inducing points. As shown in Fig. 3 the SVGP drastically smooths the data, whereas ST-SVGP fits the short-term structure well. The mean-field approach is significantly more efficient, especially when using the parallel algorithm, but we do observe a slight reduction in prediction quality.

heading: 6 Conclusions and Discussion
text: We have shown that variational inference and spatio-temporal filtering can be combined in a principled manner, introducing an approach to GP inference for spatio-temporal data that maintains the benefits of variational GPS, whilst exhibiting favourable scaling properties. Our experiments confirm that ST-SVGP outperforms baseline methods because the effective number of inducing points grows with the temporal horizon, without introducing a significant computational burden. Crucially, this leads to improved predictive performance, because fast varying temporal information is captured by the model. We demonstrated how to apply parallel filtering and smoothing in the non-conjugate GP case, but our empirical analysis identified a maximum state dimension of around d ≈ 50, after which the sub-linear temporal scaling is lost. However, our proposed spatial mean-field approach alleviates this issue somewhat, making the combined algorithm extremely efficient even when both the number of time steps and spatial points are large. The resemblance of our framework to the VGP approach suggests many potential extensions, such as multi-task models [4] and deep GPs [16]. We provide JAX code for all methods at https://github.com/AaltoML/spatio-temporal-GPs.

heading: 6.1 Limitations and Societal Impact
text: We believe our work takes an important step towards allowing sophisticated GP models to be run on both resource constrained CPU machines and powerful GPUs, greatly expanding the usability of such models whilst also reducing unnecessary consumption of resources. However, when using predictions from such methods, the limitations of the model assumptions and potential inaccuracies when using approximate inference should be kept in mind, especially in cases such as crime rate monitoring where actions based on biased or incorrect predictions can have harmful societal consequences.
Acknowledgments and Disclosure of Funding
This work was supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence (FCAI). Parts of this project were done while OH was visiting Finland as part of the FCAI–Turing initiative, supported by HIIT/FCAI. OH acknowledges funding from The Alan Turing Institute PhD fellowship programme. AS and WJW acknowledge funding from the Academy of Finland (grant numbers 324345 and 339730). NAL contributed under the NVIDIA AI Technology Center (NVAITC) Finland program. TD acknowledges support from EPSRC (EP/T004134/1), UKRI Turing AI Fellowship (EP/V02678X/1), and the Lloyd’s Register Foundation programme on Data Centric Engineering. We acknowledge the computational resources provided by the Aalto Science-IT project and CSC – IT Center for Science, Finland.

Review summaries:
This paper presents an efficient algorithm for variational inference of spatio-temporal processes using Gaussian processes as the variational family. The paper leverages so-called conjugate-computation variational inference to compute/represent distributions in the approximating variational family. By leveraging Markov, separability of the covariance kernel and prior sparsity assumptions on the spatio-temporal process, the paper demonstrates that the cubic computational complexity can be improved to linear in terms of the time variable. 

The paper presents a novel approach to variational inference in Gaussian Processes for spatio-temporal data (including non-gaussian likelihood), using state-space representation for linear scaling in the temporal dimension. The authors show how using natural gradients and conjugate-computation VI approach leads to decomposition of the ELBO which, together with linear-time filtering and smoothing of the state-space representation, results in significant computational gains.   The results are supported experimentally on one synthetic and 2 real datasets (conjugate and non-conjugate data). The proposed models outperform the baselines (sparse variational GP).  Overall this is a very solid paper.

The authors propose a novel combination of existing strategies for scalable inference for spatio-temporal GPs. The techniques utilized were Kernel separability for convenient Kronecker structure, inference by filter-smoother in the  time domain, inducing point methods, and what appear to be 2 original contributions, a mean-field approximation to their spatial latent processes, and some computational strategies leveraging parallelization of the optimization. The authors make some camparisons to other methods in terms of computation time and accuracy.

The paper presents an approach for GP-Regression in the case of spatio-temporal data. The paper discusses a new scalable approach in the temporal domain by utilizing Bayesian filtering and smoothing in conjugation with conjugate-computation variational inference. Further scalability in the spatial domain is established by the use of an inducing point method with and without a mean field assumption.