SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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, we counted the number of examples in each domain and multiplied by the number of epochs that domain specified in Gao et al. [17]. We then normalized these counts to obtain the baseline domain weights.
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 1e-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 GLaM dataset. Models under 1B parameters were trained with TPUv3 accelerators, while 1B and 8B models were trained with TPUv4.
Model architectures. Table 9 shows the architecture hyperparameters for the model sizes used in the paper. All the models we use are vanilla Transformer decoder-only models with a 256k vocab size.
D Simple Example Where Data Reweighting Has No Tradeoff
Motivated by the findings in Section 3.2, we present a 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^*$  is a mixture over  $k$  domains, where each domain  $z \in \{1, \dots, k\}$  is defined by a different unigram distribution  $p^*(x \mid z)$  over  $m$  tokens. Given a
Table 9: Architecture hyperparameters for various model scales used in the paper. All models are vanilla Transformer decoder-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
1B	16	32	64	2048	8192
8B	32	32	128	4096	24576
budget of  $n$  training samples, the goal is to choose domain weights  $p(z)$  ( $k$  scalars that add to 1) to sample training data with such a set of parameters. We use the unigram distributions  $p^{*}(\cdot | z)$  well for all  $z$  from 1 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_z$  per domain  $z$  by drawing from a multinomial distribution over  $k$  possibilities with probabilities defined by  $p(z)$  and  $n$  total trials. Then, for each domain  $z$ , we sample  $n_z$  tokens from  $p^*(\cdot | z)$ , forming a vector of tokens  $X_z$  with length  $n_z$ .
Model. For each domain  $z$ , we consider a Bayesian model of the unigram distribution  $p(x|z;\theta)$  with a Dirichlet prior  $p(\theta | z; \beta)$  over the unigram distribution parameters  $\theta \in \Delta^m$ . The Dirichlet prior has hyperparameters  $\beta \in \mathbb{R}^m$ , which can be viewed as a "pseudo-count" for each token. For each domain  $z$ , we estimate the parameters  $\hat{\theta}_z$  by computing the mean of the posterior distribution conditioned on the data:
$$
\hat {\theta} _ {z} (x) = \frac {1}{n _ {z} + s _ {z}} \left[ \lambda_ {z} (x) + \sum_ {i = 1} ^ {n _ {z}} \mathbf {1} \left[ X _ {z} [ i ] = x \right] \right] \text {f o r a l l} x \in \{1, \dots , m \} \tag {2}
$$
where  $s_z = \sum_x\lambda_z(x)$  is the sum of pseudocounts.
For a domain  $z$ , we can write the parameter error of this estimator as a function of the "difficulty"  $H_{z}$  of predicting the next token and the "quality" of the prior  $\Delta_z$ , defined below.
Lemma D.1. For domain index  $z$  with  $n_z$  samples, the parameter error is
$$
\sum_ {x} \mathbb {E} \left[ \left(\hat {\theta} _ {z} (x) - p ^ {*} (x \mid z)\right) ^ {2} \right] = \frac {n _ {z} H _ {z} + s _ {z} ^ {2} \Delta_ {z}}{\left(n _ {z} + s _ {z}\right) ^ {2}} \tag {3}
$$
$$
H _ {z} = \sum_ {x} p ^ {*} (x \mid z) \left(1 - p ^ {*} (x \mid z)\right) \tag {4}
$$
$$
\Delta_ {z} = \sum_ {x} \left(p ^ {*} (x \mid z) - \frac {\lambda_ {z} (x)}{s _ {z}}\right) ^ {2}. \tag {5}
$$
Proof. The parameter error is
$$
\sum_ {x} \mathbb {E} [ (\hat {\theta} _ {z} (x) - p ^ {*} (x \mid z)) ^ {2} ] = \sum_ {x} \mathbb {E} [ \hat {\theta} _ {z} (x) ^ {2} ] - 2 \mathbb {E} [ \hat {\theta} _ {z} (x) ] p ^ {*} (x \mid z) + p ^ {*} (x \mid z) ^ {2}. \tag {6}
$$
Evaluating the terms separately,
$$
\begin{array}{l} \mathbb {E} \left[ \hat {\theta} _ {z} (x) \right] = \frac {1}{n _ {z} + s _ {z}} \left[ \lambda_ {z} (x) + \sum_ {i = 1} ^ {n _ {z}} \mathbf {1} \left[ X _ {z} [ i ] = x \right] \right] (7) \\ = \frac {1}{n _ {z} + s _ {z}} \left(\lambda_ {z} (x) + n _ {z} p ^ {*} (x \mid z)\right) (8) \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} \left[ \hat {\theta} _ {z} (x) ^ {2} \right] = \frac {1}{(n _ {z} + s _ {z}) ^ {2}} \mathbb {E} \left[ \left(\lambda_ {z} (x) + \sum_ {i = 1} ^ {n _ {z}} \mathbf {1} \left[ X _ {z} [ i ] = x \right]\right) ^ {2} \right] (9) \\ = \frac {1}{\left(n _ {z} + s _ {z}\right) ^ {2}} \left[ \lambda_ {z} (x) ^ {2} + 2 \lambda_ {z} (x) n _ {z} p ^ {*} (x \mid z) + n _ {z} p ^ {*} (x \mid z) + \left(n _ {z} ^ {2} - n _ {z}\right) p ^ {*} (x \mid z) ^ {2} \right] (10) \\ \end{array}
$$
Putting it all together, the parameter error can be written as
$$
\begin{array}{l} \sum_ {x} \mathbb {E} \left[ \left(\hat {\theta} _ {z} (x) - p ^ {*} (x \mid z)\right) ^ {2} \right] = \sum_ {x} \frac {\left(s _ {z} ^ {2} - n _ {z}\right) p ^ {*} (x \mid z) ^ {2} + \lambda_ {z} (x) ^ {2} + \left(n _ {z} - 2 s _ {z} \lambda_ {z} (x)\right) p ^ {*} (x \mid z)}{\left(n _ {z} + s _ {z}\right) ^ {2}} (11) \\ = \sum_ {x} \frac {n _ {z} p ^ {*} (x \mid z) \left(1 - p ^ {*} (x \mid z)\right) + s _ {z} ^ {2} \left(p ^ {*} (x \mid z) - \frac {\lambda_ {z} (x)}{s _ {z}}\right) ^ {2}}{\left(n _ {z} + s _ {z}\right) ^ {2}} (12) \\ = \frac {n _ {z} H _ {z} + s _ {z} ^ {2} \Delta_ {z}}{\left(n _ {z} + s _ {z}\right) ^ {2}}. (13) \\ \end{array}
$$
No-tradeoff example. Suppose there are 3 domains  $z \in \{1,2,3\}$  and  $m = 3$  vocabulary tokens  $x \in \{1,2,3\}$ . We use a symmetric Dirichlet prior (preferring a uniform token distribution) where  $\lambda_z(x) = 1/3$  for all tokens  $x$  and domains  $z$ . Here,  $s_z = \sum_x \lambda_z(x) = 1$ . In this setting, we show that there is a 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_z$  are equal for all domains  $z$ .
Suppose the ground truth paramaters for the unigram distributions are
$$
\left[ \begin{array}{c c c} 1 & 0 & 0 \\ 0. 7 & 0. 2 & 0. 1 \\ 1 / 3 & 1 / 3 & 1 / 3 \end{array} \right], \tag {14}
$$
where row  $z$  contains the parameters for domain  $z$ . For example, token 1 has probability 1 under domain 1's unigram distribution.
For domain  $z = 1$  (non-noisy domain), we have  $H_{1} = 0$  so the parameter error (according to Lemma D.1) is
$$
\frac {s _ {1} ^ {2} \Delta_ {1}}{\left(n _ {1} + s _ {1}\right) ^ {2}} \tag {15}
$$
which is strictly decreasing in the number of samples  $n_1$ .
For domain  $z = 3$  (noisy domain), we have  $\Delta_3 = 0$  so the parameter error is
$$
\frac {n _ {3} H _ {3}}{\left(n _ {3} + s _ {3}\right) ^ {2}}, \tag {16}
$$
by Lemma D.1. This error is minimized to zero at  $n_3 = 0$  (no samples). This means that we can allocate samples elsewhere while still reducing error.
For  $z = 2$  (intermediate entropy domain), we have  $\Delta_2 = 0.207$  and  $H_2 = 0.46$ . The derivative of the parameter error with respect to the number of samples  $n_2$  is
$$
\frac {\partial}{\partial n _ {2}} \frac {n _ {2} H _ {2} + s _ {2} ^ {2} \Delta_ {2}}{(n _ {2} + s _ {2}) ^ {2}} = \frac {H _ {2} (s _ {2} - n _ {2}) - 2 s _ {2} ^ {2} \Delta_ {2}}{(n _ {2} + s _ {2}) ^ {3}} \tag {17}
$$
which is negative when
$$
n _ {2} > s _ {2} - \frac {2 s _ {2} ^ {2} \Delta_ {2}}{H _ {2}}. \tag {18}
$$
This inequality holds in this case since  $\frac{2\Delta_2}{H_2} < 1$  and  $s_2 = 1$ . Therefore the parameter error is decreasing in the number of samples  $n_2$ .
Thus, any domain weights that reallocate the examples from domain 3 to domains 1 and 2 reduces the parameter error for all domains.
What kind of domains are downweighted? Intuitively, we can downweight the very noisy (high entropy/difficulty) domain 3 because the initialization perfectly matches the ground truth. This allows us to reallocate samples to the other domains 1 and 2. Between these, domain 1 requires less additional samples since the parameter error decreases very quickly with the number of samples  $n_1$  (the difficulty  $H_1$  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 running DoReMi on the above no-tradeoff 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 a loop over  $T$  iterative model updates, while the estimator from Equation 2 is computed in 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 1 over the  $n = 500$  training examples with domain weight update rate  $\eta = 0.5$ . For the model update at step  $t$  on an example  $x$  from domain  $z$ , we increase the pseudo-count  $\hat{\theta}_z(x)$  by the current domain weight  $\alpha_t$  corresponding to domain  $z$ . Instead of using the examples in the minibatch (which is only size 1 and doesn't represent all domains), we compute the per-domain excess log-perplexities in Algorithm 1 using a fixed, independent evaluation set of 30 examples.
We compare DoReMi against a model trained with baseline domain weights, which are uniform over the 3 domains. All models are trained on  $n = 500$  training examples. We evaluate the log-perplexity of a model on each domain in closed form using the ground truth unigram distribution parameters.
On this simple example, DoReMi returns domain weights [0.39, 0.61, 0.0] after rounding to 2 decimal places. These weights correspond to our intuitions — the first domain (non-noisy) is increased by a 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 a new dataset of 500 examples. The model trained with this new dataset improves over the baseline model in perplexity on all domains.