Distributionally robust optimization. Within DRO methods for deep learning (Ben-Tal et al.,
2013, Oren et al., 2019, Sagawa et al., 2020, Sinha et al., 2018), we target a restricted form of shift
called group shifts (Duchi et al., 2019, Oren et al., 2019, Sagawa et al., 2020), where the test distri-
bution can be an unknown mixture of groups (domains). We follow DRO-LM (Oren et al., 2019),
which employs DRO for LMs in the group shift setting. DRO-LM also uses a baselined loss, but
with a simple bigram reference model. DoReMi uses a reference model of the same size and archi-
tecture as the proxy model to ensure that losses can be compared. During optimization, DRO-LM
takes a worst-case subset of each minibatch to update the model on, while we use the Group DRO
optimizer (Sagawa et al., 2020) which doesn’t require subselection. Subselection is expensive since
it requires evaluating the model on a large minibatch, while only selecting a small fraction to up-
DoReMi- Optimizing Data Mixtures Speeds Up Language Model Pretraining