SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sang Michael Xie $^{1,2}$ , Hieu Pham $^{1}$ , Xuanyi Dong $^{1}$ , Nan Du $^{1}$ , Hanxiao Liu $^{1}$ , Yifeng Lu $^{1}$ , Percy Liang $^{2}$ , Quoc V. Le $^{1}$ , Tengyu Ma $^{2}$ , and Adams Wei Yu $^{1}$
$^{1}$ Google DeepMind
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.
Datasets for training language models (LMs) are typically sampled from a mixture of many domains [17; 13; 10; 9]. For example, The Pile [17], a 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; 55]. However, it is unclear how much of each domain to include to produce a 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 of downstream tasks. For example, The Pile uses heuristically-chosen 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 a 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 a set of downstream tasks, our approach aims to 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 al. [35]; Mindermann et al. [30]. The excess loss is the loss gap between the model being evaluated and a 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 a small reference model (e.g., 280M parameters)
<sup>1</sup>The domain weights, which are based on token count in this paper, varies by tokenizer; see Appendix C.
Figure 1: Given a dataset with a set of 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 a reference model (Step 1). The reference model is used to guide the training of a 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 a large model (Step 3).
Figure 2: DoReMi optimizes domain weights with a small model (280M params) and uses these domain weights to train a much larger model (8B params, 30x larger). Here, optimizing the domain weights (training a small model twice) takes  $8\%$  of 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, DoReMi trains a small distributionally robust language model (DRO-LM) [35], which minimizes the worst-case excess loss (relative to the reference'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 a 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 training objective, instead of sub-selecting examples from a 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 to train an 8B parameter LM (over  $30\mathrm{x}$  larger). On The Pile, DoReMi reduces perplexity on all domains over baseline domain weights, even when it downweights a 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 to tuning domain weights on downstream task performance.2
2 Domain Reweighting with Minimax Optimization (DoReMi)
In this section we define DoReMi, an algorithm for using a small proxy model to optimize the domain weights of a language modeling dataset, which then improves the training of a large model.
Setup. Suppose that we have  $k$  domains (e.g., Wikipedia, GitHub), where for each domain  $i$ , we have a set of examples  $D_{i}$ . Domain weights  $\alpha \in \Delta^{k}$  specify a probability distribution over the  $k$  domains, and consequently a distribution over the training data:  $P_{\alpha} = \sum_{i=1}^{k} \alpha_{i} \cdot \mathrm{unif}(D_{i})$  where  $\mathrm{unif}(D) = \frac{1}{|D|} \sum_{x \in D} \delta_{x}$  is the uniform distribution over the examples in  $D$  and  $\delta_{x}(x')$  is 1 if  $x' = x$  and 0 otherwise.
DoReMi. The inputs of DoReMi are the data  $D_{1},\ldots ,D_{k}$ , reference domain weights  $\alpha_{\mathrm{ref}}$  (e.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  $\bar{\alpha}$  and ultimately, a large model trained on  $P_{\bar{\alpha}}$ .
Step 1: Obtain a small reference model. We first train a model  $p_{\mathrm{ref}}$  on some reference domain weights  $\alpha_{\mathrm{ref}}$  (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 a baseline level of difficulty of each example/domain. The reference model can be a relatively small model (280M parameters in our experiments).
Step 2: Train proxy model with Group DRO to obtain domain weights. To obtain domain weights, we train a small proxy model  $p_{\theta}$  in the distributionally robust language modeling (DRO-LM) [35] framework with the Group DRO optimizer [43], where  $\theta$  are the weights of the proxy model. This framework trains a robust model by optimizing the worst-case loss over domains, which is equivalent to the following minimax objective:
$$
\min  _ {\theta} \max  _ {\alpha \in \Delta^ {k}} L (\theta , \alpha) := \sum_ {i = 1} ^ {k} \alpha_ {i} \cdot \left[ \frac {1}{\sum_ {x \in D _ {i}} | x |} \sum_ {x \in D _ {i}} \ell_ {\theta} (x) - \ell_ {\text {r e f}} (x) \right] \tag {1}
$$
where the losses  $\ell_{\theta}(x) = -\log p_{\theta}(x)$  and  $\ell_{\mathrm{ref}}(x) = -\log p_{\mathrm{ref}}(x)$  are the negative log-likelihoods 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  $\alpha$  puts all the weight on the domain with the highest excess loss.
Intuitively, the excess loss  $(\ell_{\theta}(x) - \ell_{\mathrm{ref}}(x))$  measures the headroom for the proxy model to improve, with respect to the reference model, on example  $x$ . Examples with higher excess loss are those where the reference model achieves low loss (such that the example is "learnable") 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 is high) or very low entropy (i.e., easy to learn, and thus the proxy loss is low). The Group DRO optimizer works by interleaving exponentiated gradient ascent updates on domain weights  $\alpha_{t}$  with gradient updates on the proxy model weights  $\theta_{t}$  over training steps  $t$ . The optimizer updates  $\alpha_{t}$  to upweight domains with high excess loss, which scales up the proxy model's gradient update on examples from these domains. Following Nemirovski et al. [34], we return the average weights over the training trajectory  $\bar{\alpha} = \frac{1}{T}\sum_{i=1}^{T}\alpha_{t}$  as the optimized domain weights to use in step 3.
Step 3: Train large model with new domain weights. The tuned domain weights  $\bar{\alpha}$  define a new training distribution  $P_{\bar{\alpha}}$ . We resample the data from this new distribution to train a main model (larger than the reference/proxy models), using a standard training procedure.
Details for Step 2. Algorithm 1 provides the pseudocode for Step 2. The main structure of Algorithm 1 is a training loop which updates the proxy model over  $T$  steps. At each step, we follow Sagawa et al. [43] and sample a minibatch with uniform domain weights (regardless of the reference domain
2A public re-implementation of DoReMi and optimized domain weights for The Pile can be found at https://github.com/sangmichaelxie/doremi.
Algorithm 1 DoReMi domain reweighting (Step 2)
Require: Domain data  $D_{1},\ldots ,D_{k}$  , number of training steps  $T$  , batch size  $b$  , step size  $\eta$  , smoothing parameter  $c\in [0,1]$  (e.g.,  $c = 1$  e-3 in our implementation). Initialize proxy weights  $\theta_0$    
Initialize domain weights  $\alpha_0 = \frac{1}{k}\mathbf{1}$    
for  $t$  from 1 to  $T$  do Sample minibatch  $B = \{x_{1},\dots,x_{j}\}$  of size  $b$  from  $P_{u}$  , where  $u = \frac{1}{k}\mathbf{1}$  Let  $|x|$  be the token length of example  $x(|x|\leq L)$    
Compute per-domain excess losses for each domain  $i\in \{1,2,\dots,k\}$ $(\ell_{\theta ,j}(x)$  is  $j$  -th token-level loss):  $\lambda_t[i]\gets \frac{1}{\sum_{x\in B\cap D_i}|x|\sum_{x\in B\cap D_i}\sum_{j = 1}^{|x|}\max\{\ell_{\theta_{t - 1},j}(x) - \ell_{\mathrm{ref},j}(x),0\}}$  Update domain weights (exp is entrywise):  $\alpha_{t}^{\prime}\gets \alpha_{t - 1}\exp (\eta \lambda_{t})$  Renormalize and smooth domain weights:  $\alpha_{t}\gets (1 - c)\frac{\alpha_{t}^{\prime}}{\sum_{i = 1}^{k}\alpha_{t}^{\prime}[i]} + cu$  Update proxy model weights  $\theta_t$  for the objective  $L(\theta_{t - 1},\alpha_t)$  (using Adam, Adafactor, etc.)   
end for   
return  $\frac{1}{T}\sum_{t = 1}^{T}\alpha_{t}$
weights  $\alpha_{\mathrm{ref}}$ , 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 to update the domain weights  $\alpha_{t}$  at each step. We first compute the per-domain excess loss at a per-token level and then aggregate, where the token-level losses at index  $j$  are  $\ell_{\theta_{t-1}, j}(x) = -\log p_{\theta_{t-1}}(x_{j} \mid x_{1}, \ldots, x_{j-1})$  and  $\ell_{\mathrm{ref}, j}(x) = -\log p_{\mathrm{ref}}(x_{j} \mid x_{1}, \ldots, x_{j-1})$ . Since the Group DRO optimizer [43] requires a non-negative loss, we clip the per-token excess loss at 0. Finally, we update the proxy model for the objective  $L(\theta_{t-1}, \alpha_{t})$  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 to  $\eta = 1$  and the smoothing parameter to  $c = 1\mathrm{e}-3$  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  $\alpha_{\mathrm{ref}}$  for the next round to be  $\bar{\alpha}$  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  $\| \bar{\alpha} -\alpha_{\mathrm{ref}}\|_{\infty}$  is less than 1e-3. Empirically, this takes only 3 rounds on the GLaM dataset (Section 3.2).
3 DoReMi Improves LM Training Efficiency and Performance
In this section, we use DoReMi domain weights optimized with a 280M-parameter proxy model to train a 8B-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 GLaM dataset where domain weights tuned on downstream datasets are available, DoReMi finds domain weights with comparable performance to downstream-tuned domain weights.
The Pile dataset. The Pile [17] is a 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  $\alpha_{\mathrm{ref}}$  in DoReMi (see Appendix C).
GLaM dataset. The GLaM dataset [13] (also used in training PaLM [10]) includes text from 8 domains (Table 2). For comparison, the GLaM domain weights (downstream-tuned) 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 in our evaluation set. We use uniform domain weights both for training the baseline and the reference domain weights  $\alpha_{\mathrm{ref}}$  for DoReMi.
Training setup. We train Transformer [51] decoder-only LMs with the standard next-token language modeling loss. We conduct a controlled comparison by equalizing the amount of compute, measured by the number of tokens processed during training. For The Pile, we train each model for 200k steps; for the GLaM 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 in Appendix C).
Evaluation. We use held-out validation data to measure the perplexity on each domain. For downstream evaluation, we use 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 on these datasets (particularly TriviaQA) has been shown to correlate well with model scale even at the 100M-1B range [9].
Compute used for optimizing domain weights. We train two 280M models (the reference and proxy models) to optimize the domain weights. This is  $8\%$  of the FLOPs required to 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 are always the same size in our experiments) and the size of the main model trained with DoReMi domain weights as "DoReMi (size of reference/proxy  $\rightarrow$  size of main model): for example, DoReMi  $(280\mathrm{M}\rightarrow 8\mathrm{B})$ . When we are 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.
(a) The Pile
(b) GLaM dataset
Figure 3: Average one-shot downstream accuracy (exact match) on 5 tasks, with 8B parameter models trained on 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 to oracle domain weights tuned with downstream tasks that are in our evaluation set.
Figure 4: Per-domain log-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	DoReMi (280M)	Difference
Pile-CC	0.1121	0.6057	+0.4936	DM Mathematics	0.0198	0.0018	-0.0180
YoutubeSubtitles	0.0042	0.0502	+0.0460	Wikipedia (en)	0.0919	0.0699	-0.0220
PhilPapers	0.0027	0.0274	+0.0247	OpenWebText2	0.1247	0.1019	-0.0228
HackerNews	0.0075	0.0134	+0.0059	Github	0.0427	0.0179	-0.0248
Enron Emails	0.0030	0.0070	+0.0040	FreeLaw	0.0386	0.0043	-0.0343
EuroParl	0.0043	0.0062	+0.0019	USPTO Backgrounds	0.0420	0.0036	-0.0384
Ubuntu IRC	0.0074	0.0093	+0.0019	Books3	0.0676	0.0224	-0.0452
BookCorpus2	0.0044	0.0061	+0.0017	PubMed Abstracts	0.0845	0.0113	-0.0732
NIH ExPorter	0.0052	0.0063	+0.0011	StackExchange	0.0929	0.0153	-0.0776
OpenSubtitles	0.0124	0.0047	-0.0077	ArXiv	0.1052	0.0036	-0.1016
Gutenberg (PG-19)	0.0199	0.0072	-0.0127	PubMed Central	0.1071	0.0046	-0.1025
Table 2: Domain weights in the GLaM dataset. Iterated DoReMi (280M) converges within 3 rounds, with a similar overall pattern to domain weights tuned on downstream tasks.
Round 1	Round 2	Round 3	Downstream-tuned
Wikipedia	0.09	0.05	0.05	0.06
Filtered webpages	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  $(280\mathrm{M}\rightarrow 8\mathrm{B})$  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 is 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 a 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 a 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 3 rounds. We find that the second and third round domain weights are almost identical (Table 2). Figure 3 (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 to 22 in The Pile.
Inspecting the DoReMi domain weights. Tables 1 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 (e.g., TriviaQA, Appendix Table 5). Domain weights for a 1B proxy model (Appendix 8)
Figure 5: Average one-shot downstream accuracy across 4 model scales (280M, 510M, 760M, 1B) where the reference/proxy models for DoReMi are the same size as the main model trained with DoReMi domain weights. DoReMi consistently improves downstream accuracy across scales, with a similar  $3\%$  accuracy gap at 200k steps at most scales (except for 510M). DoReMi achieves the baseline accuracy 4x faster on average across scales.
shows a different trend, where OpenWebText is the mostly upweighted instead of Pile-CC. This suggests that there may be multiple possible local minima in the domain weight space. On the GLaM dataset, the DoReMi weights have the same general pattern as the downstream-tuned domain weights. DoReMi is able to recover a similar set of domain weights by starting from uniform initial reference domain weights, without any use of downstream data.
4 Ablations and Analysis Across Scales
Previously in Section 3, we showed that DoReMi finds domain weights using 280M models that can improve training of 8B models. In this section, we conduct an 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 in 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 in practice. In particular, we run DoReMi  $(\mathrm{X} \rightarrow \mathrm{X})$  where X is 280M, 510M, 760M, or 1B on The Pile. Figure 5 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 a rough picture of how much is lost when using a smaller proxy model; our DoReMi  $(280\mathrm{M} \rightarrow 8\mathrm{B})$  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 to train a proxy model, which reweights the objective with the domain weights. In contrast, the main model is 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