Inverse-Free Sparse Variational Gaussian Processes

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 13:58:01

Abstract: Gaussian processes (GPs) are a powerful prior over functions, but they require to invert/decompose the kernel matrix to perform inference, making them poorly suited to modern hardware. To address this, variational bounds that require only matmuls by introducing an additional variational parameter $\mathbf T \in \mathbb{R}^{M\times M}$ were proposed. However, in practice, the optimisation of $\mathbf{T}$ with typical deep learning optimisers is challenging, limiting the practical utility of these bounds. In this work, we solve this by introducing a preconditioner for a variational parameter in the bound, a tailored update for $\mathbf T$ based on natural gradients, and a stopping criterion to determine the number of updates. This leads to an inverse-free method on-par with existing approaches on an iteration basis, with low-precision computation and wall-clock speedups being the next step.

Main paper: 
 introduction: 
 Gaussian processes (GPs) [22] are a prior over functions with many desirable properties, but also with many computational difficulties. Despite progress, the popular variational method [18, 8] is still complicated by the need to invert/decompose the kernel matrix, since this is 1) difficult in low-precision, 2) less parallelisable than matmuls, and 3) incurs cubic time cost in model capacity.Van der Wilk et al. [21, 20] proposed variational bounds that required only matmuls by introducing an additional variational parameter T ∈ R M ×M that becomes the required inverse at the optimum. This turns the inversion into an optimisation. While theory proved that this would not introduce performance gaps, practical success required the optimisation of T to be fast enough to keep up with the changing kernel. Unfortunately, typical deep learning optimisers converged unusably slowly.We solve this through 1) a preconditioner for a variational parameter, 2) a tailored update for T based on natural gradients or Newton's method, and 3) a stopping criterion to determine the number of updates. This leads to an inverse-free method on-par with with existing approaches on an iteration basis, with low-precision computation and wall-clock speedups being the next step. 
sparse variational gaussian process (svgp) parameterisations: 
 We consider learning a function f : R D → R through Bayesian inference, with a prior f (•) ∼ GP(0, k θ ) and an arbitrary factorised likelihood with density p(y|f ) = N n=1 p(y n |f (x n ), θ). To make predictions, we are need to approximate the maximum marginal likelihood hyperparameters θ * = argmax θ log p(y|θ), and the posterior p(f |y). The predominant approach is the sparse variational Gaussian processes (SVGP) [8] , which selects an approximate posterior distribution q(f ) by minimising the KL-divergence to the true posterior p(f |y). The approximate posterior is constructed by conditioning the prior on M ≪ N inducing points [14] , at input locations Z ∈ R M ×D and output values u = f (Z), resulting in a variational posterior q(f ) = p(f ̸ =u |u)q(u)du (see [19, 12] ). Following variational inference [6] , we minimise the KL[q(f )||p(f |y)] by maximising the evidence lower bound (ELBO) with respect to Z and parameters ofq(u). The ELBO becomes [8] EQUATIONwhere µ n and σ n are, respectively, the predictive mean and variance at the input location x n . The choice of the parameterisation of the Gaussian distribution q(u) can have a significant impact on the optimisation stability and the tightness of the ELBO. Below, we briefly review two common parameterisations, as well as their relationship with the inverse-free variational bound of [20] . We highlight with red colour the terms in µ n , σ 2 n and KL[q(u)||p(u)] that involve matrix decompositions.Marginal parameterisation (M-SVGP). The most common freeform q(u) = N (m, S) results inEQUATIONSubmitted to Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (BDU at NeurIPS 2024). Do not distribute.EQUATIONwhereK uu = k(Z, Z), k un = k ⊤ nu = k(Z, x n ), and k nn = k(x n , x n ).Likelihood parameterisation (L-SVGP). Panos et al. [13] suggest to reparameterise q(u) = N (m, S) with m = K uu m and S = (K -1 uu + S-1 ) -1 , where S is PD diagonal. This results inEQUATIONEQUATIONwhere K = K uu + S. Inverting K instead of K is more stable and does not need jitter terms, due to the lower-bounded minimum eigenvalue.Inverse-free likelihood parameterisation (R-SVGP). Van der Wilk et al. [20] develop an inversefree bound for L-SVGP, starting by upper bounding the predictive variance asEQUATIONwhere T ∈ R M ×M matrix, and with equality when T = K-1 . Working backwards, they then find that if the marginal parameterisation is reparameterised asS = K uu -K uu (2T -T KT)K uu ,the predictive variance matches the upper bound σ 2 n = U n . This leads to an improvement over the inverse-free bound for M-SVGP [21] , by allowing a closed-form inverse-free bound for the KLKL [q(u)||p(u)] ≤ 1 2 tr((T KT -2T)K uu ) + m⊤ K uu m + tr( KT) -M -log |T| -log | S| ,with equality when T = K-1 . Putting all this together leads to a new ELBO that 1) depends on the additional parameter T, 2) is free of matrix decompositions, 3) is always a lower bound to the existing ELBOs and 4) recovers the original variational solutions when T = K-1 . 
