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