SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 6: Average downstream accuracy for models trained on The Pile. (Left) Increasing the size of the reference/proxy models from 70M to 280M in DoReMi improves downstream accuracy for a 8B main model, but the trend does not continue for the 1B proxy model. We hypothesize that the Group DRO optimizer is worse for larger proxy models. (Right) Optimizing for the hardest or easiest domains rather than excess loss (which combines both) do not achieve the same average downstream accuracy as DoReMi (280M models).
Table 3: Summary of per-domain log-perplexities on The Pile (22 total domains). Average log-perplexity is an unweighted average of the per-domain log-perplexities.
(a) Varying the size of the proxy/reference model and training at 8B.
Worst-case log-ppl	Avg log-ppl	#domains beating baseline
Baseline (8B)	1.71	1.64	0/22
DoReMi (70M->8B)	1.63	1.53	22/22
DoReMi (150M->8B)	1.56	1.52	22/22
DoReMi (280M->8B)	1.46	1.40	22/22
DoReMi (1B->8B)	1.58	1.54	22/22
(b) Perplexity of the DoReMi main model and proxy model of the same size. Although the 1B proxy model is relatively poor quality, the resulting domain weights still improve the main model.
Worst-case log-ppl	Avg log-ppl	#domains beating baseline
Baseline (280M)	2.39	2.32	0/22
DoReMi (280M->280M)	2.19	2.13	22/22
Proxy (280M)	2.33	2.27	19/22
Baseline (1B)	1.94	1.87	0/22
DoReMi (1B->1B)	1.92	1.83	19/22
Proxy (1B)	2.11	2.02	0/22
model typically underperforms the main model in this case. The gap between the proxy and main model increases with scale, as the 1B proxy model not only underperforms the 1B main model but also the 1B baseline model, while the 280M proxy model achieves better perplexity than the 280M baseline model on 19/22 domains. Despite the relatively poor quality of the 1B proxy model, the domain weights still allow the 1B main model to achieve the baseline performance over  $2\mathrm{x}$  faster. This suggests that DoReMi can succeed even if the proxy model is not trained well. However, we hypothesize that the mismatch between the proxy and main model training (loss reweighting vs. resampling) explains their performance difference and therefore a resampling-based Group DRO optimizer may improve DoReMi for larger proxy models.
Effect of proxy model scale on larger main model's performance. We consider 70M, 150M, 280M, and 1B scales for the DoReMi proxy model while fixing the main model size at 8B (DoReMi  $(\mathrm{X}\rightarrow 8\mathrm{B})$ ). From 70M to 280M, increasing the proxy model size improves downstream accuracy at 8B (Figure 6 left). We hypothesize that this trend does not continue for the 1B proxy model because the Group DRO optimizer is worse at larger scales (Table 3b). While DoReMi  $(280\mathrm{M}\rightarrow 8\mathrm{B})$  results in the most improvement at 8B, DoReMi  $(150\mathrm{M}\rightarrow 8\mathrm{B})$  and DoReMi  $(1\mathrm{B}\rightarrow 8\mathrm{B})$  still achieve the baseline accuracy almost  $2x$  faster. This suggests that DoReMi is robust to the proxy model scale. In practice, we suggest choosing a relatively small proxy model size (280M) to save compute.
Choosing the easiest or hardest domains do not suffice. We ablate the components of the excess loss metric  $\ell_{\theta}(x) - \ell_{\mathrm{ref}}(x)$  by running DoReMi using only the loss of the proxy model  $p_{\theta}$  on example  $x$ , i.e.  $\ell_{\theta}(x)$  (prefer hardest domains for the proxy model) or only the negative loss of the reference  $-\ell_{\mathrm{ref}}(x)$  (prefer easiest domains for the reference model). Figure 6 (right) shows that neither of the components of the excess loss alone are sufficient to achieve the gains of DoReMi.
**Curating pretraining data for LMs.** Most closely related is the GLaM dataset [13] (also used for training PaLM [10]), which has domain weights that are tuned using downstream data. Optimizing domain weights for downstream tasks can be expensive and could require search/zero-order optimization [48], RL [56], or heuristic assumptions on how positive/negative transfer between domains work. Example-level filtering also brings benefits for LM training. The C4 dataset [39] shows gains over CommonCrawl via heuristic data cleaning methods. Du et al. [13]; Xie et al. [55] show that filtering the data at an example level for high-quality text that look like Wikipedia and books can significantly improve downstream performance for LMs. In contrast to these works, DoReMi sets domain weights automatically with only two small LM training runs and does not make assumptions about the type of data to prefer (Wikipedia-like, etc.).
General data selection methods. Moore-Lewis selection [32; 3; 15] selects examples with high cross-entropy difference (similar to excess log-perplexity) between language models trained on target and raw data. In contrast, DoReMi reweights the data without a target distribution. Coleman et al. [11] select examples based on the uncertainty of a small proxy model for active learning, while DoReMi uses DRO on the excess loss with respect to a reference model, and focuses on data mixture reweighting. Mindermann et al. [30] select examples in an online fashion by taking the top  $k$  examples in a minibatch according to excess loss. DoReMi optimizes the data mixture before training, allowing the larger main model to train in a standard way. Many other works on data selection are in vision [49; 22; 24; 23; 25; 53; 54; 38; 31; 45] and mainly focus on example-level subset selection with metrics such as gradient matching. Overall, these methods do not address data selection for pretraining, where the downstream data distribution may be very different from the pretraining distribution. DoReMi aims to address the pretraining/downstream distribution shift with a robust optimization approach. To the best of our knowledge, we are the first to show that reweighting the data according to losses of a small proxy LM can improve the training efficiency of much larger LM.
Distributionally robust optimization. Within DRO methods for deep learning [4; 47; 35; 43], we target a restricted form of shift called group shifts [14; 35; 43], where the test distribution can be an unknown mixture of groups (domains). We follow DRO-LM [35], which employs DRO for LMs in the group shift setting. DRO-LM also uses a baselined loss, but with a simple bigram reference model. DoReMi uses a reference model of the same size and architecture as the proxy model to ensure that the losses are on a similar scale. During optimization, DRO-LM takes a worst-case subset of each minibatch to update the model on, while we use the Group DRO optimizer [43] which doesn't require online subselection. If we equalize the number of examples in each minibatch used for gradient updates, online subselection is more expensive than Group DRO since it requires running forward passes on a larger minibatch (e.g., double the minibatch size) before selecting a subset to update the model with. In comparison, the Group DRO optimizer updates the model on all examples in a weighted fashion. Overall, in contrast to these DRO methods which aim to produce robust models, we use DRO to optimize the data for training larger models more efficiently.
Data-centric AI. Large-scale datasets and benchmarks have driven much of the recent progress in AI, including vision, NLP, and multimodal models [12; 42; 52; 40; 39; 17; 44; 16]. However, most datasets are still painstakingly created with human-generated data, manual work, and heuristics [12; 39; 17; 44; 16]. DoReMi is a principled data-centric method that aims to improve language model training efficiency. We hope that DoReMi can provide a starting point for a general data-centric framework for language modeling via robust optimization.
6 Discussion and Limitations
Saving compute in DoReMi with extrapolation. In Section 2, we run DoReMi for the number of training steps that will be used to train the final model, which could be unnecessarily expensive. A future direction for saving compute would be to stop running DoReMi at an early step and extrapolate the domain weights for the desired number of steps, since we found that most of the variation in the domain weights during a DoReMi run seems to occur in the beginning of training (Appendix Figure 8).
Choice of reference model. The choice of reference model can affect the domain weights found by DoReMi. For example, iterated DoReMi (Section 3) improves performance by using a reference model trained on the tuned domain weights from a previous round of DoReMi. Further directions include varying the reference model size and using specialized reference models to optimize domain weights for a specific application area.
What is a domain? We define a domain by data provenance in our experiments, but this only enables coarse-grained control. Using fine-grained domains could improve the gains from DoReMi. For example, DoReMi is more effective on The Pile (22 domains) than the GLaM dataset (8 domains). Open directions include automatically finding fine-grained domains (e.g., via clustering as in DRO-LM [35]) and reweighting the data at an example level. When domains are very fine-grained, it will be important to control the pessimism of DRO (e.g., DRO can put all the weight on a small set of worst-case examples).
Transferability of domain weights across scales. We optimized the domain weights with a small proxy model (280M) and directly used these domain weights to improve training at a larger scale (8B). Understanding why the domain weights can be transferred across scales and the limits of how far these domain weights transfer are important questions to answer in future work.
Broader impacts. Large language models are We hope to improve training efficiency and reduce the environmental impact of training large LMs [50; 28; 37; 29]. In particular, by reducing the training time by  $2\mathrm{x}$ , we can halve the cost and energy consumption of training large language models. Since such efficiency improvements may be used to develop even larger models, there may be no absolute improvement in energy consumption. Ultimately, we hope to improve the training efficiency and cost of developing future language models relative to existing methods.
Large LMs have also been well-documented to have risks and biases [1; 33; 7; 6; 18]. For example, GPT-3 tends to have an anti-Muslim bias, where Muslims are frequently related to violence or terrorism in analogy and completion tasks [1]. As large language models are increasingly relied upon in applications, the magnitude of the risks increases [8]. Distributionally robust optimization (DRO), which is used in DoReMi to optimize the data mixture, can have a favorable impact on fairness [19]. While the standard approach of minimizing the average loss can lead to disparate performance on minority subgroups that do not contribute heavily to the loss [2], DRO promotes good performance on all groups via a worst-case loss. In this way, DRO-style data-centric methods such as DoReMi can improve the representation disparity between majority and minority subgroups in a dataset.
We introduced DoReMi, an algorithm reweighting data domains for training language models. DoReMi is able to run on small models and transfer the benefits to 30x larger models, resulting in a 2.6x speedup in training on the Pile just by changing the sampling probabilities on domains. We hope to instigate more research on data-centric approaches for improving language model training efficiency.
We thank Xiangning Chen, Andrew Dai, Zoubin Ghahramani, Balaji Lakshminarayanan, Paul Michel, Yonghui Wu, Steven Zheng, Chen Zhu, anonymous reviewers, and the broader Google Bard team members for insightful discussions and pointers.
[1] Abubakar Abid, Maheen Farooqi, and James Zou. Persistent anti-muslim bias in large language models. arXiv preprint arXiv:2101.05783, 2021.
[2] Dario Amodei et al. Deep speech 2 end to end speech recognition in English and mandarin. In International Conference on Machine Learning (ICML), pages 173-182, 2016.
[3] Amittai Axelrod. Cynical selection of language model training data. CoRR, abs/1709.02279, 2017. URL http://arxiv.org/abs/1709.02279.
[43] Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations (ICLR), 2020.
[44] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[45] Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. In International Conference on Learning Representations (ICLR), 2018.
[46] Noam Shazeer and Mitchell Stern. 2018.
[47] Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations (ICLR), 2018.
[48] Jasper Snoek, Hugo Larochelle, and Ryan P. Adams. Practical Bayesian optimization of machine learning algorithms. In Advances in Neural Information Processing Systems (NeurIPS), 2012.
[49] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari S. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. arXiv, 2022.
[50] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in NLP. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3645–3650, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1355. URL https://aclanthology.org/P19-1355.
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv:1706.03762, 2017.
[52] Alex Wang, Amapreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations (ICLR), 2019.
[53] Xinyi Wang, Hieu Pham, Paul Michel, Antonios Anastasopoulos, Jaime Carbonell, and Graham Neubig. Optimizing data usage via differentiable rewards. In International Conference on Machine Learning (ICML), 2020.
[54] Kai Wei, Rishabh Iyer, and Jeff Bilmes. Submodularity in data subset selection and active learning. In International Conference on Machine Learning (ICML), 2015.
[55] Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03169, 2023.
[56] Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv:1611.01578, 2016.
(a) 280M
(b) 510M
(c) 760M
(d) 1B
Figure 7: Average one-shot downstream accuracy across 4 model scales, where the reference/proxy models for DoReMi are the same size as the final model trained with DoReMi domain weights. All models in this figure are trained on the GLaM dataset. DoReMi consistently improves downstream accuracy across scales.
A Results Across Scales on the GLaM dataset
Figure 7 presents results across different scales (280M, 510M, 760M, 1B) on the GLaM dataset, where the proxy/reference models are the same size as the main model trained with DoReMi domain weights. Across all scales, DoReMi is comparable or better than both the baseline (uniform) domain weights and downstream-tuned domain weights. Interestingly, for iterated DoReMi at the 280M scale, the second round weights achieve slightly worse downstream accuracy than the round 1 weights when used to train 280M models, but transfer better to training 8B models.
B Detailed Results for The Pile
Per-domain perplexities for 8B models. Table 4 shows per-domain perplexities for 8B models trained on the Pile. The reference/proxy models in this case are 70M, 150M, 280M, and 1B. DoReMi improves the perplexity on each domain compared to the baseline domain weights.
Per-task accuracies for 8B models. Table 5 shows the accuracies on one-shot generative tasks for various reference/proxy model sizes from 70M to 1B. All DoReMi models improve downstream performance significantly over the baseline.
Summary of perplexity results across scales. Table 6 shows a summary of per-domain perplexities for DoReMi across 4 scales (280M, 510M, 760M, 1B). Here, the reference/proxy models are the same size as the main model trained with DoReMi domain weights. On average, DoReMi improves perplexity on 18.25 out of 22 domains from The Pile. The worst-case perplexity is always reduced (or comparable in the 510M case) with respect to the baseline domain weights.
Table 4: Per-domain log-perplexities for 8B models trained on The Pile where the reference/proxy models are or smaller sizes (70M, 150M, 280M, 1B). Models trained with DoReMi domain weights have lower perplexity on all domains than the baseline weights.
Baseline (8B)	DoReMi (70M->8B)	DoReMi (150M->8B)	DoReMi (280M->8B)	DoReMi (1B->8B)
Pile-CC	1.64	1.51	1.48	1.41	1.55
PubMed Central	1.60	1.58	1.54	1.46	1.56
Books3	1.65	1.52	1.50	1.42	1.57
OpenWebText2	1.66	1.48	1.54	1.36	1.58
ArXiv	1.64	1.56	1.53	1.38	1.51
Github	1.65	1.55	1.54	1.42	1.53
FreeLaw	1.64	1.55	1.54	1.45	1.55
StackExchange	1.61	1.52	1.54	1.39	1.55
USPTO Backgrounds	1.70	1.53	1.50	1.41	1.56
PubMed Abstracts	1.61	1.56	1.51	1.44	1.55
Gutenberg (PG-19)	1.70	1.56	1.54	1.35	1.52
OpenSubtitles	1.58	1.56	1.52	1.40	1.55
Wikipedia (en)	1.66	1.49	1.53	1.35	1.56
DM Mathematics	1.63	1.50	1.56	1.38	1.48
Ubuntu IRC	1.71	1.53	1.49	1.42	1.48
BookCorpus2	1.64	1.57	1.54	1.43	1.57
EuroParl	1.59	1.52	1.51	1.37	1.53
HackerNews	1.66	1.50	1.55	1.45	1.55
YoutubeSubtitles	1.67	1.63	1.55	1.42	1.53
PhilPapers	1.67	1.55	1.49	1.39	1.53
NIH ExPorter	1.63	1.51	1.48	1.36	1.52
Enron Emails	1.62	1.48	1.52	1.44	1.56
Table 5: Per-task exact-match accuracies for generative one-shot tasks. All DoReMi models improve downstream performance significantly over the baseline domain weights.
Baseline	DoReMi (1B->8B)	DoReMi (280M->8B)	DoReMi (150M->8B)	DoReMi (70M->8B)
LAMBADA	20.10	22.55	29.19	20.59	26.20
NaturalQuestions	4.35	6.01	7.73	6.26	5.10
SQuADv2	44.43	42.22	51.89	46.53	40.99
TriviaQA	24.55	32.25	34.86	30.01	26.30
WebQuestions	6.74	8.71	9.15	9.15	6.99
Average	20.03	22.35	26.56	22.51	21.11
Table 6: Summary of per-domain log-perplexities for 280M, 510M, 760M, and 1B models trained on The Pile, where the reference/proxy models are the same size. DoReMi improves the worst-case and average perplexity of the baseline domain weights in all cases. On average, DoReMi improves perplexity on 18 out of 22 domains.
Worst-case log-ppl	Avg log-ppl	# domains beating baseline
Baseline (280M)	2.39	2.32	0/22
DoReMi (280M->280M)	2.19	2.13	22/22
Proxy (280M)	2.33	2.27	19/22
Baseline (510M)	2.14	2.08	0/22
DoReMi (510M->510M)	2.14	2.06	15/22
Proxy (510M)	2.23	2.18	0/22
Baseline (760M)	2.05	1.97	0/22
DoReMi (760M->760M)	2.00	1.94	17/22
Proxy (760M)	2.15	2.10	0/22
Baseline (1B)	1.94	1.87	0/22
DoReMi (1B->1B)	1.92	1.83	19/22
Proxy (1B)	2.11	2.02	0/22
Table 7: Summary of perplexity results for ablations on the DRO objective (excess loss). The individual components (which prefer hardest and easiest domains respectively) do not reduce perplexity over the baseline.
Worst-case log-ppl	Avg log-ppl	# domains beating baseline
Baseline (280M)	2.39	2.32	0
DoReMi (280M->280M)	2.19	2.13	22/22
Hardest (280M->280M)	2.66	2.62	0/22
Easiest (280M->280M)	4.27	4.18	0/22
——Pile-CC	……	OpenSubtitles
——PubMed Central	——	Wikipedia (en)
——Books3	——	DM Mathematics
——OpenWebText2	——	Ubuntu IRC
——ArXiv	——	BookCorpus2
——Github	——	EuroParl
——FreeLaw	——	HackerNews
——StackExchange	——	YoutubeSubtitles
——USPTO Backgrounds	——	PhilPapers
——PubMed Abstracts	——	NIH ExPorter
——Gutenberg (PG-19)	——	Enron Emails
(b) 1B
Figure 8: Exponential moving average of domain weights throughout a DoReMi run for 280M and 1B reference/proxy models. In the beginning of the run, the domain weights change quickly and then become more stable after 50k steps. This suggests that 1) smaller compute budgets may require drastically different domain weights, and 2) we may be able to save compute by extrapolating the domain weights after 50k steps.
Perplexity results for ablations. Table 7 shows the perplexities for ablations on the DRO objective. We change the DRO objective and use these to tune domain weights on 280M reference/proxy models. These tuned domain weights are then used to train a main 280M model. Hardest refers to optimizing the domain-level log-perplexity without baselining with a reference model. Easiest refers to optimizing for the domains with lowest log-perplexity under the reference model. Both ablations do not improve perplexity on any domain over the baseline. Optimizing for the "hardest" domain does not actually result in improving worst-case perplexity, supporting the results of Oren et al. [35], which also employs DRO for language modeling with a baselined loss.
Trajectory of domain weights. Figure 8 shows the exponential moving average (smoothing parameter 0.99) of domain weights during a run of DoReMi. In both cases, there are domains with very high weight initially and decrease in weight very quickly (within 50k steps). Since we compute the final domain weights by integrating these curves over steps and normalizing, this suggests that if we have a smaller compute budget, these domains could become more important — this highlights the dependence of the mixture weights on the compute budget. At the same time, the domain weights tend to quickly stabilize after 50k steps, suggesting that the optimal domain weights should be similar for larger compute budgets. We may also be able to take advantage of this stability after 50k steps to run DoReMi for a smaller number of steps and extrapolate the domain weights to save compute.
Comparison of domain weights for 280M and 1B. Table 8 presents the DoReMi domain weights for The Pile at 280M and 1B proxy models. Different proxy model sizes can result in different domain weights, which suggests that there may be multiple local minima in domain weight space. With a 280M proxy model, most of the weight is put on the Pile-CC web text domain, while DoReMi with a
Table 8: Domain weights on The Pile. Baseline domain weights are computed from the default Pile dataset. With different proxy model sizes, DoReMi (280M) and DoReMi (1B) result in different domain weights. Despite the differences, the qualitative patterns are similar other than the which web domain has the most weight.
Baseline	DoReMi (280M)	DoReMi (1B)
Pile-CC	0.1121	0.6057	0.1199
PubMed Central	0.1071	0.0046	0.0149
Books3	0.0676	0.0224	0.0739
OpenWebText2	0.1247	0.1019	0.3289
ArXiv	0.1052	0.0036	0.0384
Github	0.0427	0.0179	0.0129
FreeLaw	0.0386	0.0043	0.0148
StackExchange	0.0929	0.0153	0.0452
USPTO Backgrounds	0.0420	0.0036	0.0260
PubMed Abstracts	0.0845	0.0113	0.1461
Gutenberg (PG-19)	0.0199	0.0072	0.0250
OpenSubtitles	0.0124	0.0047	0.0017
Wikipedia (en)	0.0919	0.0699	0.0962
DM Mathematics	0.0198	0.0018	0.0004
Ubuntu IRC	0.0074	0.0093	0.0044
BookCorpus2	0.0044	0.0061	0.0029
EuroParl	0.0043	0.0062	0.0078
HackerNews	0.0075	0.0134	0.0058
YoutubeSubtitles	0.0042	0.0502	0.0159
PhilPapers	0.0027	0.0274	0.0063
NIH ExPorter	0.0052	0.0063	0.0094
Enron Emails	0.0030	0.0070	0.0033
1B proxy model puts most of the weight on OpenWebText2. The overall pattern of the domain weights for the rest of the domains are similar.
Data preprocessing. For all datasets, we preprocessed the data by chunking into length 1024 examples with respect to a SentencePiece tokenizer with 256k vocabulary size. The examples are separated by domain to facilitate hierarchical sampling (first sample a domain according to some domain weights, then sample an example from that domain at random). To reduce the amount of padding tokens, we made an effort to pack examples (possibly from different domains) together into the same sequence. When doing such a packing, we compute the domain perplexities on a per-token level in DoReMi.