PaperlayerView exact revision on arXiv

Source-linked AI summary

DoGE: Domain Reweighting with Generalization Estimation

Simin Fan, Matteo Pagliardini, Martin Jaggi

arXiv:2310.15393v2cs.LGcs.AIcs.CL

TL;DR

Pretraining domain composition matters for LLM generalization, yet domain weights are often selected heuristically or through costly downstream-task tuning. DoGE learns domain weights through generalization estimation and a proxy-model-to-base-model training pipeline. It improves perplexity and few-shot reasoning across six tasks and consistently outperforms uniform weighting for out-of-domain targets.

  • Problem

    LLM pretraining domain composition is underexplored, while current domain weights are mostly chosen heuristically or tuned across downstream tasks, which can be sub-optimal and costly.

  • Method

    DoGE optimizes domain weights for generalization to target domains, using a proxy model and a fully first-order bi-level update that reuses inner-loop gradients.

  • Results

    DoGE outperforms strong baselines in average perplexity and few-shot reasoning across 6 tasks, and consistently outperforms uniform domain weights for out-of-domain targets.

  • Takeaways & Limitations

    DoGE provides robust domain reweighting across proxy-model scales and accelerates learning of target domains and low-resource languages by selecting related data sources.

  • Takeaways & Limitations

    The proxy model consistently performs worse in validation perplexity than a same-scale 82M base model, and scaling to larger models and datasets remains future work.

Abstract

from arXiv · show

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.

1. Introduction

DOGE addresses the costly, heuristic selection of pretraining domain weights by explicitly optimizing generalization to target domains. It learns weights with a proxy model, then trains a larger language model using the resulting domain mixture.

  • Current domain weighting commonly relies on heuristics or downstream-task tuning, which can be sub-optimal and costly.
  • DOGE learns domain weights by explicitly optimizing proxy-model generalization to specified target domains.
  • The method uses one proxy model to obtain optimized weights, then samples training data according to those weights for a potentially larger base model.
  • DOGE outperforms strong baselines including DOREMI in average perplexity and few-shot reasoning across 6 tasks.
  • For target domains absent from training, DOGE handles out-of-domain generalization and consistently outperforms uniform domain weighting.

2. Domain Reweighting with Generalization Estimation

DOGE formulates domain reweighting as a bi-level optimization problem that selects training mixtures for universal or out-of-domain generalization. Its first-order updates favor domains whose gradients align with target learning, after which averaged weights guide base-model sampling.

  • Generalization objectives: DOGE reweights training domains to improve generalization either across all source domains or toward a specific unseen target domain.
  • Setup & notation: The final mixture samples a domain according to α and then samples instances uniformly within that domain.
  • Motivation: Uniform sampling ignores lexical, syntactic, and semantic overlap and differing domain difficulty, and can hinder generalization compared with other methods.
  • Bi-level optimization: At each step, a proxy-model update uses domain-weighted gradients while a first-order outer update adapts weights to the updated model.
  • Generalization estimation: A domain receives higher weight when its gradient contributes to other domains or, for universal generalization, when its own gradient magnitude indicates insufficient learning.
  • Training the base model: DOGE estimates domain weights over proxy training, averages them, and trains the base model by sampling according to the final weights.

3. DOGE Improves Generalization

