Improving Linear System Solvers for Hyperparameter Optimisation in Iterative Gaussian Processes
Scaling hyperparameter optimisation to very large datasets remains an open problem in the Gaussian process community. This paper focuses on iterative methods, which use linear system solvers, like conjugate gradients, alternating projections or stochastic gradient descent, to construct an estimate of the marginal likelihood gradient. We discuss three key improvements which are applicable across solvers: (i) a pathwise gradient estimator, which reduces the required number of solver iterations and amortises the computational cost of making predictions, (ii) warm starting linear system solvers with the solution from the previous step, which leads to faster solver convergence at the cost of negligible bias, (iii) early stopping linear system solvers after a limited computational budget, which synergises with warm starting, allowing solver progress to accumulate over multiple marginal likelihood steps. These techniques provide speed-ups of up to $72\times$ when solving to tolerance, and decrease the average residual norm by up to $7\times$ when stopping early.
Paper
Similar papers
Peer review
Summary
This paper focuses on solving linear systems arising from large-scale Gaussian process hyperparameter optimisation. The first technique proposed in this paper is a reformulation of the linear system for log determinant gradient estimation, which is called the "pathwise estimator". The authors argue that this reformulation makes the all-zero initialization closer to the solution. This reformulation also allows them to reuse the linear system solutions in the prediction time by pathwise conditioning (Wilson et al., 2021). The second technique proposed in this paper is warm starting the linear solvers. Though, warm starts require using a set of **fixed** probe vectors, which necessarily introduces biases to the hyperparameter optimization. The authors demonstrate empirically that the bias does not hurt in practice.
Strengths
- The writing is very clear. - By combining several techniques, this paper achieves large speed-up factors over a wide range of datasets. Besides the speed-up, it appears that the pathwise estimator with warm start also stabilizes the linear systems solvers in some cases. - The empirical evaluations are extensive with datasets ranging from 10 thousand data points to 1.8 million data points. The extensive experiments with recent scalable GP methods could provide tips and tricks to practitioners.
Weaknesses
The main weakness is that the techniques are not super novel per se. 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}(\mathbf{0}, \mathbf{P})$ with $\mathbf{P} \approx \mathbf{H}_{\mathbf{\theta}}$. Though, the original motivation in Gardner et al. (2018) is different from this paper. Nevertheless, they do appear to be similar at least on a conceptual level. In addition, warming starting the linear solvers has been used in practice by Artemev et al. (2021) and Antoran et al. (2023). However, the main contribution of this paper is combining all techniques together and show their effectiveness on a wide range of datasets. Future work could benefit from the the insights and empirical evaluations in this paper. Minor Weaknesses: - 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. My impression is that the authors wrote this theorem to make the paper looks more "theoretical". 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$. - 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. I suggest the authors think about alternative names for this estimator. An possible option is to simply rephrase Section 3 as "a reformulation of the linear systems". - Line 172: footnote should be after the period.
Questions
- 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.
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
NA
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!
Summary
The paper considers the problem of simultaneously fitting a Gaussian process (GP) to data along with determining the hyperparameters (kernel width, noise variance) for the GP. The overall algorithm is standard, consisting of an outer loop, which is a simple gradient update to the hyperparameters, and an inner loop that fits the GP to the data. Fitting the data of course requires a system solve with the kernel matrix. Computing the gradient (or an approximation thereof) for the hyperparamters requires multiple system solves. The paper gives some numerical experiements and some rough reasoning to argue that tweaking the trace estimator used for the gradient estimate and warm starting the linear solver saves a tangible amount of computation in this process.
Strengths
Hyperparameter optimization for fitting GPs is an important and prevalent problem in machine learning. The numerical experiments are thorough and appear to be carefully done, and lend good supporting evidence to the efficacy of the method.
Weaknesses
It is hard for me to see what the really new ideas are in this paper. The the authors say, the most critical part of this whole process is the linear system solver. They do not actually propose any new type of solver, but rather show that by modifying the way the gradient of the hyperparameters is estimated and using a warm start, the linear solver they happen to be using will take fewer iterations. The idea of warm starting a linear solver using previous solutions is pervasive throughout numerical analysis, so it is no surprise that it gives you some gains. The main novelty in the paper seems to be changing the trace estimator; that is, using equation (11) in place of (7). 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 this works, and why it does not have a significant effect on the number of computations needed, needs to be spelled out more clearly. The numerical experiements are on interesting problems, but the paper would be stronger if the mathematical arguments were more precise.
Questions
How are the $\xi$ in (11) generated and what is the computational cost relative to the system solves?
Rating
5
Confidence
2
Soundness
2
Presentation
2
Contribution
2
Limitations
none
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!
Thank you for your response and detailed defense of your work. Given the re-iteration of the magnitude of the performance gain in the numerical results, I am increasing the overall score.
Summary
The paper presents approaches that can speed up solving linear systems arising in the GP regression problem. The two basic ideas are a warm start and limiting the computation budget. Another novelty is the pathwise gradient estimator, which leads to fewer iterations needed for convergence. The detailed numerical experiments confirm the performance of the suggested approaches and raise questions on the connection between the accurate solving of linear systems and the performance of the GP regression method.
Strengths
1. The novel approach for gradient estimation is proposed and demonstrates significant speed-up. 2. Different techniques to accelerate solvers for linear systems are proposed, and their efficiency is confirmed via extensive experimental evaluation. This evaluation includes target solvers CG, AP, and SGD. 3. The manuscript is well-prepared, the motivation and suggestions are clear, and the results are well-explained. 4. The presented techniques can be easily incorporated into the existing packages for GP regression and used to solve practically important problems.
Weaknesses
In this paper, the authors use some heuristics to speed up linear solvers and suggest a new gradient estimator based on the transformation of the normal distribution. Although the suggested tricks demonstrate improvement, their theoretical interpretation and the possibility of obtaining even better results remain unclear. So, a more detailed discussion of further research directions related to the presented approaches would be interesting. The theoretical bounds for the potential gain obtained from linear solvers while preserving test error can make the work more solid and fundamental.
Questions
1. 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. 2. Please add bold labeling the top timing and test log-likelihood values in Table 1 to simplify parsing such a large number of values 3. 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.
Rating
7
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors do not provide explicit limitations of the suggested approach.
Summary
This paper investigates several iterative techniques for solving linear systems when applied to the problem of finding GP hyperparameters. Specifically, the following modifications to the method are suggested: - A "pathwise" sampling estimator for the Hutchinson trace estimator - Warm starting linear systems solvers After this, a thorough investigation is made into the effect of these choices on training GP hyperparameters.
Strengths
The method follows an existing schema for training GP hyperparameters, and performs an in-depth investigation into how different choices affect the algorithm.
Weaknesses
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. For datasets of these sizes (<50k), it is feasible on modern hardware to run a Cholesky on the full dataset. This should be run with BFGS optimisation. Reporting these results will give a ground-truth comparison for how good the resulted test metrics are. If these are close, then readers will know how much to value these speedups. If they are distant, readers will be able to determine how much they value a faster approximation that does not get close to optimal predictor. Without a comparison to any other method, the value of the overall method cannot be established. I'm happy to discuss this and significantly increase my score if this is addressed sufficiently.
Questions
- Is there any information in the paper that gives a comparison to another method, or ideally a ground truth evaluation of an exact GP?
Rating
3
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
- No evaluations to any other methods, other than methods with iterative solvers.
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!
Dear Reviewer NcMH, We are reaching out to you one last time before the deadline of the author-reviewer discussion period to kindly encourage you to consider our latest update in your evaluation of our work. In particular, all other reviewers are now leaning towards acceptance of our submission, and we think that our new results are likely to address your concerns. Thank you!
I was referring to reference [3] in your paper. The "controlled setting" you mention means that you only compare approximations to approximations. In deep learning, continuous progress on benchmark datasets is shown. Each paper gets a better accuracy on benchmark datasets like Imagenet, and are compared with the strongest previous methods. These experiments don't do that, and that worries me. Overall, I am happy to accept that you have made an improvement for iterative methods. I am not convinced that your experiments show where your approximation stands in relation to other approximation methods, and what the best possible performance is. It also worries me that important baselines are being done during the rebuttal phase, in a rush, with mistakes made along the way, when they should have been core to the paper. Overall, this leaves me undecided. Perhaps the paper could be accepted on the basis of a claim of improvement only compared to other iterative methods, on limited datasets. I will discuss with the other reviewers, and AC.
Yes, the focus on truly large-scale datasets would make running a full Cholesky impossible. It is also fine to run these experiments on small datasets, where you can make the comparison, to give an indication of how your proposed method works across many datasets. (Provided you can discuss how the tuning of the method in the small-scale experiments is related to how the method can be tuned in the large-scale experiments where it cannot be compared to a ground truth.) It seems wrong that memory is a problem on Prot and Kegg. BFGS stores a Hessian approximation of size H x H, where H are the number of parameters to be optimized. Given the information in lines 58-62, Kegg should have H = 29, resulting in a minimal influence on memory? Your own results show that 100 iterations is not enough for some datasets (bike) to converge. Should you not then optimise for longer, since this is a trivial way to improve performance? In the neural networks community, comparing to a network that isn't trained for long enough isn't considered a proper baseline. The comparison to Adam is very helpful here. However, if bfgs can't be used with your proposed method, and this is a way to get better performance, this certainly would be a disadvantage of your method that should be discussed clearly. I had a look at some of the papers you cite that benchmark similar datasets, and it indeed seems like the rmse of bike can be improved by ~10x (see your reference [3]). This is a huge difference that is very impactful in places where GPs are used, e.g. in finance applications where a tiny improvement in signal from noise can be very valuable. In looking into this, I noticed that [3] raises an additional point about bike this paper does not investigate: The (CG) solver based method became unstable in the low noise setting, which is exactly what the dataset needs to perform well. This seems related to restricting the number of optimisation steps to 100, to avoid this behaviour from appearing. I'm happy to hear your thoughts on this. Overall, my current thinking is that this paper is empirical in nature (as is acknowledged by other reviewers). However, this makes good experiments crucial, and currently it seems that there are limitations in the baselines, and limitations in the method that are clear from other papers, but not properly discussed here.
*It seems wrong that memory is a problem on Prot and Kegg. BFGS stores a Hessian approximation of size H x H, where H are the number of parameters to be optimized. Given the information in lines 58-62, Kegg should have H = 29, resulting in a minimal influence on memory?* After careful investigation, we realised that our implementation based on jaxopt [1] was calculating the objective value and its gradient at the same time, using more memory than necessary in order to save computations. In particular, the high memory consumption was not due to the Hessian approximation of BFGS, as you pointed out correctly. Instead, simultaneous computation of the objective value and its gradient effectively instantiated multiple matrices of size $n \times n$, where $n$ is the amount of training data, leading to exhaustive memory consumption on the PROT and KEGG datasets. **We made the change to calculate the objective value and its gradient sequentially, thus saving memory** (at the cost of introducing more computations), and added the **new BFGS results for the PROT and KEGG datasets** to the table below (while replicating some results from the previous post to ease comparison). We conclude that, on the PROT and KEGG datasets, **BFGS achieves negligible improvements compared to Adam** in terms of test root-mean-square-error (RMSE) and test log-likelihood (LLH), similar to the POL and ELEV datasets (we discuss BIKE below). | | | | RMSE | | | | | | LLH | | | |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| | | POL | ELEV | BIKE | PROT | KEGG | | POL | ELEV | BIKE | PROT | KEGG | | Exact (BFGS) | 0.0714 | 0.3495 | 0.0317 | 0.5023 | 0.0832 | | 1.2765 | -0.3693 | 3.2685 | -0.5864 | 1.0818 | | Exact (Adam) | 0.0750 | 0.3550 | 0.0326 | 0.5025 | 0.0837 | | 1.2683 | -0.3856 | 2.1501 | -0.5868 | 1.0813 | --- *The comparison to Adam is very helpful here. However, if bfgs can't be used with your proposed method, and this is a way to get better performance, this certainly would be a disadvantage of your method that should be discussed clearly.* We agree that BFGS cannot be used with our proposed method and acknowledge this as a limitation. However, this is (currently) a **general limitation of scalable iterative methods for GPs** and Adam is commonly used instead, for example, Gardner et al. (2018) [4], Wang et al. (2019) [5] (who perform 10 iterations of L-BFGS and 10 iterations of Adam on a subset of size 10k followed by 3 iterations of Adam on the full dataset, and compare those results to performing 100 iterations of Adam on the full dataset), and Wu et al. (2024) [6] (who perform 50 or 100 iterations of Adam on the full dataset, depending on the dataset). There are several obstructions to using BFGS in this setting: stochastic gradients, evaluation of the log-marginal likelihood itself (not just its gradient), and computationally expensive linesearch. Arguably, stochastic gradients can be turned into deterministic gradients by using fixed probe vectors, which is required for warm starts anyways. However, evaluating the log-marginal likelihood in a deterministic way, in particular, **calculating the log-determinant of the kernel matrix, is intractable in the large-scale setting without low-rank approximations or lower bounds**. Even if it were possible to evaluate the log-determinant, the iterative linesearch in each BFGS iterations might become prohibitively expensive. Furthermore, on the $n$ < 50k datasets, where Cholesky factorisation + BFGS is still feasible, **the performance between BFGS and Adam, as reported in the table above, is small and arguably negligible** (except for bike, which we will discuss below). Given that Adam only requires (stochastic) evaluations of the gradient and does not need to perform computationally expensive linesearch, and thus scales to larger datasets, we believe this is a reasonable trade-off.
*I had a look at some of the papers you cite that benchmark similar datasets, and it indeed seems like the rmse of bike can be improved by ~10x (see your reference [3]). This is a huge difference that is very impactful in places where GPs are used, e.g. in finance applications where a tiny improvement in signal from noise can be very valuable.* Lin et al. (2024) [3] report a best RMSE value of 0.04 for the BIKE dataset (using SDD or CG), which is actually slightly worse than the values which we report for our iterative methods (ranging from 0.0324 - 0.0332). **Could it be that you accidentally compared the wrong numbers or columns?** Please correct us if we misunderstood your statement. --- *Your own results show that 100 iterations is not enough for some datasets (bike) to converge. Should you not then optimise for longer, since this is a trivial way to improve performance? In the neural networks community, comparing to a network that isn't trained for long enough isn't considered a proper baseline.* *In looking into this, I noticed that [3] raises an additional point about bike this paper does not investigate: The (CG) solver based method became unstable in the low noise setting, which is exactly what the dataset needs to perform well. This seems related to restricting the number of optimisation steps to 100, to avoid this behaviour from appearing.* **We performed additional experiments on the BIKE dataset using 1000 iterations of Adam** and report the test RMSE, test LLH, and noise variance $\sigma^2$ every 100 iterations in the table below. We observe that RMSE and LLH both continue to improve and do not appear to be fully converged even after 1000 iterations. After 10x the amount of iterations, RMSE values decreased by less than 15%, LLH values increased by more than 60%, while $\sigma^2$ decreases exponentially. **We conclude that the BIKE dataset is essentially noiseless, such that the test LLH can likely be increased indefinitely by reducing $\sigma^2$ until reaching numerical limits.** In particular, after 1000 iterations, $\sigma^2$ reaches ~$10^{-6}$ for iterative methods and ~$10^{-8}$ for exact optimisation, which is far below typical noise values used in this setting (Wu et al. (2024) [6] enforce $\sigma^2 > 10^{-4}$ and Wang et al. (2019) [5] use $\sigma^2 > 10^{-1}$ for large datasets), and close to the amount of jitter which is often added to the diagonal of a kernel matrix to ensure numerical stability when performing Cholesky factorisation. Notably, iterative methods (with warm starts and using the pathwise estimator) lead to larger $\sigma^2$ values than exact optimisation, which is expected because the pathwise estimator makes use of random features which introduce a small amount of noise. **Together with our conclusion above, the difference in $\sigma^2$ explains the performance gap in terms of test LLH.** This is supported by the fact that iterative methods actually achieve marginally better test RMSE than exact optimisation. We initially chose to perform 100 iterations of Adam to match previous work [5, 6] and because we mainly cared about providing a controlled environment, where all settings such as initial conditions, optimiser, learning rate, etc. are identical, such that the evaluation focuses on comparing the estimated gradients and the resulting optimisation trajectory. | Iterations | | 100 | 200 | 300 | 400 | 500 | 600 | 700 | 800 | 900 | 1000 | |:---:|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| | RMSE | CG + ws + pw | 0.0329 | 0.0311 | 0.0304 | 0.0299 | 0.0293 | 0.0289 | 0.0286 | 0.0283 | 0.0281 | 0.0279 | | | AP + ws + pw | 0.0330 | 0.0311 | 0.0304 | 0.0299 | 0.0294 | 0.0289 | 0.0286 | 0.0283 | 0.0281 | 0.0279 | | | Exact (Adam) | 0.0326 | 0.0312 | 0.0314 | 0.0312 | 0.0307 | 0.0301 | 0.0295 | 0.0290 | 0.0285 | 0.0281 | | | | | | | | | | | | | | | LLH | CG + ws + pw | 2.0615 | 2.4447 | 2.6774 | 2.8406 | 2.9655 | 3.0676 | 3.1544 | 3.2306 | 3.2991 | 3.3594 | | | AP + ws + pw | 2.0616 | 2.4439 | 2.6758 | 2.8387 | 2.9634 | 3.0652 | 3.1515 | 3.2268 | 3.2948 | 3.3554 | | | Exact (Adam) | 2.1501 | 2.6182 | 2.9202 | 3.1397 | 3.3092 | 3.4466 | 3.5621 | 3.6617 | 3.7492 | 3.8273 | | | | | | | | | | | | | | | $\sigma^2$ | CG + ws + pw | 3.52e-4 | 1.00e-4 | 4.69e-5 | 2.91e-5 | 2.08e-5 | 1.59e-5 | 1.24e-5 | 9.92e-6 | 7.99e-6 | 6.52e-6 | | | AP + ws + pw | 3.52e-4 | 9.88e-5 | 4.67e-5 | 2.88e-5 | 2.05e-5 | 1.56e-5 | 1.23e-5 | 9.87e-6 | 7.99e-6 | 6.53e-6 | | | Exact (Adam) | 2.38e-4 | 4.89e-5 | 1.22e-5 | 3.10e-6 | 9.13e-7 | 3.34e-7 | 1.47e-7 | 7.44e-8 | 4.16e-8 | 2.51e-8 |
Finally, we thank you for being a responsive and engaging reviewer, and hope that you find our additional results and discussion convincing. We are happy to include any of the results and discussion above in the camera-ready version. Please reach out to us again if there are further questions or concerns! Otherwise we kindly ask you to reconsider your score. --- [4] Gardner et al. (2018), "GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration", *Neural Information Processing Systems*. [5] Wang et al. (2019), "Exact Gaussian Processes on a Million Data Points", *Neural Information Processing Systems*. [6] Wu et al. (2024), "Large-Scale Gaussian Processes via Alternating Projection", *Artificial Intelligence and Statistics*.
Dear authors, Thanks for the detailed response to my questions! However, a theoretical analysis of the introduced combination technique is still missing. Such theoretical analysis is important to understand why, for example, pathwise + warm start is the best combination for AP and SGD but not for CG. Please comment on what approaches can help to model the effect of the selected combinations of the subroutines. This analysis may improve your work from purely empirical to more theoretical-based and introduce more novelty. Currently, I keep my score the same.
Dear Reviewer Xgts, thank you for responding to our rebuttal. While we did not have enough time to compose a rigorous mathematical proof about the effectiveness of warm starts + pathwise estimator, here is some intuition about their efficacy for different solvers: Both warm starts and the pathwise estimator reduce the number of required solver iterations until reaching the tolerance by reducing the distance which the solver has to travel in the quadratic optimisation landscape. Since CG performs an exact line search during each iteration, the actual distance to the solution is less impactful than the direction of descent (for example, starting arbitrarily far away from the solution but choosing an eigenvector as direction of descent still converges in one iteration). Nonetheless, the distance to the solution still matters, because each CG iteration monotonically decreases the energy norm of the residual, such that a smaller initial residual in general translates to a smaller number of iterations until convergence. For SGD, arguably, the initial distance is more important due to the constant step size / learning rate. A "bad" initial direction of descent can be compensated via accumulated momentum. For AP, the effects of warm starts are most profound, and we believe this is the case because AP chooses the descent subspace (the next "block") based on the corresponding residual. Therefore, if a warm start initialisation is "generally close" to the solution and only "far away in certain dimensions", the residual-based selection of blocks will converge quickly. Interestingly, if the problem is ill-conditioned enough, such as in the case of continued optimisation on the BIKE dataset (as requested by Reviewer NcMH), eventually warm starts + pathwise also becomes the best combination for CG, because (i) the standard estimator becomes too sensitive to the low noise (as explained in lines 116-127 and illustrated Figure 3), whereas the pathwise estimator is (more) robust, and (ii) warm starts seem to be generally helpful. Empirically, we observe the following behaviour: | Adam iterations | | 100 | 200 | 300 | 400 | 500 | 600 | 700 | 800 | 900 | 1000 | |:---:|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| | Solver iterations | CG + ws | 75 | 159 | 296 | 558 | 944 | 1401 | 2032 | 2861 | 3954 | 5251 | | until tolerance | CG + ws + pw | 80 | 100 | 148 | 185 | 189 | 214 | 221 | 376 | 389 | 491 | Thank you again for your time!
Thanks for the response. I think the paper will be in good shape provided that the authors add the promised discussions on 1. The pathwise estimator and its connection with the reparameterization trick, 2. The similarity and difference compared to existing work using probe vectors $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{P})$. I will maintain my score for now. I think the paper presents some solid contributions, though I do understand the concerns raised by other reviewers regarding the novelty. The review scores are quite divided, and it seems that addressing the concerns in the negative reviews is in order.
We thank you again for your effort put into reviewing our paper. Since there are only a few working days left for the discussion period, we would like to ask if our response resolved your concerns. If there are any remaining questions or concerns, we are happy to discuss them here. Otherwise we kindly invite you to raise your score.
Decision
Accept (poster)