Scaling Gaussian Processes for Learning Curve Prediction via Latent Kronecker Structure

Authors: andreas  jihao
  lin
sebastian  ament
maximilian  balandat
eytan  bakshy
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 21:08:03

Abstract: A key task in AutoML is to model learning curves of machine learning models jointly as a function of model hyper-parameters and training progression. While Gaussian processes (GPs) are suitable for this task, naive GPs require $\mathcal{O}(n^3m^3)$ time and $\mathcal{O}(n^2 m^2)$ space for $n$ hyper-parameter configurations and $\mathcal{O}(m)$ learning curve observations per hyper-parameter. Efficient inference via Kronecker structure is typically incompatible with early-stopping due to missing learning curve values. We impose $\textit{latent Kronecker structure}$ to leverage efficient product kernels while handling missing values. In particular, we interpret the joint covariance matrix of observed values as the projection of a latent Kronecker product. Combined with iterative linear solvers and structured matrix-vector multiplication, our method only requires $\mathcal{O}(n^3 + m^3)$ time and $\mathcal{O}(n^2 + m^2)$ space. We show that our GP model can match the performance of a Transformer on a learning curve prediction task.

Main paper: 
 introduction: 
 Most modern machine learning (ML) models are trained with iterative methods, giving rise to learning curves which enable an analysis of the model quality as the training progresses. Being able to predict learning curves accurately based on results from partial training facilitates decisions about whether to continue training or to stop early, such that compute resources can be used more efficiently. This can substantially accelerate model exploration and improve the efficiency of AutoML algorithms.Existing work on learning curve prediction considered Gaussian processes (GPs) [Swersky et al., 2014 , Klein et al., 2020 , Wistuba et al., 2022] , parametric functions [Domhan et al., 2015 , Klein et al., 2017 , Kadra et al., 2023] , and Transformers [Adriaensen et al., 2023 , Rakotoarison et al., 2024] (see Appendix A for details). In this paper, we revisit the idea of a joint GP over hyper-parameters and learning curves. In principle, it is possible to define a GP on the product space of hyper-parameters and learning curves, but, for n hyper-parameter configurations and O(m) learning curve observations per hyper-parameter, this naïve approach requires O(n 3 m 3 ) time and O(n 2 m 2 ) space, which quickly becomes infeasible. An efficient approach is to impose Kronecker structure in the joint covariance matrix [Bonilla et al., 2007 , Stegle et al., 2011 , Zhe et al., 2019] . However, this requires complete learning curves for each hyper-parameter configuration, which is incompatible with early-stopping.Herein, we leverage latent Kronecker structure, which enables efficient inference despite missing entries by obtaining the joint covariance matrix of observed values from the full Kronecker product via lazily evaluated projections. With suitable posterior inference techniques, this reduces the asymptotic time complexity from O(n 3 m 3 ) to O(n 3 +m 3 ), and space complexity from O(n 2 m 2 ) to O(n 2 +m 2 ) compared to the naïve approach. We demonstrate this superior scalability empirically, and show that our GP with generic stationary kernels and 10 free parameters can match the performance of a specialized Transformer model with 14.69 million parameters on a learning curve prediction task. Ground Truth Posterior SamplesFigure 1 : Learning curve predictions on the Fashion-MNIST data from the LCBench benchmark [Zimmer et al., 2021] . The GP is fit to 16 partially observed learning curves (black). Their ground truth continuations (orange) are contained within the spread of posterior samples (blue). A typical learning curve which is observed close to convergence is predicted with confidence (left). Observing a smaller fraction of the learning curve leads to increased uncertainty in the prediction (left middle).The model also adapts well to less common noisy and spiky learning curves (right).2 Gaussian Processes with Latent Kronecker StructureWe consider the problem of defining a GP on the product space of hyper-parameter configurations x ∈ R d and learning curve progression t ∈ R, namely f : R d × R → R, where f ∼ GP(µ, k) with mean function µ and kernel k. Throughout this paper, we suppress the dependence on kernel hyperparameters for brevity of notation, set µ = 0 without loss of generality, and assume homoskedastic observation noise σ 2 . For an introduction to GPs, we refer to Rasmussen and Williams [2006] .The simplest way to address this problem is to define a kernel directly on the product space, which results in a joint covarianceCov(f (x, t), f (x ′ , t ′ )) = k((x, t), (x ′ , t ′ )). However, this quickly runs into scalability issues. Assuming we evaluate n hyper-parameter configurationsX := {x i } n i=1 at m progressions t = {t 1 , . . . , t m }, observing learning curves Y := {y i ∈ R m } n i=1, the joint covariance matrix requires O(n 2 m 2 ) space, and computing its Cholesky factorization takes O(n 3 m 3 ) time.Latent Kronecker Structure A common way to improve the scalability of GPs on product spaces is to introduce Kronecker structure [Bonilla et al., 2007 , Stegle et al., 2011 , Zhe et al., 2019] . In particular, one may define a product kernel k((x, t), (x′ , t ′ )) = k 1 (x, x ′ ) k 2 (t, t ′ ),where k 1 only acts on hyper-parameter configurations x and k 2 only considers learning curve progressions t. When applied to the observed data, the resulting joint covariance matrix K has Kronecker structure,K nm×nm = k((X × t), (X × t)) = k 1 (X, X) ⊗ k 2 (t, t) = K 1 n×n ⊗ K 2 m×m ,which can be exploited by expressing the eigenvalue decomposition of K in terms of K 1 and K 2 . This reduces the asymptotic time complexity to O(n 3 + m 3 ) and space complexity to O(n 2 + m 2 ).Unfortunately, the joint covariance matrix only has Kronecker structure if each x i is evaluated at every t j , that is, each learning curve must be fully observed, which conflicts with early-stopping. However, the joint covariance matrix over partially observed learning curves is a submatrix of the latent Kronecker product, and the former can be selected from the latter using a projection matrix P,K joint = P K latent P T = P(K 1 ⊗ K 2 )P T ,where P is constructed from the identity matrix by removing rows corresponding to missing values.In practice, projections can be implemented efficiently via slice indexing without instantiating P.Latent Kronecker Projected Kronecker Joint CovarianceFigure 2 : Selecting the joint covariance matrix using projections of the latent Kronecker product after observing {(x 1 , t 1 ), (x 1 , t 2 ), (x 2 , t 1 ), (x 2 , t 2 ), (x 2 , t 3 )}, two learning curve values from a first hyper-parameter configuration (blue) and three values from a second configuration (orange).Efficient Inference with Iterative Methods As a result of introducing projections, the eigenvalues and eigenvectors of K joint cannot be expressed in terms of eigenvalues and eigenvectors of K 1 and K 2 anymore, which prevents the use of Kronecker structure for efficient matrix factorization. However, despite projections, the Kronecker structure can still be leveraged for fast matrix multiplication,(A ⊗ B) vec(C) = vec(BCA T ) → P(A ⊗ B)P T vec(C) = P vec(B vec -1 (P T vec(C))A T ),where, in practice, vec and vec -1 correspond to reshaping, P T vec(C) amounts to zero padding, and P is slice indexing. This facilitates efficient GP inference via iterative methods, which only rely on matrix-vector multiplication (MVM) to compute solutions to systems of linear equations [Gardner et al., 2018 , Wang et al., 2019] . Leveraging latent Kronecker structure and lazy kernel evaluations, iterative methods require O(n 2 m + nm 2 ) time and O(n + m) space. This is similar to structured kernel interpolation (SKI) [Wilson and Nickisch, 2015] , an inducing point approximation that creates Kronecker structure by placing inducing points on regular grids. Compared to SKI, our approach only applies a product structure to separate hyper-parameters x from progressions t, and permits exact MVMs as long as progressions are logged on a shared grid, such as epochs.Posterior Samples via Matheron's Rule Maddox et al. [2021] proposed to draw posterior samples efficiently by exploiting Kronecker structure via Matheron's rule [Wilson et al., 2020 [Wilson et al., , 2021]] , which expresses a posterior sample in terms of a transformed prior sample,(f |Y)(• x , • t ) = f (• x , • t ) + (k 1 (• x , X) ⊗ k 2 (• t , t))(K 1 ⊗ K 2 + σ 2 I) -1 (vec(Y) -f ((X × t)) -ϵ),where (f |Y) is the posterior sample, f is the prior sample, f ((X × t)) is its evaluation at the training data, and ϵ ∼ N (0, σ 2 I). To support latent Kronecker structure, we introduce projections,f (• x , • t ) + (k 1 (• x , X) ⊗ k 2 (• t , t))P Tzero padding and slice indexing(P K 1 ⊗ K 2 P T + σ 2 I) -1 (vec(Y) -f ((X × t)) -ϵ)solve inverse matrix-vector product using iterative methodsIn combination with iterative methods, latent Kronecker structure can be exploited to compute the inverse matrix-vector product using fast MVMs. Further, computations can be cached or amortized to accelerate Bayesian optimization or marginal likelihood optimization [Lin et al., 2023 [Lin et al., , 2024a,b],b] . Drawing a single posterior sample takes O(n 3 + m 3 + n 3 * ) time and O(n 2 + m 2 + n 2 * ) space. 