optimising the inverse-free variational bound: 
 The effectiveness of the R-SVGP parameterisation relies on the ability to optimise T close to its optimal value. In their experiments, Van der Wilk et al. [20] parameterise through its Cholesky factor L as T = LL ⊤ , which they update jointly with the other model parameters using Adam [9] .However, they report that this results in poor optimisation performance even with simple datasets.Here, we solve this problem with a tailored optimisation procedure.Given a symmetric PD matrix A, the Cholesky factor L A -1 of its inverse A -1 can be written asEQUATIONNatural gradient (NG) updates are known to converge quickly on KL objectives [3] . In particular, we descend along the natural gradient ∇ℓ A = F -1 (∇ℓ A ), which is found by preconditioning the standard gradient by the inverse Fisher information matrix F = -E N (x;0,LL ⊤ ) ∇ 2 L log N (x; 0, LL ⊤ ) .For Eq. (7), the natural gradient admits (see Appendix A) a simple closed-form expression.Proposition 1. Let ℓ(L) be as in Eq. (7) . Then,EQUATIONwhere tril(•) and diag(•) return the lower triangular part and the diagonal of a matrix.With this, we can optimise the R-SVGP bound by alternating between two steps:1) NG step: for the current K, update L by applying L t+1 = L t -γ ∇L ℓ K, t * times.2) Adam step: given the current value of L, update the other model parameters by applying an Adam update to maximise the R-SVGP bound.Crucially, the expression for the natural gradient in Eq. ( 8) does not involve matrix decompositions, thereby maintaining the overall optimisation procedure inverse-free. Interestingly, such an update for L is related to the well-known Newtonian iteration for computing the inverse of a matrix [5] .The R-SVGP parameterisation introduces some slack to the L-SVGP objective through the upper bound on the predictive variance σ2(L) nin Eq. ( 6). For Gaussian likelihoods, we can monitor such a bias to choose the number t * of NG updates for T. This approach is similar in spirit to several stopping criteria used for conjugate gradient approximations of Gaussian processes [10, 4] . We start by deriving the following lower bound on σ2(L) n , which is proved in Appendix B. Proposition 2. Let σ 2(L) nbe as in Eq. (4). Since S is diagonal and PD, without loss of generality, we can rewrite it as S = S′ + σ 2 I, where S′ is diagonal and PD. Then,EQUATIONwhere K-= K uu + S′ , so that K = K-+ σ 2 I.By subtracting the upper and lower bounds on σn , we find the quantityEQUATION)Moreover, for Gaussian likelihoods, the predictive variance σ is at most ϵ.When T is optimised perfectly, the R-SVGP bound coincides with the L-SVGP bound. Hence, the ability for L-SVGP to reach competitive performance is a necessary condition for R-SVGP to do the same. Unfortunately, we find across a range of experiments that the L-SVGP bound is more difficult to optimise than the standard M-SVGP bound, often leading to worse predictive performance. We identify the cause of this issue in the inducing mean parameterisation m = K uu m and we propose to address it by pre-multiplying m with the matrix K-1 = (K uu + S) -1 . The intuition behind this change is that, in the case where the data points X and the inducing inputs Z coincide, the resulting posterior mean µ n = k nu (K uu + S) -1 m is precisely in the form of the true posterior mean of a GP with Gaussian likelihood, for which S = σ 2 obs I and m = y [13, 2] . We cannot pre-multiply with this matrix without defeating the purpose of the inverse-free bound, so we use T as an approximation. In doing so, we also hardcode the custom gradient of T m wrt to K as if T was exactly equal to K-1 .Our hope is that when T is sufficiently well optimised, this will lead to similar benefits. 
