Scaling Gaussian Processes with Derivative Information Using Variational Inference

1 Introduction . Gaussian processes ( GPs ) are a popular tool for probabilistic machine learning , widely used in scenarios where uncertainty quantification for regression is necessary [ 27 , 38 , 14 ] . When used for Bayesian optimization ( BO ) [ 18 , 30 ] , or in some regression settings found in the physical sciences like estimation of arterial wall stiffness , derivative information may be available [ 37 , 34 ] . In these settings , we have not only noisy function values y = f ( x ) + but also noisy gradients ∇y = ∇xf ( x ) + at some set of training points X ∈ RN×D . On paper , GPs are ideal models in these settings , because they allow for training on both labels y and gradients∇y in closed form . Though analytically convenient , Gaussian process inference with derivative information scales poorly : computing the marginal log likelihood and predictive distribution for an exact GP in this setting requires O ( N3D3 ) time and O ( N2D2 ) memory . Recent work has addressed this scalability in certain settings , e.g . for many training points in a low-dimensional space [ 5 ] or for few training points in a high-dimensional space [ 3 ] . Despite these advances , training and making predictions for a GP with derivatives remains prohibitively expensive in regimes where both N and D are on the order of hundreds or even thousands . We introduce a novel method to scale Gaussian processes with derivative information using stochastic variational approximations . We show that the expected log likelihood term of the Evidence Lower 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Bound ( ELBO ) decomposes as a sum over both training labels and individual partial derivatives . This lets us use stochastic gradient descent with minibatches comprised of arbitrary subsets of both label and derivative information . Just as variational GPs with inducing points replace the training label information with a set of learned inducing values , we show how to sparsify the derivative information with a set of inducing directional derivatives . The resulting algorithm requires only O ( M3p3 ) time per iteration of training , where M N and p D. We demonstrate the quality of our approximate model by comparing to both exact GPs with derivative information and DSKI from [ 5 ] on a variety of synthetic functions and a surface reconstruction task considered by [ 5 ] . We then demonstrate the full scalability of our model on a variety of tasks that are well beyond existing solutions , including training a graph convolutional neural network [ 16 ] on Pubmed [ 28 ] with Bayesian optimization and regression on a large scale Stellarator fusion dataset with derivatives . We then additionally show that , surprisingly , our variational Gaussian process model augmented with inducing directional derivatives can achieve performance improvements in the regression setting even when no derivative information is available in the training set . 2 Background . In this section we review the background on Gaussian processes ( GP ) ( Section 2.1 ) , Gaussian processes with derivative observations ( Section 2.2 ) , and variational inference inducing point methods for training scalable Gaussian processes ( Section 2.3 ) . Derivative notation . Throughout this paper for compactness we abuse notation slightly and use ∂jyi to refer to the jth element of ∇yi . In this particular case , this would correspond to the partial derivative observation in dimension j for training example xi . We also use ∂vyi to refer to the directional derivative in the direction v , i.e . ∇y > i v . 2.1 Gaussian processes . A Gaussian process ( GP ) is a distribution over functions f ∼ GP ( µ ( x ) , k ( x , x′ ) ) specified by mean and covariance function µ , k [ 23 ] . Given data points X = { x1 , ... , xN } and function observations f = { f ( x1 ) , ... , f ( xN ) } , placing a GP prior assumes the data is normally distributed with f ∼ N ( µX , KXX ) where µX is the vector of mean values atX andKXX ∈ RN×N is a covariance matrix . Conditioning on noisy observations y = f + where ∼ N ( 0 , σ2I ) induces a posterior distribution p ( f∗|y ) over the value of f at points x∗ , which is Gaussian with mean µ ( x∗ ) − Kx∗X ( KXX + σ2I ) −1 ( f − µX ) and covariance k ( x∗ , x∗ ) − Kx∗X ( KXX + σ2I ) −1KXx∗ . Thus , standard GP inference takes O ( N3 ) time . Hyperparameters such as σ , θ are generally estimated by Maximum Likelihood . The log marginal likelihood L ( X , θ , σ|y ) = −1 2 ( y−µX ) T ( KXX +σ2I ) −1 ( y−µX ) − 1 2 log |KXX +σ2I| − n 2 log ( 2π ) ( 1 ) can be optimized with methods like BFGS [ 19 ] at a complexity of O ( N3 ) flops per iteration . 2.2 Gaussian processes with derivatives . GPs can leverage derivative information to enhance their predictive capabilities . Notably , as differentiation is a linear operator , the derivative of a GP is a GP [ 20 ] . Derivative observations can then be naturally included in a GP by defining a multi-output GP over the tuple of function observations and partial derivative observations ( y , ∇y ) [ 24 ] . The GP has mean and covariance functions µ∇ ( x ) = [ µ ( x ) ∇xµ ( x ) ] , k∇ ( x , x′ ) = [ k ( x , x′ ) ( ∇x′k ( x , x′ ) ) T ∇xk ( x , x′ ) ∇2k ( x , x′ ) ] . ( 2 ) While including partial derivative observations can enhance prediction of f , and vice versa , a price is paid in the computational cost , as training and inference of GPs with derivatives scale as O ( N3D3 ) . This scalability issue has been addressed in the low D setting , and is discussed in section 3 . 2.3 Stochastic Variational Gaussian Processes . Inducing point methods [ 29 , 22 , 33 , 10 ] achieve scalability by introducing a set of inducing points : an “ artificial data set ” of points Z = [ zj ] Mj=1 with associated inducing values , u = [ uj ] M j=1 . Stochastic Variational Gaussian Processes ( SVGP ) [ 9 ] augment the GP prior p ( f | X ) → p ( f | u , X , Z ) p ( u | Z ) and then learn a variational posterior q ( u ) = N ( m , S ) . Inference for an observation y∗ at x∗ takes time O ( M3 ) : q ( y∗ ) = N ( y∗ ; Kx∗ZK−1ZZm , σf ( x ∗ ) 2 + σ2 ) ( 3 ) where σf ( x ) 2 = Kxx − KxZK−1ZZKZx + KxZK −1 ZZSK −1 ZZKZx is the data-dependent variance . Using Jensen ’ s inequality and the variational ELBO [ 10 , 11 ] , SVGP develops a loss that is separable in the training data and amenable to stochastic gradient descent ( SGD ) [ 25 ] , as the Kullback-Leibler ( KL ) divergence regularization only depends on u ELBOSVGP = N∑ i=1 { logN ( yi|µf ( xi ) , σ2 ) − σf ( xi ) 2 2σ2 } − KL [ q ( u ) ||p ( u ) ] . ( 4 ) In equation ( 4 ) , µf ( xi ) , σf ( xi ) 2 are the predicted mean and variance , respectively . The ELBO is maximized over the variational parameters m , S and the GP hyperparameters θ . Training with SGD on mini-batches of B data points brings the time per iteration to O ( BM2 +M3 ) . While SVGP scales well , its predictive variances are often dominated by the likelihood noise [ 13 ] . Modeling derivatives necessarily involves heteroscedastic noise , or at least different noise for the function values and gradients , which may make SVGP with a Gaussian likelihood ill-suited to the task . The Parametric Gaussian Process Regressor ( PPGPR ) achieves heteroscedastic modeling by using the latent function variances without modifying the likelihood by symmetrizing the dependence of the loss on the data-dependent variance σf ( xi ) 2 term ELBOPPGPR = N∑ i=1 logN ( yi|µf ( xi ) , σ2 + σf ( xi ) 2 ) − KL [ q ( u ) ||p ( u ) ] . ( 5 ) In Section 5 , we evaluate our approach as an extension to both SVGP and PPGPR , and find that PPGPR gives significant performance gains . 3 Related Work . DSKI and DSKIP [ 5 ] , derivative extensions of SKI [ 36 ] and SKIP [ 7 ] , are among the first methods to address scaling Gaussian processes with derivative information in a low dimensional setting . DSKI and DSKIP approximate derivative kernels by differentiating interpolation kernels ∇k ( x , x′ ) ≈∑ i∇wi ( x ) k ( xi , x′ ) where wi ( x ) are interpolation weights used in SKI . Like SKI , DSKI suffers from the curse of dimensionality , and matrix-vector products cost O ( ND6D + M logM ) time . DSKIP improves the dependence on D , but still costs O ( D2 ( N +M logM + r3N logD ) ) to form the approximate kernel matrices , where r N is the effective rank of the approximation . Thus while these methods exhibit high model fidelity , they are limited to low dimensional settings . Recently [ 3 ] introduced an exact method for training GPs with derivatives in time O ( N2D+ ( N2 ) 3 ) , which improves on the naive O ( N3D3 ) when N < D. This method is not applicable as N grows moderately large , while our paper chiefly focuses on the high-N and high-D setting . Bayesian optimization with derivatives was considered in [ 37 ] . Here , the authors consider conditioning on directional derivatives to achieve some level of scalability , but the dataset sizes considered never exceed N of around 200 or D of around 8 . Distinct from their consideration of directional derivative information , we will be equipping each inducing point in a sparse GP model with its own set of distinct directional derivatives , allowing the model to learn derivatives in many directions in regions of space where there are a large number of inducing points . 4 Methods . Our goal is to enable training and inference on data sets with large N and D when derivatives are available . We will present our method in three steps . First , we describe a naive adaptation of stochastic variational Gaussian processes to the setting with derivatives . Second , we argue that this adaptation again scales poorly in D. Finally , we show that using additional sparsity on the derivatives gives us scalability in both N and D . 4.1 Variational Gaussian processes with derivatives .. As described in Section 2.3 , SVGP creates a dataset of inducing points Z = [ zj ] Mj=1 with labels ( or inducing values ) u = [ uj ] Mj=1 . Assume we are given a dataset X = [ xi ] N i=1 with labels y = [ yi ] N i=1 and derivative observations ∇y = [ ∇yi ] Ni=1 . A natural extension of SVGP to this data is to augment the standard inducing dataset with inducing derivatives , ∇u = [ ∇uj ] Mj=1 , each of length D , so that each inducing point becomes a triple ( zj , uj , ∇uj ) . This corresponds to a new augmented GP prior : p ( f , ∇f | X ) → p ( f , ∇f | u , ∇u , X , Z ) p ( u , ∇u | Z ) . ( 6 ) Analogous to SVGP , we introduce a variational posterior : q ( u , ∇u ) = N ( m∇ , S∇ ) = N ( [ m ∇m ] , [ S ∇S ∇S > ∇2S ] ) . ( 7 ) Here , m and∇m are trainable parameters learned by maximizing the ELBO . We abuse notation and call the second portion of the vector∇m because these variational mean parameters correspond to the M ×D inducing derivative values . This also holds for the matrices∇S and ∇2S . With this augmented variational posterior , the ELBO becomes : Eq ( f , ∇f ) [ log p ( y , ∇y | f , ∇f ) ] − KL ( q ( u , ∇u ) ||p ( u , ∇u ) ) . ( 8 ) Assuming the typical iid Gaussian noise likelihood for regression and expanding the first term further : Eq ( f , ∇f ) [ log p ( y , ∇y | f , ∇f ) ] = N∑ i=1 Eq ( fi ) [ log p ( yi | fi ) ] + N∑ i=1 D∑ j=1 Eq ( ∂jfi ) [ log p ( ∂jyi | ∂jfi ) ] . ( 9 ) Here , we have used linearity of expectation and the conditional independence between yi and ∂jfi given fi to show that the term of the ELBO that depends on training data decomposes as a sum over labels yi and partial derivatives ∂jyi . Thus , minibatches can contain an arbitrary subset of labels ( xi , yi ) and partial derivatives ( xi , ∂jyi ) , and the minibatch size B remains independent of N and D. The moments of q ( f , ∇f ) = ∫ p ( f , ∇f | u , ∇u ) p ( u , ∇u ) du d∇u are similar to those in SVGP , but the kernel matrices have been augmented with derivatives ( i.e. , using the kernel k∇ ( x , x′ ) ) : µf , ∇f = K ∇ XZ ( K∇ZZ ) −1 m∇ , Σf , ∇f = K ∇ XX +K ∇ XZK ∇−1 ZZ ( S ∇ −K∇ZZ ) ( K∇ZZ ) −1 K∇ZX . ( 10 ) Here , K∇XX is a B ×B matrix that corresponds to a randomly sampled subset of label and partial derivative information . K∇XZ is B ×M ( D + 1 ) , and both S∇ and K∇ZZ are M ( D + 1 ) ×M ( D + 1 ) . Similarly , the KL divergence KL ( q ( u , ∇u ) ||p ( u , ∇u ) ) involves multivariate Gaussians with covariance matrices of size M ( D + 1 ) ×M ( D + 1 ) . As a result , the running time complexity of an iteration of training under this framework is O ( M3D3 ) which , grows rapidly with dimension .

Summary:    The manuscripts proposes a method to accelerate GP inference in situations where derivative observations are available and the number of data points N and the data dimension D tends to be large. Extending previous work on stochastic variational approximations, the central idea of the paper is to use a number p of directional derivatives as variational parameters rather than full derivatives which greatly reduces the computational footprint.     Strengths    + Computational methods to tackle high D, high N GP inference problems with derivative information available are not fully scalable and hence a relevant research subject.    + The proposed approach is clear and builds on established techniques.    + The computational benefits demonstrated in an extensive set of experiments are substantial and convincing.      Weaknesses    - Experimental results are without error bars.