Research map / Large language models / Language, safety and interpretability
GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining
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 GRAPE, a multi-source-multi-target domain reweighting method that jointly adapts domain weights and task weights via a Group DRO minimax over a rate-of-improvement metric, improving multi-task reasoning and low-resource multilingual LM.
Problem
Pretraining data mixture optimization for multiple target tasks simultaneously; existing domain reweighting methods target a single task or average performance, causing overfitting and imbalance across benchmarks.
Approach
Interleaved minimax optimization: an inner Group DRO step upweights target tasks with the lowest normalized rate-of-improvement, and an outer step updates domain weights to maximize loss reduction on those prioritized tasks, using multiplicative (mirror descent) updates computed from gradient alignments.
Key results
- On 125M ClimbLab models, GRAPE reaches 44.16 average 5-shot accuracy vs 40.89 for uniform and 42.45 for DoGE.
- On 0.7B ClimbLab models, GRAPE reaches 49.56 average vs 48.11 uniform and 48.72 DoGE, with 6 of 6 tasks above uniform.
- On multilingual low-resource targets, GRAPE accelerates log-perplexity learning by no less than 60% relative to baselines (claimed).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use GRAPE instead of DoGE | 125M model pretrained on ClimbLab · ClimbLab reasoning | average 5-shot exact match accuracy: 42.45 -> 44.16 | Helped |
| Use GRAPE instead of DoGE | 0.7B model pretrained on ClimbLab · ClimbLab reasoning | average 5-shot exact match accuracy: 48.72 -> 49.56 | Helped |
| Replace RoI with EMA-Rate-of-Improvement (ablation) | GRAPE · 12 task combinations | average test log-perplexity: 3.53 vs. 3.68; 3.37 vs. 3.39 | Mixed |
| Replace RoI with Gap-of-Improvement (ablation) | GRAPE · 12 task combinations | average test log-perplexity: 3.53 vs. 3.58; 3.68 vs. 3.61 | Mixed |
| Use unnormalized Gap-of-Improvement instead of normalized RoI (ablation) | GRAPE task-weight updates · six reasoning tasks | task prioritization | Hurt |
| Normalize progress by an EMA of past losses instead of current loss (ablation) | GRAPE task-weight updates · six reasoning tasks | task-weight adaptation | Hurt |
Surprising findings
- DoGE with uniform task weights struggles to deliver competitive performance when facing multiple target languages.
- Offline methods RegMix and CRISP show strongly biased performance across target languages, sacrificing some targets.
Limitations
- Strong convexity assumption in theory does not hold in non-convex LLM pretraining; convergence guarantees are weak.
- Hyperparameters (mu, update frequencies) not exhaustively tuned due to compute cost.
- Gains diminish on coarse source-defined domains such as SlimPajama; results depend on data quality and domain granularity.
Experimental setup
Datasets: ClimbLab SlimPajama wiki40b ARC-Easy ARC-Challenge SciQ PIQA LogiQA HellaSwag GSM8K Kodcode
Benchmarks: ARC-Easy ARC-Challenge LogiQA PIQA SciQ HellaSwag GSM8K Kodcode wiki40b low-resource languages log-perplexity
Metrics: 5-shot exact match accuracy average test log-perplexity
Baselines: Uniform DoGE DoGE-PCGrad RegMix CRISP ClimbMix GRAPE-ClimbMix
Builds on: DoGE Group DRO RegMix ClimbMix PCGrad
Closest papers by meaning
- Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.89)
- Rethinking Data Curation in LLM Training: Online Reweighting Offers Better Generalization than Offline Methods (ICLR 2026 · similarity 0.88)
- Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima (NeurIPS 2026 · similarity 0.87)
- Scaling Laws for Optimal Data Mixtures (NeurIPS 2025 · similarity 0.87)
- TANDEM: Bi-Level Data Mixture Optimization with Twin Networks (NeurIPS 2025 · similarity 0.87)
- Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning (ICML 2025 · similarity 0.87)
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.