experiments: 
 We conducted two empirical experiments. In our first experiment, we illustrate the superior scalability of our Latent Kronecker GP (LKGP) compared to naïve Cholesky factorization of the joint covariance matrix. In our second experiment, we predict the final values of partially observed learning curves and evaluate the quality of predictions via mean-square-error (MSE) and log-likelihood (LLH).We compared the time and memory requirements for training and prediction of LKGP, which uses iterative methods to exploit latent Kronecker structure, to naïve Cholesky factorization of the joint covariance matrix. To this end, we generated random training data with n = m ∈ {16, 32, ..., 512} and d = 10 (see Appendix C for details).2 4 2 5 2 6 2 7 2 8 2 9Training Data Size Figure 3 visualizes the time and memory consumption as a function of training data size. In general, empirical trends roughly match asymptotic complexities (see Section 2). In particular, LKGP is easily scalable to n = m = 512, taking about 6 seconds for training while using less than 2 GB of memory, whereas naïve Cholesky factorization already takes more than 3.5 minutes for training at n = m = 128 and runs out of memory at n = m = 256. Empirically, in terms of time, LKGP seems to scale better than its asymptotic complexity, which is likely due to the iterative solver converging in fewer iterations than would be mathematically required for an exact solution.Learning Curve Prediction Quality We replicated the experiment from Rakotoarison et al.[2024], Section 5.1, who used data from LCBench [Zimmer et al., 2021] to define a learning curve prediction task. In particular, the final validation accuracy is predicted given partially observed learning curves. We used the experimental setup and baseline results from Rakotoarison et al. [2024] , which are publicly available on GitHub. 2 The baselines include DPL [Kadra et al., 2023] , a neural network ensemble which makes predictions based on power laws; DyHPO [Wistuba et al., 2022] , a Gaussian process model equipped with a deep kernel; FT-PFN [Rakotoarison et al., 2024] , a Transformer model which is pre-trained on synthetic learning curves and makes predictions via in-context learning; and FT-PFN (no HPs), a version of FT-PFN which does not consider correlations across hyper-parameters.Figure 4 illustrates that LKGP achieves better or similar MSE compared to other methods and slightly worse LLH compared to FT-PFN. We consider this an impressive result, because LKGP only has access to the partially observed learning curves which are passed to FT-PFN as context at prediction time, uses basic stationary kernels without any inductive bias for modeling learning curves, and only has 10 parameters; whereas the most competitive model, FT-PFN, is pre-trained on millions of samples drawn from a specifically designed learning curve prior, uses a flexible non-Gaussian posterior predictive distribution, and has 14.69 million parameters [Rakotoarison et al., 2024] . 
