DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

Summary

This paper proposes to pretrain language models (LMs) by first automatically learning domain weights using a small proxy model and then pretraining a large model under the learned weights. The proposed method, Domain Reweighting with Minimax Optimization (DoReMi), improves pretraining perplexity across all domains and results in better downstream task accuracy with better efficiency.

Strengths

* Originality: The paper studies an interesting direction in seeking the optimal pretraining data mixture/weighting. Both the angle and the proposed method are novel. * Quality: The method is generally well-designed to accomplish the goal discussed and there have been plenty of experiment results that analyze the effects of DoReMi, but I feel that the paper has not sufficiently motivated the necessity and benefits of learning domain weights (see weaknesses below). * Clarity: The paper is overall clear. * Significance: The problem tackled in this paper (i.e., automatic data selection/weighting in pretraining) is important, and the paper shows that the method can improve pretraining perplexity and certain downstream task performance, which can be considered moderately significant. However, there are concerns regarding whether the evaluations are comprehensive w.r.t. tasks and model scales (see weaknesses below).

Weaknesses

* Insufficient motivation: While it's well acknowledged that the data quality is variable across different domains, it's unclear to me whether learning domain weights and using them to construct a "better" corpus is the appropriate way. First, domains are quite coarse partitions of the data, and I'm not very convinced that assigning a single scalar weight to each domain as a whole is an ideal setup. As the authors mentioned, some data might be noisy that should be down-weighted, but shouldn't this be done at an instance level? For example, some code snippets from Github might be erroneous but others are correct. If the entire Github domain is down-weighted, it doesn't really tell apart the clean vs. noisy data, but instead puts a lower priority on learning all code-related data. Second, the authors do not seem to mention or compare with a naive baseline that directly removes noisy domains from the pretraining corpus. For example, what if only high-quality data are used? Gunasekar et al. (I'm aware that this paper came out after the paper submission deadline, but it seems relevant) showed that textbook-quality data with only 6B tokens are sufficient to train good models. * Unclear generalization ability to large model scales: Although the authors conduct experiments across different model scales, the largest model size tested is 8B. I understand that it's very expensive to train even larger models, but for a paper studying pretraining, I'd expect to see results on larger scales (e.g., 65B), considering that the model performance has a strong correlation with model sizes. This concern appears imminent given that the proxy model does not seem to scale to larger models. * Choice of evaluation tasks: The paper mainly evaluates on pretraining perplexity and QA-related tasks. For a paper that aims at "models that perform well on all domains", I believe the evaluation should have a more comprehensive coverage of tasks, such as coding and reasoning, especially considering that these tasks are relevant to certain domains in the pretraining corpora (e.g., Github). I'd be curious to know if the performance on code completion will be still higher than the baseline if the Github domain is down-weighted. * (Minor) As I understand, the appendix should be submitted separately from the main paper. Reference: Gunasekar et al. “Textbooks Are All You Need.” 2023.

Questions

* How to set the size of the proxy model given the main model size? * How do the learned domain weights by different sizes of proxy models look like? Please also clarify any misunderstandings in my review.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Please refer to the Weaknesses section.

Summary

This paper introduces a significant advancement by exploring the topic of data mixture proportions during pre-training, which holds great importance. Determining how to sample pre-trained data from diverse sources to achieve balanced results is a fundamental question. Previous approaches have often relied on intuitive-based weights or extensive experimentation to select a setting. However, these approaches either require extensive computational resources or lack generalization across different settings. Therefore, finding the optimal mixture settings using a smaller model is an intriguing and crucial question. To address this, the authors propose a group-DRO-based method for determining the weights. The weight is dynamically adjusted during the learning process, and the final weight is selected as the sampling weight. Additionally, the authors conducted extensive experiments, providing substantial evidence to support the effectiveness of the proposed approach.

Strengths

The direction is really important and the general framework is valuable that uses a small network to get the best practice and apply it on larger models. The proposed method is clearly-written, and the improvements it offers have been convincingly substantiated through extensive experiments. The general idea is also very elegant. If my understanding is correct, the method generally assigns larger weights to the data that can be learned in the future, and assigns smaller weights to those too easy or difficult data. The general idea is cool.

Weaknesses

The selected weights vary significantly across different models with varying scales. I still have some concerns about the final implementation. Following the optimization objectives, the authors use the learned weight as the re-sampling weight. It is a little bit strange.

Questions

For instance, the selection of the best weight differs among models with different scales. Does this imply the existence of multiple sub-optimal weight candidates? Moreover, I still do not comprehend why the authors do not apply the final model during the learning process but instead utilize the final weights for re-weighting purposes.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

n/a

Summary

