SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(I) The Prior Region. This corresponds to points sufficiently distant from the observed data. Here, for kernels that decay over space, the canonical basis functions  $k(x_{i},\cdot)$  go to zero. Thus, both the true posterior and any approximations formulated pathwise revert to the prior. More precisely, let  $X = \mathbb{R}^d$ , let  $k$  satisfy  $\lim_{c\to \infty}k(x',c\cdot x) = 0$  for all  $x'$  and  $x$  in  $X$ , and let  $(f\mid \pmb {y})(\cdot)$  be given by  $(f\mid \pmb {y})(\cdot) = f(\cdot) + \mathbf{K}_{(\cdot)\pmb{x}}\pmb{v}$ , with  $\pmb {v}\in \mathbb{R}^{N}$ . Then, by passing the limit through the sum, for any fixed  $\pmb{v}$ , it follows immediately that  $\lim_{c\to \infty}(f\mid \pmb {y})(c\cdot x) = f(c\cdot x)$ . Therefore, SGD cannot incur error in regions which are sufficiently far away from the data. This effect is depicted in Figure 4.
(II) The Interpolation Region. This includes points close to the training data. We characterize this region through subspaces of the RKHS, where we show SGD incurs small error.
Let  $\mathbf{K}_{xx} = \mathbf{U}\boldsymbol{\Lambda}\mathbf{U}^T$  be the eigendecomposition of the kernel matrix. We index the eigenvalues  $\boldsymbol{\Lambda} = \mathrm{diag}(\lambda_1,\dots ,\lambda_N)$  in descending order. Define the spectral basis functions as eigenvector-weighed linear combinations of canonical basis functions
$$
u ^ {(i)} (\cdot) = \sum_ {j = 1} ^ {N} \frac {U _ {j i}}{\sqrt {\lambda_ {i}}} k (x _ {j}, \cdot). \tag {14}
$$
These functions—which also appear in kernel principal component analysis—are orthonormal with respect to the RKHS inner product. To characterize them further, we lift the Courant-Fischer characterization of eigenvalues and eigenvectors to the RKHS  $H_{k}$  induced by  $k$ , obtaining the expression
$$
u ^ {(i)} (\cdot) = \underset {u \in H _ {k}} {\arg \max } \left\{\sum_ {i = 1} ^ {N} u \left(x _ {i}\right) ^ {2}: \| u \| _ {H _ {k}} = 1, \langle u, u ^ {(j)} \rangle_ {H _ {k}} = 0, \forall j <   i \right\}. \tag {15}
$$
This means in particular that the top spectral basis function,  $u^{(1)}(\cdot)$ , is a function of fixed RKHS norm—that is, of fixed degree of smoothness, as defined by the kernel  $k$ —which takes maximal values at the observations  $x_{1},\ldots ,x_{n}$ . Thus,  $u^{(1)}$  will be large near clusters of observations. The same will be true for the subsequent spectral basis functions, which also take maximal values at the observations, but are constrained to be RKHS-orthogonal to previous spectral basis functions. A proof of the preceding expression is given in Appendix E.3. Figure 4 confirms that the top spectral basis functions are indeed centered on the observed data.
Empirically, SGD matches the true posterior in this region. We formalize this observation by showing that SGD converges quickly in the directions spanned by spectral basis functions with large eigenvalues. Let  $\mathrm{proj}_{u^{(i)}}(\cdot)$  be the orthogonal projection onto the subspace spanned by  $u^{(i)}$ .
x
Figure 4: SGD error and spectral basis functions. Top-left: SGD (blue) and exact GP (black, dashed) fit to a  $N = 10\mathrm{k}$ ,  $d = 1$  toy regression dataset. Top-right: 2-Wasserstein distance (W2) between both processes' marginals. The W2 values are low near the data (interpolation region) and far away from the training data. The error concentrates at the edges of the data (extrapolation region). Bottom: The low-index spectral basis functions lie on the interpolation region, where the W2 error is low, while functions of index 10 and larger lie on the extrapolation region where the error is large.
Proposition 1. Let  $\delta >0$ . Let  $\pmb{\Sigma} = \sigma^2\mathbf{I}$  for  $\sigma^2 >0$ . Let  $\mu_{\mathrm{SGD}}$  be the predictive mean obtained by Polyak-averaged SGD after  $t$  steps, starting from an initial set of representative weights equal to zero, and using a sufficiently small learning rate of  $0 < \eta < \frac{\sigma^2}{\lambda_1(\lambda_1 + \sigma^2)}$ . Assume the stochastic estimate of the gradient is  $G$ -sub-Gaussian. Then, with probability  $1 - \delta$ , we have for  $i = 1,\dots,N$  that
$$
\left\| \operatorname {p r o j} _ {u ^ {(i)}} \mu_ {f | \boldsymbol {y}} - \operatorname {p r o j} _ {u ^ {(i)}} \mu_ {\mathrm {S G D}} \right\| _ {H _ {k}} \leq \frac {1}{\sqrt {\lambda_ {i}}} \left(\frac {\| \boldsymbol {y} \| _ {2}}{\eta \sigma^ {2} t} + G \sqrt {\frac {2}{t} \log \frac {| \mathcal {I} |}{\delta}}\right). \tag {16}
$$
The proof, as well as an additional pointwise convergence bound and a variant that handles projections onto general subspaces spanned by basis functions, are provided in Appendix E.1. In general, we expect  $G$  to be at most  $\mathcal{O}(\lambda_1^2\| \boldsymbol {y}\|_{\infty})$  with high probability.
The result extends immediately from the posterior mean to posterior samples. As a consequence, SGD converges to the posterior  $GP$  quickly in the data-dense region, namely where the spectral basis functions corresponding to large eigenvalues are located. Since convergence speed on the span of each basis function is independent of the magnitude of the other basis functions' eigenvalues, SGD can perform well even when the kernel matrix is ill-conditioned. This is shown in Figure 3.
(III) The Extrapolation Region. This can be found by elimination from the input space of the prior and interpolation regions, in both of which SGD incurs low error. Consider the spectral basis functions  $u^{(i)}(\cdot)$  with small eigenvalues. By orthogonality of  $u^{(1)}, \ldots, u^{(N)}$ , such functions cannot be large near the observations while retaining a prescribed norm. Their mass is therefore placed away from the observations. SGD converges slowly in this region, resulting in a large error in its solution in both a Euclidean and RKHS sense, as seen in Figure 3. Fortunately, due to the lack of data in the extrapolation region, the excess test error incurred due to SGD nonconvergence may be low, resulting in benign nonconvergence [52]. Similar phenomena have been observed in the inverse problems literature, where this is called iterative regularization [22, 27]. Figure 4 shows the Wasserstein distance to the exact GP predictions is high in this region, as when initialized at zero SGD tends to return small representative weights, thereby reverting to the prior.
Figure 5: Test RMSE and NLL as a function of compute time on a TPUv2 core for CG and SGD.
We now turn to empirical evaluation of SGD GPs, focusing on their predictive and decision-making properties. We compare SGD GPs with the two most popular scalable Gaussian process techniques: preconditioned conjugate gradient (CG) optimization [19, 46] and sparse stochastic variational inference (SVGP) [43, 23]. For CG, we use a pivoted Cholesky preconditioner of size 100, except in cases where this slows down convergence, where we instead report results without preconditioning. We employ the GPJax [32] SVGP implementation and use  $M = 4096$  inducing points for all datasets, initializing their locations with the  $K$ -means algorithm. Full experimental details are in Appendix A.
4.1 Regression Baselines
We first compare SGD-based predictions with baselines in terms of predictive performance, scaling of computational cost with problem size, and robustness to the ill-conditioning of linear systems. Following Wang et al. [46], we consider 9 datasets from the UCI repository [16] ranging in size from  $N = 15\mathrm{k}$  to  $N \approx 2\mathrm{M}$  datapoints and dimensionality from  $d = 3$  to  $d = 90$ . We report mean and standard deviation over five  $90\%$ -train  $10\%$ -test splits for the small and medium datasets, and three splits for the largest dataset. For all methods, we use a Matérn-3/2 kernel with a fixed set of hyperparameters obtained using maximum marginal likelihood [2], as described in Appendix A.1.
We run SGD for 100k steps, with a fixed batch size of 512 for both the mean function and samples. For CG, we run a maximum of 1000 steps for datasets with  $N \leq 500\mathrm{k}$ , and a tolerance of 0.01. On the four largest datasets, CG's per-step cost is too large to run 1000 steps: instead, we run 100 steps, which takes roughly 9 hours per function sample on a TPUv2 core. For SVGP, we learn the variational parameters for  $M = 4096$  inducing points by maximizing the ELBO with Adam until convergence. For all methods, we estimate predictive variances for log-likelihood computations from 64 function samples drawn using pathwise conditioning (2).
Predictive performance and scalability with the number of inputs. Our complete set of results is provided in Table 1, including test RMSE, test negative log-likelihood (NLL) and compute time needed to obtain the predictive mean on a single core of a TPUv2 device. Drawing multiple samples requires repeating this computation, which we perform in parallel. In the small setting ( $N \leq 20\mathrm{k}$ ), taking 100k steps of SGD presents a compute cost comparable to running CG to tolerance, which usually takes around 500-800 steps. Here, CG converges to the exact solution, while SGD tends to present increased test error due to non-convergence. In the large setting ( $N \geq 100\mathrm{k}$ ), where neither method converges within the provided compute budget, SGD achieves better RMSE and NLL and results. SVGP always converges faster than CG and SGD, but it only performs best on the BUZZ dataset, which can likely be summarised well by  $M = 4096$  inducing points.
From Figure 5, we see that SGD makes the vast majority of its progress in prediction space in its first few iterations, improving roughly monotonically with the number of steps. Thus, early stopping after  $100\mathrm{k}$  iterations incurs only moderate errors. In contrast, CG's initial steps actually increase test error, resulting in very poor performance if stopped too early. This interacts poorly with the number of CG steps needed, and the per-step cost, which generally grow with increased amounts of data [41].
Table 1: Regression task mean and std-err for GP predictive mean RMSE, low-noise RMSE  $(\dagger)$ , TPUv2 node hours used to obtain the predictive mean, and negative-log-likelihood (NLL) computed with variances estimated from 64 function samples. SVGP is omitted for the low noise setting, where it fails to run. Metrics are reported for the datasets normalized to zero mean and unit variance.
DatasetN	POL15000	ELEVATORS16599	BIKE17379	PROTEIN45730	KEGGDIR48827	3DROAD434874	SONGS515345	BUZZ583250	HOUSEELECC2049280
RMSE	SGD	0.13±0.00	0.38±0.00	0.11±0.00	0.51±0.00	0.12±0.00	0.11±0.00	0.80±0.00	0.42±0.01	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.15±0.01	0.85±0.03	1.41±0.08	0.87±0.14
SVGP	0.10±0.00	0.37±0.00	0.07±0.00	0.57±0.00	0.09±0.00	0.49±0.01	0.81±0.00	0.33±0.00	0.11±0.01
RMSE†	SGD	0.13±0.00	0.38±0.00	0.11±0.00	0.51±0.00	0.12±0.00	0.11±0.00	0.80±0.00	0.42±0.01	0.09±0.00
CG	-0.16±0.01	0.68±0.09	0.05±0.01	3.03±0.23	9.79±1.06	0.34±0.02	0.83±0.02	5.66±1.14	0.93±0.19
SVGP	-	-	-	-	-	-	-	-	-
Minutes	SGD	3.51±0.01	3.51±0.01	5.70±0.02	7.10±0.01	15.2±0.02	27.6±11.4	220±14.5	347±61.5	162±54.3
CG	2.18±0.32	1.72±0.60	2.81±0.22	9.07±1.68	12.5±1.99	85.2±36.0	195±2.31	351±48.3	157±0.41
SVGP	21.2±0.27	21.3±0.12	20.5±0.02	20.8±0.04	20.8±0.05	21.0±0.12	24.7±0.05	25.6±0.05	20.0±0.03
NLL	SGD	-0.70±0.02	0.47±0.00	-0.48±0.08	0.64±0.01	-0.62±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±0.00	-2.62±0.06	0.62±0.01	-0.92±0.10	16.27±0.45	1.36±0.07	2.38±0.08	2.07±0.58
SVGP	-0.71±0.01	0.43,±0.00	-1.27±0.02	0.86±0.01	-0.70±0.02	0.67±0.02	1.22±0.00	0.25±0.04	-0.89±0.10
Robustness to kernel matrix ill-conditioning. GP models are known to be sensitive to kernel matrix conditioning. We explore how this affects the algorithms under consideration by fixing the noise variance to a low value of  $\sigma^2 = 10^{-6}$  and running them on our regression datasets. Table 1 shows the performance of CG severely degrades on all datasets and, for SVGP, optimization diverges for all datasets. SGD's results remain essentially unchanged. This is because the noise only changes the smallest kernel matrix eigenvalues substantially and these do not affect convergence for the top spectral basis functions. This mirrors results presented previously for the ELEVATORS dataset in Figure 3.
Regression with large numbers of inducing points. We demonstrate the inducing point variant of our method, presented in Section 3.3, on HOUSEELECTRIC, our largest dataset ( $N = 2\mathrm{M}$ ). We select varying numbers of inducing points with a  $K$ -nearest-neighbor algorithm, described in Appendix A.2. Figure 2 shows the time required for  $100\mathrm{k}$  SGD steps scales roughly linearly with inducing points. It takes  $68\mathrm{m}$  for full SGD and  $50\mathrm{m}$ ,  $25\mathrm{m}$ , and  $17\mathrm{m}$  for  $M = 1099\mathrm{k}$ ,  $728\mathrm{k}$ , and  $218\mathrm{k}$ , respectively. Performance in terms of RMSE and NLL degrades less than  $10\%$  even when using  $218\mathrm{k}$  points.
4.2 Large-scale Parallel Thompson Sampling
A fundamental goal of scalable Gaussian processes is to produce uncertainty estimates useful for sequential decision making. Motivated by problems in large-scale recommender systems, where both the initial dataset and the total number of users queried are simultaneously large [36, 17], we benchmark SGD on a large-scale Bayesian optimization task.
We draw a target function from a GP prior  $g \sim \mathrm{GP}(0, k)$  and optimize it on  $X = [0, 1]^d$  using parallel Thompson sampling [25]. That is, we choose  $x_{\mathrm{new}} = \arg \max_{x \in X} (f \mid y)(\cdot)$  for a set of posterior function samples drawn in parallel. We compute these samples using pathwise conditioning with each respective scalable GP method. For each function sample maximum, we evaluate  $y_{\mathrm{new}} = g(x_{\mathrm{new}}) + \varepsilon$  with  $\varepsilon \sim \mathrm{N}(0, 10^{-6})$  and add the pair  $(x_{\mathrm{new}}, y_{\mathrm{new}})$  to the training data. We use an acquisition batch size of 1000 samples, and maximize them with a multi-start gradient descent-based approach described in Appendix A.3. We set the search space dimensionality to  $d = 8$ , the largest considered by Wilson et al. [50], and initialize all methods with a dataset of 50k observations sampled uniformly at random from  $X$ . To eliminate model misspecification confounding, we use a Matérn-3/2 kernel and consider length scales of  $(0.1, 0.2, 0.3, 0.4, 0.5)$  for both the target function and our models. For each length scale, we repeat the experiment for 10 seeds.
In large-scale Bayesian optimization, training and posterior function optimization costs can become significant, and predictions may be needed on demand. For this reason, we consider two variants of our experiment with different levels of compute. In the small-compute setting, SGD is run for 15k steps, SVGP is given  $M = 1024$  inducing points and 20k steps to fit the variational parameters,
Figure 6: Maximum function values (mean and std. err.) obtained by Thompson sampling with our approximate inference methods as a function of acquisition steps and of compute time on an A100 GPU. The latter includes time spent drawing function samples and finding their maxima. All methods share a starting dataset of 50k points and we take 30 Thompson steps, acquiring 1000 points in each.
and CG is run for 10 steps. In the large-compute setting, all methods are run for 5 times as many optimization steps.
Mean results and standard errors, across length scales and seeds, are presented in Figure 6. We plot the maximum function value achieved by each method. In the small-compute setting, the  $\sim 1.5$  (A100 GPU) hours required for 30 Thompson steps with SVGP and SGD are dominated by the algorithm used to maximise the models' posterior samples. In contrast, CG takes roughly twice the time, requiring almost  $3\mathrm{h}$  of wall-clock time. Despite this, SGD makes the largest progress per acquisition step, finding a target function value that improves upon the initial training set maximum twice as much as the other inference methods. VI and CG perform comparably, with the latter providing slightly more progress both per acquisition step and unit of time. Despite their limited compute budget, all methods outperform random search by a significant margin. In the large-compute setting, all inference methods achieve a similar maximum target value by the end of the experiment. CG and SGD make similar progress per acquisition step but SGD is faster per unit of time. SVGP is slightly slower than CG per both. In summary, our results suggest that SGD can be an appealing uncertainty quantification technique for large-scale GP-based sequential decision making.
In this work, we explored using stochastic gradient algorithms to approximately compute Gaussian process posterior means and function samples. We derived optimization objectives with linear and sublinear cost for both. We showed that SGD can produce accurate predictions—even in cases when it does not converge to an optimum. We developed a spectral characterization of the effects of non-convergence, showing that it manifests itself mainly through error in an extrapolation region located away—but not too far away—from the observations. We benchmarked SGD on regression tasks of various scales, achieving state-of-the-art performance for sufficiently large or ill-conditioned settings. On a Thompson sampling benchmark, where well-calibrated uncertainty is paramount, SGD matches the performance of more expensive baselines at a fraction of the computational cost.
We are grateful to David R. Burt for suggesting a few tricks which helped us complete the development of our theory and thank Bruno Mlodozeniec for identifying and correcting a small mistake in the proof of Appendix E.1. JAL and SP were supported by the University of Cambridge Harding Distinguished Postgraduate Scholars Programme. JA acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC. JMHL acknowledges support from a Turing AI Fellowship under grant EP/V023756/1. AT was 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. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (http://www.hpc.cam.ac.uk) funded by an EPSRC Tier-2 capital grant. This work was also supported with Cloud TPUs from Google's TPU Research Cloud (TRC).
[47] V. Wild, M. Kanagawa, and D. Sejdinovic. Connections and Equivalences between the Nyström Method and Sparse Variational Gaussian Processes. arXiv:2106.01121, 2023. Cited on pages 5, 18.
[48] W. J. Wilkinson, P. E. Chang, M. R. Andersen, and A. Solin. State Space Expectation Propagation: Efficient Inference Schemes for Temporal Gaussian Processes. In International Conference on Machine Learning, 2020. Cited on page 1.
[49] A. G. Wilson and H. Nickisch. Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP). In International Conference on Machine Learning, 2015. Cited on page 1.
[50] J. T. Wilson, V. Borovitskiy, A. Terenin, P. Mostowsky, and M. P. Deisenroth. Efficiently Sampling Functions from Gaussian Process Posteriors. In International Conference on Machine Learning, 2020. Cited on pages 2, 9.
[51] 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.
[52] D. Zou, J. Wu, V. Braverman, Q. Gu, and S. M. Kakade. Benign Overfitting of Constant-Stepsize SGD for Linear Regression. In Conference on Learning Theory, 2021. Cited on pages 1, 3, 7.
We employ the CG implementation of the JAX SciPy module and follow Wang et al. [46] in using a pivoted Cholesky preconditioner of size 100. Our preconditioner implementation resembles the implementation of the TensorFlow Probability library. For a small subset of datasets, we find the preconditioner to lead to slower convergence, and we report the results for conjugate gradients without preconditioning instead.
In addition to the experiment hyperparameters described in Section 4, for all methods, we use 2000 random Fourier features to draw each prior function used for computing posterior function samples via pathwise conditioning in (4).
For SGD, at each step, we draw 100 new Fourier features to estimate the regularizer term. In all SGD experiments, we use a Nesterov momentum value of 0.9 and Polyak averaging, following Antorán et al. [3]. For all regression experiments we use a learning rate of 0.5 to estimate the mean function representative weights, and a learning rate of 0.1 to draw samples. For Thompson sampling, we use a learning rate of 0.3 for the mean and 0.0003 for the samples. In both settings, we perform gradient clipping using optax.clip_by_global_norm with max_norm set to 0.1.
Code available at: HTTPS://GITHUB.COM/CAMBRIDGE-MLG/SGD-GP.
A.1 Hyperparameter Selection for Regression
We use a zero prior mean function and the Matérn-3/2 kernel, and share hyperparameters across all methods, including baselines. For each dataset, we choose a homoscedastic Gaussian noise variance, kernel variance and a separate length scale per input. For datasets with less than 50k observations, we tune these hyperparameters to maximize the exact GP marginal likelihood. The cubic cost of this procedure makes it intractable at a larger scale: instead, for datasets with more than 50k observations, we obtain hyperparameters using the following procedure:
1. From the training data, select a centroid data point uniformly at random.
2. Select a subset of 10k data points with the smallest Euclidean distance to the centroid.
3. Find hyperparameters by maximizing the exact GP marginal likelihood using this subset.
4. Using 10 different centroids, repeat the preceding steps and average the hyperparameters.
This approach avoids aliasing bias [1, 5] due to data subsampling and is tractable for large datasets.
A.2 Inducing Point Selection
For the inducing point SGD experiment shown in Figure 2, we choose inducing points as a subset of the training points. Due to the large number of datapoints in our dataset, we found  $k$ -means to converge very slowly. Instead, we develop an ad-hoc point elimination algorithm based on  $k$ -nearest-neighbors (KNN). We use the KNN implementation ANNOY, which we first run on the HOUSEELECTRIC dataset with num_trees set to 50. We then iterate through our dataset, retrieving 100 nearest neighbors for each original point. Of these 100 nearest neighbors, only the ones closer to the original point than some length scale parameter, in terms of  $\ell_2$  distance, are selected. If the number of points selected is larger than one, we eliminate the original point from our dataset and we also eliminate other points which are within the length scale neighborhood of the original and selected points simultaneously. The selected points are kept. We vary the number of points eliminated by our algorithm through the modification of the length scale parameter.
A.3 Maximizing Posterior Functions for Thompson Sampling
We maximize a random acquisition function sampled from the GP posterior in a three step process:
1. Evaluate the posterior function sample at a large number of nearby input locations. We find nearby locations using a combination of exploration and exploitation. For exploration, we sample locations uniformly at random from  $[0, 1]^d$ . For exploitation, we first subsample the training data with probabilities proportional to the observed objective function values and then add Gaussian noise  $\varepsilon_{\mathrm{nearby}} \sim \mathrm{N}(0, \sigma_{\mathrm{nearby}}^2)$ , where  $\sigma_{\mathrm{nearby}} = l/2$  and  $l$  is the kernel
Figure 7: Illustration of a single Thompson sampling acquisition maximization step on a 1D problem.
length scale. Throughout experiments, we find  $10\%$  of nearby locations using the uniform exploration strategy and  $90\%$  using the exploitation strategy.