results: 
 To assess the effectiveness of the NG updates for T (Section 3.1) and the inducing mean preconditioning (Section 3.2), we compare versions of the R-SVGP bound that take advantage of either or both of these tools. The main baseline is the R-SVGP bound without preconditioning and where T is trained with Adam, as in van der Wilk et al. [20] . In this context, the optimisation of the R-SVGP bound is successful when the resulting model matches the performance of one trained with the corresponding L-SVGP bound. In all comparisons, we include the M-SVGP parameterisation, which is the de facto standard baseline for stochastic variational GP approximations. The setup of all the experiments discussed below is presented in Appendix C.The benefits of the tools proposed in Section 3 are apparent even in toy datasets. Fig. 1 shows training loss traces for each bound optimised on the SNELSON regression and BANANA classification datasets. Conversely, the bounds trained only with Adam suffer from either slower convergence (R-SVGP) or worse results (R-SVGP-P) than their NG-counterparts. Furthermore, preconditioning is crucial for the L-SVGP and R-SVGP bounds to match the performance of M-SVGP. However, it is only beneficial for R-SVGP when coupled with the NG updates. This is explained by the fact that, as discussed in Section 3.2, preconditioning is expected to help only when T is a good enough approximation to K.As shown in Appendix D.1, this is usually only the case when we use NG updates.Larger datasets make the performance gap between RSVGP-NP and the other versions of the RSVGP bound even more significant. To show this, we evaluate our method on four UCI regression datasets.As shown in Appendix D.2 and Appendix D.3, the R-SVGP-NP bound recovers the performance of L-SVGP-P on all datasets, regardless of whether the inducing locations Z are fixed and trained. WhenZ is fixed, we find that a single NG update with step size one at every iteration is enough. On the other hand, when Z is trained, multiple NG updates are required, which represents a use case for the stopping criterion derived in Section 3.1. In particular, at each iteration, we employ an exponentially increasing schedule for the NG step size and stop optimising T when G falls below a predetermined threshold. In all cases, RSVGP-NP is the only version of the RSVGP bound that is able to match the performance of M-SVGP and LSVGP-P, demonstrating the effectiveness of our approach for optimising the inverse-free bound. 
discussion: 
 We introduced new and effective techniques to optimise the inverse-free bound of van der Wilk et al. [20] . For the first time, our method allows to successfully train a sparse variational GP model on realistic datasets without the computation of matrix decompositions. However, there is still room for improvement. In particular, while we find that multiple NG updates for T are required when training Z, we only experiment with a basic exponential schedule for the NG step size. We expect that a more tailored approach, possibly using a log-linear schedule [16] , may further improve stability and robustness. Furthermore, since our method only employs matmuls, it should be possible to implement it using only single-precision arithmetic. A careful runtime benchmark of our method, both using single and double precision floats, is required to assess the computational benefit of using the inverse-free bound to train variational GPs. Lastly, our method could be extended to derive and optimise inverse-free versions of other SVGP parameterisations, as well as other GP architectures based on the SVGP framework, such as deep GPs [15] . All in all, we believe that our results, while preliminary, hint at the exciting possibility of building complex GP models that are more scalable and better suited to modern hardware.Proof. Given the decomposition for our variational parameter T = LL ⊤ , our goal is to find the optimal L so that T = (K uu + S) -1 , without explicitly computing the inverse. We consider the problem of finding the Cholesky factor of the inverse of a matrix A via:EQUATIONThe above optimisation problem describes the update we take in the direction of the natural gradient:EQUATIONwhere g and g are column vectors denoting the gradient and the natural gradient of ℓ A wrt L, while we denote with F the Fisher information matrix of the distribution p(x; L) = N (x|0, LL ⊤ ) parameterised by L:EQUATIONwith vec(•) the column-wise vectorisation operator.We start by computing the gradient of the loss ∇ℓ A wrt L. As in Eq. ( 7),EQUATIONwhere c 1 is a constant that does not depend on L. The first differential of ℓ A with respect to L is given byEQUATIONTerm (a) can be computed asEQUATIONwhere the last equality holds because A is symmetric. Term (b) is given byEQUATION)Putting it all together, the differential of ℓ A takes the formEQUATIONwhich implies thatEQUATION)We now move to computing the Fisher information matrix F of p(x; L) wrt L. We have thatEQUATION)where c 2 is a constant that does not depend on L. The quantity of interest is the expected value of the second differential of log p(x; L) wrt p, which takes the formEQUATIONbecause log |LL ⊤ | does not depend on x. By taking advantage of Eq. ( 17), term (c) is given byEQUATIONwhere Eq. ( 22) follows from tr(X ⊤ Y) = vec(X) ⊤ vec(Y), andC = C ⊤ is the commutator matrix such that C vec(X) = vec(X ⊤ ) and C(X ⊗ Y) = (Y ⊗ X)C.Similarly, term (d) can be computed starting from the first differentialEQUATIONThen, the second differential takes the formEQUATION)By taking the expectation wrt p, we have thatEQUATIONEQUATIONwhere the first component of Eq. ( 26) follows from the expression already derived in Eq. ( 23). Overall, the expectation in Eq. ( 21) becomesEQUATIONwhich implies that the Fisher information matrix F takes the formEQUATIONBy rearranging the terms of Eq. ( 12), we have that the natural gradient may be found without explicitly computing the inverse F by solving the equationEQUATIONfor g, under the constraint that the latter is the vectorisation of a lower triangular. Before we proceed, let us introduce the matrix form of the the natural gradient vector as G = unvec(g). Then we computeEQUATIONThen, Eq. ( 30) is equivalent toEQUATIONBy left-multiplying both sides by L ⊤ , we obtainEQUATIONwhich implies thatEQUATIONFinally, by left-multiplying both sides by L, we find thatEQUATIONNotice that a similar procedure as the one above can be used to derive the natural gradient of the loss ℓ A with respect to different parameterisations of the covariance matrix of p(x). This includes the marginal parameterisation p(x; T) = N (0, T), as well as a general square root parameterisation p(x; B) = N (0, BB). In particular, it can be shown that, when using the marginal parameterisation, natural gradient descent with unit step size recovers the well-known Newtonian iteration for computing the inverse of a matrix [5] .Proof. As in the statement of Proposition 2, we use the shorthand K-= K uu + S′ , so that K = K-+ σ 2 I. Then, proving the inequalityEQUATIONis equivalent to showing thatEQUATIOND Additional resultsWe aim to evaluate the effectiveness of Adam and NG in keeping the matrix T close to its optimal value throughout training. To do so, we consider the KIN40K dataset with M = 1000 and initialise T to its optimal value K-1 . Then, we train the R-SVGP and R-SVGP-N bounds for 100000 iterations. updates are able to consistently keep T close to its optimal value, whereas this is not the case for Adam. In particular, the R-SVGP bound exhibits the same erratic behaviour observed by van der Wilk et al. [20] , where the quality of T repeatedly improves until a point where it suddenly gets worse, hampering effective training.We evaluate the ability of our method to optimise the R-SVGP bound beyond toy datasets by testing it on four larger UCI regression datasets. In particular, we use each of the bounds discussed in Section 4to train a sparse variational GP model with fixed inducing locations Z for 20000 iterations following the specifics described in Appendix C. As reported by van der Wilk et al. [20] , we find that the R-SVGP bounds trained only with Adam (i.e.R-SVGP and R-SVGP-P) exhibit erratic performance, and we report their results separately in Fig. 4 .We repeat the experiments in Appendix D.2, while also training the inducing locations Z as part of the Adam step. In this case, the optimisation of T in the inverse-free bounds is more challenging because, when Z moves, we expect the optimal value of T, K-1 = (K uu + S) -1 , to change more between iterations compared to when Z is fixed. For this reason, for R-SVGP-N and R-SVGP-NP, we perform multiple NG updates at each iteration according to the schedule described in Appendix C.3. where Z is fixed. We believe that using a more refined scheduling for the step-size of the NG updates, such as commonly used log-linear scheduling [16] , may help to reduce it. In any case, the main takeaway remains the same: the R-SVGP-NP bound achieves competitive performance compared to the standard M-SVGP parameterisation even when Z is trained. https://archive.ics.uci.edu/datasetshttps://www.openml.org/search?type=data 
 
 Appendix:  



