3

Step 1: Obtain a small reference model. We ﬁrst train a model pref on initial domain weights α0
(e.g., uniform domain weights as a default) for T steps and batch size b. This model serves as the
reference model for step 2 and captures a baseline level of difﬁculty of each example/domain. The
reference model can be a relatively small model (280M parameters in our experiments).

θ

θ

i=1

min

|x|

(cid:34)

x∈Di

x∈Di

αi ·

max
α∈∆k

max
α∈∆k

(cid:88)

k(cid:88)

L(θ, α) = min

1(cid:80)

Step 2: Train proxy model with Group DRO to obtain domain weights. To obtain domain
weights, we train a small proxy model pθ in the distributionally robust language modeling (DRO-
LM) (Oren et al., 2019) framework with the Group DRO optimizer (Sagawa et al., 2020), where
(cid:35)
θ 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:
max{(cid:96)θ(x) − (cid:96)ref(x), 0}
DoReMi- Optimizing Data Mixtures Speeds Up Language Model Pretraining