Improving Linear System Solvers for Hyperparameter Optimisation in Iterative Gaussian Processes

We thank Reviewer upE1 for their time to read and review our work, and are delighted to hear that our "writing is very clear", and "empirical evaluations are extensive". In the following, we want to address their specific concerns and questions:

There exists something similar to the pathwise estimator proposed in this paper. For example, see Equation (7) in Gardner et al. (2018), where the probe vectors are sampled from $\mathcal{N}(0, P)$ with $P \approx H$. Though, the original motivation in Gardner et al. (2018) is different from this paper. [...]

We agree that the probe vectors used by Gardner et al. (2018) have conceptual similarities but the motivation is different. In particular, Gardner et al. (2018) use $z \sim \mathcal{N}(0, P)$ probe vectors, where $P$ is constructed using a low-rank pivoted Cholesky decomposition, to implement the preconditioner, allowing their Lanczos quadrature and trace estimation to return the correct values. In contrast, our pathwise probe vectors are sampled using random features, and we used the pivoted Cholesky preconditioner in addition to the pathwise probe vectors.

Furthermore, we leverage the solution of the pathwise probe vector systems to construct posterior samples via pathwise conditioning [1], which has not been done by Gardner et al. (2018). We do this by adapting the standard trace estimator to share linear solves with pathwise conditioning, effectively amortising expensive computations. In contrast, Gardner at al. (2018) solve additional linear systems to obtain the posterior predictive distribution.

In addition, warming starting the linear solvers has been used in practice by Artemev et al. (2021) and Antoran et al. (2023).

Indeed, and we state exactly this in our footnote on page 6. However, we also point out in our footnote that they considered slightly different settings.

In particular, Artemev et al. (2021) only apply warm starts to the mean system and not to any probe vectors, because they do not perform stochastic trace estimation. Arguably, applying warm starts to the probe vectors is non-trivial due to the introduced correlations and bias, as discussed in Section 4.

Antoran et al. (2023) used warm starts for linear solvers in the context of finite-dimensional linear models with Laplace approximation rather than Gaussian processes.

Theorem 1 is not very interesting. This theorem is based on the fact that the gradient bias goes to zero as the number of probe vectors goes to infinity. [...] Frankly, I think removing this theorem makes the paper look cleaner. A middle ground would be just presenting how fast the bias goes to zero as $s \to \infty$.

To avoid misconceptions about Theorem 1, we would like to point out that it is not as simple as it may seem. In particular, to obtain the result, it is not sufficient to show that the error of the gradient estimator goes to zero as the number of probe vectors goes to infinity (which would indeed not be very interesting), since the rate of convergence of the error could technically be different in different parts of the optimisation landscape, leading to asymptotically biased optima. Therefore, the actual result was quite non-trivial to obtain. The complete proof in Appendix A gives a non-asymptotic result with a convergence rate. To avoid excessive notation, we only provided a simplified asymptotic statement in Section 4.

We are currently working on a simplified proof with an improved convergence rate. In general, we believe that this is a non-trivial result and would be inclined to keep Theorem 1 in the main paper. However, we are open to keeping it in Appendix A if Reviewer upE1 feels strongly about this.

The name "pathwise estimator" might be somewhat confusing to some readers. This name is usually reserved for the reparameterization trick gradient, which is evidently not the case for this paper. [...]

This is an interesting observation! In fact, our pathwise estimator inherits its name from pathwise conditioning [1], which can be interpreted as reparameterising a sample from the GP posterior as a deterministic transformation of a sample from the GP prior (and thus its name)! To make the connection to the reparameterisation trick more evident, note that the sample from the GP prior itself can again be reparameterised as an affine transformation of a standard normal random variable.

We will update our manuscript to include a brief version of the explanation above.

The conjugate gradient method estimates both the log marginal likelihood (by tridiagonalization) and its gradient. I am wondering if using the pathwise estimator for conjugate gradient produces an estimate for the log marginal likelihood. This is useful monitoring the progress of hyperparameter optimization.

Currently, the pathwise estimator does not estimate the marginal likelihood (although this could be a future research endeavor!). However, it produces posterior predictive samples via pathwise conditioning [1], which can be used to monitor the progress of hyperparameter optimisation by evaluating the predictive performance. Conventionally, evaluating predictions is expensive because it requires additional linear solves. This becomes intractable for existing iterative methods without relying on additional approximations.

