Research map / Large language models / Language, safety and interpretability
Aioli: A Unified Optimization Framework for Language Model Data Mixing
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Unifies existing LM data mixing methods under a Linear Mixing Optimization framework and proposes Aioli, an online method that estimates mixing-law parameters during training and uses them to adjust group proportions.
Problem
Learning optimal mixture proportions over data groups for language model pretraining; existing mixing methods do not consistently beat stratified sampling in test perplexity.
Approach
Shows existing offline and online data mixing methods are equivalent to minimizing average loss subject to a method-specific mixing law, finds their mixing-law parameters are inaccurate, and proposes Aioli, which fits the linear dynamic mixing law parameters online from validation loss changes on interleaved one-hot-smoothed mixtures and updates proportions via exponentiated gradient descent.
Key results
- Aioli outperforms stratified sampling on all 6 dataset settings, by an average of 0.274 test perplexity points (up to 0.439), while existing methods are worse than stratified on at least one setting by up to 6.9 points.
- In the restricted setting, adding Aioli to existing methods improves average test perplexity in 28 of 30 settings, by an average of 1.202 and up to 12.012 points.
- Log-linear static and linear dynamic mixing laws fit the loss-proportion relationship well (average R^2 0.991 and 0.947, respectively).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Aioli instead of stratified sampling to estimate mixing parameters online and adjust proportions dynamically | 160M GPT-style decoder-only LM · Six SlimPajama data settings | average test perplexity: by an average of 0.274 and up to 0.439 points | Helped |
| Use existing data mixing methods instead of stratified sampling | 160M GPT-style decoder-only LM · SlimPajama data settings | average test perplexity: up to 6.9 points worse | Mixed |
| Dynamically adjust proportions learned by GS on shorter runs using Aioli | 160M GPT-style decoder-only LM · Six SlimPajama data settings | average test perplexity: average of 1.202 and a maximum of 12.012 points | Helped |
| Dynamically adjust shortened-run proportions using Aioli | Skill-It, DoReMi, DoGE, DML, and GS · Restricted training-budget setting | average test perplexity: 28 out of 30 cases; worsening by at most 0.025 points | Mixed |
| Fit log-linear static and linear dynamic mixing laws to loss-proportion data | 160M GPT-style decoder-only LM · Six SlimPajama data settings | MSE: 8.9 × 10^-4; 1.0 × 10^-4 | Helped |
| Use linear dynamic mixing law instead of log-linear static mixing law | 160M GPT-style decoder-only LM · Six SlimPajama data settings | R^2: 0.991; 0.947 | Mixed |
| Use online methods with method-specific mixing-law parameters | Skill-It, DoReMi, and DoGE · SlimPajama data settings | correlation between parameter similarity and improvement over stratified samplin: R^2=0.491 | Helped |
| Use greedy EGD selection rather than full optimization of dynamic proportions | Online data mixing methods · Three data settings | optimality of selected dynamic proportions: 2 out of 3 data settings | No clear effect |
| Use log-linear static or linear dynamic mixing laws on Books/C4 | 160M GPT-style decoder-only LM · Books/C4 | R^2: 0.558; 0.721 | Hurt |
Surprising findings
- No existing data mixing method consistently outperforms simple stratified sampling in average test perplexity across datasets.
- Greedy exponentiated gradient descent recovers near-optimal dynamic proportions in 2 of 3 tested data settings.
Limitations
- Aioli incurs extra inference cost from repeated validation evaluations used to estimate A^t.
- Linear and log-linear mixing laws fit poorly for some data groups with nonlinear loss-proportion relationships (e.g., Books/C4).
- Experiments are limited to 160M-parameter models on SlimPajama group settings.
Experimental setup
Datasets: SlimPajama RedPajama
Metrics: test perplexity mean squared error R^2
Baselines: Stratified sampling Grid search Data Mixing Laws (DML) Skill-It DoReMi DoGE
Builds on: DoReMi DoGE Skill-It Data Mixing Laws RegMix
Closest papers by meaning
- Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance (ICLR 2025 · similarity 0.92)
- RegMix: Data Mixture as Regression for Language Model Pre-training (ICLR 2025 · similarity 0.89)
- Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training (ICML 2026 · similarity 0.89)
- Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 2025 · similarity 0.89)
- A Hitchhiker's Guide to Scaling Law Estimation (ICML 2025 · similarity 0.88)
- MixMin: Finding Data Mixtures via Convex Minimization (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.