Paper: Stochastic Gradient Descent for Gaussian Processes Done Right (ICLR-2024)

Reviewer comment (Evaluation):
In Fig. 1 you note that the primal gradient makes more progress in K^2-norm while the dual one in K-norm (with the same step size). However, in the left panel of Fig. 1 it seems that for a few iterations in the beginning of optimisation, the primal gradient was also making more progress than dual in the K-norm. Why do you think it is the case?

Author response:
Regarding Figure 1, specifically the ordering of performance right at the beginning of optimisation: we’ve thought about this, and we genuinely do not know! We think that this behaviour is relatively rare, as it occurs way out of convergence, with relatively few iterations spent in this regime, which should make it unlikely to affect behaviour in practice very much.

Reviewer comment (Experiments):
The GP hyper-parameters (e.g. observational noise variance, kernel parameters, etc.) are typically estimated by maximising the marginal log-likelihood using gradient-based methods. Do you think it could be possible run the gradient-based hyper-parameters inference jointly with the posterior inference that you discussed in this paper?

Author response:
Our method can indeed be easily combined with standard approaches to kernel hyperparameter optimisation. **We have added a demonstration of this in Appendix C.7** (we use an approach similar to [Gardner et al. 2018](https://arxiv.org/pdf/1809.11165.pdf)), and include some empirical results to that end where **we match performance of marginal likelihood optimisation with exact linear system solves (see the figure [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/mll_optim.pdf)).** Thank you for pointing out that this may be interesting to our readers to highlight!

Reviewer comment (Experiments):
More figures/tables to explicitly show the weakness/instability of the baseline methods are expected.

Author response:
**We have added a number of further comparisons where we show failure modes of the baselines we consider**, namely SGD from [Lin et al. 2023](https://arxiv.org/pdf/2306.11589.pdf), CG, and SVGP.

1. In Figures 12 and 13 in Appendix C.6, also available [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/primal_lr_divergence_pol_10.pdf) and [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/primal_lr_divergence_pol_20.pdf), you can see how the primal SGD baseline diverges very quickly at even moderately high learning rates.
2. In Figure 6 in Appendix C.1 ([here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/toy_comparison_rebuttal.pdf)), you can see failure modes for CG (on the infill asymptotics task) and SVGP (large-domain asymptotics task). Our dual SGD method performs robustly in both tasks.
3. In Figure 11 in Appendix C.5 ([here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/ill_conditioned_pol.pdf)), you can see that CG struggles in ill-conditioned settings, where the dual SGD method is quite robust.

Reviewer comment (Experiments):
Sharing more insights into the algorithm settings, such as the choice of geometric averaging, the effect/influence on the sparsity of the unbiased estimator \hat(g)(\alpha), etc, are expected.

Author response:
Sparsity of the random coordinate estimate plays no role in our recommendation of this gradient estimate. Rather we point it out to help the reader understand what this estimate looks like.

Regarding averaging, **our main point is that iterate averaging may not be necessary for convergence under multiplicative noise (Section 3.3)**, and can be completely done away with. Indeed, arithmetic iterate averaging slows down convergence empirically (see Figure 3).

However, we find that geometric averaging, which is much softer a form of averaging (places much more weight on recent iterates), is useful in practice. We thus recommend its use.

Reviewer comment (Theory):
A theoretical convergence analysis is expected (not only some figures).

Author response:
Our paper is targeted at GP researchers and practitioners. While using SGD for GP inference has recently been found to be very promising ([Lin et al. 2023](https://arxiv.org/pdf/2306.11589.pdf) and [Dai et. al. 2014](https://arxiv.org/pdf/1407.5599.pdf)), **previous work makes a number of suboptimal choices when it comes to the details of the algorithm's formulation. Our contribution is handling these details carefully, highlighting to researchers and practitioners how to make SGD inference even more effective**, particularly on harder problems like the graph kernel comparison.


In the **full-batch (non stochastic) case, the analysis of the algorithm is standard** ([Nesterov 1983](https://www.mathnet.ru/php/archive.phtml?wshow=paper&jrnid=dan&paperid=46009&option_lang=eng)), and it is known that the dual objective will result in faster convergence than its primal counterpart due to conditioning. See, for instance Section 3.7, page 289, of [Bubeck 2015](https://arxiv.org/pdf/1405.4980.pdf), a textbook on convex optimisation. 

Our full algorithm makes use of an estimator with multiplicative noise. **The convergence of Stochastic GD under multiplicative noise is currently at the forefront of theoretical research and beyond the scope of our more practical work.** We have extended our discussion of pertinent work in this space in Section 3.4. We hope that our work's demonstration of the empirical performance achievable with SGD under multiplicative noise will motivate research in this area by the optimisation community.

Reviewer comment (Evaluation):
In Figure 1, we can not see the primal gradient descent becomes unstable and diverges for $\beta n$>0.1. Please show the unstable or compare the evaluated conditional numbers. Under higher step sizes, why does the gradient descent of the primal return $NaN$ (any possible reasons)?

Author response:
**We include the requested figures showing the divergence of primal SGD with high step-sizes, [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/primal_lr_divergence_pol_10.pdf) and [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/primal_lr_divergence_pol_20.pdf), and Appendix C.6. This behaviour quickly exhausts the numerical floating point precision, leading to  overflows and thus NaNs.**

We did not include these in the original text because the loss curves diverge so quickly that they cannot be viewed simultaneusly with the optimisation trajectories of the non-diverging methods. We run our algorithm for 100k iterations. After only 10 iterations, amplifying oscillations are visible and after 20 iterations an exponential trend is already present.

Reviewer comment (Experiments):
Figure 2 shows the random coordinate estimate with a step size equal to 50. What is the performance on varied step sizes? Can any explanation of the rising part (the blue dashed line) in the middle figure in Figure 2 be given?

Author response:
We have conducted **further experiments** under the same settings with **varying step-sizes** for both random coordinate estimators (the one with additive and the one with multiplicative noise). **Please see the corresponding figure [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/lr_comparison_pol.pdf)** and more details in Appendix C.3.


The performance of the additive noise estimator plateaus earlier as the step-size increases. Smaller step-sizes converge slower but eventually reach better solutions, such that a trade-off is necessary: the step-size is proportional to the speed of convergence but inversely proportional to the final performance. In contrast, the multiplicative noise estimator consistently improves both convergence speed and final performance as the step-size is increased, as long as the step-size is not large enough to cause divergence.

Reviewer comment (Reproducibility):
What is the step size used to generate the Figure 3? It seems less than 50 and has a competitive rate compared to the random feature estimate shown in Figure 2. Extra clarification and comparison would be better.

Author response:
The step-size used to generate Figure 3 was **$\beta n = 50$. We added this to the figure caption now**. 

**Both plots use a step-size of 50**, but different batch sizes. However, we have discovered the **lines are different due to the floating point precision** used when running the experiment. Figure 2 was generated with 32-bit floating point, which becomes unstable when the error is very small (~ 1e-7). We have now replaced this with 64-bit floating point, so **this is now fixed. Please see the fixed plot [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/batch_variants_pol.pdf). To improve clarity, we also added "random features" and "random coordinates" to the labels in Figure 2.**

Reviewer comment (Experiments):
How do different batch sizes affect the overall convergence?

Author response:
**We conducted additional experiments with varying batch sizes. Please see our figures [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/batch_size_comparison_pol.pdf) and [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/batch_size_comparison_pol_per_time.pdf)** and Appendix C.4 for further details.

Smaller batch sizes degrade the final performance of the additive noise estimator because the amount of additive noise is inversely proportional to the batch size. In contrast, the multiplicative noise estimator produces nearly identical optimisation traces with different batch sizes, as long as the batch size is not small enough to cause divergence. Therefore, increasing the batch size consistently improves the final performance of the additive noise estimator while it does not impact the performance of the multiplicative noise estimator, given that the algorithm converges.

In terms of performance vs time, a larger batch size requires additional compute per step. Thus, the additive noise estimator performance is proportional to the batch size and the multiplicative noise estimator performance per time is inversely proportional to the batch size, as long as the algorithm converges. Therefore, the multiplicative noise estimator should use the smallest batch size which does not cause divergence. Selecting such a batch size a priori is, in general, non-trivial. We have included this discussion in Appendix C4.

Reviewer comment (Experiments):
It is better to add a test where samples are generated by a non-stationary kernel, to show the ability of the random coordinate estimate (to distinguish with the random Fourier features).

Author response:
Thank you for the question. **Our optimisation method does not require the existence of closed-form random features: for computing the posterior mean, we can work with any kernel**. We only introduce feature-based gradient estimators due to their use in previous literature ([Lin et al. 2023](https://arxiv.org/pdf/2306.11589.pdf) and [Dai et. al. 2014](https://arxiv.org/pdf/1407.5599.pdf)). We ultimately recommend against their use for gradient estimation due to producing additive rather than multiplicative noise.

We do use random features for efficiently sampling from the GP prior, which is needed, for instance, in the Bayesian optimisation experiments. We similarly use random hashing for the Tanimoto kernel in the graph experiment. In non-stationary kernels, efficient sampling can often be done on a case-by-case basis, but we prefer not to focus on it as it is orthogonal to our contributions.

Reviewer comment (Writing):
What is the difference between the $\beta n$ in the main texts and the $\beta$ in Algorithm 1?

Author response:
The $\beta n = \beta \times n$ is simply the learning rate $\beta$ times $n$, the size of the training data set. This formulation allows us to **keep the step-size $\beta$ times the data set size $n$ constant across data sets**. We mostly use $\beta n$ throughout our writing, since the algorithm ultimately depends on the respective product.

Reviewer comment (Presentation):
The green dashed line is missing in Figure 3.

Author response:
Thank you for pointing this out! **We have updated the caption to reflect that this refers to the olive green line.**

Reviewer comment (Theory):
This paper only provides numerical experiments to evaluate the performance of different algorithms. However, it would be good if rigorous theoretical guarantees could be proved, at least for some special cases. Besides, I think the authors stress too much on the algorithm details, which can be deferred to the appendix for a major part of them while trying to leave some room for theoretical analysis.

Author response:
**1. Theoretical guarantees**

> "This paper only provides numerical experiments... would be good if rigorous theoretical guarantees could be proved, at least for some special cases."

Our paper is targeted at GP researchers and practitioners. While using SGD for GP inference has recently been found to be very promising ([Lin et al. 2023](https://arxiv.org/pdf/2306.11589.pdf) and [Dai et. al. 2014](https://arxiv.org/pdf/1407.5599.pdf)), **previous work makes a number of suboptimal choices** when it comes to the details of the algorithm's formulation. **Our contribution is handling these details carefully, highlighting to researchers and practitioners how to make SGD inference even more effective**, particularly on harder problems like the graph kernel comparison.

In the **full-batch (non-stochastic) case, the analysis of the algorithm is standard** ([Nesterov 1983](https://www.mathnet.ru/php/archive.phtml?wshow=paper&jrnid=dan&paperid=46009&option_lang=eng)), and it is known that the dual objective will result in faster convergence than its primal counterpart due to conditioning. See, for instance section 3.7, page 289, of [Bubeck 2015](https://arxiv.org/pdf/1405.4980.pdf), a textbook on convex optimization. 

Our full algorithm makes use of an estimator with multiplicative noise. **The convergence of Stochastic GD under multiplicative noise is currently at the forefront of theoretical research and beyond the scope of our more practical work.** We have extended our discussion of pertinent work in this space in Section 3.4. We hope that our work's demonstration of the empirical performance achievable with SGD under multiplicative noise will motivate research in this area by the optimisation community.

**2. Algorithmic details.** 

> "The authors stress too much on the algorithm details. which can be deferred to the appendix for a major part of them while trying to leave some room for theoretical analysis."

Thank you for this comment. **Our aim was to show how to make the already promising SGD-based learning algorithms for Gaussian processes more effective by giving a careful treatment to algorithmic details. To do so, it is very important to precisely state the algorithm we recommend**, particularly for readers coming from the Gaussian process rather than optimisation community.

At the same time, we recognize that papers which focus more on the side of theoretical analysis are also valuable. We believe it is best to have both kinds of papers available in the literature, so that readers can more precisely find what they are looking for.

Reviewer comment (Experiments):
There are many different optimizers for the kernel ridge regression, such as AdaGrad, Adam, etc. The authors should also try these methods in the experiments.

Author response:
**3. Comparison with other optimisers.** 

>  different optimizers such as AdaGrad, Adam, etc. The authors should also try these methods.

We now include an empirical comparison on the POL dataset of AdaGrad vs RMSprop vs Adam vs our method. **We show that our method compares favourably to these methods by a significant margin: see [linked figure](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/optimiser_comparison_pol.pdf) and Appendix C.2.**

Thank you for pointing this comparison out. We originally opted to omit these algorithms, as they are usually designed to tackle problems with a non-constant curvature. GP regression has a constant curvature: it yields a quadratic objective where Nesterov-type momentum is theoretically rate-optimal ([Dieuleveut et. al. 2017](https://arxiv.org/abs/1602.05419) and [Jain et. al. 2018](https://arxiv.org/abs/1704.08227)). However, given the prevalence of other optimisation algorithms, we agree that a comparison to them is valuable to the larger community!

Reviewer comment (Novelty):
The algorithm design is a bit incremental to me, as it looks like a combination of standard existing approaches, which is tuned for the specific tasks. Then, the idea of the algorithm design may be difficult to extend to other tasks.

Author response:
**4. Significance of work**

> The algorithm design is a bit incremental... tuned for the specific tasks. ...may be difficult to extend to other tasks.

We would like to highlight that **Gaussian processes are a very important class of models, which are widely used for Bayesian optimisation in areas like computational chemistry** (for instance in protein design, see [here](https://pubs.acs.org/doi/10.1021/acscentsci.7b00572)), **hyperparameter tuning** (see for instance [Google Vizier](https://github.com/google/vizier)), and others. Designing specialised optimisation algorithms which take advantage of their specific structure is therefore a worthwhile problem, which has long been held back by cubic computational cost in data set size. Building on the initial results of [Lin et al. 2023](https://arxiv.org/pdf/2306.11589.pdf), **our algorithm is the first alternative which reliably outperforms conjugate gradient-based algorithms used in state-of-the-art Gaussian process software packages.**

Extending this, solving quadratic problems emerging from kernel regression efficiently can also be applied to inference in non-conjugate settings, like classification. For instance, [Antoran et. al. 2023](https://arxiv.org/pdf/2210.04994.pdf) combine SGD with the Laplace approximation to perform inference with the NTK on ImageNet classification. Our faster dual SGD technique is a drop-in replacement into their paper.

**5. Combination of existing approaches**

>The algorithm design is ... a combination of standard existing approaches 

We agree that our paper uses insights from the theory literature—indeed, this is highlighted throughout our text. To our knowledge, these insights have not been studied in a unified way by the kernel/GP community before. Even worse, **all applied kernel/GP work we are familiar with, which use SGD methods, make suboptimal algorithmic design choices**. On this basis we believe that a simple algorithm that combines recent theoretical advancements in a principled way and produces state-of-the-art results is valuable to the ICLR community.

We would also like to note that in machine learning, there is a long history of combinations of small improvements eventually leading to breakthrough performance: the most prominent example, arguably, is the development of deep neural networks, which required small tweaks to initialisation, activation functions, architecture, optimisation, and other seemingly incremental factors in order to result in today’s effectiveness.

Reviewer comment (Experiments):
Besides, it is not clear to me whether the variance of stochastic gradient is really a big issue from Figure 2, as the authors do not add the full-gradient version for comparison. If controlling the variance is important, the authors may also need to consider variance-reduce techniques (e.g., SVRG) and add them to the experiments.

Author response:
**6. Gradient variance**

> ...it is not clear to me whether the variance of stochastic gradient is really a big issue from Figure 2, as the authors do not add the full-gradient version for comparison.

Gradient variance can significantly harm SGD performance if proper precautions are not taken. This can be seen through the effect of geometric averaging, which aims to ameliorate the effects of noisy gradient estimates. This is illustrated **in Figure 3, where the blue line with averaging (i.e. variance reduction) is consistently below the non-averaged olive-green line.**

This can also be seen when using an estimator of the gradient with additive (instead of multiplicative) noise. With geometric averaging and our proposed multiplicative noise gradient estimator, the effects of variance are notably lessened. **We have added ablations over step-size and batch size for both additive and multiplicative estimators of the gradient (see figures [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/lr_comparison_pol.pdf) and [here](https://anonymous.4open.science/r/SDD-GPs-5486/rebuttal_plots/batch_size_comparison_pol.pdf) and Appendix C.3 and C.4 for further details).** The additive estimator's performance suffers when the step-size is increased, or batch size is decreased; that is when the gradient estimator's variance increases. In contrast, the proposed multiplicative estimator's performance increases as the step-size increases, whereas the additive estimator gets worse. Similarly, the multiplicative estimator is significantly more robust to reductions in batch size.