Line 172: footnote should be after the period.

We thank Reviewer upE1 for pointing this out and will correct this for the camera-ready version.

Finally, we hope that we successfully addressed all concerns and questions, and kindly encourage Reviewer upE1 to consider increasing their score or to reach out again with follow-up questions if there are any unresolved concerns.

[1] Wilson et al. (2021), "Pathwise Conditioning of Gaussian Processes", Journal of Machine Learning Research.

Dear Reviewer upE1,

Thank you for acknowledging our rebuttal. We will make sure to include those discussions in the updated version of our manuscript.

Thank you again for your time!

We thank Reviewer 89S9 for their time to read and review our work, and are excited to hear that our experiments are "thorough" and "carefully done". In the following, we want to address their specific concerns and questions:

It is hard for me to see what the really new ideas are in this paper. [...] They do not actually propose any new type of solver [...]. The idea of warm starting a linear solver using previous solutions is pervasive throughout numerical analysis [...]. The main novelty in the paper seems to be changing the trace estimator [...].

Indeed, we do not propose a new solver, nor do we claim to; the title of our paper is “Improving Linear System Solvers […]”.

We want to emphasise that a major (if not our main) contribution consists of an extensive empirical evaluation which

• quantifies the effectiveness of using warm starts and the pathwise estimator;
• across three different kinds of linear system solvers;
• for five small ( < 50k) datasets in the regime of running solvers until convergence;
• and four large ( up to 1.8M) datasets in the regime of limited compute budgets of 10, 20, 30, 40, and 50 epochs.

As a result, we demonstrate speed-ups of up to 72x when running solvers until reaching the tolerance, and achieve up to 7x lower average relative residual norms under a limited compute budget. We believe that this evaluation will be of great value to future researchers and practitioners, and should not be underrated when judging the novelty and contributions of our work.

While we agree that warm starts are commonly used, we also believe that warm starts have not been considered for stochastic trace estimation in the context of marginal likelihood optimisation for GPs, nor compared across various solvers before (if yes, we would appreciate a reference to the literature). Thus, we do not claim warm starts per se as a novel contribution, but rather our application and theoretical / empirical analysis in this particular setting.

Something which is not really discussed, but seems important, is that the $z$ in (7) are easier to generate than the $\xi$ in (11). The $z$ can be random vectors with independent entries, the $\xi$ need a prescribed covaraince structure that changes at every iteration. Since this covariance matrix (the kernel matrix) is large, taking a random draw of the $\xi$ seems nontrivial. [...] How are the $\xi$ in (11) generated and what is the computational cost relative to the system solves?

We agree that the $z$ are easier to generate than the $\xi$, because the latter effectively requires a matrix square root of the kernel matrix $K$ of size $n \times n$. To draw $\xi$ efficiently, we use random features, which produce a low-rank square root with an unbiased approximation to the kernel, reducing the asymptotic time complexity from $\mathcal{O}(n^3)$ to $\mathcal{O}(nm)$, where $m$ is the number of random features. Since $m$ is constant and much smaller than $n$ (e.g. we used $m$ = 2k while $n$ > 1M), this effectively becomes an $\mathcal{O}(n)$ operation, which is faster than CG with $\mathcal{O}(n^2)$ per iteration, and similar to AP / SGD with $\mathcal{O}(n)$ per iteration. Existing work used random features for this purpose and demonstrated strong empirical performance despite this approximation [1, 2]. A discussion of alternatives can be found in Section 4 of [1].

We mention random features in Section 2, lines 70-71, and discuss them further in Section 3, lines 143-149, and Figure 5. Additionally, a detailed description of how to sample random probe vectors is provided in Appendix B, lines 486-501.

The numerical experiements are on interesting problems, but the paper would be stronger if the mathematical arguments were more precise.

We thank Reviewer 89S9 for describing the problems in our experiments as "interesting". To clarify any potential misconceptions about our mathematical arguments, we summarise them here:

In Section 3, we calculate the expected squared distance between zero initialisation and the solution of the linear solver. In particular, we show that, for the standard estimator, this distance depends on the spectrum of the regularised kernel matrix, while it is constant for the pathwise estimator. Thus, for the standard estimator, this distance tends to increase as the spectrum changes during hyperparameter optimisation, leading to substantially increased number of required solver iterations. In contrast, for the pathwise estimator, the constant distance translates to a roughly constant (and significantly lower) number of required solver iterations.

