Spatio-Temporal Variational Gaussian Processes

1 Introduction . 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 . 1.1 Related Work . 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 . 2 Background . 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 ) . 2.1 State Space Spatio-Temporal Gaussian Processes . 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 ) .

Summary: 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.