DOGE reweights source domains to improve both universal and out-of-domain generalization, using proxy-model optimization followed by weighted training of larger base models. Across SlimPajama and multilingual settings, it improves perplexity and few-shot reasoning while detecting useful inter-domain affinities.

  • Universal Generalization: DOGE-10k achieves the best average perplexity among all baselines and outperforms the uniform baseline on 5 of 7 SlimPajama domains.The comparison uses 684M-parameter base models and gives DOREMI-50k a substantial iteration advantage.
  • Universal Generalization: DOGE-10k improves 5-shot reasoning accuracy over the uniform baseline by 1.7 accuracy points on average and leads on 5 of 6 tasks.DOREMI-50k slightly exceeds DOGE on PIQA despite using 40k more auxiliary-model training iterations.
  • Universal Generalization: DOGE produces a two-phase easy-to-hard curriculum, whereas DOREMI weights oscillate and vary substantially with auxiliary-model capacity and training iterations.The mean absolute difference between DOREMI-10k and DOREMI-50k final weights is 0.08.
  • Universal Generalization: DOGE is robust to proxy scale, with mean absolute weight differences below 0.015 and 0.005 for 60M and 124M proxies relative to 82M.DOGE also uses one proxy model, while DOREMI requires two auxiliary models and 5× more tokens plus 10× more floating-point operations.
  • Out-of-Domain Generalization: For unseen target domains, DOGE consistently beats uniform weighting and identifies related source domains, improving target perplexity even after target-domain fine-tuning.The method detects affinities such as C4–CommonCrawl, Stackexchange–Github, Catalan–French/Spanish, and Dutch–German.

4. Discussion and Limitations

The discussion finds that globally averaged domain weights outperform stage-wise schedules overall, while the proxy model remains weaker than an equally sized base model.

  • Proxy-model limitation: The proxy model consistently performs worse in validation perplexity than a same-scale 82M base model.This behavior is also reported for DOREMI’s auxiliary models, which do not match same-scale resampled base models.
  • Efficiency: Parameter selection is explored because DOGE’s generalization-estimation computation budget scales quadratically with model size.The approach ranks transformer blocks and embedding layers by cancellation effect and uses selected-module gradients to compute the estimate.

5. Related Work

Related work spans corpus quality and selection, scalable data filtering, instance-level reweighting, and domain-level reweighting for language-model pretraining.

  • Out-of-domain generalization: The figure’s out-of-domain experiments compare DOGE with uniform weighting on unseen Catalan and Dutch targets and show faster target-loss reduction for DOGE.The figure also displays coherent cross-domain weights for SlimPajama out-of-domain targets.
  • Data selection: Selected high-quality corpora can improve downstream performance with fewer training tokens, but traditional selection methods often face scalability limitations for pretraining.Examples include quality-controlled corpora and PHI-1, while classifier filtering and KL-based selection address corpus construction with different computational trade-offs.
  • Data reweighting: Data reweighting preserves access to the full dataset while rescaling contributions under target tasks, with prior methods operating at instance level or through weighting networks.Domain reweighting instead rescales or resamples broader data groups, offering better scalability for language-model pretraining.

6. Conclusion

The paper concludes that DOGE learns domain weights tailored to generalization objectives and improves language-model and reasoning generalization, including to unseen targets through related data sources.

  • Conclusion: DOGE finds domain-weight distributions tailored to different generalization objectives through generalization estimation.The conclusion presents this as the framework’s central contribution.
  • Conclusion: DOGE improves universal language-model generalization and downstream few-shot reasoning ability on SlimPajama.For out-of-domain objectives, it accelerates target-domain and low-resource-language learning by selectively using related sources.
  • Conclusion: DOGE produces robust reweighting across proxy-model scales, supporting the use of small proxy models to guide larger-model training.The authors identify scaling experiments with larger models and datasets as an important future direction.

A. Model Architectures

The experiments use vanilla decoder-only Transformer models across several scales, with training hyperparameters varying for the largest model.

  • Training hyperparameters: The largest 684M model uses a lower learning-rate range than the other models, while all models share weight decay 0.01 and gradient clipping 1.0.The learning rates use cosine scheduling.
  • Model architectures: The paper uses vanilla Transformer decoder-only architectures across multiple model scales.Table 4 lists the architecture hyperparameters for those scales.

B. Derivation of Domain Weights Update Rule