In Section 4, we investigate the concern that warm starts require deterministic probe vectors which introduce correlations into subsequent gradient estimates. Despite individual estimates being unbiased, these correlations introduce bias into the optimisation. Theorem 1 shows that, under certain assumptions, the marginal likelihood after optimisation with deterministic probe vectors converges in probability to the marginal likelihood of a true maximum as the number of probe vectors goes to infinity. The proof is provided in Appendix A, and we are working on an updated, simpler proof with better constants.

After the clarifications above, we are happy to provide further details upon request and would appreciate it if Reviewer 89S9 has specific suggestions for further mathematical arguments which they think might strengthen the paper.

Finally, we hope that we successfully addressed all concerns and questions, and kindly encourage Reviewer 89S9 to increase their score or to reach out again with follow-up questions if there are any unresolved concerns.

[1] Wilson et al. (2021), "Pathwise Conditioning of Gaussian Processes", Journal of Machine Learning Research.

[2] Lin et al. (2023), "Sampling from Gaussian Process Posteriors using Stochastic Gradient Descent", Neural Information Processing Systems.

Dear Reviewer 89S9,

Given that the discussion phase is soon coming to an end , this is a gentle reminder to react to our rebuttal before the deadline.

Thank you!

We thank Reviewer Xgts for their time to read and review our work, and are grateful to hear that our "motivation and suggestions are clear", our manuscript is "well-prepared", and our experimental evaluation is "extensive". In the following, we want to address their specific concerns and questions:

So, a more detailed discussion of further research directions related to the presented approaches would be interesting.

At the end of Section 6, we mention that an interesting direction for future research would be to investigate why low relative residual norms do not always translate to good predictive performance. To elaborate on that, we believe that interesting research questions are:

• Which metric, norm, or quantity actually corresponds to predictive performance (other than predictive performance itself)?
• Can we develop a stopping criterion based on this particular quantity (instead of relative residual norm) to optimise for predictive performance?

We are happy to include this discussion at the end of Section 6 in the camera-ready version of our manuscript.

The authors mentioned in the Appendix that they use precondition based on the Cholesky factorization. Please provide more details since, typically, such preconditioners, their structure, and properties significantly affect the performance of CG.

We will update our manuscript to provide the following details (and references) about the pivoted Cholesky preconditioner in Appendix B:

The pivoted Cholesky factorisation [1] constructs a low-rank approximation to the full (but intractable) Cholesky factorisation by permuting rows and columns to prioritise the entries with the highest variance in a greedy way. Given the extensive and successful use of this preconditioner in the context of iterative Gaussian processes [2, 3, 4], we follow and refer to existing work.

Please add bold labeling the top timing and test log-likelihood values in Table 1 to simplify parsing such a large number of values.

We will update our manuscript to include a summary of the following explanation in the caption of Table 1:

In Table 1, we compare different iterative methods when solving until the tolerance is reached. This results in all methods performing (nearly) identical in terms of log-likelihood (except for a few cases where the bias of the pathwise estimator due to random features leads to a slightly different outcome). Therefore, we did not bold these log-likelihood values because almost every value in the table would be bold. While the runtimes indeed differ, the comparison mainly evaluates the effects of warm starts and the pathwise estimator per solver. Thus, we bold the relative speed-ups per solver instead of the absolute runtimes.

Why are large datasets (lines 219-235) not included in Table 1 or a similar table where the runtime gain can be observed? Please add these results, too. The large datasets can illustrate the gain from the introduced approach even better.

We do provide runtimes for large datasets in Tables 7-10 in Appendix C, and we are happy to include these in main paper for the camera-ready version.

Table 1 considers small datasets and presents the predictive performances of runtimes of solving until reaching the tolerance. In this setting, the final performance is basically fixed and gains in runtime are the main focus. In contrast, for the large datasets, solvers run for a fixed maximum number of epochs, implying a (roughly) fixed maximum amount of time. While in some cases solvers eventually reach the tolerance in fewer epochs and less time, this is not the main emphasis of these experiments. Instead, the main goal is to show that warm starts and the pathwise estimator lead to better performance under a fixed compute budget due to accumulation of solver progress, which is visualised in Figure 10.

