Research map / Label-efficient and robust learning / Domain adaptation and generalization
DOGE: Domain Reweighting with Generalization Estimation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DoGE, a first-order bi-level domain reweighting method that trains one small proxy model and uses gradient alignment with target domains to set pretraining domain sampling weights, outperforming DoReMi on SlimPajama perplexity and 5-shot reasoning.
Problem
Choosing pretraining domain weights for LLMs relies on heuristics or costly downstream tuning; the prior DoReMi method depends heavily on a well-trained reference model and cannot handle out-of-domain targets.
Approach
Reweight domain gradients in a proxy model's single-step inner update and update domain weights via multiplicative weights on a generalization estimation score (gradient inner products with the average or target-domain gradient), giving a first-order, one-proxy-model method that also supports out-of-domain targets.
Key results
- On the 684M base model, DoGE-10k achieves average validation perplexity 15.806 vs 16.526 for the uniform baseline and 17.172 for DoReMi-10k.
- DoGE-10k improves average 5-shot accuracy across 6 tasks to 52.29 vs 50.59 for uniform baseline, outperforming DoReMi-50k (51.49).
- In out-of-domain SlimPajama experiments, DoGE without target access beats the uniform baseline on all 7 domains; e.g. Book 63.89 vs 82.57 and C4 63.96 vs 89.56 perplexity.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DoGE domain reweighting instead of uniform domain weights | 684M Transformer trained on SlimPajama · SlimPajama 7 domains | average perplexity: 16.526 -> 15.806 | Helped |
| Use DoGE domain reweighting instead of DoReMi-10k | 684M Transformer trained on SlimPajama · SlimPajama 7 domains | average perplexity: 17.172 -> 15.806 | Helped |
| Use DoGE domain reweighting instead of DoReMi-50k | 684M Transformer trained on SlimPajama · SlimPajama 7 domains | average perplexity: 16.124 -> 15.806 | Helped |
| Use DoGE domain reweighting instead of uniform domain weights | 684M Transformer · Six 5-shot reasoning tasks, average | exact-match accuracy (%): 50.59 -> 52.29 | Helped |
| Use DoGE instead of DoReMi-10k domain weights | 684M Transformer · Six 5-shot reasoning tasks, average | exact-match accuracy (%): 49.79 -> 52.29 | Helped |
| Use DoGE instead of DoReMi-50k domain weights | 684M Transformer · Five of six 5-shot reasoning tasks | exact-match accuracy (%) | Mixed |
| Use DoGE domain reweighting instead of uniform source-domain weights | 124M Transformer · SlimPajama OoD target domains | perplexity | Helped |
| Use DoGE instead of uniform domain weights | 124M Transformer trained on SlimPajama · SlimPajama universal, 124M | average perplexity: 18.566 -> 18.066 | Helped |
| Use DoGE instead of uniform domain weights | 210M Transformer trained on SlimPajama · SlimPajama universal, 210M | average perplexity: 17.218 -> 16.661 | Helped |
| Use DoGE-low10 instead of full-gradient DoGE (ablation) | 124M Transformer trained on SlimPajama · SlimPajama universal | average perplexity: 18.066 -> 18.848 | Hurt |
| Use DoGE-low30 instead of full-gradient DoGE (ablation) | 124M Transformer trained on SlimPajama · SlimPajama universal | average perplexity: 18.066 -> 18.442 | Hurt |
| Use DoGE-low50 instead of full-gradient DoGE (ablation) | 124M Transformer trained on SlimPajama · SlimPajama universal | average perplexity: 18.066 -> 18.151 | Hurt |
Surprising findings
- The proxy model itself performs worse than a same-scale base model trained with the learned domain weights, and even worse than the uniform baseline.
- Stage-wise (time-varying) domain weights give no clear improvement over globally averaged domain weights; K=10 stages hurts performance substantially.
Limitations
- Base models are small (124M to 684M) and experiments are not scaled to larger models or datasets.
- Proxy model performs worse than same-scale base model, so the proxy's own performance is not a good indicator of final quality.
- Generalization estimation cost scales quadratically with the number of domains and model size, motivating parameter selection heuristics that did not improve results.
Experimental setup
Datasets: SlimPajama Wiki40b
Benchmarks: COPA SciQ PIQA LogiQA WiC WinoGrande SlimPajama validation perplexity per domain Wiki40b Catalan/Dutch perplexity
Metrics: perplexity exact-match accuracy (5-shot)
Baselines: DoReMi Uniform baseline Oracle
Builds on: DoReMi influence functions Group DRO
Closest papers by meaning
- Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning (ICML 2025 · similarity 0.89)
- Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.88)
- Rethinking Data Curation in LLM Training: Online Reweighting Offers Better Generalization than Offline Methods (ICLR 2026 · similarity 0.87)
- QT-DoG: Quantization-Aware Training for Domain Generalization (ICML 2025 · similarity 0.87)
- GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining (NeurIPS 2025 · similarity 0.86)
- Auto-GDA: Automatic Domain Adaptation for Efficient Grounding Verification in Retrieval-Augmented Generation (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.