Can Small Training Runs Reliably Guide Data Curation? Rethinking Proxy-Model Practice

TL;DR

This work reveals that using fixed hyperparameters for proxy-model assessments of data recipes can yield unreliable conclusions that may flip with small hyperparameter changes. It introduces a simple yet effective patch: train proxy models with a tiny learning rate $\eta_{tiny}$, which, both theoretically (for random-feature models) and empirically across 23 data recipes, preserves dataset rankings when transferred to optimally-tuned large-scale targets. The authors prove a transferability theorem linking tiny-LR proxy rankings to the infinite-width optimum, and demonstrate dramatic improvements in cross-scale reliability and reduced decision regret. Practically, this approach enables data teams to identify high-potential data recipes with far greater confidence before committing to expensive full-scale pretraining, while pointing to future directions in joint data-training hyperparameter optimization and multi-epoch proxy protocols.

Abstract

Data teams at frontier AI companies routinely train small proxy models to make critical decisions about pretraining data recipes for full-scale training runs. However, the community has a limited understanding of whether and when conclusions drawn from small-scale experiments reliably transfer to full-scale model training. In this work, we uncover a subtle yet critical issue in the standard experimental protocol for data recipe assessment: the use of identical small-scale model training configurations across all data recipes in the name of "fair" comparison. We show that the experiment conclusions about data quality can flip with even minor adjustments to training hyperparameters, as the optimal training configuration is inherently data-dependent. Moreover, this fixed-configuration protocol diverges from full-scale model development pipelines, where hyperparameter optimization is a standard step. Consequently, we posit that the objective of data recipe assessment should be to identify the recipe that yields the best performance under data-specific tuning. To mitigate the high cost of hyperparameter tuning, we introduce a simple patch to the evaluation protocol: using reduced learning rates for proxy model training. We show that this approach yields relative performance that strongly correlates with that of fully tuned large-scale LLM pretraining runs. Theoretically, we prove that for random-feature models, this approach preserves the ordering of datasets according to their optimal achievable loss. Empirically, we validate this approach across 23 data recipes covering four critical dimensions of data curation, demonstrating dramatic improvements in the reliability of small-scale experiments.

Figures (29)

• Figure 1: Overview of the industrial practice of data recipe ablation using proxy models. Left: Data teams assess each candidate dataset using small-scale proxy models trained with identical hyperparameters, and recommend the data recipe with the best performance. Right: Model training teams perform extensive hyperparameter tuning to optimize performance on the large target model.
• Figure 2: Validation loss rankings of 23 data recipes evaluated on proxy models (GPT2-125M) and target models (Pythia-1B), where target models undergo extensive dataset-specific hyperparameter tuning. Rankings are determined by the pretrained models' loss on Pile's validation split gao2020pile. Left: When proxy models are trained with a standard learning rate ($3 \times 10^{-4}$), data recipe rankings exhibit severe disagreement between proxy and target scales, with many dramatic reorderings that would lead to suboptimal data recipe ablation. Right: When proxy models are trained with a tiny learning rate ($1 \times 10^{-6}$), dataset rankings remain highly consistent across scales.
• Figure 3: Performance comparison of DCLM variants by training GPT2-Small with two similar learning rates in a near-optimal regime. DCLM-dedup-GS is a variant of the DCLM dataset constructed with more stringent deduplication thresholds (see Table \ref{['table:dataset-recipes']} for details). (a) Validation loss on the Pile dataset. (b) Average accuracy across 5 downstream benchmarks (see Section \ref{['sec:eval-settings']}). Star markers indicate the superior recipe in each setting. The shaded regions represent the standard errors calculated across 3 random seeds.
• Figure 4: (a) Correlation between losses of GPT2-Small trained with a small learning rate $1\times 10^{-6}$ versus optimally tuned hyperparameters, evaluated on Pile's validation split. Each point represents one dataset, demonstrating that tiny learning rate performance is strongly correlated with the optimal performance for the same model architecture. (b) Heatmap showing Spearman rank correlation between dataset rankings from proxy (GPT2-Small) and target (GPT2-Large) models trained with varying learning rate combinations. High correlation (darker area) in the lower-left indicates that when both models use tiny learning rates, dataset rankings are preserved across model scales.
• Figure 5: Average rank correlation between proxy (GPT2-Small) and target (GPT2-Large) for the loss computed over a variety of validation domains (from Pile) and downstream benchmarks. Compared to the standard learning rate used in the literature karpathy2022nanogpt, a tiny learning rate consistently improves the proxy model's rank transferability, often by a significant margin.

Theorems & Definitions (24)

• Theorem 1: Informal
• Remark
• Remark : Computational challenges in evaluating proxy model transferability
• Remark : Architecture and scale considerations.
• Theorem 2: Main Theorem
• Lemma 3
• proof : Proof of \ref{['thm:main']}
• Lemma 4
• proof
• Corollary 5