DoGE chooses domain weights at each step to minimize future average loss, using first-order gradient estimates and regularization; for out-of-domain targets, it uses gradient alignment with the target domain.

  • DoGE optimizes each step’s domain weights to minimize average cross-entropy across all domains at the final training stage.The objective uses domain-specific next-token prediction losses and a greedy stepwise approximation.
  • The update estimates first-order loss changes with a remainder term and groups per-domain generalization estimates into a score vector.The remainder is treated as higher-order error in the Taylor approximation.
  • A Bregman-divergence regularizer stabilizes the domain-weight optimization and yields a log-domain update from the estimated generalization scores.The derivation uses h(α) = Σ α_i ln α_i and its gradient to obtain the update condition.
  • Out-of-domain Generalization: For an out-of-domain target, DoGE selects weights that minimize the target domain’s next-step loss using source-target gradient alignment.The generalization gain for a source domain is defined through the inner product between its gradient and the out-of-domain target gradient.
  • Evaluation: The learned weights produce better average perplexity than baseline methods, while DOGE accelerates early learning across all seven domains.The reported comparison covers 124M and 210M base models, and the early-stage result includes domains that receive fewer training tokens.

C.4. Ablation on the Scale of Proxy Model.

The proxy model’s scale has little effect on the learned DOGE weights, while the resulting base-model comparisons show consistent DOGE advantages over baselines.

  • Proxy models at 60M, 82M, and 124M produce consistent domain weights, with only 1.45% and 0.04% MAE across seven domains.Because the differences were negligible, the authors did not retrain the base model.
  • The proxy model performs worse than both the resampled DOGE base model and the uniformly sampled baseline.This comparison separates proxy-model gradient rescaling from the larger model trained on the resampled corpus.
  • During the first 2500 steps, DOGE’s validation perplexity drops faster than every baseline across all seven domains.The acceleration includes domains that DOGE down-weights and therefore exposes to fewer tokens.

D. Out-of-Domain Generalization Evaluation

DOGE improves out-of-domain validation perplexity without access to the target domain, but gains are limited where source domains lack domain-specific knowledge.

  • DOGE outperforms the uniform baseline without the target domain on all seven evaluated domains.On Book, Github, C4, and CC, DOGE is comparable to or better than the oracle baseline that accesses the target domain.
  • On Arxiv, Stackexchange, and Wikipedia, both DOGE and the target-free uniform baseline remain substantially worse than the oracle.The result indicates that source-domain reweighting has limited benefit when the target requires domain-specific knowledge unavailable in the source domains.

E. Stage-wise Curriculum

Stage-wise curricula preserve DOGE’s performance with two or three stages but substantially hurt performance with ten fine-grained stages; parameter-selection analyses reveal distinct domain-weighting patterns.

  • Experimental setup: The stage-wise and cancellation-effect experiments evaluate alternative sampling weights derived from proxy-model training and parameter-module selection.The stage-wise setup averages weights within K stages, while cancellation-effect selection computes estimates from selected parameter modules after an initial measurement period.
  • E. Stage-wise Curriculum: With K = 2 or 3 stages, stage-wise curriculum maintains performance comparable to global DOGE averaging.The shorter curricula slightly improve hard domains such as C4, CC, and Book while sacrificing performance on Arxiv and Github.
  • E. Stage-wise Curriculum: With K = 10 stages, the fine-grained curriculum severely degrades performance across all domains.The result suggests that training order matters even when the curriculum uses the same diverse data coverage.
  • F. Cancellation Effect: Parameter-selection strategies based on cancellation effect do not outperform full-gradient DOGE, although they produce distinct domain-weight patterns.The corresponding evaluations compare domain weights and validation perplexity by domain.
  • F. Cancellation Effect: Low-cancellation modules tend to up-weight domains with more domain-specific knowledge, whereas high-cancellation modules tend to up-weight broadly informative domains.Examples include Wikipedia, Arxiv, and Stackexchange for unique knowledge, versus CC and C4 for broader coverage.
Loading 2310.15393v2…