# GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining

arXiv: https://arxiv.org/abs/2505.20380 · HTML: https://mlsanity.com/paper/2505.20380/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 30 references that are in the map.
- First posted on arXiv 26 May 2025.

## Builds on (most cited first, 24 of 24)

- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Retrieve, Program, Repeat: Complex Knowledge Base Question Answering via Alternate Meta-learning](https://mlsanity.com/paper/2010.15875/) (2020)
- [Multi-Task Learning as a Bargaining Game](https://mlsanity.com/paper/2202.01017/) (2022)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Skill-it! A Data-Driven Skills Framework for Understanding and Training Language Models](https://mlsanity.com/paper/2307.14430/) (2023)
- [DoGE: Domain Reweighting with Generalization Estimation](https://mlsanity.com/paper/2310.15393/) (2023)
- [Efficient Online Data Mixing For Language Model Pre-Training](https://mlsanity.com/paper/2312.02406/) (2023)
- [Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training](https://mlsanity.com/paper/2504.13161/) (2025)
- [Learning to Combine: Knowledge Aggregation for Multi-Source Domain Adaptation](https://mlsanity.com/paper/2007.08801/) (2020)
- [An Online Method for A Class of Distributionally Robust Optimization with Non-Convex Objectives](https://mlsanity.com/paper/2006.10138/) (2020)
- [Task-Adaptive Pretrained Language Models via Clustered-Importance Sampling](https://mlsanity.com/paper/2410.03735/) (2024)
- [Dynamic Gradient Alignment for Online Data Mixing](https://mlsanity.com/paper/2410.02498/) (2024)
- [Task Weighting through Gradient Projection for Multitask Learning](https://mlsanity.com/paper/2409.01793/) (2024)

## Built on by (most cited first, 4 of 4)

- [Multi-Task GRPO: Reliable LLM Reasoning Across Tasks](https://mlsanity.com/paper/2602.05547/) (2026)
- [DynaMiCS: Fine-tuning LLMs with Performance Constraints using Dynamic Mixtures](https://mlsanity.com/paper/2605.10770/) (2026)
- [ADAPT: Learning Task Mixtures for Budget-Constrained Instruction Tuning](https://mlsanity.com/paper/2512.04555/) (2025)
- [Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time](https://mlsanity.com/paper/2605.15220/) (2026)

## Abstract (arXiv)

The performance of large language models (LLMs) across diverse downstream applications is fundamentally governed by the quality and composition of their pretraining corpora. Existing domain reweighting algorithms primarily optimize data mixtures for a single target task, thereby resulting in models that overfit to specialized objectives while exhibiting substantial performance degradation on other benchmarks. This paper introduces Group Robust Multi-target Adaptive PrEtraining (GRAPE), a novel multi-source-multi-target domain reweighting framework designed to calibrate pretraining data mixtures for robust performance across multiple target tasks simultaneously. GRAPE dynamically adjusts sampling weights across source domains (domain weights) while concurrently modulating task weights that quantify the relative importance of each individual target task. This adaptive process prioritizes tasks based on their learning difficulty throughout training. We formulate this interleaved reweighting mechanism as a minimax optimization problem: The inner maximization adjusts task weights leveraging group distributed-robust-optimization (DRO), where those tasks demonstrating the least improvement under the current data mixture are prioritized with higher weights; The outer minimization then optimizes domain weights to maximize loss reduction on the prioritized tasks. Experiments on ClimbLab and SlimPajama datasets demonstrate that GRAPE consistently outperforms baseline methods in terms of reasoning performance across 6 benchmarks. Furthermore, when applied to multilingual targets, GRAPE effectively identifies optimal training mixtures from mainstream languages, achieving superior language modeling capabilities across 8 low-resource target languages.