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: In a spirit similar to the SVGPs which sparsify the training data, the authors propose to augment the SVGP framework with derivative information by sparsifying the derivative data with *inducing directional derivatives*. By virtue of these modeling choices, the resulting algorithm scales favorably independent of data size as in SVGPs, and now additionaly also w.r.t. input dimension. Empirically, the authors show much better scalability in both high-N and high-D than existing work.