SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Steven Adriaensen*
adriaens@cs.uni-freiburg.de
Herilalaina Rakotoarison*
rakotoah@cs.uni-freiburg.de
Samuel Müller
muellesa@cs.uni-freiburg.de
Frank Hutter
fh@cs.uni-freiburg.de
Learning curve extrapolation aims to predict model performance in later epochs of training, based on the performance in earlier epochs. In this work, we argue that, while the inherent uncertainty in the extrapolation of learning curves warrants a Bayesian approach, existing methods are (i) overly restrictive, and/or (ii) computationally expensive. We describe the first application of prior-data fitted neural networks (PFNs) in this context. A PFN is a transformer, pre-trained on data generated from a prior, to perform approximate Bayesian inference in a single forward pass. We propose LC-PFN, a PFN trained to extrapolate artificial right-censored learning curves generated from a parametric prior proposed in prior art using MCMC. We demonstrate that LC-PFN can approximate the posterior predictive distribution over learning curves more accurately than MCMC, while being over 10 000 times faster. We also show that the same LC-PFN achieves competitive performance extrapolating a total of 20 000 real learning curves from four learning curve benchmarks (LCBench, NAS-Bench-201, Taskset, and PD1) that stem from training a wide range of model architectures (MLPs, CNNs, RNNs, and Transformers) on 53 different datasets with varying input modalities (tabular, image, text, and protein data). Finally, we investigate its potential in the context of model selection and find that a simple LC-PFN based predictive early stopping criterion obtains  $2 - 6 \times$  speed-ups on 45 of these datasets, at virtually no overhead.
Learning curve extrapolation [Mohr and van Rijn, 2022] aims to predict how much a machine learning model will improve with more training, e.g., to determine how much more training data to collect [Cortes et al., 1993, Frey and Fisher, 1999, Leite and Brazdil, 2004, Kolachina et al., 2012], or to define an early stopping criterion in online learning [Yao et al., 2007]. Learning curve extrapolation has recently been widely studied to speed up automated machine learning (AutoML) and hyperparameter optimization (HPO) of deep neural networks, by discarding non-promising configurations early [Swersky et al., 2014, Domhan et al., 2015, Klein et al., 2017, Baker et al., 2017, Chandrashekaran and Lane, 2017, Gargiani et al., 2019, Wistuba et al., 2022].
* Equal Contribution.
Despite these efforts, learning curve extrapolation is not yet widely adopted in practice, e.g., state-of-the-art multi-fidelity hyperparameter optimization techniques, such as BOHB [Falkner et al., 2018], still rely on successive halving [Li et al., 2017], i.e., the crude heuristic that learning curves mostly do not cross each other.
One reason for this is that, while many learning curves are well-behaved, some exhibit chaotic behavior and are intrinsically difficult to predict accurately [Choi et al., 2018]. In this setting, Bayesian approaches [Swersky et al., 2014, Domhan et al., 2015, Klein et al., 2017, Wistuba et al., 2022], which also quantify the reliability of their extrapolation, show great potential. However, existing methods for Bayesian inference either (i) put strong restrictions on the prior, and are incapable of modeling the variable nature of learning curves, or (ii) are too computationally expensive, limiting their practical applicability. Furthermore, most of this related work focuses on demonstrating the potential that learning curve extrapolation has to accelerate downstream AutoML/HPO tasks, yet fails to fully investigate the quality of the extrapolations themselves, and to quantify the approach's ability to handle the heterogeneity of real-world learning curves, e.g., varying performance metrics, curve shapes, divergence, heteroscedastic noise, etc.
In this work, we investigate the potential of learning curve extrapolation using prior-data fitted networks (PFNs), a meta-learned approximate Bayesian inference method recently proposed by Müller et al. [2022]. PFNs combine great flexibility with efficient and accurate approximation of the posterior predictive distribution (PPD) in a single forward pass of a transformer [Vaswani et al., 2017] trained on artificial data from the prior only. As PFNs are a promising alternative to Markov Chain Monte Carlo (MCMC) for approximating Bayesian inference, we compare our approach (LC-PFN) to the MCMC approach for learning curve extrapolation of Domhan et al. [2015], taking into account both the quality and the cost of PPD approximation.
In summary, our contributions are as follow:
- We are the first to apply PFNs to an extrapolation task, introducing LC-PFN, the first PFN for learning curve extrapolation.
- We demonstrate that LC-PFN can be more than  $10000 \times$  faster than MCMC while still yielding better probabilistic extrapolations.
- We show that LC-PFN does not only yield better probabilistic extrapolations on prior samples, but also on real learning curves of a wide range of architectures (MLPs, CNNs, RNNs, Transformers) on varying input modalities (tabular, image, text and protein data).
- We demonstrate the practical usefulness of LC-PFN to construct an early stopping criterion that achieves  $2 - 6 \times$  speedups over baselines.
- To facilitate reproducibility and allow others to build on our work, we open-source all code, data, and models used in our experiments at https://github.com/automl/lcpfn.
Learning curves and how to use them for decision-making has been an active research area, as recently surveyed by Mohr and van Rijn [2022]. Most related work considers point estimates of the curve or a specific property thereof [Cortes et al., 1993, Frey and Fisher, 1999, Kolachina et al., 2012, Baker et al., 2017, Kaplan et al., 2020], or follows a non-Bayesian approach to quantify uncertainty [Chandrashekaran and Lane, 2017, Gargiani et al., 2019].
Only a few works have explored Bayesian learning curve extrapolation. For example, the Freeze-Thaw Bayesian optimization method Swersky et al. [2014] used a Gaussian process (GP) as a joint model of learning curves and hyperparameters to decide what learning run to continue for a few epochs (or whether to start a new one). The model is then dynamically updated to fit the partial learning curve data. Training data grows quickly since each performance observation of the curve is treated as a datapoint, making exact GPs intractable, and thus the work relies on approximate GPs. Furthermore, their approach makes strong (prior) assumptions. On top of the standard GP assumptions, they used a specialized kernel assuming exponential growth to improve extrapolation. Domhan et al. [2015] proposed a less restrictive parametric prior (see Section 3.2 for more details) and used the gradient-free MCMC method from Foreman-Mackey et al. [2013] as approximate inference method. While MCMC is a very general approach, it can be sensitive to its hyperparameters (e.g.,
burn-in period, chain length, etc.) and, as we will show in Section 4, generating sufficient samples to reliably approximate the PPD may impose significant overhead. Klein et al. [2017] extended this parametric prior to also capture the effect of hyperparameter settings. In particular, they used a Bayesian neural network with a specialized learning curve layer, and trained this network using gradient-based MCMC on learning curve data from previously tested hyperparameter settings. While this approach is able to predict learning curves of previously unseen configurations, conditioning on the current partial learning curve requires retraining the Bayesian neural network online, which is costly. Recently, DyHPO [Wistuba et al., 2022] followed a similar dynamic HPO setup as Swersky et al. [2014], but used deep GPs [Damianou and Lawrence, 2013]. While deep GPs relax some of the standard GP assumptions, extrapolation abilities were not thoroughly analyzed, and DyHPO only predicts one epoch into the future. Finally, it is worth noting that, except for Domhan et al. [2015], all the aforementioned probabilistic approaches [Swersky et al., 2014, Klein et al., 2017, Chandrashekaran and Lane, 2017, Gargiani et al., 2019, Wistuba et al., 2022] utilize meta-learning across the learner's hyperparameter settings. While this is an interesting line of work, it limits applicability, and introduces confounding factors. We will therefore consider a simpler and more general setting in this work (see Section 3.1). Indeed, akin to the approach of Domhan et al. [2015], we operate without the assumption of access to data from previous runs employing different hyperparameter settings, nor do we assume the ability to generalize across these settings. Our results highlight that prior-data fitted networks (PFNs) offer a significantly more efficient and practical alternative to Markov Chain Monte Carlo (MCMC) methods. As categorized by Mohr and van Rijn [2022], Domhan et al. [2015] is the only comparable prior work within this category. This underlines the novelty and importance of our approach in the context of advancing current methodologies.
While we are the first to apply PFNs [Müller et al., 2022] to learning curve extrapolation, PFNs have previously been applied in different settings: Hollmann et al. [2023] used them to meta-learn a classifier for tabular data; Müller et al. [2023] as a surrogate model for Bayesian optimization; and most recently concurrent work by Dooley et al. [2023] as a zero-shot time series forecaster.
3.1 Bayesian learning curve extrapolation
Let  $y_{t} \in [0,1]$  represent the model performance (e.g., validation accuracy) at training step  $t \in \{1, \dots, m\}$ . The problem we consider in this paper can be formulated as follows: Given a partial learning curve  $y_{1}, \ldots, y_{T}$  up to some cutoff  $T$ , and a prior distribution  $p(\pmb{y})$  over learning curves, approximate the posterior predictive distribution (PPD)  $q(y_{t'} \mid y_{1}, \ldots, y_{T})$  for  $T < t' \leq m$ . We will further assume that we can calculate the relative probability density of  $p(\pmb{y})$ , a requirement for MCMC, and that we can generate samples from  $p(\pmb{y})$ , a requirement for PFNs. Figure 1 provides an illustration of Bayesian learning curve extrapolation, showcasing the posterior predictive distributions (PPDs) of the extrapolated curves generated by LC-PFN and MCMC, along with a few representative curves sampled from the prior distribution  $p(\pmb{y})$ .
Figure 1: (Left) Visualization of Bayesian learning curve extrapolation. The plot shows the median and the  $90\%$  confidence interval of the PPDs inferred using MCMC and LC-PFN, given two partial empirical learning curves of 10 and 20 epochs, respectively, and the prior described in Section 3.2. (Right) Example of learning curves sampled from the prior  $p(\mathbf{y})$ .
While the fixed ranges for  $y_{t}$  and  $t$  are well-suited for modeling particular learning curves (e.g., accuracy over epochs), they also are restrictive. In Appendix A, we discuss the invertible normalization procedure we apply to support extrapolating, possibly diverging, iteration-based learning curves across a broad range of performance metrics (e.g., log loss).
3.2 Learning curve prior
Following Domhan et al. [2015], we model  $\pmb{y}$  as a linear combination of  $K$  basis growth curves  $f_{k}$ , each parameterized by  $\theta_{k}$ , and i.i.d. additive Gaussian noise with variance  $\sigma^2$ , i.e.,
$$
y _ {t} \sim \mathcal {N} (f _ {\mathrm {c o m b}} (t | \boldsymbol {\xi}), \sigma^ {2}) \quad \text {w i t h} \quad f _ {\mathrm {c o m b}} (t | \boldsymbol {\xi}) = \sum_ {k = 1} ^ {K} w _ {k} \cdot f _ {k} (t | \boldsymbol {\theta} _ {k}),
$$
where we assume our model parameters
$$
\boldsymbol {\xi} = \left(w _ {1}, \dots , w _ {K}, \boldsymbol {\theta} _ {1}, \dots , \boldsymbol {\theta} _ {K}, \sigma^ {2}\right)
$$
to be random variables with prior  $p(\pmb{\xi})$ . Here, Domhan et al. [2015] assumed an uninformative prior (i.e.,  $p(\pmb{\xi}) \propto 1$ ), with the exception of some hard constraints. We adopt a strictly more informative prior, because (i) the original prior puts almost all probability mass on parameterizations yielding invalid learning curves, e.g.,  $y_{t} \notin [0,1]$ ; and (ii) we cannot practically sample from this prior having unbounded support, a requirement for PFNs.<sup>1</sup> Specifically, to mimic realistic learning curves we use bounded uniform weight priors  $w_{k} \sim \mathcal{U}(0,1)$ , a low-noise prior  $\log(\sigma^2) \sim \mathcal{N}(-8,2)$ , only allow curves with values in [0, 1] and, like Domhan et al. [2015], only accept curves whose last value is higher than its first. Putting all of these together, our prior distribution thus takes the form:
$$
p (\boldsymbol {\xi}) \propto \left(\prod_ {k = 1} ^ {K} p (w _ {k}) \cdot p (\boldsymbol {\theta} _ {k})\right) \times p (\sigma^ {2}) \times \mathbb {1} \left(f _ {\operatorname {c o m b}} (1 | \boldsymbol {\xi}) <   f _ {\operatorname {c o m b}} (m | \boldsymbol {\xi})\right) \times \left(\prod_ {t = 1} ^ {m} \mathbb {1} \left(f _ {\operatorname {c o m b}} (t | \boldsymbol {\xi}) \in [ 0, 1 ]\right)\right).
$$
Finally, we limit ourselves to three parametric families of learning curves ( $K = 3$ , see Table 1). These basis curves were chosen to capture a variety of growth trends and convergence behavior. We show examples of curves sampled from this prior in Figure 1 (right).
Table 1: Formulas of the three parametric basis curves and priors over their parameters.
Reference name	Formula f_k(t)	Prior p(θ_k)
pow3	c - at^-α	c ∼ U(0,1.25)	a ∼ U(-0.6,0.6)	log(α) ∼ N(0,4)
Janoschek	α - (α - β)e^-κt^δ	α ∼ U(0,1)	β ∼ U(0,2)	log(κ) ∼ N(-2,1)	log(δ) ∼ N(0,0.25)
ilog2	c - a/√(t+1)	c ∼ U(0,1)	a ∼ U(-0.5,0.5)
3.3 Prior-data fitted networks (PFNs)
In this paper, we propose to use prior-data fitted networks (PFNs, Müller et al., 2022) instead of MCMC for learning curve extrapolation. PFNs are neural networks trained to perform approximate Bayesian prediction for supervised learning settings. That is, PFNs are trained to predict some output  $y \in \mathbb{R}$ , conditioned on an input  $t$  and a training set  $D_{train}$  of given input-output examples. The PFN is trained for this task with samples obtained from a prior over datasets  $p(\mathcal{D})$ . The loss function for training a PFN  $q_{\theta}$  with parameters  $\theta$  is the cross entropy  $\ell_{\theta} = \mathbb{E}_{(t,y) \cup D_{train} \sim p(\mathcal{D})}[-\log q_{\theta}(y|t, D_{train})]$  for predicting the hold-out example's label  $y$ , given  $t$  and  $D_{train}$ . Müller et al. [2022] proved that minimizing this loss over many sampled tasks  $(t,y) \cup D_{train}$  directly coincides with minimizing the KL divergence between the PFN's predictions and the true PPD. In essence, the PFN meta-learns to perform approximate posterior inference on (meta-train) synthetic tasks sampled from the prior, and at inference time also does so for a (meta-test) real task.
<sup>1</sup> Note that the probability density of  $p(y_{t})$  is well-defined, a requirement for MCMC, but not for PFNs.
2We found that adding additional curves did not improve predictive performance significantly, and made the MCMC baseline even more expensive and less stable. Klein et al. [2017] also limited themselves to five basis curves for similar reasons.
Figure 2: A visualization of our LC-PFN model on the task of predicting an accuracy over epochs curve.  $D$  represents the epoch accuracies up to epoch  $3 (= T')$ . Attention between test and training positions is shown using red and blue arrows, respectively. Plots based on Müller et al. [2022].
3.4 PFNs for Learning Curve Extrapolation (LC-PFNs)
To apply PFNs to learning curve extrapolation, we train them on learning curves sampled from a given prior over learning curves (in our case, the prior defined in Section 3.2). Specifically, the training set we condition on is the available partial learning curve up to some varying cutoff point  $T'$ , i.e.,  $D_{train} = \{(t', y_{t'})\}_{t'=1}^{T'}$ , the input we condition on is the epoch  $t \in \{T' + 1, \dots, m\}$  to predict for, and the desired output  $y$  is the value of the learning curve at epoch  $t$ . During training, we randomly sample the cutoff points  $T' \sim \mathcal{U}(0, m - 1)$  for every batch in order to learn to predict for initial learning curves of varying sizes. Figure 2 illustrates the information flow in a LC-PFN during learning curve extrapolation.
LC-PFN architecture and hyperparameters We use the PFN architecture proposed by Müller et al. [2022] and visualized in Figure 2 (a). That is, we use a sequence Transformer [Vaswani et al., 2017] and treat each pair  $(t,y)$  (for train) and  $t$  (for test) as a separate position/token. We encode these using a simple linear layer. We do not use positional encoding such that we are permutation invariant. Furthermore, the attention matrix is masked such that every position only attends to the training positions. This way training examples can attend to each other, but the test examples do not influence each other's predictions. Note that the output of the PFN with parameters  $\theta$  is a distribution  $q_{\theta}(y|t,D_{train})$ . Following Müller et al. [2022], we discretize  $q_{\theta}$  in a finite number of bins whose probability mass is predicted by the PFN, as is shown in Figure 2 (b). The size of each bin is set such that, under the prior,  $y_{t}$  is equally likely to fall in each bin. The number of bins is a hyperparameter that we set to 1000. The LC-PFN model further inherits hyperparameters from the Transformer, including the number of layers (nlayers), number of heads (nhheads), embedding size (emsize), and hidden size (nhidden). We use four heads, a hidden size of 1024, and conduct a thorough ablation study to investigate the effects of the number of layers and embedding size on the final performance, exploring a grid of values (see Table 2). We use a standard training procedure for all experiments, employing the Adam optimizer [Kingma and Ba, 2015] (learning rate 0.0001, batch size 100) with cosine annealing [Loshchilov and Hutter, 2017] with a linear warmup over the first  $25\%$  epochs of the training. Finally, we set  $m = 100$ , implying LC-PFN is trained for extrapolating sequences of up to 100 training steps (e.g., epochs). We found that most curves are shorter in practice, and when longer sequences are encountered, we subsample them as described in Appendix B.
Our experiments aim to test the hypothesis that PFNs present a practical Bayesian approach to learning curve extrapolation. To this end, we first compare our LC-PFN approach against the MCMC approach of Domhan et al. [2015], using the same prior on samples generated from it (Section 4.1). Then, we extend the comparison to four real-world learning curve benchmarks (Section 4.2). Finally, we look beyond the quality of individual extrapolations and evaluate the potential of LC-PFN in the context of predictive early stopping to accelerate model selection (Section 4.3).
4.1 Extrapolating samples of the prior
The goal of this first experiment is to assess the ability of LC-PFNs and MCMC to approximate the true posterior predictive distribution (PPD). To avoid artifacts due to out-of-distribution data, in this experiment, we use curves sampled from the prior (defined in Section 3.2). Furthermore, we modified the original implementation of MCMC [Domhan et al., 2015], to use the curve prior we proposed in Section 3.2. In the following, we refer to this MCMC variant as MCMC-PP and denote the one using the original prior [Domhan et al., 2015] as MCMC-OP (used in Section 4.2). Since the LC-PFN and MCMC-PP methods are both (approximate) Bayesian inference methods, using the same prior, they aim to approximate the same true target PPD, given a partial learning curve.
As a performance metric, we evaluate the log-likelihood (LL) of the unseen data (right-censored curve) under the inferred PPD. We use this metric, also known as logarithmic score, to assess a model's ability to infer the remaining part of the curve based on the initial observed values. A benefit of this metric is that it measures the quality of the PPD as a whole (rather than merely focusing on the error associated to a specific PPD statistic) and, assuming data is generated by the prior, the exact PPD maximizes this metric.
Importantly, we vary the cutoff, i.e., the percentage of the observed curve used as input, to better assess the model's performance across different amounts of available information. Furthermore, to allow a more comprehensive comparison, we vary the hyperparameters of both LC-PFN and MCMC-PP. For LC-PFN, we vary the embedding size (emsize), the number of layers (nlayers), and the total number of learning curves used during training (nb_data). For MCMC-PP, we vary the number of chains generated by the emcee [Foreman-Mackey et al., 2013] ensemble sampler (nwalkers), the length of each chain
Table 2: Grid of hyperparameter values evaluated for MCMC-PP and LC-PFN.
Hyperparameters
MCMC-PP	nsamples ∈ [100, 250, 500, 1000, 2000, 4000]
nwalkers ∈ [26, 50, 100]
burn-in ∈ [0, 50, 100, 500]
thin ∈ [1, 10, 100]
Lc-PEN	nb_data ∈ [100k, 1M, 10M]
emsize ∈ [128, 256, 512]
nlayers ∈ [3, 6, 12]
(burn-in + nsamples), the part of the chain omitted to account for mixing (burn-in), and the sub-sample frequency (thin). The considered values for each hyperparameter are summarized in Table 2.
We conducted the comparison on 10 000 sampled curves. Figure 1 shows a few curve examples, as well as inferences using LC-PFN and MCMC-PP given the data of the first 10 - 20 epochs (cutoff). We observe that both predicted median and uncertainties are indeed similar. More inference examples, with different cutoffs can be found in Appendix C.4.
Results Figure 3 displays the average log-likelihood across MCMC-PP / LC-PFN inferences for varying hyperparameters and a  $10\%$  cutoff. The log-likelihood is shown w.r.t. runtime, which is measured as the average wall-clock time for a single inference on a single Intel(R) Xeon(R) Gold 6242 CPU. Note that this inference time includes both the fit and prediction times for MCMC variants. Table 3 provides results on higher cutoffs  $(20\%, 40\%, 80\%)$  and corresponding runtimes for three variants of each method (M1-3, P1-3), labeled in Figure 3. Generally, the LC-PFN variants (left side of figure 3) are significantly faster than MCMC-PP variants (right side). LC-PFN always ran in less than 0.1 seconds while the best MCMC-PP (M3) took over 100 seconds. Figure 3 also offers insights into the importance of LC-PFN and MCMC-PP hyperparameters. For MCMC-PP, both the cost and quality of inference increase with longer chain lengths and higher cutoffs. For LC-PFN, the inference cost increases with the model complexity which is closely related to the number of trainable parameters. Inference quality positively correlates with model size ("larger is better"), and the number of data LC-PFN was trained on. Among the hyperparameter grid we examined (Table 2), except for the smallest model (P1), all LC-PFN variants that were trained on 10M samples produce higher log-likelihood than the best MCMC-PP variant (M3). In particular, an LC-PFN (P2) with 3 layers, embedding size 256, trained on 10M samples achieved better performance (log-likelihood of PPD) than the best MCMC-PP, but more than 15 000 times faster. We also find that while the runtime of the best MCMC-PP can be reduced (with minor loss of quality) by using thinning (M2), the better LC-PFN is still approximately 7 000 times faster. Finally, it is important to note that training the largest LC-PFN (P3, 10M samples with 26M parameters) on the prior took approximately eight hours (single CPU, single RTX2080 GPU), but this cost is incurred only once for all of our experiments.