Princeton Researchers Fix Costly AI Training Mistakes With 100x Smaller Learning Rates
A new study shows that proxy models used to pick pretraining data can flip their verdicts with tiny learning rate changes, with a simple fix.
- Standard proxy-model data ablation can flip dataset rankings from tiny hyperparameter tweaks, misleading full-scale training decisions.
- Fix: train proxy models with learning rates 1-2 orders of magnitude smaller than standard, roughly 1e-5 to 1e-6 for LLMs.
- Spearman rank correlation jumped above 0.95 across 253 dataset pairs between GPT2-125M proxy and tuned Pythia-1B target.
- Validated across 23 data recipes, three architectures (GPT2, Pythia, OPT), and sizes from 70M to 1B.
- Theoretical proof for random feature models shows tiny LR preserves dataset ranking at the infinite-width optimum.
- Note on DataDecide: fixed-hyperparameter benchmarks may themselves overfit to specific training configurations.
Tiny learning rates make data proxies more reliable
Teams training large language models often compare candidate corpora with small, inexpensive proxy models, then scale the highest-ranked recipe into a larger target run. In Can Small Training Runs Reliably Guide Data Curation?, a preprint from researchers at Princeton, Tsinghua, Stanford, UC Berkeley, and Virginia Tech, the authors report that dataset rankings can reverse after modest learning-rate changes. They recommend running proxies at learning rates 10 to 100 times below the usual setting.
A bad proxy ranking can send an expensive target run toward the wrong corpus. The study finds that tiny-learning-rate proxies align much more closely with results from larger models whose hyperparameters were tuned separately for each dataset.
Fixed settings can flip the winner
Standard curation experiments hold the architecture, compute budget, and optimizer settings constant across candidate datasets. That design isolates the dataset under one configuration, although the best learning rate and related settings can vary by corpus. Many production pipelines then tune the target model after selecting its data, creating a mismatch between proxy evaluation and final training.
The authors demonstrate the instability by comparing the DataComp-LM corpus, known as DCLM, with a more aggressively deduplicated variant on GPT-2 Small. At a lower learning rate, DCLM achieved lower validation loss and stronger downstream benchmark scores. A modest increase in the rate put the deduplicated variant ahead on both measures.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
