# Aioli: A Unified Optimization Framework for Language Model Data Mixing

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 50 of its 70 references that are in the map.
- First posted on arXiv 8 Nov 2024.

## Builds on (most cited first, 30 of 50)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [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)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Super-NaturalInstructions: Generalization via Declarative Instructions on 1600+ NLP Tasks](https://mlsanity.com/paper/2204.07705/) (2022)
- [Deduplicating Training Data Makes Language Models Better](https://mlsanity.com/paper/2107.06499/) (2021)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [Cross-Task Generalization via Natural Language Crowdsourcing Instructions](https://mlsanity.com/paper/2104.08773/) (2021)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms](https://mlsanity.com/paper/1905.13319/) (2019)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [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 Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [Abductive Commonsense Reasoning](https://mlsanity.com/paper/1908.05739/) (2019)
- [On The Power of Curriculum Learning in Training Deep Networks](https://mlsanity.com/paper/1904.03626/) (2019)
- [Learning Word Vectors for 157 Languages](https://mlsanity.com/paper/1802.06893/) (2018)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [A Survey on Data Selection for Language Models](https://mlsanity.com/paper/2402.16827/) (2024)
- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)

## Built on by (most cited first, 30 of 34)

- [Cartridges: Lightweight and general-purpose long context representations via self-study](https://mlsanity.com/paper/2506.06266/) (2025)
- [Organize the Web: Constructing Domains Enhances Pre-Training Data Curation](https://mlsanity.com/paper/2502.10341/) (2025)
- [Between Circuits and Chomsky: Pre-pretraining on Formal Languages Imparts Linguistic Biases](https://mlsanity.com/paper/2502.19249/) (2025)
- [TritonRL: Training LLMs to Think and Code Triton Without Cheating](https://mlsanity.com/paper/2510.17891/) (2025)
- [Learning to See Before Seeing: Demystifying LLM Visual Priors from Language Pre-training](https://mlsanity.com/paper/2509.26625/) (2025)
- [Language Models Improve When Pretraining Data Matches Target Tasks](https://mlsanity.com/paper/2507.12466/) (2025)
- [Diversity as a Reward: Fine-Tuning LLMs on a Mixture of Domain-Undetermined Data](https://mlsanity.com/paper/2502.04380/) (2025)
- [Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models](https://mlsanity.com/paper/2508.11953/) (2025)
- [Compute Optimal Scaling of Skills: Knowledge vs Reasoning](https://mlsanity.com/paper/2503.10061/) (2025)
- [SampleMix: A Sample-wise Pre-training Data Mixing Strategey by Coordinating Data Quality and Diversity](https://mlsanity.com/paper/2503.01506/) (2025)
- [DUET: Optimizing Training Data Mixtures via Feedback from Unseen Evaluation Tasks](https://mlsanity.com/paper/2502.00270/) (2025)
- [DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models](https://mlsanity.com/paper/2603.26164/) (2026)
- [Evaluating Sample Utility for Efficient Data Selection by Mimicking Model Weights](https://mlsanity.com/paper/2501.06708/) (2025)
- [Mixtera: A Data Plane for Foundation Model Training](https://mlsanity.com/paper/2502.19790/) (2025)
- [R&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training](https://mlsanity.com/paper/2505.00358/) (2025)
- [Towards Next-Generation LLM Training: From the Data-Centric Perspective](https://mlsanity.com/paper/2603.14712/) (2026)
- [TANDEM: Bi-Level Data Mixture Optimization with Twin Networks](https://mlsanity.com/paper/2606.04401/) (2026)
- [Rethinking Data Mixture for Large Language Models: A Comprehensive Survey and New Perspectives](https://mlsanity.com/paper/2505.21598/) (2025)
- [Hubs or Fringes: Pretraining Data Selection via Web Graph Centrality](https://mlsanity.com/paper/2606.11499/) (2026)
- [RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories](https://mlsanity.com/paper/2606.18663/) (2026)
- [Domain-Aware Scaling Laws Uncover Data Synergy](https://mlsanity.com/paper/2607.11052/) (2026)
- [MaD-Mix: Multi-Modal Data Mixtures via Latent Space Coupling for Vision-Language Model Training](https://mlsanity.com/paper/2602.07790/) (2026)
- [GEM: Geometric Entropy Mixing for Optimal LLM Data Curation](https://mlsanity.com/paper/2605.26121/) (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)
- [The Chicken and Egg Dilemma: Co-optimizing Data and Model Configurations for LLMs](https://mlsanity.com/paper/2602.08351/) (2026)
- [Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time](https://mlsanity.com/paper/2605.15220/) (2026)
- [Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning](https://mlsanity.com/paper/2606.24133/) (2026)
- [CausalMix: Data Mixture as Causal Inference for Language Model Training](https://mlsanity.com/paper/2607.01104/) (2026)

## Abstract (arXiv)

Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over training runs to dynamically updating proportions throughout training. Surprisingly, we find that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity. To understand this inconsistency, we unify existing methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law -- an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of a method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we observe. Using this insight, we derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.27 test perplexity points, whereas existing methods fail to consistently beat stratified sampling, doing up to 6.9 points worse. Moreover, in a practical setting where proportions are learned on shorter runs due to computational constraints, Aioli can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.012 test perplexity points.