Zero-Shot Performance Prediction for Probabilistic Scaling Laws

AI summary

Overview

• Research area: Machine learning performance prediction — specifically zero-shot forecasting of learning curves and scaling laws for NLP models.
• Technical level: Intermediate. Familiarity with Gaussian processes, Bayesian models, and scaling laws helps, but the paper's central ideas are described conceptually here.
• Scope: The paper proposes modelling NLP learning-curve prediction as a multitask problem with a two-layer hierarchical structure, using latent variable multi-output Gaussian Processes to enable zero-shot curve and probabilistic scaling-law prediction.

What This Paper Is About

Predicting how NLP models will perform as dataset size or compute grows is expensive, because it normally requires actually training many models across many configurations. This paper asks whether those learning curves contain a hidden two-layer hierarchical structure (for example, models sharing embedding parameters but differing in layer count) that can be exploited to predict the curves of unseen model sizes without training them. The goal is to produce accurate "probabilistic scaling laws" — scaling-law estimates with uncertainty — at substantially lower computational cost.

Key Contributions

1. A hierarchical multitask formulation of learning-curve prediction. Each task's data is organized in a two-layer hierarchy, and the paper shows this modelling choice is effective on three small-scale NLP datasets with up to 30 learning curves.
2. Use of latent variable multi-output Gaussian Processes (Ma GP) to capture both within-task hierarchical structure and correlations across tasks, enabling zero-shot prediction of learning curves.
3. A probabilistic scaling-law procedure in which the Ma GP zero-shot predictions feed into Monte Carlo simulation over multiple runs to estimate the coefficients of the scaling law, with uncertainty.
4. An active learning query strategy for selecting which learning curves to query, shown to reduce predictive uncertainty and yield more reliable predicted scaling laws than largest-first, smallest-first, or random ordering.

Main Findings

• Ma GP outperforms all baselines on nanoGPT. On the nanoGPT large dataset, Ma GP achieved MSE 0.02 ± 0.01 / MAE 0.10 ± 0.02 (Tri split) versus DHGP at 0.07 ± 0.00 / 0.19 ± 0.00, BNN (LC) at 11.02 ± 0.47 / 2.82 ± 0.04, and BNN (orig) at 13.67 ± 0.82 / 3.03 ± 0.06.
• Smaller training data degrades performance. Reducing to nanoGPT small worsened Ma GP results (T1 small: MSE 0.08 ± 0.00, MAE 0.21 ± 0.00 versus 0.04 ± 0.05 / 0.12 ± 0.08 for T1 on nanoGPT large). BNN results were not reported for T1 small.
• The data is extremely small compared to prior work. The largest dataset in the study comprises only 29 learning curves with 11 data points each, whereas the smallest dataset in Klein et al. [18] contained 256 configurations. The abstract describes datasets with up to 30 learning curves.
• Hierarchies are exchangeable. When the hierarchy was swapped (layer count as the task level, embedding parameters as the second level), Ma GP still outperformed DHGP across Quad, Tri, and T1 splits, supporting the hypothesis that the hierarchy ordering can be exchanged.
• Predicted scaling laws approach ground truth. Against a ground-truth law of −0.056c + 3.51, Ma GP produced (−0.043 ± 0.002)c + (2.957 ± 0.074) for Quad with AbC 0.521 ± 0.069, (−0.052 ± 0.005)c + (3.352 ± 0.202) for Tri with AbC 0.160 ± 0.172, and (−0.059 ± 0.006)c + (3.636 ± 0.245) for T1 with AbC 0.111 ± 0.222.
• Lower compute cost for obtaining scaling laws. Reported costs for obtaining training data were 1.28 · 10^5 PetaFLOPs (Quad), 2.06 · 10^5 (Tri), and 5.15 · 10^5 (T1), compared with 6.081 × 10^5 PetaFLOPs for the full nanoGPT large dataset.
• A cost-accuracy trade-off exists. The Quad split is the most compute-efficient for training-data acquisition but yields the highest AbC; Tri and T1 achieve lower AbC at the cost of higher predictive uncertainty.
• Active learning gives the most confident scaling laws. Across query counts, active learning AbC remained consistently below (or comparable to) the initial training set (query number 0), a property not shared by the other strategies. Largest First achieved similar AbC values but at higher computational cost and with larger standard deviations, while random querying introduced substantial variability.
• Hierarchical modelling also helps bilingual translation. Ma GP had higher predictive uncertainty for BLEU but consistently outperformed all baselines in mean RMSE; the naive baseline (NBL, averaging over source and target language curves) showed substantial variability, and DHGP improved over NBL but still displayed heightened uncertainty.
• Few-shot extrapolation works for multilingual translation. For the M2M100 1.2B model, Ma GP consistently outperformed the naive regression baseline (NRBL) in RMSE when extrapolating the final nine missing points, despite higher predictive uncertainty; the highest uncertainty was observed for Tagalog (tl) as a target language.
• Prediction error depends on sample density. Per Appendix E, Figure E.3, all error metrics improve once the sample size exceeds approximately 100 data points.

