Stochastic Gradient Descent for Gaussian Processes Done Right
Authors: Jihao Andreas Lin
Shreyas Padhy
Javier Antoran
Austin Tripp
Alexander Terenin
Csaba Szepesvari
José Miguel Hernández-Lobato
David Janz
Venue: ICLR 2024 poster (2024)
Keywords: Gaussian process
stochastic gradient descent
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Abstract: As is well known, both sampling from the posterior and computing the mean of the posterior in Gaussian process regression reduces to solving a large linear system of equations. We study the use of stochastic gradient descent for solving this linear system, and show that when done right---by which we mean using specific insights from the optimisation and kernel communities---stochastic gradient descent is highly effective. To that end, we introduce a particularly simple stochastic dual descent algorithm, explain its design in an intuitive manner and illustrate the design choices through a series of ablation studies. Further experiments demonstrate that our new method is highly competitive. In particular, our evaluations on the UCI regression tasks and on Bayesian optimisation set our approach apart from preconditioned conjugate gradients and variational Gaussian process approximations. Moreover, our method places Gaussian process regression on par with state-of-the-art graph neural networks for molecular binding affinity prediction.

STOCHASTIC GRADIENT DESCENT FOR GAUSSIAN PROCESSES DONE RIGHT
Jihao Andreas ! Shreyas Padhy Javier Antorán" Austin Tripp'
Alexander Terenin 1.3 Csaba Szepesvári'
José Miguel Hernández-Lobato David Janz'
University of Cambridge
Max Planck Institute for Intelligent Systems
CComell University
*University of Alberta
ABSTRACT
As is well known, both sampling from the posterior and computing the mean of the posterior in Gaussian process regression reduces solving a large linear system of equations. We study the use of stochastic gradient descent for solving this linear system, and show that when done right- by which we mean using specific insights from the optimisation and kernel communities caaamunitess-stohhsstttt gradient descent is highly effective. To that end, we introduce a particularly simple stochastic dual descent algorithm, explain its design in an intuitive manner and illustrate the design choices through a series of ablation studies. Further experiments demonstrate that our new method is highly competitive. In particular, our evaluations the UCI regression tasks and on Bayesian optimisation set our approach apart from preconditioned conjugate gradients and variational Gaussian process approxima- tions. Moreover, our method places Gaussian process regression on par with state-of-the-art graph neural networks for molecular binding affinity prediction.
INTRODUCTION
Gaussian process regression the standard modelling choice Bayesian optimisation and other applications where uncertainty-aware decision decision-making is required to gather data efficiently. The main limitation of Gaussian process models that their fitting requires solving large linear system of equations which, using direct methods, has cost cubic in the number of observations.
Standard approaches to reducing the cost of fitting Gaussian process models either apply approxi- mations to reduce the complexity of the linear system, such as the Nyström and related variational approximations (Williams and Seeger, 2000; Titsias, 2009; Hensman et al., 2013), or use carefully engineered iterative solvers, such preconditioned conjugate gradients (Wang al., 2019), or employ a combination of both (Rudi et al., 2017). An alternative approach that we focus on this work is the use of stochastic gradient descent minimise an associated quadratic objective.
Multiple works have previously investigated the use of SGD with Gaussian process and related kernel models (Lin et al., 2023; Dai et al., 2014; Kivinen al., 2004), with Lin et al. (2023) pointing out, in particular, that SGD may be competitive with conjugate gradients (CG) when the compute budget is limited, in terms of both the mean and uncertainty predictions that produces. In this work, we go step further, and demonstrate that when done right, SGD can outperform CG.
To that end, we propose simple SGD-based algorithm, which we call stochastic dual descent (SDD). Our algorithm is an adaptation of ideas from the stochastic dual coordinate ascent (SDCA) algorithm of Shalev-Shwards and Zhang (2013) the large-scale dcep-learning-tyy terriinngtyppeeee gradient descent setting, combined with insights on stochastic approximation from Dieuleveut et al. (2017) and Varre et al. (2021). We provide the following evidence supporting the strength of SDD:
On standard UCI regression benchmarks with up to 2 million observations, stochastic dual descent either matches or improves upon the performance of conjugate gradients.
'Equal contribution.
Code available at: HTTes///eTHHIn..CMMCAAAMEEEEEEE

---