conclusion: 
 Motivated by AutoML problems, we proposed a joint Gaussian process model over the product space of hyper-parameters and learning curves, which leverages latent Kronecker structure to accelerate computations and reduce memory requirements in the presence of partial learning curve observations. In contrast to existing Gaussian process models with Kronecker structure, our approach deals with missing entries by combining projections and iterative methods. Empirically, we demonstrated that our method has superior scalability than the naïve approach, and that it can match a specialized Transformer model in terms of prediction quality. Future work could investigate specialized kernels and heteroskedastic noise models, which would still be efficient due to the use of iterative methods. 
 
 Appendix: https://github.com/automl/ifBO/tree/icml-2024/src/section5.1 



Reviews: review: title: title: Official Review
main_review: main_review: The overall quality of the paper is good, utilizing latent Kronecker structure to over come the complexity issues in traditional Gaussian processes. The paper is generally clearly written, but would benefit from more detailed mathematical presentations. The idea of latent Kronecker structure is interesting, although not so ground-breaking. The paper shows its strengths in learning curve prediction tasks, but its scalability and applications to other practical tasks are unknown.

Pros:
* Propose latent Kronecker structure, which is simple but effective.

Cons: 
* Limited exploration of specialized kernels or heteroskedastic noise models.
* More diverse datasets could strengthen the empirical validation.
* Lacks discussions on scalability and applications to other practical tasks in addition to cure prediction.
title: title: Official Review of Submission79 by Reviewer UaYR
main_review: main_review: Summary:
----------
The authors propose to leverage a projection of a latent Kronecker structure to preserve the advantages brought by using kernels, but also handle missing values in inference of joint models of learning curves and hyperparameters for AutoML tasks. This roughly halves the order of time and storage space needed.

Main Review:
----------
Quality and clarity: The paper is well-written, exposition and results are well-explained.

Figure 1 shows good results.
Figure 3 is convincing, as it shows a clear advantage of LKGP over the naive approach.
I would be nice to have at least one line (in the appendix) mentioning how the choice of d impacts the results, if at all.

It would be interesting how the kernel choice affects the results and what motivated the choices of RBF and Matérn kernel.
The enable reproducibility, the authors code for LKGP should be made available online.

I would like to know the motivation for illustrating the performance per task instead of reporting the median over all tasks in Figure 4.

Typos
--------
In l. 211 the citation is formatted differently that e.g. in l. 196.

Citations
--------
The "a." in l. 127 should be removed.
The reference in ll. 142-143 is missing a venue.
In ll. 160 and 162 it should be "Gaussian" and "Bayesian".
The capitalization of titles is inconsistent: for most titles, almost all words are capitalized, but not for the reference in ll. 159-161.
The line break in ll. 178-179 is unfortunate, one should probably define a custom hyphenation for that word.