Research map / Large language models / Language, safety and interpretability
RegMix: Data Mixture as Regression for Language Model Pre-training
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.
Proposes RegMix, which fits a regression model on losses of 512 ultra-small (1M-param) proxy models trained on random Pile domain mixtures, then uses the predicted best mixture to train larger models.
Problem
Choosing an effective data mixture for LLM pre-training is hard; prior proxy-model methods (e.g., DoReMi) are expensive, unstable, and sequential.
Approach
Formulate data mixture selection as regression: train many tiny proxy models on diverse mixtures (Dirichlet-sampled), regress the target validation loss on mixture weights (linear ridge or LightGBM), simulate millions of mixtures, and train the large model on the top-predicted mixture, relying on a rank-invariance assumption across model scales.
Key results
- LightGBM regression fit on 512 1M-param models predicts rank of unseen mixtures for 1B models (25B tokens) with Spearman rho 97.12%.
- RegMix's selected mixture achieves average 47.3 across 13 downstream tasks vs. 45.1 for Human and 46.8 for DoReMi, using about 3.5e18 FLOPs versus 3.7e19 for DoReMi.
- Validation loss on Pile-CC correlates most strongly with downstream performance, more than Wikipedia loss; the best predicted mixture assigns 0.87 weight to Pile-CC.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Select the training data mixture with RegMix instead of human-curated weights | 1B-parameter pre-training on the Pile · Pile downstream benchmarks | Average Performance: 45.1 -> 47.3 | Helped |
| Select the training data mixture with RegMix instead of human-curated weights | 1B-parameter pre-training on the Pile · HellaSwag | accuracy: 37.4 -> 44.2 | Helped |
| Select the training data mixture with RegMix instead of human-curated weights | 1B-parameter pre-training on the Pile · Pile downstream benchmarks (LightEval) | Average Performance: 39.5\text{\, $\pm$\,0.3} -> 41.5\text{\, $\pm$\,0.2} | Helped |
| Use LightGBM regression rather than linear regression | RegMix regression rank prediction · Unseen Pile data mixtures, 1M models | Spearman rank correlation ρ: 90.08 -> 98.45 | Helped |
| Use LightGBM regression rather than linear regression | RegMix regression rank prediction · Unseen Pile data mixtures, 60M models | Spearman rank correlation ρ: 89.26 -> 98.64 | Helped |
| Use LightGBM regression rather than linear regression | RegMix regression rank prediction · Unseen Pile data mixtures, 1B models | Spearman rank correlation ρ: 88.01 -> 97.12 | Helped |
| Increase the number of proxy models rather than increasing proxy training tokens beyond a threshold | RegMix proxy-model search · Regression rank prediction | Spearman rank correlation ρ | Helped |
| Train 512 proxy models on 0.2B tokens rather than 128 proxy models on 0.8B tokens | RegMix proxy-model search · Regression rank prediction | performance | Helped |
| Use a RegMix-selected mixture rather than training only on Pile-CC | 1B-parameter pre-training on the Pile · Average Performance | Average Performance: 46.8 -> 47.3 | Helped |
Surprising findings
- Web corpus (Pile-CC) validation loss correlates with downstream tasks more strongly than Wikipedia loss, contrary to the common assumption that WikiText is most representative.
- Domain interactions are counterintuitive; e.g., PhilPapers appears to help all other domains under linear regression.
Limitations
- Rank invariance is only verified for model scales from 1M to 1B parameters; larger scales untested.
- Evaluated only up to 1B-parameter models for the main selection and assumes unlimited data per domain.
- Requires domain labels and shared tokenizer between proxy and target models.
Experimental setup
Datasets: The Pile (17 uncopyrighted subsets) The Stack C4 (C4100Domain validation) FineWeb (for 100-domain extension)
Benchmarks: Social IQA HellaSwag PiQA OpenBookQA Lambada SciQ ARC Easy COPA RACE LogiQA QQP WinoGrande
Metrics: Spearman rank correlation mean squared error average accuracy (0-shot to 5-shot)
Baselines: Human (Pile default weights) DoReMi PPL filtering Online Data Mixing (ODM) Pile-CC only Best of 64 random 1B models
Builds on: DoReMi Pile Chinchilla scaling laws
Closest papers by meaning
- Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance (ICLR 2025 · similarity 0.91)
- Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training (ICML 2026 · similarity 0.91)
- Aioli: A Unified Optimization Framework for Language Model Data Mixing (ICLR 2025 · similarity 0.89)
- Programming Every Example: Lifting Pre-training Data Quality Like Experts at Scale (ICML 2025 · similarity 0.88)
- Data Mixture Inference Attack: BPE Tokenizers Reveal Training Data Compositions (NeurIPS 2024 · similarity 0.88)
- Scaling Laws for Optimal Data Mixtures (NeurIPS 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.