2. On the large-scale Bayesian optimisation task considered by Lin et al. (2023), stochastic dual descent is shown to be superior to their stochastic gradient descent method and other baselines, both against the number of iterations and against wall-clock time.
3. On molecular binding affinity prediction task, the performance of Gaussian process regres- sion with stochastic dual descent matches that state-of-the-ar graph neural networks.
In short, we show that a hemple-hut-well-essiie can be very competitive with other approaches for Gaussian processes, and may make Gaussian processes competitive with graph neural networks on tasks on which the latter are state-of-the-art.
stochastic gradient method for Gaussian processes
2 GAUSSIAN PROCESS REGRESSION
We consider Gaussian process regression over a domain C R assuming a Gaussian process prior induced by continuous, bounded, positive-definite kernel k: X' X R. Our goal is to, given some observations, either (i) sample from the Gaussian process posterior, or (ii) compute its mean.
To formalise these goals, let some observed inputs -1,...,--i X' be collected in matrix X E Raad and let IR" denote the corresponding observed real-valued targets. We say that a random function fn: X R is sample from the posterior Gaussian process associated with the kernel data set (X, and likelihood variance parameter À > 0, if the distribution of any finite set of its marginals is jointly multivariate Gaussian, and the mean and covariance between the marginals of fn at any a and a' in X' are given by
(a) B[fx(a)] = + and (or( = k(a, k(a, X)( + A
respectively, where E R=x== the n x n matrix whose (i.j) entry is given by k(2,, xÿ). I the n x n identity matrix, k(a, X) is row vector in R" with entries k( ........ a column vector defined likewise (per the notation of Rasmussen and Williams, 2006).
......,((a,, k(a, xn,) and k(X, a')
Throughout, for two vectors a, b E R2, with N+, we write a for their usual inner product and ||a| for the 2-norm a. For a symmetric positive semidefinite matrix M E R we write ||a|| for the M M-weighted 2-norm of a, given by [al]; M a Ma. For a symmetric matrix S, we write ||| for the operator norm of 5, and (S) for the ith largest eigenvalue of 5, such that (S) = ||||||
STOCHASTIC DUAL DESCENT FOR REGRESSION AND SAMPLING
We show our proposed algorithm, stochastic dual descent, in Algorithm 1. The algorithm can be used for two purposes: regression, where the goal is to produce good approximation the posterior mean function m,, and sampling, that is, drawing a sample from the Gaussian process posterior.
The algorithm takes input kernel matrix K, the entries of which can be computed 'on demand' using the kernel function h and the observed inputs Ii 2i,..., Fn as needed, a likelihood variance 0, and a vector b R., It produces vector of coefficients Tr E R which approximates
(b) (K + Xb)
To translate this into goals of mean estimation and sampling, given vector R", let
ha = Σ (ak((,)))
Then, har(s) gives the mean function ma, and thus running the algorithm with the targets y as the vector b can be used estimate the mean function. Moreover, given sample fo from the Gaussian process prior associated with i, and noise ( (O, AT), we have that
fo + hat (g- (X)(+()))
is a sample from the Gaussian process posterior (Wilson et al., 2020, 2021). In practice, one might approximate fo using random features (as done in Wilson al., 2020, 2021; Lin et al., 2023).
2

---

Algorithm Stochastic dual descent for approximating a² (K + X1
Require: Kernel matrix K with rows K,...., Kn E _,", targets E IR", likelihood variance À > 0, number of steps T E N+ batch size B E f1..... n), step size O,
momentum parameter p E [0, 1). averaging parameter r E (0, 1]
1: 00 = 0; <ig = dig
all l
2: for / E , do
3:
Sample I, = (if
(if..... . (g)) Uniform((L . n.) independently
random coordinates gradient estimate b velocity update parameter update iterate averaging
4:
9t i LieT, (Ki + )ei) (at-1 + por-1) b;)e:
5:
Ut = p 3gt
6: 04 04-1 1 12
7: T = ra; + (1-11 (1-r)
8:
return OT
The SDD algorithm is distinguished from variants of SGD used in the context of GP regression by the following features: (i) SGD uses a dual objective in place of the usual kernel ridge regression objective, (ii) it uses stochastic approximation entirely via random subsampling of the data instead of random features, (iii) it uses Nesterov's momentum, and (iv) it uses geometric, rather than arithmetic, iterate averaging. the following subsections, we examine and justify each of the choices behind the algorithm design, and illustrate these the UCI data set POL (Dua and Graff, 2017), chosen for its small size, which helps us to compare against less effective alternatives.
3.1 GRADIENT DESCENT: PRIMAL VERSUS DUAL OBJECTIVES
For any b R" computing the vector a (b) of equation (1) using Cholesky factorisation takes on the order of n operations. We now look at how a (b) may be approximated using gradient descent.
As is well known, the vector a
a²(b) (b) is the minimiser of the kernel ridge regression objective,
[(a)= = 2 1,1bb Kall++ + 2 ||all
R" (Smola and Schölkopf, 1998). We will refer to L as the primal objective. Consider using gradient descent to minimise L(a). This entails constructing a sequence (ax), of elements in R" which initialise the standard but otherwise arbitrary choice og 0, and setting
at-1 = a4 LV
where B > O is step-size and VL is the gradient function of T.. Recall that the speed at which approaches (b) determined by the condition number of the Hessian: the larger the condition number, the slower the convergence speed (Boyd and Vandenberghe, 2004). The intuitive reason for this correspondence is that, to guarantee convergence, the step-size needs to scale inversely with the largest eigenvalue of the Hessian, while progress in the direction of an eigenvector underlying eigenvalue is governed by the step-size multiplied with the corresponding eigenvalue. With that in mind, the primal gradient and Hessian
(((a)= = K(Aa ++ Ka) and V2L(a) = K(K + Al)
respectively, and therefore the relevant eigenvalues are bounded as
An (K AQ) S (K + Af)) S Kn(Kn A),
where λ = sup k(x, x is finite by assumption. These bounds only allow for a step-size 3 on the order of (Km(Rn -1 Moreover, since the minimum eigenvalue not bounded away from zero, we do not have priori guarantees for the performance of gradient descent using L.
Consider, instead, minimising the dual objective
L (a) = |aal Tb.
The dual L* has the same unique minimiser L. namely a² (b), and the two are, up rescaling, strong dual pair (in a sense made precise in Appendix A). The dual gradient and Hessian are given by
g(a) VLA (a) Xa
Ka and (LI K N.
3

---

1. ||a
a² (1)
1|@ a² (v)ll
Test RMSE
0.60
0.20
0.48
0.45
0.15
0.36
0.30
0.10
0.24
0.15
0.05
0.12
0.00
0.00
0.00
25k
75k
100k
25k
75k
100k
25k
50k
75k
100k
Iterations
Iterations
Iterations
VL with 8n = 0.1
with 3n = 0.1
with In = 5
with On = 50
Figure 1: Comparison of full-batch primal and dual gradient descent POL with varying step-sizes. Primal gradient descent becomes unstable and diverges for 3m greater than 0.1. Dual gradient descent is stable with larger step-sizes, allowing for markedly faster convergence than the primal. For 38 0.1, the dual method makes more progress in the K Knnrrm whereas the primal in the K -orm.
Examining the eigenvalues the above Hessian, we see that gradient descent on L (dual gradient descent) may
step-size of up order of en higher than that L (primal gradient descent), leading to faster convergence. Moreover, since the condition number of the dual satisfies cond K + provide priori bound on the number of iterations
XI) S 1 + sn/A and À is positive, required for dual gradient descent to convergence to any fixed error level for a length n sequence of observations.
In Figure 1, we plot the results of an experimental comparison of primal and dual gradient descent on the UCI POL regression task. There, gradient descent with the primal objective is stable up to 3n = 0.1 but diverges for larger step-sizes. In contrast, gradient descent with the dual objective stable with step-size as much as 500x higher, converges faster and reaches a better solution; see also Appendix B for more detailed plots and recommendations on setting step-sizes. We show this on three evaluation metries: distance to a (yy) measured in || l1 the K-norm (squared) and in || 1lka,, the normorm (squared), and test set root-mean-square error (RMSE). To understand the difference between the two norms, note that the K-norm error bounds the error of approximating (b) uniformly. Indeed, shown in Appendix A, we have the bound
hha (1)l S V - a² a (b)lkK
(5)
where recall that i = supzex k( 2). Uniform norm guarantees of this type are crucial for sequential decision-making tasks, such as Bayesian optimisation, where test input locations may be arbitrary. The norm metric, on the other hand, reflects the training error.
Examining the two gradients, its immediate that the primal gradient optimises for the K²-norm, while the dual for the K-norm. And indeed, we see in Figure that when both methods use Bn 0.1, up to 70k iterations, the dual method is better on the K-norm metric and the primal on K² Later, the dual gradient method performs better on all metrics. This, too, is to be expected, as the minimum eigenvalue of the Hessian of the dual loss is higher than that of the primal loss.
3.2 RANDOMISED GRADIENTS: RANDOM FEATURES VERSUS RANDOM COORDINATES
To compute either the primal or the dual gradients, presented in equations (2) and (4), we need to compute matrix-vector products of the form Ka, which requires order n² computations. We now introduce and contrast two types stochastic approximation for our dual gradient g(a) that reduce the cost order n per-step, and carefully examine the noise they introduce into the gradient.
For the sake of exposition (and exposition only, this is not an assumption our algorithm). assume that we in an m-dimensional (finite) linear model setting, such that K is of the form ΣØ= 2,2]. where 2j E Rn collects the values of the jth feature of the observations Z1,..., Zn. Then, for j Uniform((1,.. .,,,m)) m)), we have that E|mzz ==] K. and therefore
π(a) Aa b m2,2Jaa
4

---

1 ||a
a" (v)ll
x l|a
a" 21)))
Test RMSE
10
10
1.00
10
10 4
10
0.50
10-6
10-6
0.25
10
0.00
25k
100k
25k
50k
75k
100k
0
25k
50k
75k
100k
Iterations
Iterations
Iterations
with On = 5 x 10
ne:e! (Ka) da with 8n
50
with An = 5 x 10-3
neçe, (Ka y + An) = g(a) with 3n =
Random features
Random coordinates
Figure 2: A comparison of dual stochastic gradient descent on the POL data set with either random Fourier features random coordinates, using batch size B 512, momentum p= 0.9 and averaging parameter r' = 0.001 (see Section 3.3 for explanation of latter two). Random features converge with 3n 5 10 but perform poorly, and diverge with a higher step-size. Random coordinates stable with 3n 50 and show much stronger performance on all metrics. We include a version of random coordinates where only the Ka term subsampled: this breaks the multiplicative noise property, and results in an estimate which is worse on both the K K-norm and the Kornorm metric.
an unbiased estimate of g(x):: we call ğ the random feature estimate of g. The alternative we consider is random coordinates, where we take i (1iiform(11....... and use
ī(a) ne.e g(a) = ne; (X;; b, + Kça)).
Observe that g(a) zeros all but the ith (of n) coordinate of g(a), and then scales the result by n; since E [ = I, g(a) is also an unbiased estimate of g(a). Note that the cost of calculating either g(a) or ((x) is linear in n, and therefore both achieve our goal of reduced computation time.
However, while ÿ and ğ may appear similarly appealing, the nature of the noise introduced by these, and thus their qualities, are quite different. In particular, one can show that
(xx(x) g(a)l s II(M (K + XT)IIIaa - a² (b)1ll
where .M ne,e; the random coordinate approximation the identity matrix. As such, the noise introduced by F(a)) is proportional to the distance between the current iterate and the target a*, making it a so-called multiplicative noise gradient estimate (Dieuleveut al., 2017). Intuitively, multiplicative noise estimates automatically reduce the amount of noise injected as the iterates get closer to their target. On the other hand, for @(a), letting K m2ya], we have that
(ī(a) g()) || 11(K K)al1,
and thus the error in ğ(a) is not reduced as approaches a (b). The behaviour of g is that of an additive noise gradient estimate. Algorithms using multiplicative noise estimates, when convergent, often yield better performance (see Varre al., 2021).
Another consideration is that when m, the number of features, is larger than n. individual features 2j may also be less informative than individual gradient coordinates egc g(a). Thus. picking features uniformly at random, as in y. may yield poorer estimate. While this can be addressed by introducing an importance sampling distribution for the features zj (as per Li et al., 2019), doing so adds implementation complexity, and could be likewise done to improve the random coordinate estimate.
In Figure 2, we compare variants of stochastic dual descent with either random (Fourier) features or random coordinates. We see that random features, which produce high-variance additive noise, can only be used with very small step-sizes and have poor asymptotic performance. We test two versions of random coordinates: δ(a), where, presented, alternative, ne,e, (Ka) much higher sizp-sizes than random features, the latter has worse asymptotic performance. This is
subsample the whole gradient, and an λa, where only the Ka term is subsampled. While both are stable with
5

---

1 ||
a² a>)il
1 ||@ a a²(v)l(2>
Test RMSE
10
10a
0.11
10
10-
0.10
m
10-3
10-6
0.09
10-12
0
25k
50k
75k
100k
25k
50k
75k
100k
0
25k
50k
75k
100k
Iterations
Iterations
Iterations
AdaGrad
RMSprop
Adam
Nesterov
Figure 3: Comparison of dual stochastic gradient descent on the POL data set with different accelera- tion methods, using batch size B 512, geometric averaging parameter 0.001, and step-sizes tuned individually for each method (AdaGrad 3n 10; RMSprop & Adam 3n = 0.05; Nesterov's momentum Bn 50). Both Adam and Nesterov's momentum perform well Test RMSE, but the latter performs better on the K and K2 norms.
kind of Rao-Blackwellisation trap: introducing the known value of -y + λa in place of its estimate ne;e? (-y + λo:) destroys the multiplicative property of the noise, making things worse, not better.
Algorithm combines a number of samples into a minibatched estimate to further reduce variance. We discuss the behaviour of different stochastic gradient estimates under minibatching in Appendix B.
There are many other possibilities for constructing randomised gradient estimates. For example, Dai et al. (2014) applied the random coordinate method on top of the random feature method, in an effort further reduce the order cost per iteration of gradient computation. However, based on the above discussion, we generally recommend estimates that produce multiplicative noise, and our randomised coordinate gradient estimates as the simplest of such estimates.
3.3 NESTEROV'S S MOMENTUM AND POLYAK- RUPPERT ITERATE AVERAGING
Momentum, acceleration, range of modifications the usual gradient descent updates that aim to improve the rate of convergence, in particular with respect to its dependence on the curvature of the optimisation problem (Polyak, 1964). Many schemes for momentum exist, with AdaGrad (Duchi et al., 2011), R.SPProp (Tieleman and Hinton, 2012) and Adam (Kingma and Ba, 2015) particularly popular the deep learning literature. These, however, are designed to adapt to changing curvatures. As our objective has a constant curvature, we recommend the use of the simpler Nesterov's momentum (Nesterov, 1983; Sutskever al., 2013), as supported by our empirical results in Figure with further results in Figure 4, showing that momentum is vital for the POL data set. The precise form of the updates used is shown Algorithm 1; we use momentum of p = 0.9 throughout.
Polyak-Rupper averaging returns the average of the (tail of the) stochastic gradient iterates, rather than the final iterate, so to reduce noise in the estimate (Polyak, 1990; Ruppert, 1988; Polyak and Juditsky, 1992). While Polyak-Ruppert averaging is necessary with constant step-size and additive noise, not under multiplicative noise (Varre al., 2021), and indeed can slow convergence, For our problem, we recommend using geometric averaging instead, where we let Tig og and, at each step compute
T2 = ra, + -1- r- for an averaging parameter r E (0,1)),
and return as our estimate ür.. Geometric averaging an anytime approach, in that does not rely on fixed lail-window size; it can thus be used in combination with early stopping, and the value of can be tuned adaptively. Figure 4 shows that geometric averaging outperforms both arithmetic averaging and returning the last iterate, ar. Here and throughout, we set 100/T.
3.4 CONNECTIONS TO THE LITERATURE
The dual formulation for the kernel ridge regression objective was first pointed out in the kernel literature in Saunders al. (1998). It features in textbooks the topic (see equations (10.95) and

---

1
|| a² k
1 || a a"(v)li2>
Test RMSE
10a
11
10-2
0.088
10 -
10-4
0.086
10
10 6
0.084
10
10
0.082
10-10
0.080
0
25k
50k
75k
100k
25k
50k
75k
100k
0
25k
50k
75k
100k
Iterations
Iterations
Iterations
No momentum
Momentum
Momentum
Momentum
No averaging
No averaging
Arithmetic avg.
Geometric avg.
Figure 4: Comparison of optimisation strategies for random coordinate estimator of the dual objective on the POL data set, using momentum p = 0.9, averaging parameter r 0.001, batch size B = 128, and step-size 3n = 50. Nesterov's momentum significantly improves convergence speed across all metries. The dashed olive line, marked arithmetic averaging, shows the regular iterate up until 70k steps, at which point averaging commences and the averaged iterate is shown. Arithmetic iterate averaging slows down convergence in K- K-norm once enabled. Geometric iterate averaging, on the other hand, outperforms arithmetic averaging and unaveraged iterates throughout optimisation.
(10.98) in Smola and Schölkopf, 1998), albeit with no mention of its better conditioning, Gradient descent on the dual objective equivalent applying the kernel adatron update rule of Frie et al. (1998), stated there for the hinge loss, and to the kernel least-mean-squar algorithm of Liu al. (2008), analysed theoretically in Dieuleveut et (2017). Shalev-Shwartz and Zhang (2013) introduce stochastic dual coordinate ascent (SDCA), which uses the dual objective with random coordinates and analytic line search, and provide convergence results. Tu et al. (2016) implement block (minibatched) version of SDCA. Bo and Sminchiseseu (2008) propose method similar to SDCA, but with coordinates chosen by active learning. Wu et al. (2023) reimplement the latter two for Gaussian process regression, and link the algorithms to the method of alternating projections.
From the perspective of stochastic gradient descent, the work closest to ours is that of Lin et al. (2023); which we refer to SGD, contrast with SDD. SGD uses the gradient estimate
(I(a) nK, b.)
232 0
where i ~ Uniforr((11... . n)) is random index and Σ 2j2 is random Fourier feature approximation of K. This can be seen targeting the primal loss, with mixed multiplicative- additive objective. Like SDD, SGD uses geometric iterate averaging and Nesterov's momentum; unlike SDD, the SGD algorithm requires gradient clipping to control the gradient noise. Note that Lin et al. (2023) also consider modified notion of convergence which involves subspaces of the linear system coefficients. This is of interest in situations where the number of iterations is sufficiently small relative to data size, and generalization properties are important. While our current work does explicitly cover this setting, could be studied in similar manner.
Stochastic gradient descent approaches were also used in similar contexts by, amongst others, Dai al. (2014) for kernel regression with the primal objective and random Fourier features; Kivinen et al. (2004) for online learning with the primal objective with random sampling of training data; and Antorán et al. (2023) for large-scale Bayesian linear models with the dual objective.
4 EXPERIMENTS AND BENCHMARKS
We present three experiments that confirm the strength of our SDD algorithm standard benchmarks. The first two experiments, on UCI regression and large-scale Bayesian optimisation, replicate those of Lin et al. (2023), and compare against SGD (Lin et al., 2023), CG (Wang et al., 2019) and SVGP (Hensman et al., 2013). Unless indicated otherwise, we use the code, setup and hyperparameters of Lin et al. (2023). Our third experiment tests stochastic dual descent five molecule-protei binding
7

---

Table 1: Koot-mean-square
(RMSE), compute time (on an A100 GPU), and negative log-
likelihood (NLL), for 9 UCI regression tasks for all methods considered. We report mean values and standard error across five 90%-train 10%-test splits for all data sets, except the largest, where three splits are used. Targets are normalised to zero mean and unit variance. This work denoted by SDD".
Data Size
POL
ELEVATORS 17k
BIKE
PROTEIN
KEGGDIR
3DROAD
SONG 515k
BUZZ 583k
HOUSBELEC
15k
17k
46k
49k
435k
2M
SDD
0.08 1: 0.00
0.35 :: 0.00
0.04
0.00
0.50 1: 0.01
0.08 : 0.00
0.04 1 0.00
0.00
0.28 1: 0.00
0.04 :: 0.00
RMSS
SGD
0.38 0.00
0.11
= 0.00
0.51
0.00
0.12 0,00
0.11 0.00
0.00
0.00
0.09 + 0,00
CG
0.08
+: 0.00
0.35
0.00
0.04 = 0.00
0.50
0.00
0.08
0.00
0.18 0.02
0.87
L88
0.19
0.87
0.14
SVGP
0.10 +: 0.00
0.37 +: 0.00
0.08 0.00
0.57
+: 0.00
0.10
0.00
0.47 :: 0.01
0.80
+: 0.00
0.32 0.00
0.12 0.00
Time (min))
SDD"
1.88 0.01
1.13 0.02
1.15 0.02
1,00 0.06
3.02
0.56
207
0.10
0.02
SGD
2.80 1:
2.07 1: 0.03
2.12 = 0.04
2.87 0.01
3.30 0.12
6,68 1: 0.02
: 0.61
0.15
69.5 1.0.
CG
1.17
0.04
+ 0.00
0.11 = 0.01
0.16
0.17
=: 0.00
13.4
244 0.77
157 0.01
SVGP
11.5
1: 0.01
11.3 :: 0.06
11.1 = 0.02
11.1 1: 0.02
11.5 1 0.04
152
0.15
0.13
209 0.37
154 1 0.12
SDD"
-1.18
100000000 0.01
0.38 0.01
-2.99
0.63 + 0.02
-0.92
0.11
-1.70
0.01
0.17 1.03
0,06
-1.46 0.10
NLL
SGD
-0,70
0.47 0.02
0.00
40.48
0.08
0.64
-0622
0.07
-0.60
0.00
1.21
0.00
0.83
+: 0.07
-1.09 0.04
CG
.1.17
= 0.01
0.38
1: 0.00
-2.62
1: 0.06
0.62
1: 0.01
40.92
1: 0.10
16.3
0.45
1.36
0.07
2.38
1 0.08
2.07 0.58
SVGP
-0.67
= 0.01
0.43 0.00
-1.21
+ 0.01
0.85
0.01
0.60 +0.54
0.00
1.21 0.00
0.22
0.03
-0.61 0.01
Table 2: Test set R² scores obtained for each target protein the DOCKSTRING molecular binding affinity prediction task. Results with (-) are from García-Ortegón et al. (2022), those with (.)+ are from Tripp al. (2023). SVGP uses 1000 inducing points. SDD* denotes this work.
Method
ESR2
F2
KIT
PARPI
PGR
Method
ESR2
PARP1
PGR
Attentive FP!
0.627
0,880
0.806
0,910
0.678
SDD*
0.627
0.880
0.790
0.907
0.626
MPNN
0.506
0.798
0.755
0.815
0.324
SGD
0.526
0.832
0.697
0.857
0.408
XGBoost
0.497
0.688
0.674
0.723
0.345
SVGP
0.533
0.839
0.696
0.872
0.477
affinity prediction benchmarks of García- PrteeónOrtegónnneo al. (2022). We include detailed descriptions of all three experimental set-ups and some additional results in Appendix C.
4.1 UCI REGRESSION BASELINES
We benchmark on the 9 UCI repository regression data sets (Dua and Graff, 2017) used by Wang et al. (2019) and Lin et al. (2023). We run SDD for 100k iterations, the same number used by Lin et al. (2023) for SGD, but with step-sizes 100x larger than Lin et al. (2023), except for ELEVATORS, KEGGDIRECTED, and BUZZ, where this causes divergence, and We run CG to a tolerance of 0.01, except for the 4 largest data sets, where we stop CG after 100 heraaiios--the still provides CG with larger compute budget than first-order methods. CG uses a pivoted Cholesky preconditioner of rank 100. For SVGP, we use 3, 000 inducing points for the smaller five data sets and 9, 000 for the larger four, so to match the runtime of the other methods.
10x larger step-sizes instead.
The results, reported in Table 1, show that SDD matches or outperforms all baselines on all UCI data sets in terms of root-mean-square error of the mean prediction across test data. SDD strictly outperforms SGD on all data sets and metrics, matches CG on the five smaller data sets, where the latter reaches tolerance, and outperforms CG on the four larger data sets. The same holds for the negative log-liked ligllikel metric (NLL)., computed using 64 posterior function samples, except on BIKF, where CG marginally outperforms SDD. Since SDD requires only one matrix-vectho multiplication per step, as opposed to two for SGD, it provides about 30% wall-clock time speed-up relative to SGD. We run SDD for 100k iterations to match the SGD baseline, SDD often converges earlier than that.
4.2 LARGE-SCALE THOMPSON SAMPLING
Next, replicate the synthetic large-scale black-box function optimisation experiments of Lin et al. (2023), which consist of finding the maxima of functions mapping [0. 1 + IR sampled from Matémn33/2 Gaussian process priors with 5 different length-scales and 10 random functions per
8

---

Small Compute Budget
Large Compute Budget
6.0
Max value
6.5
5.0
4.5
4.0
50k
60k
70k
80k
0h
1h
2h
3h
50k
60k
70k
80k
0h
5h
10h
15h
Acquisitions
Time elapsed
Acquisitions
Time elapsed
Random Search
SDD
SGD
CG
SVGP
Figure 5: Results for the Thompson sampling task. Plots show mean and standard error of the maximum function values identified, across 5 length scales and 10 seeds, against both the number of observations acquired and the corresponding compute time on an A100 GPU. The compute time includes drawing posterior function samples and finding their maxima. All methods share initial data set of 50k points, and take 30 steps of parallel Thompson sampling, acquiring 1 points at each.
length-scale, using parallel Thompson sampling (Hernández-Lohato et al., 2017). We set the kernel used by each model to match that of the unknown function. All methods are warm-started with the same 50k points chosen uniformly at random the domain. We then run 30 iterations of the parallel Thompson sampling. acquiring 1000 points at each iteration. We include two variants of the experiment, one with a small compute budget, where SGD and SDD are run for 15k steps, SVGP is given 20k steps and CG is run for 10 steps, and one with a large budget, where all methods are run for times as many steps. We present the results on this task in Figure 5, averaged over both length-scales and seeds, and detailed breakdown in Appendix C. In both large and small compute settings, SDD makes the most progress, in terms of maximum value found, while using the least compute. The performance of SDD and SGD degrades gracefully when compute budget is limited.
4.3 MOLECULE-PROTEIN BINDING AFFINITY PREDICTION FOR DRUG DISCOVERY
In this final experiment, show that Gaussian processes with SDD competitive with graph neural networks in predicting binding affinity, a widely used filter drug discovery (Pinzi and Rastelli, 2019; Yang et al., 2021). We use the DOCISTRING regression benchmark of Garcfa-Ortegón et al. (2022), which contains five tasks, corresponding to five different proteins. The inputs are the graph structures of 250k candidate molecules, and the targets are real-valued affinity scores from the docking simulator AutoDock Vina (Trott and Olson, 2010). For each protein, we use a standard restt splits of 210k and 40k molecules, respectively. We use Morgan fingerprints of dimension 1024 (Rogers and Hahn, 2010) to represent the molecules, and use Gaussian process model based on the Tanimoto kernel Ralaivola et al. (2005) with the hyperparameters of Tripp et (2023).
In Table 2, following García-Ortegón et al. (2022), we report R² values. Alongside results for SDD and SGD, we include results from Carcfa-Ortegón et al. (2022) for XGBoost, and for two graph neural networks, MPNN (Gilmer et al., 2017) and Attentive FP (Xiong al., 2019), the latter of which is the state-of-the-an for this task. We also include the results for SVGP reported by Tripp et al. (2023). These results show that SDD matches the performance of Attentive FP the ESR2 and FP2 proteins, and comes close on the others. To the best of our knowledge, this is the first time Gaussian processes have been shown to be competitive large-scale molecular prediction task.
CONCLUSION
We introduced stochastic dual descent, specialised first-order stochastic optimisation algorithm for computing Gaussian process mean predictions and posterior samples. We showed that stochastic dual descent performs very competitively standard regression benchmarks and on large-scale Bayesian optimisation, and matches the performance of state-of- ofte-of-thee the-hartt graph neural networks on a molecular binding affinity prediction task.

---

ACKNOWLEDGMENTS
JAL and SP were supported by the University of Cambridge Harding Distinguished Postgraduate Scholars Programme, AT supported by Cornell University. jointly via the Center for Data Science for Enterprise and Society, the College of Engineering, and the Ann S. Bowers College of Computing and Information Science. JMHL acknowledges support from Turing AI Fellowship under grant EP/VO23756/1. CS gratefully acknowledges funding from the Canada CIFAR Al Chairs Program, Amii and NSERC.
REFERENCES
J. Antorán, S. Padhy, R. Barbano, E. T. Nalisnick, D. Janz, and J. M. Hernández-Lobato. Sampling- based inference for large linear models, with application linearised Laplace. In International Conference on Learning Representations, 2023. Cited on pages 7, 16.
L. Bo and C. Sminchisescu. Greedy Block Coordinate Descent for Large Scale Gaussian Process Regression. In Uncertainty Artificial Intelligence, 2008. Cited on page
P. Boyd and L.. Vandenberghe. Convex Optimization. Cambridge University Press, 2004. Cited on pages 3, 13.
J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, van der Plas, Wanderman-Milne and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. Cited page 14.
B. Dai, B. Xie, N. He, Y. Liang, A. Raj. M.-F. F. Balcan, and L. Song. Scalable Kernel Methods via Doubly Stochastic Gradients. Advances in Neural Infommaion Processing Systems, 2014. Cited on pages 1, 6, 7.
A. Dieuleveut, N. Flammarion, and F. Bach. Harder, Better, Faster, Stronger Convergence Rates for Seast-Squares Regression. Journal of Machine Learning Research, 2017. Cited on pages 1, 5, 7.
D. Dua and C. Graff. UCI Machine Learning Repository, 2017. Cited on pages 3, 8.
J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011. 1. Cited page 6.
T.-T. Frie, N. Cristianini, and Campbell. The Kernel-Adatron Algorithm: Fast and Simple Learn ing Procedure for Support Vector Machines. In International Conference on Machine Learning, 1998. Cited on page 7.
M. Garcia-Ortegón, N. C. Simm, A. J. Tripp, M. Hernández-Lobato, A. Bender, and Bacallado. DOCKSTRING: Easy Molecular Docking Yields Better Benchmarks for Ligand Design. Journal of Chemical Information and Modeling, 2022. Cited on pages 8, 9, 16.
Gilmer, S. Schoenholz, P. F. Riley, Vinyals, and G. E. Dahl. Neural Message Passing for Quantum Chemistry. In International Conference on Machine Learning, 2017. Cited on page 9.
J. Hensman, N. Fusi, and N. D. Lawrence. Gaussian Processes for Big Data. In Uncertainty in Artificial Intelligence, 2013. Cited on pages 7.
M. Hernández-Lobato, J. Requeima, E. O. Pyzer-Knapp, and A. Aspuru-Guzik, Parallel and Distributed Thompson Sampling for Large-scale Accelerated Exploration of Chemical Space. In International Conference on Machine Learning, 2017. Cited on pages 9, 16.
loffe. Improved Consistent Sampling, Weighted Minhash and LI Sketching. International Conference Data Dining, 2010. Cited page 18.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015. Cited on page 6.
Kivinen, A. J. Smola, and R. C. Williamson. Online Learning with Kernels. IEEE Transactions on Signal Processing, 2004. Cited pages 7.
10

---

Z. J.-FF Ton, D. Oglic, and D. Sejdinovic. Towards Unified Analysis of Random Fourier Features. In International Conference on Machine Learning, 2019. Cited on page 5.
J. A. Lin, J. Antorán, S. Padhy, D. Janz, J. M. Hernández-Lobato, and A. Terenin. Sampling from Gaussian Process Posteriors using Stochastic Gradient Descent. In Advances in Neural Information Processing Systems, 2023, Cited on pages 1, 2, 7, 8, 14, 16.
W. Liu, P. P. Pokharel, and J. C. Principe. The Kernll Least-Mean-Squee -Suaree Algorithm. IEEE Transac- tions on Signal Processing, 2008. Cited on page
Y. Nesterov. A method for unconstrained convex minimization problem with the rate convergence O(1/k²). In Doklady Akademii Nauk SSSR, 1983. Cited page 6.
T. Pinder and D. Dodd. GPJax: A Gaussian Process Framework in JAX. Journal of Open Source Software, 2022. Cited on page 16.
L. Pinzi and G. Rastelli. Molecular Docking: Shifting Paradigms in Drug Discovery. International Journal of Molecular Sciences, 2019. Cited on page 9.
B. T. Polyak. New stochastic approximation type procedures. Avtomatika Telemekhanika, 1990. Cited on page 6.
B. T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 1964. Cited on page 6.
B. T. Polyak and A. B. Juditsky. Acceleration of Stochastic Approximation by Averaging, SIAM Journal on Control and Optimization, 1992. Cited on page 6.
L. Ralaivola, Swamidass, H. Saigo, and Baldi. Graph kernels for chemical informatics. Neural Networks, 2005. Cited on page 9.
C. E. Rasmussen and C. K. Williams. Gaussian Processes for Machine Learning. MIT Press, 2006. Cited on page 2.
D. Rogers and M. Hahn. Extended-Conneettin Fingerprints. Journal of Chemical Information and Modeling, 2010. Cited on pages 9, 16.
A. Rudi, L. Carratino, and I.. Rosasco. Falkon: Optimal Large Scale Kernel Method. Advances in Neural Information Processing Systems, 2017. Cited on page 1.
D. Ruppert. Efficient Estimations from Slowly Convergent Sobbins-Monne Process. Technical report, Cornell University, 1988. Cited on page 6.
C. Saunders, A. Gammerman, and V. Vovk. Ridge Regression Learning Algorithm in Dual Variables. In International Conference on Machine Learning, 1998. Cited on page 6.
Shalev-Shwartz and Zhang. Stochastic Dual Coordinate Ascent Methods for Regularized Loss Minimization. Journal of Machine Learning Research, 2013. Cited on pages 1, 7.
A. J. Smola and B. Schölkopf. Learning with Kernels. MIT Press, 1998. Cited pages 3, 7.
I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning, 2013. Cited on page 6.
T. Tieleman and G. Hinton. Lecture 6.5-RMSprop, coursera: Neural networks for machine learning, 2012. Cited on page
M. K. Titsias. Variational Learning of Inducing Variables Sparse Gaussian Processes. In Artificial Intelligence and Statistics, 2009. Cited on page 1.
A. Tripp, S. Bacallado, S. Singh, and J. M. Hernández-Lobaoe Tanimoto Random Features for Scalable Molecular Machine Learning. In Advances in Neural Information Processing Systems, 2023. Cited on pages 8, 9, 14, 16, 18.
11

---

O. Trott and A. J. Olson. AutoDock Vina: Improving the Speed and Accuracy of Docking with New Scoring Function, Efficient Optimization, and Multithreading. Journal of Computational Chemistry, 2010. Cited on page 9.
S. Tu, R. Roelofs, S. Venkataraman, and B. Recht. Large Scale Kernel Learning using Block Coordi- nate Descent. 1702:::600.053100 2016. Cited on page 7.
A. V. Varre, L. Pillaud- Vivien, and N. Flammarion. Last iterate convergence of SGD for Least-Squares in the Interpolation regime. Advances Neural Information Processing Systems, 2021. Cited on pages 1, 5, 6.
K. A. Wang, G. Pleiss, J. R. Gardner, S. Tyree, K. Q. Weinberger, and A. G. Wilson. Exact Gaussian Processes on Million Data Points. In Advances in Neural Information Processing Systems, 2019. Cited on pages 1, 8.
C. Williams and M. Seeger. Using the Nyström Method to Speed Up Kernel Machines. Advances in Neural Information Processing Systems, 2000. Cited on page 1.
J. T. Wilson, V. Borovitskiy, A. Terenin, P. Mostowsky, and M. P. Deisenroth. Efficiently Sampling Functions from Gaussian Process Posteriors. In International Conference on Muchine Learning, 2020. Cited on page 2.
J. T. Wilson, V. Borovitskiy. A. Terenin, P. Mostowsky, and M. P. Deisenroth. Pathwise Conditioning of Gaussian Processes. Journal of Machine Learning Research, 2021. Cited on page 2.
K. Wu, J. Wenger, H. Jones, G. Pleiss, and J. R. Gardner. Large-Scale Gaussian Processes via Alternating Projection, arXiv preprint arXiv:2310.17137, 2023. Cited on page 7.
Z. Xiong, D. Wang. X. Liu, F. Zhong, X. Wan, X. Li, Li, X. Luo, K. Chen, H. Jiang, et al. Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism. Journal of Medicinal Chemistry, 2019. Cited on page 9.
Y. Yang, K. Yao, M. P. Repasky, K. Leswing, R. Abel, B. K. Shoichet, and S. V. Jerome. Efficient Exploration of Chemical Space with Docking and Deep Learning. Journal of Chemical Theory and Computation, 2021. Cited on page 9.
12

---

A CONVEX DUALITY AND UNIFORM APPROXIMATION BOUNDS
Claim (Strong duality). We have that
min oER
L(a) = - min [a)a), oER"
for L, L defined per equations (2) and (4) respectively, with a (b)b) minimising both sides.
Proof. That o" (b) minimises both L and L can be established from the first order optimality conditions. Now, for the duality, observe that we can write min L(a) equivalently as the constrained optimisation problem
min min 2 =Ill11++ +2 |Ialli subject to Ka b. HER* oER"
Note that this is quadratic in both and Introducing Lagrange multipliers / R2. in the form A8, where we recall that À > O, the solution of the above is equal to that of
min min sup |ull 2 11lllli + (b-Kaauu
-)-
"ER" deR" BER 2
This is finite-dimensional quadratic problem, and thus we have strong duality (see, e.g., Examples 5.2.4 in Boyd and Vandenberghe, 2004). We can therefore exchange the order of the minimum operators and the supremum, yielding the again equivalent problem
sup
min KER
|/|l2 AB u
lall A3 28" Ka + b aE3"
min
BER"
Noting that the two inner minimisation problems are quadratic, we solve these analytically using the B and u AB, to obtain that the above is equivalent to first order optimality conditions, that is
sup -> (x 110 bTb) À min L'(3). BER"
BER*
The result follows by chaining the above equalities and relabelling 3 + a.
To show equation (5), the uniform approximation bound, we first need to define reproducing kernel Hilbert spaces (RKHSes). Let H be Hilbert space of functions and corresponding norm 1/ We say H is the RKHS associated with bounded kernel l: if the reproducing property holds:
with inner product (:, -)
E X, vh E H, That is, k(x, -) is the evaluation functional (at x E X') on H. For observations X, let t: H. - R² be the linear operator mapping h + h(X), where h(X) (h(z1)....,h(a) h(x,)) We will write 4* for the adjoint of 4 and observe that K is the matrix of the operator " with respect to a standard basis on R (that used implicitly throughout).
(k(z, )) = h(x).
Claim (Uniform approximation). For any a, a' E RR,
ha - hao S VKlla a'Ilk,
where k = supzex k(x,æ).
Proof. First, observe that,
I|aa ha /loo sup ha(z)) - ha))())) (z))
(defn. of sup norm)
sEx
sup (k(I, :), ha har)l
(reproducing property)
=EX
< sup ||k(,,, )11ze1/aa xEX
Mhaellha-aaa - tar
(Cauchy-Schwarz)
S k Na - har l|H
(defn. of xt)
Now, observe that hn 4² and hw 80 a'' and
have the equalities
lha had 112 ($'(e ((a-a)), 8' a'))
(defn. norm) (defn. adjoint) (defn. K)
(a a', 40 (a a'))
||a a lix
where for the final equality, we note that is the matrix of the operator 4$* with respect to the standard basis. Combining the above two displays yields the claim.
n
13

---

1 +|@
a² a'(1))
a² ay(y)l
Test RMSE
10
10
0.086
me;e; (Ka)
B + Xa
10
10
0.084
10
10
0.082
10
0.080
10
10-2
0.083
10-2
0.082
10-6
10
10
0.081
10-6
10-11
0.080
0
25k
50k
75k
100k
0
25k
50k
75k
100k
25k
50k
75k
100k
Iterations
Iterations
Iterations
Bn
20
8n
40
8n 50
Figure 6: Comparison dual stochastic gradient descent on the POL with a momentum parameter P 0.9, averaging parameter 0.001, batch size B 512, and varying step-sizes Bn. For the mixed additive-multipiicttie noise gradient estimator in the top row, higher step-size leads to faster convergence, but worse asymptotic result. For recommended multiplicative noise estimator in the bottom row, higher step-size improves both the speed of convergence and the asymptotic result.
B EFFECTS OF VARYING STEP-SIZE AND BATCH-SIZE SIZE
In Figures we examine the trade-offs related to step-size and batch-size for stochastic dual gradient descent with gradient estimators either of the mixed additive-multiplictivv type, used by SGD, or the purely multiplicative type, recommended in this work. In short, trade-offs exist in the aoditive-multiplicaatee case, and we cannot make clear recommendations. For purely multiplicative noise, the picture clearer:
Step-size should be chosen large possible while avoiding divergence, improve both the rate of convergence and the asymptotic quality of the result.
Batch-size should chosen as small as possible while avoiding divergence, since larger batch incurs larger per-step computational cost without affecting the convergence of the algorithm.
Of course, this then leaves step-size versus batch-size trade-off for SDD, translating into trade-off between the rate of convergence and the asymptotic behaviour versus wall-clock time. This trade-off' will need to be addressed case-by-case basis.
C ADDITIONAL DETAILS ON EXPERIMENTAL SETUPS AND RESULTS
We the implementation and hyperparameter configurations from Lin al. (2023) for the SGD, SVGP and CG baselines. which uses the jax library (Bradbury et al., 2018). Some more detail on the implementations is follows.
CG We implement CG using the jax.scipy library, with pivoted Cholesky preconditioner of size 100 computed using the TensorFlow Frobability library.
SGD We minibatches of 512 estimate the fit term of equation (6), and 100 random features for the regulariser. For UCI and Bayesian optimisation experiments, random Fourier features are used. For the molecule task, random hashes, described in Tripp et al. (2023). SGD uses
14

---

(1())
(2)l
Test RMSE
10
10
ne;e; (Ko)
-V + Ao
10
10-1
10
10-3
10-4
10-3
10a
10
0.083
10
10-5
0.082
g)
10
10"8
0.081
10
10
0.080
0
25k
50k
75k
100k
25k
50k
75k
100k
25k
50k
75k
100k
Iterations
Iterations
Iterations
64 B 128 B 256
B 512
1024
Figure Dual stochastic gradient descent the POL with momentum parameter p 0.9, averaging parameter r = 0.001, step-size 3n = 50, and varying batch sizes B, with metrics plotted against the number of iterations. For the mixed additive-multiples muliii-auulipiicaiiiee noise gradient estimator in the top row, higher batch size improves the final performance. For our preferred multiplicative noise estimator in the bottom row, batch sizes has little effect on performance, so long as it not so low that it causes the optimisation to diverge.
(2)
a
Test RMSE
10a
n1
10a
mexe; (
39 + Aa
10
10
10
0.10
10-1
10-3
0.08
101
10
0.083
10
10
0.082
()))
10-5
10-8
0.081
10
10
11
0.080
100
150
200
100
150
200
50
100
150
200
Time in sec
Time sec
Time sec
64
# = 128
B
256
512
....
1024
Figure 8: Dual stochastic gradient descent the POL with momentum parameter 0.9, averaging parameter = 0.001, step-size 3n. = 50, and varying batch sizes B, with metrics plotted against wall- clock time. For the mixed noddiive-muliilicaaiv noise gradient estimator, conclusions match those in Figure 7. For our preferred multiplicative noise estimator in the bottom row, since performance for batch sizes that lead to convergence is similar, performance is best for the smallest batch-size that does not lead to divergence. as batch-size increases the per-step computation cost.
15

---

Nesterov's momentum of 0.9 and geometric iterate averaging, as implemented in optax, with r = 100/T, where T the total number SGD iterations. We clip the 2-norm of the gradient estimates to 0.1. These settings are replicated from Lin al. (2023), which in turn takes these from Antorán et al. (2023).
SVGP We initialise inducing points using k-means and k-means++. We then optimise all variational parameters, including inducing point locations, by maximising the ELBO with the Adam optimiser until convergence, using the GP Jax library (Pinder and Dodd, 2022).
C.I UCI REGRESSION
For all methods, we use a Matérn-3/2 kernel with a fixed set of hyperparameters chosen by Lin al. (2023) via maximum likelihood. To ease reproducibility, we make the full hyperparameter set available as Python file in our source code HERE. SGD uses 33 0.5 to estimate the mean function weights, and fn 0.1 to draw samples. For SDD, we use step-sizes 3n which are 100 times larger, except for ELEVATORS, KEGGDIR and BUZZ, where this causes divergence; there, we a 10x larger step-size instead. SDD and SGD are both run for 100k steps with batch size B = 512 for both the mean function and posterior samples. For CG, with N' < 50k, and a tolerance of 0.01. On the four largest data sets, the per step cost of CG is too large to 1000 steps, and we run 100 steps instead. For SVGP, the number of inducing points is chosen such that the compute cost approximately matches that of other methods: 3000 for the smaller five data sets and 9000 for the larger four. Negative log-likelihood computations are done by estimating the predictive variances using 64 posterior samples, with 2000 random Fourier features used approximate the requisite prior samples.
maximum of 1000 steps for data sets
C.2 LARGE-S SCALE THOMPSON SAMPLING
We draw target functions X' + R from Gaussian process prior with Matérn-3/2 kernel and length scales (0.1, 0.3, 0.4, 0.5), using 2000 random Fourier features. For each length scale, we repeat the same kernel that used to generate the data.
the experiment for 10 seeds. All methods
We optimise the target functions on X' = [0, 1] using parallel Thompson sampling (Hernández- Lobato et al., 2017). That is, we choose Vnew arg m fn for set of posterior function samples drawn in parallel. We replicate the multi-start gradient optimisation maximisation strategy of Lin et al. (2023) (see therein). For each function sample maximum,
evaluate Veew = ((Fwew) ) + with drawn from a zero-mean Gaussian with variance of 10 6 We then add the pair (Inews Inew to the training data. We of 50k observations sampled uniformly at random from X.
acquisition batch size of 1000. We initialise all methods with data set
Here, SGD step-size of .3n 0.3 for the mean and Ann 0.0003 for the samples. SDD uses step-sizes that are 10x larger: .n = 3 for the mean and Bn 0.003 for the samples.
C.3 MOLECULE-PROTER BINDING AFFINITY PREDICTION
We use the data and standard train-test splits from Garefa-Ortegón et al. (2022), which were produced by structure-basse clustering to avoid similar molecules from occurring both in the train and test set. We perform all the preprocessing steps for this benehmark outlined by Garefa-Ortegón et al. (2022), including limiting the maximum docking score to 5.
Fingerprints, Tanimoto Kernel and Random Features Molecular fingerprints are a way encode the structure of molecules by indexing sets of subgraphs present in molecule. There are many types of fingerprints. Morgan fingerprints represent the subgraphs up a certain radius around each atom in molecule (Rogers and Hahn, 2010). The fingerprint can be interpreted as sparse vector of counts, analogous to a 'bag of words' representation a document. Accordingly, the Tanimoto coefficient T(x, .'), also called the Jaccard index, is way to measure similarity between fingerprints, given by
E, min(2,,
T(x,x') Σ, max(x,,
This function is a valid kernel and has a known random feature expansion using random hashes (Tripp et al., 2023). The feature expansion builds upon prior work for fast retrieval of documents
16

---

Small Compute Budget
Large Compute Budget
7.0
Max value
6.5
0.1
5.5
5.0
4.5
6.5
Max value
6.0
I= 0.2
5.5
6.5
Max value
1= 0.3
Max valtte
1= 0.4
3.5
Max value
1= 0.5
4.5
4.0
3.5
50k
60k
70k
80k
0h
1h
2h
3h
50k
60k
70k
80k
Oh
5h
10h
15h
Acquisitions
Time elapsed
Acquisitions
Time elapsed
Random Search
SDD
SGD
CG
SVGP
Figure 9: Maximum function values, with mean and standard error across 10 seeds, obtained by parallel Thompson sampling, for functions with different length-scales plotted as functions of acquisition steps and the compute time on an A 100 GPU. All methods share an initial data set of 50k points, and take 30 Thompson steps, acquiring a batch of 1000 points in each. The algorithms perform differently across the length-scales: CG performs better in settings with smaller length-scale which give better conditioning: SVGP tends perform better in settings with larger length-scale and thus higher smoothness; SGD and SDD perform well in both settings.
17

---

using random hashes that approximate the Tanimoto coefficient; that is, distribution PA. over hash functions h such that
p h(=')) ==(="))=((zzz)))
Per Tripp et al. (2023), we extend such hashes into random by using them to index a random tensor whose entries are independent Rademacher random variables, and use the random hash of loffe (2010).
Gaussian Process Setup The results for the SVGP baseline are taken from Tripp et al. (2023). As the Tanimoto kernel itself has hyperparameters, the only kernel hyperparameters scaling factor A > for the kernel and the noise variance X, and a constant GP prior mean (Ag (the Gaussian process regresses on y Po in place of y). These are chosen by using an exact GP randomly chosen subset of the data and held constant during the optimisation of the inducing points. The values of these are given Table 3. The same values are also used for SGD and SDD to ensure that the differences in accuracy are solely due to the quality of the GP posterior approximation. The SGD method uses 100-dimensional random features for the regulariser.
constant
Table 3: Hyperparameters for all Gaussian process methods used in the molecule-proted binding affinity experiments of Section 4.3.
Data ESR2
F2
KIT
PARP1
PGR
0.497
0.385
0.679
0.560
0.630
A
0.373
0.049
0.112
0.024
0.332
140
-6.79
-6.33
-6.39
-6.95
-7.08
18


Reviews:
[{"review_id": "37ahVnNmx6", "reviewer": "Reviewer_WMvZ", "summary": "This paper proposes a stochastic gradient descent method for solving the kernel ridge regression problem. In particular, three aspects are covered: (1) a dual objective that allows a larger learning rate; (2) a stochastic approximation that brings in effective utilization of stochastic gradients; (3) momentum and geometric iterate averaging. By combining these aspects, the algorithm is demonstrated be faster compared to baselines in experiments.", "questions": "See the weakness part.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* This paper proposes a new method for the kernel ridge regression problem.\n* Experimental results show that the proposed algorithms can achieve better performance than baselines. When combined with the Gaussian process, the method can also achieve comparable performance to that of graph neural networks.", "weaknesses": "* 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.\n* 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.\n* 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.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FmrcoTvImX", "reviewer": "Reviewer_dz2H", "summary": "This paper introduces a stochastic dual gradient descent method for optimizing the Gaussian process posterior computation.", "questions": "See weakness.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present a novel \"dual\" formulation for the Gaussian process regression problem. After studying the condition number of new and old formulations, the authors observe that the \"dual\" formulation allows for the use of larger learning rates, indicating its potential to converge faster. They then propose the stochastic dual gradient descent method, leveraging various optimization techniques based on the \"dual\" formulation, including feature and coordinate sampling (or minibatch) [1], Nesterov's acceleration [2], and Polyak averaging. Notably, the authors introduce a new averaging scheme called geometric averaging.\n\nThe paper is overall well-structured, clear, logically presented, and readable. It contains minimal typos and lacks theoretical flaws. Moreover, the authors conduct sufficient numerical experiments to validate the effectiveness of their proposed optimizer.\n\n[1] \"Sampling from Gaussian Process Posteriors using Stochastic Gradient Descent\"\n[2] Y. Nesterov, \"A method for unconstrained convex minimization problems with a convergence rate of O(1/k^2)\"\n[3] B. T. Polyak, \"New stochastic approximation type procedures,\" Avtomatika i Telemekhanika, 1990.", "weaknesses": "The authors do not provide a theoretical justification to verify the convergence of the proposed method. Nevertheless, it is likely that convergence can be ensured under mild conditions, as the optimization techniques employed are standard and well-established in the community and literature.\n\nFrom my perspective, the primary contribution of this paper lies in the introduction of the \"dual\" formulation, as presented on page 4 after Equation (2). This formulation allows for the use of larger step sizes, which suggests the potential for faster convergence. While the remaining studies and techniques are also important, they are somewhat incremental and standard. Consequently, I am uncertain about whether the paper's contribution alone justifies its publication in ICLR. As a result, I have assigned a boundary score and defer to the area chair's judgment for the final decision on acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CmR2R4IAIh", "reviewer": "Reviewer_eaDf", "summary": "This paper uses insights drawn from the application of gradient descent in the kernel and optimisation communities to develop a stochastic gradient descent approach for Gaussian processes. In particular this method is useful in regresssion to approximate the posterior mean and to draw samples from the GP posterior. This method, stochastic dual descent, is compared to conjugate gradient, stochastic gradient descent and stochastic variational Gaussian processes on regression benchmarks and molecular binding affinity\nprediction.", "questions": "What are the implications of limiting the kernel to the form $\\sum_{j=1}^mz_jz_j^T$?\n\nHow does ill conditioning affect the performance of the method?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a well written paper that considers an interesting problem. The use of several benchmarks in the experimental section and comparison with recent work is a plus.\n\nThe justification for use of the dual objective as well as the illustrative example is clear.\n\nThe reason behind the choice of random coordinate estimates is well done.", "weaknesses": "It would be useful to emphasise that this work is useful when the Kernel is already known. Comments on whether these methods would be useful in hyperparameter estimation would be useful.\n\nThe claim that the method can be implemented in a few lines of code should be demonstrated. The repo given does not clearly illustrate this using a simple example.\n\nThe paper would benefit from a visualisation comparing samples from a GP using SDD to an exact GP fit to show that the samples lie within the confidence interval.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HrEjW464cV", "reviewer": "Reviewer_1x4h", "summary": "This paper introduced a stochastic dual gradient descent algorithm for kernel ridge regression and sampling. The stochastic dual descent algorithm admits better-conditioned gradients and a faster convergence rate compared to the SGD proposed by Lin et al. (2023). With the selected kernels, experimental results showed competitive performance with a number of SOTA methods on UCI regression/Bayesian optimization/ molecular binding affinity prediction tasks. Overall, the paper is easy to follow and well-written, while technical contributions seem to be below the bar of ICLR.", "questions": "1. 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)?\n2. 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?\n3. 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.\n4. How do different batch sizes affect the overall convergence?\n5. 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)\n6. what is the difference between the $\\beta n$ in the main texts and the $\\beta$ in Algorithm 1?\n7. The green dashed line is missing in Figure 3.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The strengths are: \n(1) Some fresh insights from the optimization and kernel communities were explored. \n(2) Uniform approximation bound and duality of objectives were both analyzed. \n(3) Different randomized gradients and convergence performance were compared.", "weaknesses": "Some suggestions on improving the weakness points are: \n(1) More figures/tables to explicitly show the weakness/instability of the baseline methods are expected. \n(2) 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.  \n(3) A theoretical convergence analysis is expected (not only some figures).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OVJ1IC0qzH", "reviewer": "Reviewer_NBgz", "summary": "The authors consider the problem of computing a Gaussian process posterior, specifically its mean and random draws from it. While the naive computation scales cubically in the number of observations, the authors propose a iterative solver with linear cost per iteration. The idea behind this solver is that the expensive quantity in the GP posterior (kernel matrix inverse) can be thought of as a minimiser of a particular regression problem, which can be solved iteratively with gradient-based methods. The authors consider two formulations of such a regression problem (primal and dual), study their convergence properties, as well as discuss randomised gradients computations to achieve linear computational cost. The proposed algorithm is shown to perform competitively on a number of benchmarks.", "questions": "- 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?\n\n- 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "+ The paper is clearly written and is easy to follow\n+ The differences to the closely related work of Liu et al. (2023) are clearly discussed\n+ I think the results are quite significant for the community. I was especially interested to see that the proposed algorithm performs competitively in comparison to a neural network in Table 2.", "weaknesses": "I didn't notice any significant weaknesses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)