Finally, we hope that we successfully addressed all concerns and questions, and kindly encourage Reviewer Xgts to consider increasing their score or to reach out again with follow-up questions if there are any unresolved concerns.

[1] Harbrecht et al. (2012), "On the low-rank approximation by the pivoted Cholesky decomposition", Applied Numerical Mathematics.

[2] Gardner et al. (2018), "GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration", Neural Information Processing Systems.

[3] Wang et al. (2019), "Exact Gaussian Processes on a Million Data Points", Neural Information Processing Systems.

[4] Wu et al. (2024), "Large-Scale Gaussian Processes via Alternating Projection", Artificial Intelligence and Statistics.

We thank Reviewer NcMH for their time to read and review our work. In the following, we want to address their specific concerns and questions:

The main weakness of this paper is that it does not compare the results of the approximate methods to a reliable exact method. This makes it impossible to determine from the paper whether the predictive performances are actually good.

We thank Reviewer NcMH for mentioning BFGS. Indeed, this is the preferred algorithm for GP hyperparameter optimisation when its use is computationally tractable. We initially did not consider it due to our focus on the large-scale setting. We conducted additional experiments using the Cholesky factorisation + BFGS optimiser on the POL, ELEVATORS, and BIKE datasets. We were not able to run BFGS on larger datasets because our A100 80GB GPUs did not support the necessary memory requirements. In particular, we used jaxopt [1], and performed optimisation for a maximum of 100 iterations, a maximum of 30 linesearch steps per iteration, and used a stopping tolerance of 1e-5.

Additionally, we provide results using the Adam optimiser + Cholesky factorisation (instead of stochastic gradient estimation) with all other settings being identical to the iterative methods in our paper. This can be considered a pure comparison of the quality of estimated gradients. For both BFGS and Adam, all hyperparameters were initialised at the same values as for the iterative methods in our paper (namely 1.0) and also used the softplus transformation to enforce positive value constraints.

Furthermore, we also include results for SVGP [2], a popular variational approximation, taken from Lin et al. [3], who used the same datasets and train / test splits, and performed optimisation until convergence using the Adam optimiser and 3000 inducing points.

The table below lists the mean over 10 splits of test root-mean-square-errors (RMSE) and test log-likelihoods (LLH), with "Exact" referring to Cholesky factorisation + backpropagation. In summary, BFGS achieves marginally better results on the POL and ELEVATORS datasets and significantly better results on the BIKE dataset. However, we conclude that the discrepancy on the BIKE dataset is due to Adam not converging in 100 iterations, because Adam with Cholesky factorisation + backpropagation achieves nearly identical performance compared to the iterative methods (which also use Adam). Therefore, the performance gap cannot be linked to the gradient estimation via iterative methods, but rather to the different optimisers. However, in the large-scale setting, memory requirements and stochastic gradient estimation make BFGS infeasible. Therefore, we conducted all our experiments with Adam. In comparison, SVGP generally yields worse results, particularly in terms of log-likelihood.

Is there any information in the paper that gives a comparison to another method, or ideally a ground truth evaluation of an exact GP?

Currently, Figures 5 and 8 visualise the influence of the pathwise estimator and warm starts by comparing their optimisation trajectories to exact optimisation using Cholesky factorisation + backpropagation (and otherwise identical configurations).

Additionally, Figures 11-13 in Appendix C also compare iterative methods to exact optimisation with Cholesky factorisation + backpropagation by visualising the optimisation trajectories of selected hyperparameters. Furthermore, we are happy to include the quantitative results from the table above in the camera-ready version.

Finally, we hope that we successfully addressed all concerns and questions, and kindly encourage Reviewer NcMH to increase their score or to reach out again with follow-up questions if there are any unresolved concerns.

[1] Blondel et al. (2021), "Efficient and Modular Implicit Differentiation", arXiv.

[2] Hensman et al. (2013), "Gaussian Processes for Big Data", Uncertainty in Artificial Intelligence.

[3] Lin et al. (2024), "Stochastic Gradient Descent for Gaussian Processes Done Right", International Conference on Learning Representations.

Dear Reviewer NcMH,

Given that the discussion phase will end very soon, this is a friendly reminder to respond to our latest update before the deadline. We remain excited to share these new results and clarifications with you, and we are confident they are likely to address your concerns!

Thank you!