DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining
Authors: Sang Michael Xie
Hieu Pham
Xuanyi Dong
Nan Du
Hanxiao Liu
Yifeng Lu
Percy Liang
Quoc V Le
Tengyu Ma
Adams Wei Yu
Venue: NeurIPS 2023 spotlight (2023)
Keywords: language models
pretraining
domain reweighting
data curation
TL;DR: We present an algorithm that reweights how much of each domain/data source is in a language modeling dataset (e.g. The Pile), resulting in faster LM training and improvements in perplexity and downstream accuracy.

Abstract: 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.

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining
Sang Michael Xie! Hieu Pham', Xuanyi Dong', Nan Du', Hanxiao Liu', Yifeng Lu',
Percy Liang", Quoe V. Le', Tengyu Ma¹, and Adams Wei Yu!
Google DeepMind 2 Stanford University
Abstract
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 downstream tasks. We then resample a dataset with these domain weights and train larger, full-sized model. In experiments, we DoReMi parameter proxy model to set the domain weights for training an 8B- 18-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 baseline model trained using The Pile's default domain weights and reaches the bascline accuracy with 2.6x fewer training steps. On the GLaM dataset, DoReMi, which has no knowledge downstream tasks, even matches the performance of using domain weights tuned on downstream tasks.
280M-
Introduction
Datasets for training language models (LMs) are typically sampled from a mixture of many domains [17; 13; 10; 9]. For example, The Pile [17], large publicly available dataset, is composed of 24% web data, 9% Wikipedia, 4% GitHub, etc.' The composition of the pretraining data greatly affects the effectiveness of an LM [13; 20; 55J. However, it is unelear how much of each domain to include to produce model that performs well for a wide variety of downstream tasks.
Existing works determine domain weights (the sampling probabilities for each domain) by using intuition or a set downstream tasks. For example, The Pile uses heuristically-choser domain weights, which could be suboptimal. On the other hand, existing LMs such as PaLM [10] and GLaM [13] tune the domain weights based on set of downstream tasks, but requires training potentially thousands of LMs on different domain weights and risks overfitting to the particular set of downstream tasks.
Instead of optimizing domain weights based on set downstream tasks, our approach aims find domain weights which lead to models that perform well on all domains by minimizing the worst-case excess loss over domains, following Oren et [35]; Mindermann et al. [30]. The excess loss the loss gap between the model being evaluated and pretrained reference model.
This motivates our algorithm, Domain Reweighting with Minimax Optimization (DoReMi), which leverages distributionally robust optimization (DRO) to tune the domain weights without knowledge of downstream tasks (Figure 1). First, DoReMi trains small reference model (c.g.. 280M parameters)
'The domain weights, which are based on token count in this paper, varies by tokenizer; see Appendix C.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Stop 1
Stop
Stop
Books
Books
News
News
language model with
Code
reference
model
to get domalh
weights
Code
Web
noweighttd
Law
Law
cotaset
Med
Med
Reference domain
Small reference model
Optimized domain weights define reweighted dataset
Large language model
weights
Small procy model
Figure 1: Given dataset with a domains, Domain Reweighting with Minimax Optimization (DoReMi) optimizes the domain weights to improve language models trained on the dataset. First, DoReMi uses some initial reference domain weights to train reference model (Step 1). The reference model is used to guide the training of small proxy model using group distributionally robust optimization (Group DRO) over domains [35; 43; 34), which we adapt to output domain weights instead of a robust model (Step 2). We then use the tuned domain weights to train large model (Step 3).
Avg One-shot Acc on 8B LMs
25
6.5% better
Exact match (%)
2.6x faster
20-
15
Tune domain weights
with 30x smaller model
Baseline (8B)
DoReMi (280M->8B)
50000
100000
150000
200000
Steps
Figure 2: DoReMi optimizes domain weights with small model (280M params) and uses these domain weights to train much larger model (8B params, 30x larger). Here, optimizing the domain weights (training small model twice) takes 8% the compute of training the large model. DoReMi improves average one-shot downstream accuracy by 6.5% points and reaches the baseline accuracy 2.6x faster when pretraining on The Pile.
in a standard way. Second, DoRcMi trains small distributionally robust language model (DRO-LM) [35], which minimizes the worst-case excess loss (relative to the reference's 's model's loss) across all domains. Notably, rather than using the robust LM, we take the domain weights produced by DRO training. Finally, we train a large (8B) LM on a new dataset defined by these domain weights.
Our approach adapts the DRO-LM framework [35] to optimize domain weights instead of producing robust model. To do this, DoReMi uses the online learning-based optimizer from Group DRO [43; 34], which dynamically updates domain weights according to the loss on each domain for rescaling the train- ing objective, instead sub-selecting examples from minibatch as in Oren et al. [35]; Mindermann et al. [30]. Finally, DoReMi takes the averaged domain weights over DRO training steps.
In Section 3, we run DoReMi on 280M proxy and reference models to optimize domain weights on The Pile [17] and the GLaM dataset [13] (used in PaLM [10]). The DoReMi domain weights are used train an 8B parameter LM (over 30x larger). On The Pile, DoReMi reduces perplexity on all domains over baseline domain weights, even when it downweights domain. DoReMi improves average downstream accuracy over a baseline model trained on The Pile's default domain weights by 6.5% points on generative few-shot tasks and achieves the baseline downstream accuracy 2.6x faster (Figure 2). In Section 4, we find that DoReMi consistently improves LM training when varying the sizes of the proxy model and the main model trained with optimized domain weights. On the

---

GLaM dataset where domain weights tuned on downstream tasks are available, DoReMi even performs comparably tuning domain weights on downstream task performance."
2 Domain Reweighting with Minimax Optimization (DoReMi)
In this section we define DoReMi, an algorithm for using small proxy model to optimize the domain weights of a language modeling dataset, which then improves the training of large model.
Setup. Suppose that we have & domains (e.g.. Wikipedia, GitHub), where for each domain i, we have a set of examples D,. Domain weights a A specify a probability distribution over the domains, and consequently a distribution over the training data: Pa = Σi= Og unif( D;) where unif( D)- = DD Ey is the uniform distribution over the examples in and (x') is if = x and otherwise.
DoReMi.
The inputs of DoReMi are the data D1.... , Dx. reference domain weights Ore (c.g., uniform or based on raw token count of each domain), and training hyperparameters for the large, full-size model (number of training steps T and batch size b). DoReMi returns optimized domain weights â and ultimately, a large model trained on Pa.
Step 1: Obtain a small reference model. We first train model Prer on some reference domain weights Ore (e.g., based on raw token count as a default) for T steps, batch size b. This model serves as the reference model for step 2 and captures baseline level of difficulty of each example/domain. The reference model can be a relatively small model (280M parameters our experiments).
Step 2: Train proxy model with Group DRO to obtain domain weights. To obtain domain weights, we train a small proxy model po in the distributionally robust language modeling (DRO- LM) [35] framework with the Group DRO optimizer [43], where 0 are the weights of the proxy model. This framework trains robust model by optimizing the worst-case loss over domains, which equivalent to the following minimax objective:
Elo(z)- Lo(x) - (x) max L(0,a) Σ aF' EmED, SED;
(1)
a
where the losses ()(x) log (o(x) and (x)(T) = log Pre((x) are the negative log- liglliielsooods of the proxy and reference models respectively in this paper, and |x|| is the number of tokens in an example x. The objective aims to minimize the worst-case excess loss across domains because the inner maximization over puts all the weight on the domain with the highest excess loss.
Intuitively, the excess loss (x) measures the headroom for the proxy model to improve, with respect to the reference model, on example .r. Examples with higher excess loss are those where the reference model achieves low loss (such that the example is "learnalle") but the proxy model still has high loss. Examples with low excess loss may be very high entropy (i.e. optimal loss is high, and thus the reference loss high) or very low entropy (i.e., easy to learn, and thus the proxy loss low). The Group DRO optimizer works by interleaving exponentiated gradient ascent updates on domain weights a with gradient updates on the proxy model weights 0, over training steps f. The optimizer updates o, upweight domains with high excess loss, which scales up the proxy model's 's gradient update on examples from these domains. Following Nemirovski et al. [34], we return the average weights over the training trajectory á= E O4 as the optimized domain weights to use in step 3.
Step 3: Train large model with domain weights. The tuned domain weights define new training distribution Pa. We resample the data from this new distribution to train a main model (larger than the reference/proxy models), using standard training procedure.
Algorithm 1 provides the pseudocode for Step 2. The main structure of Algo- rithm 1 is training loop which updates the proxy model over T steps. each step, follow Sagawa et al. [43] and sample a minibatch with uniform domain weights (regardless of the reference domain
Details for Step 2.
'A public re- implemennation of DoReMi and optimized domain weights for The Pile can be found at https:/ /tt/tttttttttt thu.. www.aaawwwawaiieeaaaaaaaaaaaaaea
3

---

Algorithm 1 DoReMi domain reweighting (Step 2)
Require:
Domain data Di,...,Da Dr. number of training steps T, batch size b, step size smoothing parameter E [0.1] (e.g.. c=le-3 in our implementation).
Initialize proxy weights 0n
Initialize domain weights ag =0= 1 for f from I to I" do
(., of size from Pu. where u= 1.1 Let [z]] be the token length of example x ((ar) s L)
Sample minibatch
Compute per-domain excess losses for each domain i (fo.) (x) is j-th token-level
loss):
X E.corD, EREBOD, lzl
(
(zd (2)).0)
Update domain weights (exp is entrywise): af f- 1--00-1ee 1x-eee(((]))
Renormalize and smooth domain weights: a4 + (c) Σ -ا tcu
Update proxy model weights O, for the objective L(O,- 1,00) 1,0)) (using Adam, Adafactor, etc.)
end for return +
weights arer which only affects the reference model). We then compute the per-domain excess losses, normalized by the total number of tokens in each domain, and use them update the domain weights O₁ at each step. We first compute the per-domain excess loss at per-token level and then aggregate, where the token-level losses index j 1 (a)) = logge. (ar (=1,--.,---1) and (er.(()) -log (z) Prer Prer(z)) x1,....Zy-1). Since the Group DRO optimizer [43] requires a negatvvee loss, we clip the per-token excess loss at O. Finally, we update the proxy model for the objective L(0,-1,0t) using a standard optimizer such as Adam [26] or Adafactor [46]. All experiments in this paper use Adafactor. We set the domain weight update step size 7= and the smoothing parameter to c=le33 in all our experiments and did not extensively tune these hyperparameters.
Iterated DoReMi. We extend DoReMi by running it for multiple rounds, setting the reference domain weights Ore for the next round to be a from the previous round. We call this iterated DoReMi. The entire iterated process still only uses small models for tuning domain weights. We stop iterating when the domain weights converge, which we define as when maximum change in any domain weight ||ā [ mo is less than le-3. Empirically, this takes only rounds the GLaM dataset (Section 3.2).
3 DoReMi Improves LM Training Efficiency and Performance
In this section, we use DoReMi domain weights optimized with 280M-parameter proxy model to train a 8B 18-parameter main model (30x larger). We consider two datasets, The Pile [17] and the GLaM dataset [13]. On The Pile, DoReMi reduces perplexity significantly on every domain, improves average downstream accuracy on generative one-shot tasks by 6.5%, and achieves the baseline accuracy 2.6x faster. On the GL.aM dataset where domain weights tuned downstream datasets available, DoReMi finds domain weights with comparable performance to downstream-tunet domain weights.
3.1 Experimental setup
The Pile dataset. The Pile [17] is 800GB text dataset with 22 domains (Table 1). The default domain weights were determined heuristically. We use the default domain weights from The Pile dataset to train the baseline and as the reference domain weights O in DoReMi (see Appendix C).
GLaM dataset. The GLaM dataset [13] (also used in training PaLM [10]) includes text from were tuned according to the downstream performance of models trained on each domain and the size of each domain [13]. We consider this an oracle comparison, since these domain weights are tuned on downstream tasks that are our evaluation set. We use uniform domain weights both for training the baseline and the reference domain weights age for DoReMi.
domains (Table 2). For comparison, the GLaM domain weights (downstream-tuum)

---

Training setup. We train Transformer [51] decoder-only LMs with the standard next- tokentoken language modeling loss. We conduct controlled comparison by equalizing the amount compute, measured by the number of tokens processed during training. For The Pile, we train each model for 200k steps; for the GL dataset, we train each model for 300k steps. All models use a batch size of 512 and maximum token length of 1024. The proxy and reference models have 280M parameters. All models are trained from scratch (other hyperparameters are Appendix C).
Evaluation. We use held-out validation data to measure the perplexity on each domain. For the generative one-shot tasks from the GPT-3 paper [9]: TriviaQA [21]] NaturalQuestions [27]. WebQuestions [5], SQuADv2 [41], and LAMBADA [36]. We use the standard exact-match accuracy metric for the these datasets. The performance these datasets (particularly TriviaQA) has been shown to correlate well with model scale even at the 100M-1B range [9].
downstream evaluation,
Compute used for optimizing domain weights. We train two 280M models (the reference and proxy models) optimize the domain weights. This is 8% of the FLOPs required train the main 8B model. All FLOPs come from standard forward and backward passes.
Notation for model sizes in DoReMi. We denote the size of the reference/proxy models (which always the same size in experiments) and the size of the main model trained with DoReMi domain weights as "DoReMi (size of reference/proxy--siin of main model)": for example, DoReMi (280M~8B). When
discussing the optimized domain weights independently of the main model, we only include one number (e.g., DoReMi (280M)) which refers to the reference/proxy model size.
25
Exact match (%)
Exact Match (%)
25
20
20
15
15
Baseline (8B)
DoReMi (280M->80) round
DoReMi 2880->8BB)
10
DoReMi (280M-280) round
Downstream-tuned (88)
50000
100000
150000 200000
100000
200000
300000
Steps
Steps
(a) The Pile
(b) GLaM dataset
Figure 3: Average one-shot downstream accuracy (exact match) on 5 tasks, with 8B parameter models trained The Pile (left) and the GLaM dataset (right). On The Pile. DoReMi improves downstream accuracy by 6.5% points and achieves the baseline accuracy 2.6x faster (same plot as Figure 2). On the GLaM dataset, iterated DoReMi (round 2) attains comparable performance oracle domain weights tuned with downstream tasks that are in our evaluation set.
1.5
Baseline (88) DoReMi (280M-288)
ppl
1.0
log
0.5
Pile-CC CC
PubMed Central
Books3
OpenWeb
ArXiv
Github-
Freelan
StackExchange
USPTo Backgrounds
Pubtied Abstracts
Gutenberg (PG-
OpenSubeities
Wikipedis (en)
DM Mathemalics
Ubuntu IRC
BookCorpus2
EuroPart
Hackertiews
YoutubeSubtttte
PhilPapers
NIH ExPorter
Emails
Source
Figure Per-domain log-lerrpe perplexity of 8B models on The Pile. Despite downweighting some domains,
DoReMi improves log-perplexity on all domains.

---

Table 1: Domain weights on The Pile. Baseline domain weights are computed from the default Pile
dataset. DoReMi (280M) uses a 280M proxy model to optimize the domain weights.
Domain
Baseline
DoReMi (280M)
Difference
Domain
Baseline
DoRcMi (2SIM)
Difference
Pil-
0.1
0.6057
+0,4936
DM Mathematics
0.0198
0.0018
.0.0180
YoutubeSubiidlss
0.0012
0,0502
+0.0160
Wikipedia (en)
0.0919
0,0699
-0.0220
PhillPapers
0.0027
0,0244
OpenWeb
0.1247
0.1019
-0.0238
HackerNews
0.0075
0.0134
+0.0059
Github
0.0427
0.0179
-0.0248
Enron Emaill
0.0030
0.0070
+0.0010
FreeLaw
0.0386
0.0013
-0.0341
EunoPari
0.0043
0.0062
+0.0019
USPTO Backgrounds
0,0420
0,0036
-0.0384
Ubuntu
Books3
BookCorpus2
0.0014
0.0061
+0.0017
PubMed Abstracts
0.0815
0.0113
-0.0732
NIH ExPorter
0,0052
0.0063
+0.0011
Stnck
0,0929
0.0153
-0.0776
OpenSubtiles
0.0124
0.0047
-0.0077
ArXiv
0.1052
0,0036
-0.1016
Gutenberg (PG-199
0,0072
-0.0127
PubMed Central
0.1071
0.0016
-0.1025
Table 2: Domain weights in the GLaM dataset. Iterated DoReMi (280M) converges within rounds, with similar overall pattern to domain weights tuned on downstream tasks.
Round
Round
Round 3
Downstream- tuned
Wikipedia
0.09
0.05
0.05
0,06
Filtered wehpages
0.44
0.51
0.51
0.42
Conversations
0.10
0.22
0.22
0.27
Forums
0.16
0.04
0.04
0.02
Books
0.11
0.17
0.17
0.20
News
0.10
0.02
0.02
0.02
3.2 DoReMi improves perplexity and downstream accuracy
We show that DoReMi significantly improves both the perplexity and downstream accuracy of 8B models trained on The Pile and the GLaM dataset over their respective baseline domain weights.
Downstream accuracy improves on The Pile. Figure 3 (left) shows the average downstream performance for baseline and DoReMi (280M~88B) models on The Pile. DoReMi improves the downstream accuracy by 6.5% points and achieves the baseline accuracy within 75k steps 2.6x faster than the baseline (200k steps). Thus, DoReMi can dramatically speed up training and improve downstream performance.
DoReMi can reduce perplexity across all domains without a tradeoff. Figure 4 shows the per-domain log-perplexity of the 8B models on The Pile. DoReMi significantly reduces the perplexity over the baseline across all domains, despite allocating lower weight to some domains. How can this occur? One hypothesis that the domains with the lowest and highest entropy can be downweighted without impacting the perplexity much. The lowest entropy domains statistically require few samples to learn. The highest entropy domains have token distributions that are close to common uniform priors for example, models at random initialization tend to output uniform next token distribution. Thus, we need less samples to fit these domains. Positive transfer from allocating more samples to medium entropy domains can then improve perplexity on all domains. In Appendix D, we provide simple example where reweighting domains can improve perplexity on all domains and DoReMi finds such domain weights in simulations.
Iterated DoReMi achieves performance of downstream-tuned weights on the GLaM dataset. We employ iterated DoReMi on the GLaM dataset over rounds. We find that the second and third round domain weights are almost identical (Table 2). Figure (right) shows one-shot results for the first two rounds of iterated DoReMi. After the first round, the DoReMi main model has comparable downstream accuracy to the baseline (uniform domain weights). After the second round, the DoReMi main model achieves comparable downstream accuracy to oracle domain weights tuned on downstream tasks in our evaluation set. Overall, domain reweighting has a smaller effect on GLaM, possibly because there are only 8 domains compared 22 in The Pile.
Inspecting the DoReMi domain weights. Tables I and 2 present the DoReMi domain weights for The Pile and the GLaM dataset. When running DoReMi on a 280M proxy model (DoReMi (280M)), most weight is put on the diverse Pile-CC web text domain. Note that Wikipedia is downweighted in comparison to the baseline, but DoReMi still improves the downstream accuracy on tasks derived from Wikipedia (c.g., TriviaQA, Appendix Table 5). Domain weights for 1B proxy model (Appendix 8)

---

Exact match (%)
Exact match (%)
8
6-
Baseline (510M)
Daseline
DoReMi (510M->510M)
DoRaMi (280M:>280M)
50000
100000
150000
200000
50000
100000
150000
200000
Steps
Steps
12
14
Exact match (%)
Exact match (%)
10
10
8
Baseline (760M)
Baseline (1B)
DoReMi (760M->760M)
DoReMi (1B->1B)
50000
100000
150000 200000
50000
100000
150000 200000
Steps
Steps
Figure 5: Average one-shot downstream accuracy across model scales (280M, 510M, 760M, IB) where the reference/proxy models for DoReMi the same size the main model trained with DoReMi domain weights. DoReMi consistently improves downstream accuracy across scales, with similar 3% accuracy gap 200k steps most scales (except for 510M). DoReMi achieves the baseline accuracy 4x faster on average across scales.
shows different trend, where Open is the mostly upweighted instead of Pile-CC. This suggests that there may be multiple possible local minima dataset, the DoReMi weights have the same general pattern as the downstream-tuned domain weights. DoReMi able to recover similar set domain weights by starting from uniform initial reference domain weights, without any use of downstream data.
domain weight space. the GLaM
4 Ablations and Analysis Across Scales
Previously Section 3, we showed that DoReMi finds domain weights using 280M models that can im- prove training of 8B models. In this section, we conduct analysis of DoReMi where we vary the scale of the proxy model in relation to the main model and ablate the components of the excess loss objective.
DoReMi improves LMs consistently across scales. We consider using proxy and main models of the same size to analyze DoReMi's behavior a simple setting, without the need for the domain weights to generalize across scales. Note that this is just for scientific purposes since this does not save compute practice. particular, we run DoReMi (X-xX) where is 280M, 510M, 760M, or 1B on The Pile. Figure shows that DoReMi consistently improves downstream accuracy over the baseline by 2% and achieves the baseline accuracy 4x faster on average across scales, and this improvement does not shrink with larger model size. DoReMi improves the worst-case perplexity on all scales and improves 18 of 22 individual domain perplexities on average across scales (Appendix Table 6). These experiments give rough picture of how much is lost when using smaller proxy model; our DoReMi (280M-8B) model achieves the baseline accuracy 2.6x faster, while matching the proxy and main model sizes results in a 4x average speedup.
Proxy model underperforms main model, especially at larger sizes. Recall that DoReMi uses Group DRO train proxy model, which reweights the objective with the domain weights. In contrast, the main model trained by resampling on the domain weights from DoReMi. When the proxy model and the main model are the same size, which one is the better model? Table 3b shows that the proxy
7

---

25
8
Exact match (%)
Exact match (%)
20
6
15
4-
Baseline (88)
DoReMM 280M)
DoReMi (18->68)
Hardast (280M4x280M)
(280M)288)
2
Easiest
DoReMi (150M:288)
DoReMi (70M->88)
0
50000
100000
150000 200000
50000
100000 Steps
150000 200000
Steps
Figure 6: Average downstream accuracy for models trained The Pile. (Left) Increasing the size of the reference/proxy models from 70M to 280M in DoReMi improves downstream accuracy for 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 unweighted average of the per-domain pegppepeexisties
(b) Perplexity of the DoReMi main model and proxy model of the same size. Although the IB proxy model is relatively poor quality, the resulting domain weights still improve the main model.
(a) Varying the the proxy/reference model and training at 8B.
Warst-casc log-ppi
4 domains beating baseline
log-pp)
Warst-case log
Avg log-pp!
domains beating haseline
Baseline (8B)
1.71
1.64
(/22
DoReMi (20M->8R)
1.63
1.53
22/22
DoRcMi (150M-28B)
22/22
DuReM: (280M)280MM
2.19
2.13
22/22
DoReMi (280M-88B)
1.46
1.40
22/22
Proxy (2RIM)
2.35
2:27
19/22
DoReMi (IB-o8B)
1.58
1.54
22/22
Baseline (IB)
1:94
1.87
0/22
DoReMi (BB.oI))
1.92
188
19/2
Proxy
model typically underperforms the main model in this case. The gap between the proxy and main model increases with scale, as the IB proxy model not only underperforms the 1B main model but also the IB 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 IB main model to achieve the baseline performance over 2x 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 resampling-basee 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 IB scales for the DoReMi proxy model while fixing the main model size at 8B (DoReMi (X++B))). From 70M 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 (280M-x8 -8) results in the most improvement at 8B, DoReMi (150M- +BB) and DoReMi (IB -BB still achieve the baseline accuracy almost faster. This suggests that DoReMi is robust to the proxy model scale. practice, 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 En (x) (x)--(rte (x)(()) by running DoReMi using only the loss of the proxy model pa on example i.e. (x( (prefer hardest domains for the proxy model) or only the negative loss of the reference -bre((a) (prefer easiest domains for the reference model). Figure (right) shows that neither of the components of the excess loss alone are sufficient to achieve the gains of DoReMi.

---

5 Related Work
Curating pretraining data for LMs. Most closely related is the GLaM dataset [13] (also used for training PaLM [10D, which has domain weights that are tuned using downstream data. Optimizing domain weights for downstream tasks can be expensive and could require search/zero-ore 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-perplexipy 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 reference model, and focuses data mixture reweighting. Mindermann et al. [30] select examples in an online fashion by taking the top examples in 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 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/downstreann distribution shift with a robust optimization approach. To the best of our knowledge, we are the first 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 restricted form of shift called group shifts [14; 35; 431, 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 baselined loss, but with a simple bigram reference model. DoReMi uses reference model of the same size and architecture as the proxy model to ensure that the losses are on 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 subselelction is more expensive than Group DRO since it requires running forward passes on a larger minibatch (c.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 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-centrie 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 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 general data-centric framework for language modeling via robust optimization.
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 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 domain by data provenance in our experiments, but this only enables coarse-graine soare--rrained control. Using tine-graine 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 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 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 future work.
Broader impacts. Large language models are We hope to improve training efficiency and reduce the environmental impact training large LMs [50; 28; 37; 29]. In particular, by reducing the training time by 2x, 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 energy consumption. Ultimately, we hope 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 applications, the magnitude of the risks increases [8]. Distributionally robust optimization (DRO), which is used in DoReMi to optimize the data mixture, can have favorable impact on fairesss [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 dataset.
7 Conclusion
We introduced DoReMi, 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 2.6x speedup in training the Pile just by changing the sampling probabilities domains. We hope to instigate more research on data-centric approaches for improving language model training efficiency.
Acknowledgments
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.
References
[1] Abubakar Abid, Maheen Farooqi, and James Zou. Persistent anti-muslim bias in large language models. arXiv preprint arXiv:2/01.05783, 2021.
[2] Dario Amodei et al. Deep speech 2 end to end speech recognition 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/1 709.02279
2017. URL htth: /ttttp://rxiv...../// arxiv. wwwg/rsg/////eeeee/oooeeee abo9 02279.
10

---

[4] Aharon Ben- Tal, Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59:341-357,
[5] Jonathan Berant, Andrew Chou, Roy Frostig, and Perey Liang. Semantic parsing Freebase from question-answer pairs. In Empirical Methods in Natural Language Processing (EMNLP),
[6] Su Lin Blodgett and Brendan OConnor. Racial disparity in natural language processing: A case study of social media Arrccaa-mmricann English. arXiv preprint
2017.
[7] Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson. Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dorottya Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Ste- fano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitr, Daniel E. Ho, Jenny Hong. Kyle Hsu, Jing Huang. Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Gcoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning. Suvir Mirchandani, Eric Mitchell, Zanele Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadim- itriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose F. Wang. William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga, Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Percy Liang. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07250 2021.
[8] Rishi Bommasani, Kathleen A. Creel, Ananya Kumar, Dan Jurafsky, and Percy Liang. Picking on the same person: Does algorithmic monoculture lead to outcome homogenization? In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[9] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint 241...
[10] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, A. Rao, Parker Barnes, Yi Tay, Noam M. Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, B. Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, M. Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, S. Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Gareía, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, D. Luan, Hycontack Lim, Barret Zoph, A. Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, T. S. Pillai, Marie Pellat, Aitor Lewkowycz, E. Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele D. Eck, J. Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling language modeling with pathways. arXiv, 2022.
Catasta, Jason Wei, K.
[11] Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning. In International Conference on Learning Representations (ICLR), 2020.
11

---

[12] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognition (CVPR), pages 248-255, 203, 2009.
[13] Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, M. Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang. Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Toju Duke, Lucas Dixon, Kun Zhang. Quoe V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-experee arXiv, 2021.
[14] John Duchi, Tatsunori Hashimoto, and Hongseok Namkoong. Distributionally robust losses against mixture covariate shifts. https: wtt.ss.sttsor..ooooo/oooo..ooo.. ed//" aatau//thahhinaasstttas
wwteraaiiiin/ooddriiiiii pdf, 2019.
[15] Yukun Feng, Patrick Xia, Benjamin Van Durme, and João Sedoc. Automatic document selection for efficient encoder pretraining, 2022. URL https:////axii httpa/aaww..i..raaaaaaaaaraaaaaa wwwp/////////eeee 10951.
[16] Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, Eyal Orgad, Rahim Entezari, Giannis Daras, Sarah Pratt, Vivek Ramanujan, Yonatan Bitton, Kalyani Marathe, Stephen Mussmann, Richard Vencu, Mehdi Cherti, Ranjay Krishna, Pang Wei Koh, Olga Saukh, Alexander Ratner, Shuran Song, Hannaneh Hajishirzi, Ali Farhadi, Romain Beaumont, Sewoong Oh, Alex Dimakis, Jenia Jitsev, Yair Carmon, Vaishaal Shankar, and Ludwig Schmidt. Datacomp: In search of the next generation of multimodal datasets. arXiv preprint ar arXiv:2304. 14108, 2023.
[17] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima. Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. arXiv, 2020.
[18] Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith. Real- toxicityprompts: Evaluating neural toxic degeneration in language models. arXiv preprint arXiv:2009. /162, 2020.
[19] Tatsunori B. Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In International Conference on Machine Learning (ICML), 2018
[20] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero. Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[21] Mandar Joshi, Funsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. Association for Computational Linguistics (ACI.) 2017.
[22] Vishal Kaushal, Rishabh Iyer, Suraj Kothawade, Rohan Mahadey, Khoshrav Doctor, and Ganesh Ramakrishnan, Learning from less data: unified data subset selection and active learning framework for computer vision. IEEE/CVF Winter Conference on Applicatios of Computer Vision (WACV), 2019.
[23] Krishnateja Killamsetty, Durga S, Ganesh Ramakrishnan, Abir De, and Rishabh Iyer. GRAD- MATCH: Gradient matching based data subset selection for efficient deep model training. In International Conference on Machine Learning (ICML), 2021.
[24] Krishnateja Killamsetty, Durga Sivasubramanian. Ganesh Ramakrishnan, and Rishabh Iyer. Glister: Generalization based data subset selection for efficient and robust learning. In Association for the Advancement of Artificial Intelligence (AAAI), 2021.
[25] Krishnateja Killamsetty, Xujiang Zhao, Feng Chen, and Rishabh Iyer. Retrieve: Coreset selection for efficient and robust semi- remi-supervised learning. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
12

---

[26] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
[27] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. In Association for Computational Linguistics (ACL). 2019.
[28] Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and Thomas Dandres. Quantifying 2019.
the carbon emissions of machine learning. arXiv preprint
[29] Anne-Laure Ligozat, Julien Lefévre, Aurélie Bugeau, and Jacques Combaz. Unraveling the 2021.
hidden environmental impacts of AI solutions for environment. CoRR, URL pettpss//arxivv org/ absa 11822.
[30] Sören Mindermann, Jan Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N. Gomez, Adrien Morisot, Sebastian Farquhar, and Yarin Gal. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Conference on Machine Learning (ICMI), 2022.
[31] Baharan Mirzasoleiman, Jeff Bilmes, and Jure Leskovec. Coresets for data-efficient training of machine learning models. In International Conference on Machine Learning (ICML), 2020.
[32] Robert C. Moore and William Lewis. Intelligent selection of language model training data. In Pro- ceedings of the ACI. 2010 Conference Short Papers, pages 220-224, Uppsala, Sweden, July 2010. Association for Computational Linguistics. URL /tatttsss///carrttttttt
org/P10-2041.
[33] Moin Nadeem, Anna Bethke, and Siva Reddy. Stereoset: Measuring stereotypical bias in pretrained language models. arXiv preprint arXiv:2004.09456 2020.
[34] Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on optimization, 19(4:: 1574-1609, 2009.
[35] Yonatan Oren, Shiori Sagawa, Tatsunori Hashimoto, and Percy Liang. Distributionally robust language modeling. In Empirical Methods in Natural Language Processing (EMNLP), 2019.
[36] Denis Paperno, German Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The LAMBADA dataset: Word prediction requiring broad discourse context. In Association for Computational Linguistics (ACI.), 2016.
[37] David A. Patterson, Joseph Gonzalez, Quoc V. Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David R. So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. CoRR, abs/2104. 10350, 2021. URL https:://a/att
org/ abs//104- 10350.
[38] Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on data diet: Finding important exportantexxmpplee examples early in training. In Association for the Advancement of Artificial Intelligence (AAAI), 2021.
[39] Colin Raffel, Noam Shazeer. Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 10683, 2019.
[40] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Empirical Methods in Natural Language Processing (EMNLP), 2016.
[41] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for SQuAD. In Association for Computational Linguistics (ACL), 2018.
[42] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. ImageNet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):11-2252,200 2015.
13

---

[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 Snock, Hugo Larochelle, and Ryan P. Adams. Practical Bayesian optimization of machine learning algorithms. 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 the Association for Compu- tational Linguistics, pages 3645-3650, 3650, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.11663/v1/ 1355. URL / org/P19-1355
/
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint arXiv:1706.0376. 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 lyer, 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 drXiv:2302.03169, 2023.
[56] Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint
2016.
14

---

12-
Exact Match (%)
Exact Match (%)
10-
8
4
Baseline (280M)
Baseline (510M)
DoReMi (280M->280M)
DoReMi (510M->510M)
DoReMi (260M->260M) round 2
6
Downstream-tuned (510M)
Downstream-tuned (280M)
100000
200000
300000
100000
200000
300000
Steps
Steps
(a) 280M
(b) 510M
14
17.5
Exact Match (%)
Exact Match (%)
12
15.0
10
12.5
8
Baseline (760M)
10.0
Baseline (1B)
DoReMi (760M->760M)
DoReMi (1B->1B)
6-
Downstream-tuned (760M)
7.5
Downstream-tuned (1B)
100000
200000
300000
100000
200000
300000
Steps
Steps
(c) 760M
(d) IB
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 the GLaM dataset
Figure 7 presents results across different scales (280M, 510M, 760M, 1B) 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 comparable better than both the baseline (uniform) domain weights and downstream-tune second round weights achieve slightly worse downstream accuracy than the round weights when used to train 280M models, but transfer better to training 8B models.
domain weights. Interestingly, for iterated DoReMi at the 280M scale, the
B Detailed Results for The Pile
Per-domain perplexities for SB models. Table shows per-domain perplexities for 8B models trained on the Pile. The reference/proxy models in this case are 70M, 150M, 280M, and IB. DoReMi improves the perplexity on each domain compared the baseline domain weights.
Per-task accuracies for 8B models. Table shows the accuracies on one-shot generative tasks for various reference/proxy model sizes from 70M to IB. All DoReMi models improve downstream performance significantly over the baseline.
Summary of perplexity results across scales. Table shows a summary of per-domain perplexities for DoReMi across scales (280M, 510M, 760M, 1B). Here, the reference/proxy models are same size as the main model trained with DoReMi domain weights. On average, DoReMi improves perplexity on 18.25 out of domains from The Pile. The worst-case perplexity always reduced (or comparable in the 5 51MM case) with respect to the baseline domain weights.
15

---

Table 4: Per-domain log-perplexities for 8B models trained on The Pile where the reference/proxy models smaller sizes (70M, 150M, 280M, IB). Models trained with DoReMi domain weights have lower perplexity on all domains than the bascline weights.
(8B)
(OM-MM(0M-2SH)
(18MM)()50M-288)
DoReMi (28IM->88B)
DoReMi (IB:588B)
Pile-CC
1.64
1.51
1.18
1.41
1.55
PuhMed Central
1.60
1.58
1.54
1.46
1.56
Books3
OpenWebTest2
1.66
1.48
1.54
1.36
1.58
ArXiv
1.64
1.56
1.53
1.38
1.51
Github
1.53
Freelaw
1.64
1.55
1.54
1.45
1.55
StockEschange
1.61
1.52
1.54
1.39
1.55
USPTO Backgrounss
1.56
PubMed Abstracts
1.61
1.51
1.35
Gutenbery (PG- 19)
1.70
1.56
1.54
1.35
1.52
OpenSubities
1.58
1.56
1.52
1.40
1.55
Wikipedia
1.56
DM Mathemaics
1.63
1.50
1.56
1.38
1.48
Ubusiu IRC
1.71
1.53
1.49
1.42
1,48
BookCorpus2
EuroPast
1.59
1.52
1.51
1.37
1.53
HackerNews
1,66
1.50
1.55
1.45
1.55
YoutubeSuhiitee
1.53
Phiers
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.
(IB->8BB
DoReMi (280M:>8B)
(150M->8B)
DoReM:
LAMBADA
20.10
22.55
29.19
20.59
26.20
NaturalQuestions
SQuADv2
44.43
42.22
$1.89
46.53
40.99
TriviaQA
24.55
34.86
WebQueetions
6.74
8.71
9.15
9.15
6.99
Average
26.56
Table Summary of per-domain log-perplexities for 280M, 5 510M, 760M, and models trained on The Pile, where the referennc/proxx models are the same size. DoReMi improves the worst-case and average perplexity of the baseline domain weights in cases. On average, DoReMi improves perplexity on 18 out of 22 domains.
Worst-casc log-ppl
Avg log-ppl
domains beating bascline
Baseline (280M)
2.39
2.32
0/22
DoRcMi (280M->280M)
2.19
2.13
22/22
Proxy (280M)
2.33
2.27
19/22
Baseline (510M)
2.14
2.08
(1/2
DoRcMi (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) DoReMi (IB->1B)
1.94
1.87
0/22
1.92
1.83
19/22
Proxy (IB)
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-2280M
2.19
2.13
22/22
Hardest (280M->280M)
2,66
2.62
0/22
Easiest (280M->280M)
4.27
4.18
0/22
16

---

0.6
Pile-CC
OpenSubtitles
PubMed Central
Wikipedia (en)
Books3
DM Mathematics
Weight
0.4
OpenWebText2
Ubuntu IRC
ArXiv
BookCorpus2
Github
EuroParl
0.2
FreeLaw
HackerNews
StackExchange
YoutubeSubtittes
0.0
USPTO Backgrounds
PhilPapers
50000
100000
150000
200000
PubMed Abstracts
ExPorter
Steps
Gutenberg (PG-19)
Enron Emails
(a) 280M
Pile-CC CC
OpenSubtitles
0.6
PubMed Central
Wikipedia (en)
Books3
DM Mathematics
Weight
0.4
Mww www. wwwmmm MMMMMMM
OpenWebText2
Ubuntu IRC
ArXiv
BookCorpus2
Github
EuroParl
0.2
FreeLaw
HackerNews
StackExchange
YoutubeSubtitles
0.0
USPTO Backgrounds
PhilPapers
50000
100000 150000
200000
PubMed Abstracts
ExPorter
Steps
Gutenberg (PG-19)
Enron Emails
(b) 1B
Figure 8: Exponential moving average of domain weights throughout 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 shows the perplexities for ablations on the DRO objective. We change the DRO objective and use these 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-levee log- perplexity without baselining with 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 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. 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 smaller compute budget, these domains could become more important- this highlights dependence of the mixture weights on the compute budget. At same time, 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 take advantage of this stability after 50k steps to run DoReMi for smaller number of steps and extrapolate the domain weights to compute.
Comparison domain weights for 280M and 1B. Table presents DoReMi domain weights for The Pile at 280M and 1B proxy models. Different proxy model sizes result in different domain weights, which suggests that there may be multiple local minima domain weight space. With 280M proxy model, most of the weight is put on the Pile-CC web text domain, while DoRcMi with
17

---

Table Domain weights on The Pile. Baseline domain weights are computed from the default Pile dataset. With different proxy model sizes. DoReMi (280M) and DoReMi (IB) result in different domain weights. Despite the differences, the qualitative patterns are similar other than the which web domain has the most weight.
Bascline
DoRcM: (250M)
DoReMi((B)
Pile-CC.
0.1121
0.6157
0,1199
Central
0.1071
Books3
0.0676
0.0224
0.0739
OpenWehflex22
0.1247
0.1019
0.3289
0.1052
0,0384
Gilhub
0.0127
0.0179
0.0129
FrocLaw
0.0386
0.0043
0.0148
SuckExchange
0.0929
0.0153
0.0452
USPTO Rackgrounds
0/0420
0.0260
PubMad Abstracts
0.0845
0.0113
0.1461
Gutenberg (PG-19)
0.0199
0,0072
0.0250
OpenSuii
0:0124
0
Wikipedia (en)
0.0919
0.0699
0.0962
DM Mathemarics
0.0198
010118
0.004
Ubantu
0.0044
0,0044
BookCorpus2
0.0044
0.0061
0.0029
EaroPari
0.002
0.0066
0.0078
HackerNews
0.0075
0.0134
0.0058
YoutubeSSubiitles
010022
0,05022
0.0159
PhilPaaprrs
0.0027
0.0274
0.0063
NTH ExPorter
0.0052
0.0063
0.0094
0.0030
0.0070
0,0033
IB proxy model puts most of the weight on OpenWebTex12. The overall pattern of the domain weights for the rest of the domains are similar.
C Training Details
Data preprocessing. For all datasets, 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 domain according to some domain weights, then sample an example from that domain at random). To reduce the amount of padding tokens, made effort 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.
Baseline domain weights for The Pile. The baseline domain weights for The Pile were computed from The Pile dataset and the number of epochs for each domain given in Gao et al. [17]. After chunking into length 1024 examples, multiplied by the number of epochs that domain specified Gao al. [17]. We then normalized these counts to obtain the baseline domain weights.
counted the number of examples in each domain and
Training setup. For all training runs (including DRO runs), we train with a batch size of 512, initial learning rate of 1e-3, weight decay of 1e-2, and gradient clipping to norm 1. We decay the learning rate exponentially until it reaches a minimum of le-4 at the end of training, with a linear warmup of 6% of the total training steps. We train for 200k steps on The Pile and 300k steps on the Gl .LM dataset. Models under 1B parameters were trained with TPUv3 accelerators, while 1B and 8B models were trained with TPUv4.
Model architectures. Table shows the architecture hyperparameters for the model sizes used in the paper. All the models we use are vanilla Transformer decoder-only models with 256k vocab size.
D Simple Example Where Data Reweighting Has No Tradeoff
Motivated by the findings in Section 3.2, present simple language modeling example where reweighting the training data from different domains improves perplexity on all domains. The example shows that DoReMi downweights domains that are extremely high or low entropy.
Setup. Suppose the ground-truth distribution of text p" a mixture over i domains, where each domain 2 E {l.....k} is defined by different unigram distribution p" (xx 22) over m tokens. Given
18

---

Table Architecture hyperparameters for various model scales used in the paper. All models are vanilla Transformer decoder-only only models and use vocabulary size 256k.
Layers
Attention heads
Attention head dim
Model dim
Hidden dim
70M
3
4
64
256
1024
150M
6
8
64
512
2048
280M
12
12
64
768
3072
510M
12
16
64
1024
8192
760M
12
20
64
1280
8192
IB
16
32
64
2048
8192
8B
32
32
128
4096
24576
budget of n training samples, the goal is choose domain weights p(z) (k scalars that add to 1) to sample training data with such that we learn the parameters of the unigram distributions p (-[12)) well for all 2 from to k. Notably, we do not aim to estimate the ground truth mixture proportions across domains.
Data. Given some domain weights p(z), we sample training data hierarchically: first we determine the number of samples n per domain by drawing from multinomial distribution over & possibilities with probabilities defined by p(z) and n total trials. Then, for each domain 2, we sample n, tokens from (+1+=).. forming vector of tokens X₂ with length n2.-
Model. For each domain 2, consider Bayesian model of the unigram distribution p(r: 2;0) with a Dirichlet prior p(O) 2;;3) over the unigram distribution parameters E A The Dirichlet prior has hy- perparameters 3 E Rmm which can be viewed "pseudo-count" for each token. For each domain z, estimate the parameters Og by computing the mean of the posterior distribution conditioned on the data:
1
nx
(x(x)
{= > [x
for all x (1.....m
n2 t
where 8, - Ax) the sum of pseudocounts.
For domain 2, we write the parameter error of this estimator function of the "difficulty" H of predicting the next token and the "quality" the prior Ag. defined below.
Lemma D.1. For domain index with nz samples, the parumeter error is
n. +
Σ
(3)
(n2 +z)
where
H.
(4)
A. Σ pz(x12) \=(x)
(5)
5
Proof. The parameter error is
-200
(6)
19

---

Evaluating the terms separately,
[x), 12 82
(z) -
+ +n=pp
(8)
B[0_(x)
(x) Σ -[
(nz + 8²)
1=1
(n, + (x) +2A (z) ( (112)==)_p(((]))))(]-m)))))
(10)
Putting it all together, the parameter error can be written as
(0) n + (1 (1)-)))))){
A2
Σ [x( p" (
(11)
(n +s²)
n&p (z (1) + (x(x22) Ae(a)
(12)
(nz sz)
n. H +
(13)
(mg + 82)
No-tradeoff example. Suppose there are 3 domains 2 E {1,2,3} and m = 3 vocabulary tokens E [1,2,3). We use a symmetric Dirichlet prior (preferring a uniform token distribution) where 11/) for all tokens x and domains z. Here, = ^x = 1. In this setting, we show that there is set of domain weights that has strictly lower parameter error than the baseline where we sample the same number of tokens from each domain: n, are equal for all domains 2.
Suppose the ground truth paramaters for the unigram distributions
1
0
0
0.7
0.2
0.1
(14)
1/3
1/3
1/3
where row z contains the parameters for domain z. For example, token has probability under domain I's unigram distribution.
For domain z (non-noisy domain), we have H₂ so the parameter error (according to Lemma D.1) is
s (n1 + s1)
(15)
which is strictly decreasing in the number of samples 71.. For domain = (noisy domain), we have =3=0 so the parameter irror
n3 n3 + 8z)
(16)
by Lemma D.I. This error is minimized to zero at n3 = (no samples). This means that we can allocate samples elsewhere while still reducing error.
For (intermediate entropy domain), we have A2 0.207 and H2 0.46. The derivative of the parameter error with respect to the number samples n2
no H2 + +
H2 - m 223 A2
(17)
Ong
(n2 +12+822 22)
(n22 (12+82)*
20

---

which is negative when
28 A2
> 82
(18)
H2
This inequality holds in this case since 2 < and Therefore the parameter enror is decreasing in the number of samples 112-
Thus, any domain weights that reallocate the examples from domain 3 to domains 1 and reduces the parameter error for all domains.
What kind of domains are downweighted? Intuitively, we can downweight the very noisy (high entropy/difficutt) domain 3 because the initialization perfectly matches the ground truth. This allows us reallocate samples the other domains 1 and 2. Between these, domain requires less additional samples since the parameter error decreases very quickly with the number of samples n, (the difficulty H1 is zero). Thus, the easiest domains should also receive relatively less weight. In practice, positive transfer between domains (which is not captured here) can also contribute to scenarios where reweighting results in no tradeoff across domains.
Simulation with DoReMi. We consider rumning DoReMi on the above no- trodrordofff instance of the simple example with the ground truth unigram distributions in Equation 14. Note that DoReMi's domain reweighting step (Step 2, Algorithm 1) involves loop over 7 iterative model updates, while the estimator from Equation computed closed form. To adapt the estimator for DoReMi, we consider an iterative version where the average is computed in an online fashion. We run DoReMi for T = 500 steps using minibatch size over the n= 500 training examples with domain weight update rate 7) =05. For the model update at step t on an example x from domain 2, we increase the pseudo-count 6, (x) by the current domain weight ag corresponding to domain Instead of using the examples in the minibatch (which is only size and doesn't 't represent all domains), we compute the per-domain excess log-perllexities in Algorithm using fixed, independent evaluation set of 30 examples.
We compare DoReMi against model trained with baseline domain weights, which are uniform over the 3 domains. All models are trained on = 500 training examples. We evaluate the perreeeeyy of model on each domain closed form using the ground truth unigram distribution parameters.
On this simple example, DoReMi returns domain weights [0.39,0.6] 339.0.66,..... 0.00 after rounding to 2 decimal places. These weights correspond to our intuitions - the first domain (non-noisy) is increased by small amount, the third domain (noisy) is decreased to 0 weight, and most of the weight is allocated to the second domain. We use these domain weights to generate new dataset of 500 examples. The model trained with this new dataset improves over the baseline model in perplexity on all domains.
21


Reviews:
[{"review_id": "iskwbTkUDC", "reviewer": "Reviewer_MeFt", "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.\n\nThe 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.\n\nThe 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.", "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.\n\n2) 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?\n\nNote: The appendix should be submitted as the separated supplementary file.", "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.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "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.\n- 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.\n- 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).\nThese 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. \nWhile 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.\n\n[1] Measuring Massive Multitask Language Understanding\n\n[2] Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge\n\n[3] Complex Sequential Question Answering: Towards Learning to Converse Over Linked Question Answer Pairs with a Knowledge Graph", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. MeFt felt that the paper presented a “novel approach to learning the optimal data distribution”, provides “solid experimental evidence, exploring a variety of settings”. We address specific questions below:\n\n> “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…For example, incorporating experimental results on MMLU [1] or commonsense reasoning (ARC[2], CSQA[3], etc) could increase the credibility”\n\nWe evaluated on generative exact-match tasks since that is how the models are typically used (prompting, rather than scoring logprobs). We note that the GPT-3 paper finds that performance on TriviaQA, for example, scales smoothly with model size even at the 100M-1B range, and tracks other benchmarks well. For broad evaluation, we show that DoReMi improves the validation perplexity on all domains, which typically has a strong relationship with downstream performance.\n\n> “ 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.” \n\n- We conducted the ablations of the excess loss objective in Fig 6 using 280M models due to the high cost of running large-scale experiments. \n- With regards to the webpage-only comparison: On the GLaM dataset, we compare against GLaM domain weights that were tuned using the downstream tasks that we evaluate on as an oracle. Without any access to the downstream tasks, we are able to get a similar performance. **Part of GLaM’s process of tuning against downstream tasks is to evaluate the performance of models trained on each single domain. Thus we believe the comparison to downstream-tuned weights is a stronger comparison than the webpage-only baseline proposed by the reviewer.** We thank the reviewer for the question and will clarify this in the final revision.\n\n> “Could you clarify how the main model employs domain weight in detail (Step 3 in Section 2)? … Does the model sample instances in the mini-batch according to the domain weight during pre-training?”\n\nThe main model trains in a standard way with data sampled from the optimized domain weights. As the reviewer mentions, the data loader samples instances from each domain according to the domain weights, and then combines these instances into a minibatch on-the-fly during pre-training. This can be implemented efficiently with separate data queues for each domain. By sampling on the fly, the reweighted dataset never has to be materialized. We apologize for any confusion and will clarify in the final revision.\n\n> “Note: The appendix should be submitted as the separated supplementary file.”\n\nWe thank the reviewer for the note and will update this in the final revision.\n"}, {"review_id": "vHaLbQbUBw", "reviewer": "Reviewer_VPSZ", "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. ", "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. \n\n* **Q1.** (W1.1.) Could the authors better justify their choice of going with generative exact match for the evaluation?\n\n* **Q2.** (W1.2.) Could the authors provide scores (in 0-shot or 1-shot) using logprob-based multiple choices instead?\n\n* **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).  ", "limitations": "The authors extensively discuss limitations of their work in a dedicated section and provide interesting pointers for further research. ", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "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. \n\n* **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. \n\n* **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. \n\n* **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. \n   * **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. \n   * **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).\n\n* **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.\n\n* **W3.** (smaller nits)\n    * **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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. VPSZ felt that the method “is an incredibly valuable contribution”, “reproduces results obtained from manual tuning” from the community, and “could see wide adoption”. We address specific questions below:\n\n\n> “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… 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).”\n\n**We evaluated on generative exact-match tasks since that is how the models are typically used** (prompting, rather than scoring logprobs). We note that the GPT-3 paper finds that performance on TriviaQA, for example, scales smoothly with model size even at the 100M-1B range, and tracks other benchmarks well. For broad evaluation, we show that DoReMi improves the validation perplexity on all domains, which typically has a strong relationship with downstream performance.\n\n\n> “The poor scaling behaviour of DoReMi past 280M parameters for the proxy model is a concern for robustness.”\n\nVPSZ is referring to the scaling behavior of DoReMi with respect to proxy model size, where larger (1B) proxy models were found to be less well optimized with DRO. Overall, **we saw gains across all the proxy model sizes**, but we agree with VPSZ that improving the DRO optimization for larger proxy models is an important future step.\n\n> Nits: table presentation\n\nWe thank the reviewer for the suggestions, and will include them in the final revision.\n"}, {"review_id": "WMnUxlNNVK", "reviewer": "Reviewer_5e3R", "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.", "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?\n- 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.\n\n[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\nhttps://arxiv.org/abs/1812.05159\n", "limitations": "Yes.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n- 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.\n- 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.\n- 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].\n- 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.\n\n[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\nhttps://arxiv.org/abs/1812.05159\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. 5e3R notes that “DoReMi offers a unique and efficient way to determine domain weights” with “wide applicability”. We address specific questions below:\n\n\n> “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.”\n\n- Our paper’s comparisons reflect common ways to assess the quality/importance of domains: 1) by intuition (The Pile) or 2) by tuning on downstream tasks (GLaM dataset).\n- In particular, on the GLaM dataset, we compare against domain weights that were tuned using the downstream tasks that we evaluate on as an oracle. Without any access to the downstream tasks, we are able to get a similar performance. **We believe the comparison to downstream-tuned weights is optimal and in particular, stronger than the simple baseline proposed by the reviewer.** \n\n> “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].”\n\n- Overall, the amount of compute used to train a small proxy model is much smaller than the compute needed to train a large model (2.5% in our 280M to 8B experiment), so that **saving compute in the small proxy model training step only results in a small compute benefit.** \n- **In many practical settings, there may already exist a pretrained reference model** that can be used, reducing the number of proxy LM training processes to 1.\n\n> “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.”\n\n- We believe that we can use a fixed proxy model size (e.g. 280M) to find the domain weights for any main model size, given that the weights are able to transfer across 30x larger model scales. We thank the reviewer for the important practical question and we will add more discussion on this in the final revision.\n\n> “For the domain weight α, 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?”\n\nAlthough assigning all the weight to the domain with the highest excess loss is optimal for the inner maximization, **it can result in unstable training for the outer minimization** since it reduces the number and diversity of examples in the minibatch. Instead, we follow Sagawa et al. (https://arxiv.org/abs/1911.08731) and employ a mirror descent-based DRO optimizer that updates the domain weight α smoothly while maintaining optimality guarantees.\n"}, {"review_id": "c1PQmEQPv6", "reviewer": "Reviewer_uiyj", "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.\n\nTo 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.", "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.", "limitations": "n/a", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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. \n\n\nThe proposed method is clearly-written, and the improvements it offers have been convincingly substantiated through extensive experiments.\n\nThe 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.\n\nI 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. Overall, uiyj feels that the “the direction is really important”, “the general framework is valuable”, and “the improvements it offers have been convincingly substantiated through extensive experiments”. We address specific questions below:\n\n> “the selection of the best weight differs among models with different scales. Does this imply the existence of multiple sub-optimal weight candidates?”\n\n**We believe that there is a frontier of solutions, especially when there are domains that are similar (e.g., OpenWebText and Pile-CC).** For example, we compare the domain weights from 280M and 1B proxy models below. With a 280M proxy model, most of the weight is put on the Pile-CC web text domain, while DoReMi with a 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. We thank the reviewer for the question and will include this discussion in the final revision.\n\n|                   | Baseline | DoReMi (280M) | DoReMi (1B) |\n|-------------------|---------:|--------------:|------------:|\n| Pile-CC           |   0.1121 |        0.6057 |      0.1199 |\n| PubMed Central    |   0.1071 |        0.0046 |      0.0149 |\n| Books3            |   0.0676 |        0.0224 |      0.0739 |\n| OpenWebText2      |   0.1247 |        0.1019 |      0.3289 |\n| ArXiv             |   0.1052 |        0.0036 |      0.0384 |\n| Github            |   0.0427 |        0.0179 |      0.0129 |\n| FreeLaw           |   0.0386 |        0.0043 |      0.0148 |\n| StackExchange     |   0.0929 |        0.0153 |      0.0452 |\n| USPTO Backgrounds |   0.0420 |        0.0036 |      0.0260 |\n| PubMed Abstracts  |   0.0845 |        0.0113 |      0.1461 |\n| Gutenberg (PG-19) |   0.0199 |        0.0072 |      0.0250 |\n| OpenSubtitles     |   0.0124 |        0.0047 |      0.0017 |\n| Wikipedia (en)    |   0.0919 |        0.0699 |      0.0962 |\n| DM Mathematics    |   0.0198 |        0.0018 |      0.0004 |\n| Ubuntu IRC        |   0.0074 |        0.0093 |      0.0044 |\n| BookCorpus2       |   0.0044 |        0.0061 |      0.0029 |\n| EuroParl          |   0.0043 |        0.0062 |      0.0078 |\n| HackerNews        |   0.0075 |        0.0134 |      0.0058 |\n| YoutubeSubtitles  |   0.0042 |        0.0502 |      0.0159 |\n| PhilPapers        |   0.0027 |        0.0274 |      0.0063 |\n| NIH ExPorter      |   0.0052 |        0.0063 |      0.0094 |\n| Enron Emails      |   0.0030 |        0.0070 |      0.0033 |\n\n\n>  “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.”\n\n**We do not use DRO to directly train the large model because it is a more expensive training procedure** that evaluates the losses of two equally sized models (proxy and reference) at every training step. **Instead, we do the expensive DRO training at a small scale, transferring the benefits to the large model through the optimized domain weights.** This also preserves the standard training procedure for the large model. We apologize for the confusion and will clarify in the final revision.\n\n"}, {"review_id": "V0klVpFJkg", "reviewer": "Reviewer_ajJF", "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.", "questions": "* How to set the size of the proxy model given the main model size?\n* How do the learned domain weights by different sizes of proxy models look like?\n\nPlease also clarify any misunderstandings in my review.", "limitations": "Please refer to the Weaknesses section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n* 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).\n* Clarity: The paper is overall clear.\n* 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. \n* 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.\n* 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.\n* (Minor) As I understand, the appendix should be submitted separately from the main paper.\n\nReference:\nGunasekar et al. “Textbooks Are All You Need.” 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. **ajJF notes that “the angle and the proposed method are novel” and “generally well-designed”, with “plenty of experiment results”.** We address specific questions below:\n\n> “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?”\n\nDespite their coaseness, **we already find a significant improvement from reweighting these coarse domains, and addresses an immediate practical need for determining these domain weights**. This suggests that **a future promising direction is to define more fine-grained domains** to further increase the power of reweighting. We note this in the discussion section and will clarify in the revision.\n\n> “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.”\n\n- **Our paper’s comparisons reflect common ways to assess the quality of domains (which is a difficult problem): 1) by intuition (The Pile) or 2) by tuning on downstream tasks (GLaM dataset).**\n- In particular, on the GLaM dataset we compare against domain weights that were tuned using the downstream tasks that we evaluate on as an oracle. Without any access to the downstream tasks, we are able to get a similar performance. **In particular, part of GLaM’s tuning process is to evaluate the performance of models trained on each single domain (including books only or web only).  Thus we believe our comparison to downstream-tuned weights is optimal, and in particular a stronger comparison than the reviewer's suggestion.**\n- Furthermore, we believe that Gunasekar et al., which came out after the paper submission deadline, can focus the pretraining dataset on (coding) textbooks since they are training more specialized models for code.\n\n> “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)”\n\nAlthough we would like to train 65B models, **due to compute limitations we instead show that DoReMi brings benefits at a variety of scales from 280M to 8B, without diminishing returns**. This suggests that DoReMi will scale well to even larger models (analogous to scaling laws).\n\n> “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 … 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.”\n\nFor broad evaluation, we find that **DoReMi improves perplexity on all pretraining domains**. Even though the Github domain is downweighted, we find an improvement in perplexity / next-token prediction, which is closely related to code completion and typically tracks downstream accuracy.\n\n> “How to set the size of the proxy model given the main model size?”\n\nGiven that the weights are able to transfer across 30x larger model scales, we suggest **using a fixed proxy model size (e.g. 280M)** for any main model size. We thank the reviewer for the important practical question and will add a discussion in the final revision.\n\n> “How do the learned domain weights by different sizes of proxy models look like?”\n\nWe compare the domain weights from 280M and 1B proxy models below. With a 280M proxy model, most of the weight is put on the Pile-CC web text domain, while DoReMi with a 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. This suggests there may be multiple local minima in domain weight space, especially when there are some similar domains (such as Pile-CC vs OpenWebText2).\n\n|                   | Baseline | DoReMi (280M) | DoReMi (1B) |\n|-------------------|---------:|--------------:|------------:|\n| Pile-CC           |   0.1121 |        0.6057 |      0.1199 |\n| PubMed Central    |   0.1071 |        0.0046 |      0.0149 |\n| Books3            |   0.0676 |        0.0224 |      0.0739 |\n| OpenWebText2      |   0.1247 |        0.1019 |      0.3289 |\n| ArXiv             |   0.1052 |        0.0036 |      0.0384 |\n| Github            |   0.0427 |        0.0179 |      0.0129 |\n| FreeLaw           |   0.0386 |        0.0043 |      0.0148 |\n| StackExchange     |   0.0929 |        0.0153 |      0.0452 |\n| USPTO Backgrounds |   0.0420 |        0.0036 |      0.0260 |\n| PubMed Abstracts  |   0.0845 |        0.0113 |      0.1461 |\n| Gutenberg (PG-19) |   0.0199 |        0.0072 |      0.0250 |\n| OpenSubtitles     |   0.0124 |        0.0047 |      0.0017 |\n| Wikipedia (en)    |   0.0919 |        0.0699 |      0.0962 |\n| DM Mathematics    |   0.0198 |        0.0018 |      0.0004 |\n| Ubuntu IRC        |   0.0074 |        0.0093 |      0.0044 |\n| BookCorpus2       |   0.0044 |        0.0061 |      0.0029 |\n| EuroParl          |   0.0043 |        0.0062 |      0.0078 |\n| HackerNews        |   0.0075 |        0.0134 |      0.0058 |\n| YoutubeSubtitles  |   0.0042 |        0.0502 |      0.0159 |\n| PhilPapers        |   0.0027 |        0.0274 |      0.0063 |\n| NIH ExPorter      |   0.0052 |        0.0063 |      0.0094 |\n| Enron Emails      |   0.0030 |        0.0070 |      0.0033 |\n\n\n> (Minor) As I understand, the appendix should be submitted separately from the main paper.\n\nWe thank the reviewer for pointing this out, and will correct it in the final revision."}]

Author rebuttal:

Meta-review: This paper proposes a novel effective domain-reweighting method for LLM pretraining. The topic is very interesting and significant.

All reviewers appreciated for the contributions and gave high acceptance scores.

AC also agrees the reviewers' thoughts, so recommends accepting this paper as spotlight or oral
Decision: Accept (spotlight)