DoGE: Domain Reweighting with Generalization Estimation

The coverage and composition of the pretraining data significantly impacts the generalization ability of Large Language Models (LLMs). Despite its importance, recent LLMs still rely on heuristics and trial and error to increase or reduce the influence of data-domains. We propose DOmain reweighting with Generalization Estimation (DoGE), which optimizes the probability of sampling from each ___domain (___domain weights) in a principled way. Our approach is a two-stage process consisting of (i) training a proxy model to obtain ___domain weights using a bi-level optimization algorithm; (ii) training a larger base model by sampling training domains according to the learned ___domain weights. In our experiments, we extensively show how DoGE improves the generalization of the base model to any target data mixture. On the SlimPajama dataset, our base model gets better perplexity and few-shot reasoning accuracies across 6 tasks compared to baseline methods. Moreover, aiming to generalize to out-of-___domain target tasks, which is unseen in the pretraining corpus (OOD ___domain), DoGE can effectively identify inter-___domain dependencies, and consistently achieves better test perplexity on the target ___domain.