Source-linked AI summary
DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, Adams Wei Yu
TL;DR
Pretraining-domain mixture weights strongly affect language-model performance, but choosing them for broad downstream performance is unclear. DoReMi uses distributionally robust optimization on a small proxy model to produce domain weights without downstream-task knowledge, then transfers them to larger-model training, achieving a 2.6x speedup on The Pile and comparable performance to downstream-tuned weights on GLaM.
Problem
Choosing how much data to sample from each pretraining domain for broad downstream performance remains unclear, while existing approaches rely on intuition or downstream tasks.
Method
DoReMi trains a small distributionally robust language model to minimize worst-case excess loss across domains, then uses its produced domain weights for larger-model data resampling.
Results
2.6x speedup in training on The Pile is achieved by changing domain sampling probabilities, and DoReMi performs comparably to downstream-task-tuned weights on GLaM.
Takeaways & Limitations
Small models can produce domain weights whose benefits transfer to models 30x larger, improving training efficiency without downstream-task knowledge.
Takeaways & Limitations
The experiments define domains by data provenance, providing only coarse-grained control; finer-grained domains may improve DoReMi's gains.
Abstract
from arXiv · showhide
The mixture proportions of pretraining data domains (e.g., Wikipedia, books, web text) greatly affect language model (LM) performance. In this paper, we propose Domain Reweighting with Minimax Optimization (DoReMi), which first trains a small proxy model using group distributionally robust optimization (Group DRO) over domains to produce domain weights (mixture proportions) without knowledge of downstream tasks. We then resample a dataset with these domain weights and train a larger, full-sized model. In our experiments, we use DoReMi on a 280M-parameter proxy model to set the domain weights for training an 8B-parameter model (30x larger) more efficiently. On The Pile, DoReMi improves perplexity across all domains, even when it downweights a domain. DoReMi improves average few-shot downstream accuracy by 6.5% points over a baseline model trained using The Pile's default domain weights and reaches the baseline accuracy with 2.6x fewer training steps. On the GLaM dataset, DoReMi, which has no knowledge of downstream tasks, even matches the performance of using domain weights tuned on downstream tasks.
1 Introduction
Pretraining-domain proportions strongly affect LM performance, but choosing weights that work broadly across downstream tasks is unclear. DoReMi addresses this by using minimax optimization on a small proxy model to tune domain weights before training a larger model.
- Motivation: The Pile combines many domains, including 24% web data, 9% Wikipedia, and 4% GitHub, while domain composition substantially affects LM effectiveness.Existing domain weights are chosen heuristically or using downstream tasks.
- Approach: DoReMi uses Group DRO over domains to produce domain weights without knowledge of downstream tasks.The method adapts distributionally robust optimization to tune the data mixture rather than produce the robust proxy model itself.
- Efficiency: 280M-parameter proxy models tune weights for 8B-parameter models, with domain-weight optimization requiring 8% of large-model training compute.The large model is trained after resampling data according to the tuned weights.
- Results: 6.5% points: DoReMi improves average one-shot downstream accuracy on The Pile and reaches baseline accuracy 2.6x faster.The comparison uses a baseline trained with The Pile’s default domain weights.
- Results: DoReMi can perform comparably to tuning domain weights on downstream task performance when downstream tasks are available.This result supports the use of task-agnostic domain reweighting in the reported setting.
2 Domain Reweighting with Minimax Optimization (DoReMi)
DoReMi trains a reference model and a Group-DRO proxy to dynamically optimize domain weights, then resamples the data and trains a larger model with those weights. Iterated DoReMi repeats this process until the weights converge.
- Overview: DoReMi takes domain datasets, reference weights, and large-model training hyperparameters, then returns optimized weights for the final training distribution.The weights define a new distribution P̄α used to train the main model.
- Step 1: Reference model: Step 1 trains a small reference model on initial domain weights to establish baseline difficulty for examples and domains.The reference model can use uniform or raw-token-count weights and is 280M parameters in the experiments.
- Step 2: Proxy optimization: Step 2 trains a small proxy model with Group DRO to minimize worst-case excess loss across domains.Excess loss is measured relative to the reference model, and domains with high excess loss receive larger weights.
- Step 3: Main model: Step 3 resamples training data from the optimized distribution and trains the larger main model using standard procedures.The final model is trained on P̄α rather than by directly deploying the robust proxy model.
- Step 2: Proxy optimization: Group DRO alternates proxy-model updates with exponentiated-gradient domain-weight updates based on per-domain excess losses.The implementation clips per-token excess losses at zero, aggregates them by domain, and updates the proxy with a standard optimizer.
- Iterated DoReMi: Iterated DoReMi resets the next round’s reference weights to the previous optimized weights and stops when the maximum change is below 1e-3.On GLaM, convergence empirically takes three rounds.
3 DoReMi Improves LM Training Efficiency and Performance
DoReMi uses small proxy models to optimize domain weights, improving perplexity and downstream accuracy for larger language models on The Pile and GLaM. On The Pile, it improves accuracy over baseline weights and reaches baseline performance faster; on GLaM, iterated DoReMi approaches downstream-tuned weights.
- Experimental setup: DoReMi uses 280M-parameter proxy models to optimize domain weights for training 8B-parameter main models on The Pile and GLaM.The proxy models are 30x smaller than the 8B main models.
- Downstream performance: 6.5% points: DoReMi improves average one-shot downstream accuracy on The Pile over the baseline.The evaluation averages exact-match accuracy across five generative one-shot tasks.
- Training efficiency: 2.6x faster: DoReMi reaches baseline accuracy within 75k steps instead of the baseline’s 200k steps on The Pile.The comparison uses 8B models trained with equalized token-based compute conditions.
- Perplexity: DoReMi significantly improves perplexity over baseline domain weights on every Pile domain, despite assigning lower weight to some domains.The paper provides a hypothesis and simulation example for how reweighting can improve all-domain perplexity.
- GLaM results: On GLaM, iterated DoReMi reaches comparable downstream accuracy to oracle domain weights tuned using evaluation-set downstream tasks.The second and third rounds produce nearly identical domain weights, and the second round achieves the comparable performance.
- Domain weights: DoReMi recovers a domain-weight pattern similar to downstream-tuned weights on GLaM without using downstream data.On The Pile, Wikipedia is downweighted while downstream accuracy on Wikipedia-derived tasks still improves.
4 Ablations and Analysis Across Scales
Analyses show that DoReMi’s gains persist across model scales and proxy sizes, although larger proxy models can weaken optimization. The complete excess-loss objective outperforms variants targeting only easiest or hardest domains, while proxy models may underperform the resampled main models.
- Across scales: 2%: DoReMi improves downstream accuracy over baseline across matched proxy and main model scales from 280M to 1B.The improvement does not shrink with larger model size in the matched-scale experiments.
- Across scales: 4x faster: matched-scale DoReMi reaches baseline accuracy on average across the four evaluated model scales.The 280M→8B setup achieves a smaller 2.6x speedup, reflecting the use of a smaller proxy.
- Perplexity across scales: DoReMi improves 18 of 22 individual domain perplexities on average across matched model scales.It also improves worst-case perplexity at all evaluated scales.
- Proxy versus main model: Over 2x faster: domain weights from the relatively poor 1B proxy still let the 1B main model reach baseline performance.The 1B proxy underperforms the corresponding main model and baseline, unlike the stronger 280M proxy.
- Proxy scale: From 70M to 280M, larger proxy models improve downstream accuracy for an 8B main model, but the trend does not continue at 1B.The authors hypothesize that Group DRO performs worse for larger proxy models.
- Objective ablations: Neither proxy loss alone nor negative reference loss alone matches DoReMi’s downstream accuracy.The full excess loss combines both components and outperforms variants favoring only hardest or easiest domains.
5 Related Work
DoReMi is positioned against manual, heuristic, target-dependent, and example-level approaches to curating language-model pretraining data. It instead uses robust optimization to set domain weights automatically without a target distribution or assumptions about preferred data types.
- DoReMi addresses uncertainty about domain proportions by automatically reweighting pretraining domains without downstream-task knowledge.Its approach contrasts with domain weights tuned using downstream data.
- Unlike example-level filtering methods, DoReMi changes domain weights using only two small language-model training runs.The method does not assume that particular data types, such as Wikipedia-like text, should be preferred.
- DoReMi uses distributionally robust optimization for unknown mixtures of domain groups rather than a specified target distribution.Related DRO-LM work also considers group shifts, but DoReMi applies the framework to domain reweighting.
- DoReMi is presented as a principled data-centric method intended to improve language-model training efficiency.The authors frame it as a starting point for broader robust-optimization approaches to language-model data curation.
6 Discussion and Limitations
The discussion identifies compute cost, reference-model choice, domain granularity, cross-scale transfer, and broader-impact uncertainty as important boundaries for DoReMi. Several proposed directions aim to make weighting more efficient, finer-grained, and better understood.
- Running DoReMi for the final model’s full training duration may be unnecessarily expensive.The authors suggest early stopping and extrapolating weights because most weight variation occurs near the beginning of training.
- The reference model affects the domain weights DoReMi finds, and specialized or differently sized references remain future directions.Iterated DoReMi improves performance by using a reference model trained on weights from an earlier round.
- Data provenance provides only coarse-grained domains; finer-grained domains may improve gains but require controlling DRO pessimism.With very fine-grained domains, DRO may concentrate all weight on a small set of worst-case examples.
- The limits and reasons for transferring domain weights from a 280M proxy model to an 8B model remain unresolved.The paper identifies cross-scale transfer as an important question for future work.
- A 2x reduction in training time could halve cost and energy consumption, but larger resulting models may prevent absolute energy reductions.The broader-impact discussion also notes that large language models carry documented risks and biases.
- DRO promotes good performance across groups and may improve representation disparity between majority and minority subgroups.This potential fairness benefit is presented as a consequence of optimizing worst-case rather than average loss.
7 Conclusion
The paper concludes that DoReMi reweights language-model data domains and can transfer benefits from small models to much larger ones. On The Pile, changing domain sampling probabilities alone yields a 2.6x training speedup.
- DoReMi reweights data domains for training language models.
- 2.6x speedup in training on The Pile results from transferring weights learned on small models to models 30x larger.The speedup comes from changing domain sampling probabilities.
A Results Across Scales on the GLaM dataset
Across 280M to 1B models on GLaM, DoReMi performs comparably to or better than uniform and downstream-tuned domain weights. Iterated weights can transfer better to 8B models despite slightly worse results at 280M.
- Across 280M, 510M, 760M, and 1B scales, DoReMi is comparable or better than uniform and downstream-tuned domain weights.The proxy and main models have the same size in these comparisons.
- At 280M, second-round iterated DoReMi weights yield slightly worse downstream accuracy than round-one weights but transfer better to 8B models.
B Detailed Results for The Pile
On The Pile, DoReMi improves perplexity across domains and downstream performance relative to baseline domain weights, while proxy-model size affects the learned mixture.
- DoReMi improves perplexity on every The Pile domain for 8B models compared with baseline domain weights.
- DoReMi models significantly improve one-shot downstream performance over the baseline across proxy sizes from 70M to 1B.
- 18.25 out of 22 domains show average perplexity improvement across four model scales, while worst-case perplexity is always reduced or comparable.
- DRO ablations do not improve perplexity over baseline, and optimizing only the hardest domain fails to reduce worst-case perplexity.
- Domain weights change rapidly early in training and stabilize after 50k steps, suggesting possible compute savings through extrapolation.
- Different proxy sizes produce different domain weights, with 280M emphasizing Pile-CC and 1B emphasizing OpenWebText2.
C Training Details
The experiments use domain-separated, hierarchically sampled data and standardized optimization settings across Transformer models of several scales.
- Examples are separated by domain to enable hierarchical sampling, with optional packing across domains to reduce padding.
- DoReMi’s domain weights can differ with proxy size, although their qualitative patterns are generally similar.
- The Pile baseline weights are derived from post-chunking example counts multiplied by each domain’s specified number of epochs, then normalized.
- Training uses batch size 512, initial learning rate 1e-3, weight decay 1e-2, gradient clipping, warmup, and exponential decay.
- Models under 1B parameters use TPUv3 accelerators, while 1B and 8B models use a different accelerator setting.
- All models are vanilla decoder-only Transformer models with a 256k vocabulary.
D Simple Example Where Data Reweighting Has No Tradeoff
A simple unigram model illustrates a no-tradeoff setting: reallocating samples away from a noisy domain toward other domains can reduce parameter error for all domains.
- Setup: The example chooses domain weights to learn every domain’s unigram parameters well, rather than estimate the ground-truth mixture proportions.
- Data: Training data are sampled hierarchically by first drawing per-domain counts from the domain weights and then sampling tokens within each domain.
- Model: The Bayesian estimator uses a Dirichlet prior, whose hyperparameters act as token pseudo-counts, and estimates parameters from the posterior mean.
- Model: Domain difficulty H_z and prior quality Δ_z determine how parameter error changes with the number of samples n_z.
- Results: For the non-noisy domain, parameter error decreases with more samples; for the noisy domain, error is minimized at n_3 = 0.
- Results: Reallocating examples from domain 3 to domains 1 and 2 reduces parameter error for all domains.
- Simulation: DoReMi returns weights [0.39, 0.61, 0.0], and the resulting model improves perplexity over uniform baseline weights on every domain.