DoGE: Domain Reweighting with Generalization Estimation

TL;DR

This paper introduces DoGE, a principled method for reweighting data domains during pretraining of large language models, leading to improved generalization across various tasks and out-of-domain data.

Contribution

DoGE presents a bi-level optimization approach to learn domain sampling weights, enhancing model generalization without heuristics.

Findings

Improves perplexity and reasoning accuracy on SlimPajama dataset

Effectively identifies inter-domain dependencies for out-of-domain generalization

Achieves better test perplexity on unseen target domains

Abstract

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.