DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

By Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, Adams Wei Yu. 2023

393 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Large Language Models

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Large Language Models

Citations in the corpus

Llemma: An Open Language Model For Mathematics (2023, methodology): “When training a language model, it is common to upsample high-quality subsets of the training data according to mixture weights (Brown et al., 2020; Gao et al., 2020; Xie et al., 2023).”
Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (2023, methodology): “Inspired by recent work (Xie et al., 2023), we propose dynamic batch loading , a more efficient algorithm to simply adjust domain proportions on the fly based on the model performance.”
Qwen3 Technical Report (2025, neutral): “Unlike previous studies (Xie et al., 2023; Fan et al., 2023; Liu et al., 2024b) that optimize the data mixture at the data source or domain level, our method optimizes the data mixture at the instance-level through extensive ablation experiments on small proxy models with the fine-grained data…”
MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies (2024, background): “From the training data standpoint, various data-centric strategies have been advocated (Xie et al., 2024; Shi et al., 2023; Ye et al., 2024).”
SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore (2023, background): “More sophisticated domain weighting strategies (Xie et al., 2023) are of interest but beyond the scope of this work.”