Reviews: review: title: title: Review
main_review: main_review: The paper proposes an *inverse-free* variational approach for sparse Gaussian Processes (GPs), focusing on optimizing a new loss function that avoids matrix inversion. The method incorporates natural gradients and preconditioning techniques, demonstrating effective performance across various datasets.

**Strengths:**
- Interesting and Well-Motivated Loss Function: The newly proposed loss function is a novel contribution that significantly improves the optimization process for GPs. The combination of natural gradient updates and preconditioning provides a clear advantage over traditional methods. 

- Solid Validation: The paper includes thorough empirical validation across toy datasets and real-world datasets, demonstrating that the proposed approach is on par or even superior to existing methods in terms of performance.
title: title: Official Review for Submission125
main_review: main_review: This paper presents a novel approach to sparse variational Gaussian processes (SVGP) by introducing an inverse-free optimization method, which is a meaningful contribution. Using matrix-free natural gradient updates to avoid kernel matrix inversions is particularly relevant for large-scale GPs. The following are my comments:

1.	I would suggest adding a computational complexity analysis for the natural gradient update approach, detailing how the proposed method scales with the number of inducing points and the data size in contrast to traditional inversion-based methods.
2.	The authors mentioned that natural gradients converge faster, but is there any potential instability or sensitivity in the choice of step size for the natural gradient? For instance, are there scenarios where convergence could be slower due to poor conditioning of T?
3.	The stopping criterion for the natural gradient updates seems sound, but how did the authors choose the threshold and how sensitive is the method to different thresholds?
4.	I would suggest making more explicit comparisons to other inversion-based methods in terms of both accuracy (e.g., RMSE) and computational performance (e.g., iteration count or time to convergence).
5.	Please discuss the practical limitations of the proposed method in certain use cases (e.g., how it might perform in extremely high-dimensional or non-stationary data).