# AutoScale: Scale-Aware Data Mixing for Pre-Training LLMs

arXiv: https://arxiv.org/abs/2407.20177 · HTML: https://mlsanity.com/paper/2407.20177/

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 48 references that are in the map.
- First posted on arXiv 29 Jul 2024.

## Builds on (most cited first, 29 of 29)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models](https://mlsanity.com/paper/2206.04615/) (2022)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [CrowS-Pairs: A Challenge Dataset for Measuring Social Biases in Masked Language Models](https://mlsanity.com/paper/2010.00133/) (2020)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Scaling Data-Constrained Language Models](https://mlsanity.com/paper/2305.16264/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Investigating Bi-Level Optimization for Learning and Vision from a Unified Perspective: A Survey and Beyond](https://mlsanity.com/paper/2101.11517/) (2021)
- [Rho-1: Not All Tokens Are What You Need](https://mlsanity.com/paper/2404.07965/) (2024)
- [Skill-it! A Data-Driven Skills Framework for Understanding and Training Language Models](https://mlsanity.com/paper/2307.14430/) (2023)
- [Language models scale reliably with over-training and on downstream tasks](https://mlsanity.com/paper/2403.08540/) (2024)
- [Efficient Online Data Mixing For Language Model Pre-Training](https://mlsanity.com/paper/2312.02406/) (2023)
- [BiMix: A Bivariate Data Mixing Law for Language Model Pretraining](https://mlsanity.com/paper/2405.14908/) (2024)

## Built on by (most cited first, 22 of 22)

- [Improving Data Efficiency for LLM Reinforcement Fine-tuning Through Difficulty-targeted Online Data Selection and Rollout Replay](https://mlsanity.com/paper/2506.05316/) (2025)
- [ATLAS: Adaptive Transfer Scaling Laws for Multilingual Pretraining, Finetuning, and Decoding the Curse of Multilinguality](https://mlsanity.com/paper/2510.22037/) (2025)
- [Can Generalist Agents Automate Data Curation?](https://mlsanity.com/paper/2606.04261/) (2026)
- [InfoLaw: Information Scaling Laws for Large Language Models with Quality-Weighted Mixture Data and Repetition](https://mlsanity.com/paper/2605.02364/) (2026)
- [Rethinking Data Mixture for Large Language Models: A Comprehensive Survey and New Perspectives](https://mlsanity.com/paper/2505.21598/) (2025)
- [Rethinking Data Mixing from the Perspective of Large Language Models](https://mlsanity.com/paper/2604.07963/) (2026)
- [MixAtlas: Uncertainty-aware Data Mixture Optimization for Multimodal LLM Midtraining](https://mlsanity.com/paper/2604.14198/) (2026)
- [Repetition Mismatch: Why Data Mixture Experiments Don't Scale and How to Fix Them](https://mlsanity.com/paper/2606.07597/) (2026)
- [DataComp-VLM: Improved Open Datasets for Vision-Language Models](https://mlsanity.com/paper/2606.28551/) (2026)
- [Domain-Aware Scaling Laws Uncover Data Synergy](https://mlsanity.com/paper/2607.11052/) (2026)
- [Scaling Domain Data Repetition in LLM Pretraining](https://mlsanity.com/paper/2608.14071/) (2026)
- [Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization](https://mlsanity.com/paper/2603.08022/) (2026)
- [DataProphet: Demystifying Supervision Data Generalization in Multimodal LLMs](https://mlsanity.com/paper/2603.19688/) (2026)
- [Explaining Data Mixing Scaling Laws](https://mlsanity.com/paper/2606.08167/) (2026)
- [FastMix: Fast Data Mixture Optimization via Gradient Descent](https://mlsanity.com/paper/2606.14971/) (2026)
- [WARP: Weight-Space Analysis for Recovering Training Data Portfolios](https://mlsanity.com/paper/2607.01686/) (2026)
- [Exploring Polyglot Harmony: On Multilingual Data Allocation for Large Language Models Pretraining](https://mlsanity.com/paper/2509.15556/) (2025)
- [Characterizing Model-Native Skills](https://mlsanity.com/paper/2604.17614/) (2026)
- [How Much Is a Dataset Worth? Scaling Laws, the Vendi Score, and Matrix Spectral Functions](https://mlsanity.com/paper/2605.29448/) (2026)
- [Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning](https://mlsanity.com/paper/2606.24133/) (2026)
- [Harnessing the Potential of Optimizing Data Mixtures via Bayesian Domain Reweighting](https://mlsanity.com/paper/2607.27928/) (2026)
- [Selecting Repetition Counts Across Model Scales in Data-Constrained Pretraining](https://mlsanity.com/paper/2610.05126/) (2026)

## Abstract (arXiv)

Domain reweighting is an emerging research area aimed at adjusting the relative weights of different data sources to improve the effectiveness and efficiency of LLM pre-training. We show that data mixtures that perform well at smaller scales may not retain their advantage at larger scales, challenging the existing practice of determining competitive mixtures in small-scale experiments and directly applying them at much larger scales. To address this, we propose AutoScale, a two-stage, scale-aware data composition framework. First, AutoScale fits a parametric model that predicts the model's loss under different data compositions, then uses it to find an approximate best allocation at smaller, more manageable budgets. Next, leveraging a novel theoretical analysis of how optimal compositions evolve with scale, AutoScale extrapolates that composition to larger budgets without further retraining. Empirically, AutoScale accelerates convergence and improves downstream performance. For instance, when pre-training GPT-2 Large, it achieves a 28% faster perplexity reduction than baselines and up to a 38% speed-up over unweighted training, while yielding best-average results on various downstream tasks. Overall, our findings illustrate how domain importance shifts with training scale, underscoring the need for scale-dependent data curation in LLM training. Our code is open-sourced.