DoReMi: Optimizing Data Mixtures with Group DRO

Editorial record

Plain-language summary

The method first trains a small reference model, then trains a second small proxy model with Group Distributionally Robust Optimization that raises the sampling weight of domains where the proxy has the largest excess loss over the reference. The resulting domain weights are reused to sample data for a much larger model. This removed the need to hand-tune or grid-search corpus mixtures at full scale, and reached target accuracy in fewer training steps by reweighting domains like The Pile rather than using their default proportions.