Methodology in Plain English

The researchers treat every learning curve as a "task," and assume the curves within a dataset fall into groups and sub-groups — a two-layer hierarchy. For nanoGPT, the first layer is the number of embedding parameters and the second is the number of layers; for bilingual translation, the first layer is the source language and the second is the target language; for multilingual translation, the first layer is the source language and the second is the target language, with model sizes treated as separate tasks under the assumption that the source language and model size are the same.

They fit these hierarchies with the latent variable multi-output Gaussian process model of Ma et al. [24] ("Ma GP"), which learns a shared mean function across all tasks plus task-specific latent variables that capture correlations between tasks. This lets the model predict a curve for a model configuration it has never seen — zero-shot prediction.

To turn predicted curves into a scaling law, they repeatedly (over R runs, 10 in the experiments) train the model on the observed curves, zero-shot predict the missing curves, fit the functional form l(c) = (c/c0)^(−γ) to the compute-efficient frontier of the combined set, and average the resulting coefficients. This Monte Carlo averaging yields a probabilistic scaling law rather than a single point estimate. They evaluate against a Deep Hierarchical Gaussian Process baseline (DHGP), two Bayesian neural network variants with different basis functions (BNN orig, BNN LC), and dataset-specific naive baselines in the translation settings.

The nanoGPT data comes from curves measured over 19,072 validation steps (one pass through FineWeb-Edu, 10B), subsampled to 11 data points per curve. They use three train-test splits (Quad, Tri, T1) and a nanoGPT small subset with embedding sizes 512, 960, 1600 and layer counts 8, 12, 32, 48. Curves are z-score normalized and step/dataset sizes log-transformed for training, but evaluated in the original domain. Finally, they test four query strategies (Largest First, Smallest First, Active Learning by predictive variance, Random Order) to see which produces the most reliable scaling law per unit of compute.

Why This Matters

Impact on research. The paper argues that scaling laws — usually derived through exhaustive, expensive training runs — can instead be estimated probabilistically with far fewer trained models. It also introduces hierarchical structure as an explicit modelling assumption for learning-curve datasets, which the authors note had surprisingly not been explored in this context, and shows the assumption holds across three quite different NLP settings.

Real-world applications (as motivated by the paper):

• Model selection — choosing which model configuration to train or deploy based on predicted performance.
• Data acquisition and curation — deciding whether collecting more data for a specific source or target language is justified, based on predicted performance.
• Compute planning — estimating the compute budget needed to reach a target loss before committing resources.
• Early stopping and efficiency — informing stopping criteria and reducing the energy demands, processing times, and carbon footprint of training.

Industry relevance. The authors frame the work around the growing computational, financial, and environmental cost of large language models. Since the framework predicts from very small datasets (as few as 29 learning curves, each with 11 data points), it targets settings where practitioners cannot afford a full sweep of scaling-law experiments but still need defensible performance forecasts with uncertainty estimates.

Future Directions

• Scaling to larger, denser datasets. The paper notes that all error metrics improve once the sample size exceeds approximately 100 data points, and that its largest dataset holds only 29 learning curves, so extending to richer curve collections is an obvious next step.
• Designing better train-test split strategies. The observed trade-off between AbC and predictive uncertainty across Quad, Tri, and T1 directly motivates the question the authors themselves raise of whether a more informative split can be designed.
• Testing the hierarchy and exchangeability assumptions more broadly. Exchangeability was validated for nanoGPT, bilingual, and multilingual translation; whether other NLP learning-curve datasets share this bi-level, exchangeable structure is open.
• Beyond bi-level hierarchies. The paper uses a two-layer hierarchy throughout, while noting that hierarchical data in other domains often contains multiple layers, leaving deeper hierarchies unexplored for learning curves.

Target Audience

Researchers and practitioners working on scaling laws, learning-curve extrapolation, Bayesian optimization, and Gaussian process models for NLP; machine learning engineers who need to forecast model performance and plan compute or data budgets under tight resource constraints. The paper is most accessible to readers with some background in probabilistic modelling and scaling-law literature, though its central claims about hierarchical structure and zero-shot prediction are explained conceptually.

Authors’ abstract

The prediction of learning curves for Natural Language Processing (NLP) models enables informed decision-making to meet specific performance objectives, while reducing computational overhead and lowering the costs associated with dataset acquisition and curation. In this work, we formulate the prediction task as a multitask learning problem, where each task's data is modelled as being organized within a two-layer hierarchy. To model the shared information and dependencies across tasks and hierarchical levels, we employ latent variable multi-output Gaussian Processes, enabling to account for task correlations and supporting zero-shot prediction of learning curves (LCs). We demonstrate that this approach facilitates the development of probabilistic scaling laws at lower costs. Applying an active learning strategy, LCs can be queried to reduce predictive uncertainty and provide predictions close to ground truth scaling laws. We validate our framework on three small-scale NLP datasets with up to $30$ LCs. These are obtained from nanoGPT models, from bilingual translation using mBART and Transformer models, and from multilingual translation using M2M100 models of varying sizes.