The authors proposed DoReMi for optimizing the mixture proportions of pretraining data domains when training language models (LMs). The authors demonstrate that DoReMi, which utilizes a small proxy model trained via group distributionally robust optimization (Group DRO), can be used to determine optimal domain weights without knowledge of downstream tasks. Subsequently, these weights are used to resample a dataset for training a larger model. Experimental results indicate significant improvements in LM performance using DoReMi, including a 6.5% increase in average few-shot downstream accuracy and achieving baseline accuracy with 2.6x fewer training steps.

Strengths

- DoReMi offers a unique and efficient way to determine domain weights, speeding up LM training and improving accuracy. Moreover, the down-weighted domains are not getting a worse performance, which is surprising. - It also shows that the domain weights determined by DoReMi are transferable across a broad range of model scales, compute budgets, and other training hyperparameters, making it wide applicability. - Overall, the idea of DoReMi is simple, and the results can show its effectiveness.

Weaknesses

- Lack of the baselines: In this paper, the baselines are the LMs trained on the original data distribution of Pile, but there should be some simple but stronger baselines as well, like calculating the lexical overlap within each domain and assigning weights to maximize the lexical diversity. Although this simple baseline sounds naive, we still need to justify that their improvement would not be as much as DoReMi. - DoReMi needs two or more proxy LM training processes in order to obtain the domain weights. However, the authors need to justify that there doesn't exist simpler but equally effective heuristics that only need one proxy LM training, such as the number of example forgetting times [1]. - Based on Figure 6, it seems that the choice of proxy model size is very crucial. For 8B model, the 280M proxy is significantly better than other larger or smaller proxy models. It would be better if the author can provide explanation or principles to select the size of the proxy model, otherwise, people who want to use DoReMi actually need to empirically run different proxy models in order to get better performance. [1] An Empirical Study of Example Forgetting during Deep Neural Network Learning. Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, Geoffrey J. Gordon. ICLR 2019 https://arxiv.org/abs/1812.05159

Questions

- For the domain weight $\alpha$, I am wondering why we need gradient ascent to update it? Isn't it already optimal if we simply assign 1.0 for the domain with the highest excess loss and 0.0 for all the other domains? If the reason is that we need a single set of consistent domain weights till the end of the training, we could also collect the weights along the training process and average them at the end. Did you try this before? - I am curious whether there are simple heuristics that can achieve the same level of effects. For example, in [1], people found that the number of forgetting times of a training example is a good indicator of how important each training example is. We need this kind of baseline to prove that the 2 runs in DoReMi are necessary and its result is more effective than others. [1] An Empirical Study of Example Forgetting during Deep Neural Network Learning. Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, Geoffrey J. Gordon. ICLR 2019 https://arxiv.org/abs/1812.05159

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes.

Summary

This paper introduces DoReMi, a method for automatically deriving optimal/improved domain weights for aggregated pretraining datasets for LLMs. DoReMi works in three 3: (1) train a small reference model to use for the excess loss in DRO; (2) train a small proxy model with DRO to obtain optimised domain weights; (3) use the optimised weights with a larger model. The authors demonstrate their method leads to significant improvements on The Pile, and also matches manual tuning on GLaM. In ablations, the scaling behaviour is studied, showcasing that proxy model size improves downstream performance up to ~280M parameters, after which improvements degrade as the larger models are inadequately trained by DRO.

Strengths

* **S1.** Automatic optimisation of domain weights for pretraining datasets is an incredibly valuable contribution for the LLM community. Standard procedures have been mostly based on expensive manual tuning, and have not necessarily been principled in their finding. * **S2.** The paper is well grounded in works around DRO, and adequately position itself -- diverging from existing methods when necessary. The authors also take care in pointing out current limitations of the DRO approach, and potential for future improvements. * **S3.** The results obtained on The Pile reproduce the observations recently made by the RedPjamas & RefinedWeb datasets: some components of The Pile should ideally be downsampled, and increased web data may be beneficial. The fact that DoReMi repeatedly reproduces results obtained from manual tuning is a good validation of the method. * **S4.** The paper is well-written and presented, and easy to follow.

Weaknesses

* **W1. The evaluation setup could be broader.** The authors evaluate downstream performance: (1) in 1-shot; (2) in a generative/exact-match setting; (3) on 5 tasks. * **W1.1.** The choice of a generative exact-match setting is strange for models in the ~100M-1B range, as they consistently struggle with exact match. Instead, for small models, leveraging logprob-based evaluation of multiple choices is more common, and can deliver stronger signal. * **W1.2.** The choice of task is arbitrary. Rather than these 5 tasks, the authors could have evaluated on the full set of GPT-3 tasks, or on a popular benchmark such as HELM or BigBench (for the larger models). * **W2. The poor scaling behaviour of DoReMi past 280M parameters for the proxy model is a concern for robustness.** Notably, this might make it difficult for practitioners to apply DoReMi as a "set it and forget it" method. Some level of manual inspection & analysis is required, which holds back the method from fully delivering on its promise of completely automated domain weight optimization. However, I appreciate that the authors discuss this limitation openly and propose some potential ideas for further improvements in this direction. * **W3.** (smaller nits) * **W3.1.** The authors showcase in Table 1 the baseline & DoReMi domain weights on The Pile; the presentation of the table could be improved, to better identify which domains are upsampled and which are downsampled -- this could be as simple as sorting the domains and explicitly providing the up/downsampling value in the table.

