Sparse Gaussian Process Variational Autoencoders

1 INTRODUCTION . Increasing amounts of large , multi-dimensional datasets that exhibit strong spatio-temporal dependencies are arising from a wealth of domains , including earth , social and environmental sciences ( Atluri et al. , 2018 ) . For example , consider modelling daily atmospheric measurements taken by weather stations situated across the globe . Such data are ( 1 ) large in number ; ( 2 ) subject to strong spatio-temporal dependencies ; ( 3 ) multi-dimensional ; and ( 4 ) non-Gaussian with complex dependencies across outputs . There exist two venerable approaches for handling these characteristics : Gaussian process ( GP ) regression and deep generative models ( DGMs ) . GPs provide a framework for encoding high-level assumptions about latent processes , such as smoothness or periodicity , making them effective in handling spatio-temporal dependencies . Yet , existing approaches do not support the use of flexible likelihoods necessary for modelling complex multi-dimensional outputs . In contrast , DGMs support the use of flexible likelihoods ; however , they do not provide a natural route through which spatio-temporal dependencies can be encoded . The amalgamation of GPs and DGMs , GP-DGMs , use latent functions drawn independently from GPs , which are then passed through a DGM at each input location . GP-DGMs combine the complementary strengths of both approaches , making them naturally suited for modelling spatio-temporal datasets . Intrinsic to the application of many spatio-temporal datasets is the notion of tasks . For instance : medicine has individual patients ; each trial in a scientific experiment produces an individual dataset ; and , in the case of a single large dataset , it is often convenient to split it into separate tasks to improve computational efficiency . GP-DGMs support the presence of multiple tasks in a memory efficient way through the use of amortisation , giving rise to the Gaussian process variational autoencoder ( GP-VAE ) , a model that has recently gained considerable attention from the research community ( Pearce , 2020 ; Fortuin et al. , 2020 ; Casale et al. , 2018 ; Campbell & Liò , 2020 ; Ramchandran et al. , 2020 ) . However , previous work does not support sparse GP approximations based on inducing points , a necessity for modelling even moderately sized datasets . Furthermore , many spatio-temporal datasets contain an abundance of missing data : weather measurements are often absent due to sensor failure , and in medicine only single measurements are taken at any instance . Handling partial observations in a principled manner is essential for modelling spatio-temporal data , but is yet to be considered . Our key technical contributions are as follows : i ) We develop the sparse GP-VAE ( SGP-VAE ) , which uses inference networks to parameterise multi-output sparse GP approximations . ii ) We employ a suite of partial inference networks for handling missing data in the SGP-VAE . iii ) We conduct a rigorous evaluation of the SGP-VAE in a variety of experiments , demonstrat- ing excellent performance relative to existing multi-output GPs and structured VAEs . 2 A FAMILY OF SPATIO-TEMPORAL VARIATIONAL AUTOENCODERS . Consider the multi-task regression problem in which we wish to model T datasets D = { D ( t ) } Tt=1 , each of which comprises input/output pairsD ( t ) = { x ( t ) n , y ( t ) n } Ntn=1 , x ( t ) n ∈ RD and y ( t ) n ∈ RP . Further , let any possible permutation of observed values be potentially missing , such that each observation y ( t ) n = yon ( t ) ∪ yun ( t ) contains a set of observed values yon ( t ) and unobserved values yun ( t ) , with O ( t ) n denoting the index set of observed values . For each task , we model the distribution of each observation y ( t ) n , conditioned on a corresponding latent variable f ( t ) n ∈ RK , as a fully-factorised Gaussian distribution parameterised by passing f ( t ) n through a decoder deep neural network ( DNN ) with parameters θ2 . The elements of f ( t ) n correspond to the evaluation of a K-dimensional latent function f ( t ) = ( f ( t ) 1 , f ( t ) 2 , . . . , f ( t ) K ) at input x ( t ) n . That is , f ( t ) n = f ( t ) ( x ( t ) n ) . Each latent function f ( t ) is modelled as being drawn from K independent GP priors with hyper-parameters θ1 = { θ1 , k } Kk=1 , giving rise to the complete probabilistic model : f ( t ) ∼ K∏ k=1 GP ( 0 , kθ1 , k ( x , x ′ ) ) ︸ ︷︷ ︸ pθ1 ( f ( t ) k ) y ( t ) |f ( t ) ∼ Nt∏ n=1 N ( µoθ2 ( f ( t ) n ) , diag σ o θ2 2 ( f ( t ) n ) ) ︸ ︷︷ ︸ pθ2 ( y o n ( t ) |f ( t ) , x ( t ) n , O ( t ) n ) ( 1 ) where µoθ2 ( f ( t ) n ) and σ o θ2 2 ( f ( t ) n ) are the outputs of the decoder indexed byO ( t ) n . We shall refer to the set θ = { θ1 , θ2 } as the model parameters , which are shared across tasks . The probabilistic model in equation 1 explicitly accounts for dependencies between latent variables through the GP prior . The motive of the latent structure is twofold : to discover a simpler representation of each observation , and to capture the dependencies between observations at different input locations . 2.1 MOTIVATION FOR SPARSE APPROXIMATIONS AND AMORTISED INFERENCE . The use of amortised inference in DGMs and sparse approximations in GPs enables inference in these respective models to scale to large quantities of data . To ensure the same for the GP-DGM described in equation 1 , we require the use of both techniques . In particular , amortised inference is necessary to prevent the number of variational parameters scaling with O ( ∑T t=1N ( t ) ) . Further , the inference network can be used to condition on previously unobserved data without needing to learn new variational parameters . Similarly , sparse approximations are necessary to prevent the computational complexity increasing cubically with the size of each task O ( ∑T t=1N ( t ) 3 ) . Unfortunately , it is far from straightforward to combine sparse approximations and amortised inference in a computationally efficient way . To see this , consider the standard form for the sparse GP approximate posterior , q ( f ) = pθ1 ( f\u|u ) q ( u ) where q ( u ) = N ( u ; m , S ) with m , S and Z , the inducing point locations , being the variational parameters . q ( u ) does not decompose into a product over N ( t ) factors and is therefore not amendable to per-datapoint amortisation . That is , m and S must be optimised as free-form variational parameters . A naı̈ve approach to achieving per-datapoint amortisation is to decompose q ( u ) into the prior pθ1 ( u ) multiplied by the product of approximate likelihoods , one for each inducing point . Each approximate likelihood is itself equal to the product of per-datapoint approximate likelihoods , which depend on both the observation yon and the distance of the input xn to that of the inducing point . An inference network which takes these two values of inputs can be used to obtain the parameters of the approximate likelihood factors . Whilst we found that this approach worked , it is somewhat unprincipled . Moreover , it requires passing each datapoint/inducing point pair through an inference network , which scales very poorly . In the following section , we introduce a theoretically principled decomposition of q ( f ) we term the sparse structured approximate posterior which will enable efficient amortization . 2.2 THE SPARSE STRUCTURED APPROXIMATE POSTERIOR . By simultaneously leveraging amortised inference and sparse GP approximations , we can perform efficient and scalable approximate inference . We specify the sparse structured approximate posterior , q ( f ( t ) ) , which approximates the intractable true posterior for task t : pθ ( f ( t ) |y ( t ) , X ( t ) ) = 1 Zp pθ1 ( f ( t ) ) Nt∏ n=1 pθ2 ( y o n ( t ) |f ( t ) , x ( t ) n , O ( t ) n ) ≈ 1 Zq pθ1 ( f ( t ) ) Nt∏ n=1 lφl ( u ; y o n ( t ) , x ( t ) n , Z ) = q ( f ( t ) ) . ( 2 ) Analogous to its presence in the true posterior , the approximate posterior retains the GP prior , yet replaces each non-conjugate likelihood factor with an approximate likelihood , lφl ( u ; y o n ( t ) , x ( t ) n , Z ) , over a set ofKM ‘ inducing points ’ , u = ∪Kk=1∪Mm=1umk , at ‘ inducing locations ’ , Z = ∪Kk=1∪Mm=1 zmk . For tractability , we restrict the approximate likelihoods to be Gaussians factorised across each latent dimension , parameterised by passing each observation through a partial inference network : lφl ( uk ; y o n ( t ) , x ( t ) n , Zk ) = N ( µφl , k ( y o n ( t ) ) ; k f ( t ) nkuk K−1ukukuk , σ 2 φl , k ( yon ( t ) ) ) ( 3 ) where φl denotes the weights and biases of the partial inference network , whose outputs are shown in red . This form is motivated by the work of Bui et al . ( 2017 ) , who demonstrate the optimality of approximate likelihoods of the form N ( gn ; kf ( t ) nkuk K−1ukukuk , vn ) , a result we prove in Appendix A.1 . Whilst , in general , the optimal free-form values of gn and vn depend on all of the data points , we make the simplifying assumption that they depend only on yon ( t ) . For GP regression with Gaussian noise , this assumption holds true as gn = yn and vn = σ2y ( Bui et al. , 2017 ) . The resulting approximate posterior can be interpreted as the exact posterior induced by a surrogate regression problem , in which ‘ pseudo-observations ’ gn are produced from a linear transformation of inducing points with additive ‘ pseudo-noise ’ vn , gn = kf ( t ) nkuk K−1ukukuk + √ vn . The inference network learns to construct this surrogate regression problem such that it results in a posterior that is close to our target posterior . By sharing variational parameters φ = { φl , Z } across tasks , inference is amortised across both datapoints and tasks . The approximate posterior for a single task corresponds to the product of K independent GPs , with mean and covariance functions m̂ ( t ) k ( x ) = kf ( t ) k uk Φ ( t ) k Kukf ( t ) k Σ ( t ) φl , k −1 µ ( t ) φl , k k̂ ( t ) k ( x , x ′ ) = k f ( t ) k f ′ k ( t ) − kf ( t ) k ukK −1 ukuk kukf ′k ( t ) + k f ( t ) k uk Φ ( t ) k kukf ′k ( t ) ( 4 ) where Φ ( t ) k −1 = Kukuk + Kukf ( t ) k Σ ( t ) φl , k −1 K f ( t ) k uk , [ µ ( t ) φl , k ] i = µφl , k ( y o i ( t ) ) and [ Σ ( t ) φl , k ] ij = δijσ 2 φl , k ( yoi ( t ) ) . See Appendix A.2 for a complete derivation . The computational complexity associated with evaluating the mean and covariance functions is O ( KM2N ( t ) ) , a significant improve- ment over the O ( P 3N ( t ) 3 ) cost associated with exact multi-output GPs for KM2 P 3N ( t ) 2 . We refer to the combination of the aforementioned probabilistic model and sparse structured approximate posterior as the SGP-VAE . The SGP-VAE addresses three major shortcomings of existing sparse GP frameworks . First , the inference network can be used to condition on previously unobserved data without needing to learn new variational parameters . Suppose we use the standard sparse GP variational approximation q ( f ) = pθ1 ( f\u|u ) q ( u ) where q ( u ) = N ( u ; m , S ) . If more data are observed , m and S have to be re-optimised . When an inference network is used to parameterise q ( u ) , the approximate posterior is ‘ automatically ’ updated by mapping from the new observations to their corresponding approximate likelihood terms . Second , the complexity of the approximate posterior can be modified as desired with no changes to the inference network , or additional training , necessary : any changes in the morphology of inducing points corresponds to a deterministic transformation of the inference network outputs . Third , if the inducing point locations are fixed , then the number of variational parameters does not depend on the size of the dataset , even as more inducing points are added . This contrasts with the standard approach , in which new variational parameters are appended to m and S .

Summary: In this work generative models using a GP as prior and a deep network as likelihood (GP-DGMs) are considered. In the VAE formalism for inference, the novelty of this paper is located in the encoder: It is sparse and the posterior can be computed even when part of the observations are missing. Sparsity is obtained using inducing inputs and the missing observations are handled through the use of deep sets, i.e. the observations aren't given as a vector, but as a permutation-invariant set of (index, value) pairs.