Questions

This is an excellent paper introducing a method which could see wide adoption in the community as a way to improve aggregated pretraining datasets. The paper is well-written, and opens the door to numerous exciting follow-up works. Accordingly, I would currently rate it as a **Strong Accept (8)**. Note that should my concerns about evaluation be addressed, I would be willing to further increase my score to a 9/10 -- this paper has significant potential for the community, and my only main concern currently is regarding the robustness of the selected evaluation setup. * **Q1.** (W1.1.) Could the authors better justify their choice of going with generative exact match for the evaluation? * **Q2.** (W1.2.) Could the authors provide scores (in 0-shot or 1-shot) using logprob-based multiple choices instead? * **Q3.** (W1.2.) Could the authors provide evaluation results on additional tasks, such as the full set of GPT-3 tasks, or scores on a popular benchmark such as HELM or BigBench? (this applies mostly to the headline 8B models).

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors extensively discuss limitations of their work in a dedicated section and provide interesting pointers for further research.

Summary

This work buils upon prior studies’ empirical findings, emphasizing how the composition of pretraining data affects the performance of Language Models (LMs). To avoid reliance on heuristic or iterative performance measurements on downstream tasks, this work introduces a method that employs a trainable model, capable of assigning appropriate weights to each “domain” of pretrainng data. The proposed approach leverages the concept of Distributionally Robust Optimization (DRO) to train a small proxy model that learns to assign weights to each domain, thereby minimizing the worst-case excess losses. Subsequently, a larger model is trained using pretraining data that has been adjusted by these optimized domain weights. The authors demonstrate through experiments that their proposed method significantly accelerates the pretraining process by showing the few-shot accuracy on factual QA tasks. They show that the model trained with optimized pretraining data reaches the performance level of baseline models much faster.

Strengths

- The idea to optimize the pre-training data sounds clever, particularly their use of a proxy model trained with DRO. This offers a novel approach to learning the optimal data distribution without the costly estimation of the large pre-trained model’s task performance under different pretraining data settings. - This work provides solid experimental evidence, exploring a variety of settings including various proxy and main model sizes, two widely-used pretraining datasets (GLaM dataset and The Pile), and different domain compositions. These experiments adequately address questions that arise during a review of the paper to some extent. - This research yields some interesting observations. According to Tables 1 and 2, the webpage domain carries the largest weights across both pretraining datasets. This suggests that optimizing the model with respect to the webpage domain could potentially reduce excess loss in other domains. Consequently, pretraining large models with an optimized dataset results in lower perplexities across all domains. Interestingly, despite a decrease in weight for Wikipedia, task performance on tasks derived from it (TriviaQa, NaturalQuestions) improves (Table 5 in Appendix). These insights could prove beneficial for other researchers in the field.

Weaknesses

- The diversity of downstream tasks examined is relatively narrow, given the paper’s claim. While the experiments in this paper present solid empirical evidence that the proposed method improves language model performance, the evidence primarily focuses on language modeling and factual question answering tasks in a general domain. As such, the claim that the proposed method “speeds up language model pretraining” might mislead readers, since the evaluation of pretraining should be more rigorous. For example, incorporating experimental results on MMLU [1] or commonsense reasoning (ARC[2], CSQA[3], etc) could increase the credibility of the claim that the pre-trained model with the proposed method generally outperforms previous approaches. [1] Measuring Massive Multitask Language Understanding [2] Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge [3] Complex Sequential Question Answering: Towards Learning to Converse Over Linked Question Answer Pairs with a Knowledge Graph

Questions

1) Why did the authors only experiment with a 280M main model in Figure 6 (right)? I understand the potential cost issues, but I believe it’s essential to show that DoReMi outperforms the webpage-only setting, even with a larger model. 2) Could you clarify how the main model employs domain weight in detail (Step 3 in Section 2)? I was unable to find detailed information on this in the paper. Does the model sample instances in the mini-batch according to the domain weight during pre-training? Note: The appendix should be submitted as the separated supplementary file.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

4 excellent

Contribution

3 good

Limitations

The authors discuss limitations in Section 6. However, the limitations regarding the evaluation methods for pretraining are not thoroughly addressed in the paper. Please refer to the “Weaknesses” section for more details on my perspective.

Answer to rebuttal

First, I would like to thank the authors for taking time to write a rebuttal to each reviewer. Based on the rebuttal and the other reviews, I will maintain my score of a **Strong Accept (8)**.

Thanks for the responses

The responses well addressed my questions. I would keep my score of 7 Accept.