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

We thank the reviewer for the thoughtful questions and constructive suggestions! We provide our response as follows:

Q1: Generalizability to unseen tasks

We evaluate the model optimized using the 6 reasoning tasks from our main experiment on 6 unseen tasks and report the 5-shot reasoning accuracy as follows:

Our method outperforms both the Uniform and the strong DoGE baseline on a majority (4 out of 6) of these unseen benchmarks, which suggests that GRAPE's advantages generalize well beyond the seen tasks.

Q2: Broader applications to Post-training

GRAPE is highly applicable as a general data mixture optimization method to broader scenarios, including continued pretraining and instruction tuning. During the post-training stage like instruction tuning, a more fine-grained data selection method can be preferable like a sample-level selection to get a better control of the final model behavior. Using sample-level DRO to upweight the most difficult individual instructions is a powerful concept for the SFT stage and a promising direction for future research, as we discussed in the Discussion section (line 353-359).

We will add the above discussions to the revised version of the paper. For code availability, according to the rebuttal guideline, we cannot put external links here. We promise to open-source our codebase in the updated version.

We thank the reviewer for the thoughtful questions and constructive suggestions! We provide our response as follows:

Q1: Hyperparameter tuning & Runtime:

We perform a hyperparameter search on the reweighting intervals $\Delta T_{\alpha}, \Delta T_{z}$ = 50, 100, 200 for both GRAPE and DoGE. We train 125M model for 50k steps and report the final task accuracies and run time as follows:

These results show that GRAPE consistently outperforms DoGE across all tested update frequencies. We also observe that the most frequent updates ($\Delta T=50$) do not yield the best performance. This is likely because the Rate-of-Improvement (RoI) metric becomes more reliable after the model has trained for a sufficient number of steps on the current data mixture, making overly frequent weight adjustments suboptimal.

A side note on efficiency improvement: Regarding efficiency, while GRAPE has a theoretical FLOP overhead of ~ 28%, we note that the practical runtime is also largely influenced by the extra GPU memory occupation for storing extra gradients. A promising direction for future work is to incorporate a recently proposed "ghost gradient computation" technique [1], which could further improve the wall-clock efficiency for both GRAPE and DoGE.

We will add a more comprehensive discussion on the hyperparameter tuning in the updated version of our paper.

Q2: Benefits of the Minimax DRO Framework vs. Minimization-Only Frameworks

The primary benefit of GRAPE's minimax (min-max) framework is that it explicitly optimizes for robustness of multi-task learning performance by adaptively tuning the task weights, while the minimization-only frameworks (like RegMix, DoGE, or Aioli) only optimizes the domain weights according to a static, pre-defined target distribution. Therefore, GRAPE is able to adaptively adjust the guidance distribution (i.e. task weights) for balanced and robust improvements across a diverse suite of benchmarks.

Q3: Scalability to a large number of target tasks

We thank the reviewer for the insightful question! The scalability of gradient-based approaches is indeed a critical issue since computing gradients for all $N$ target tasks can become a bottleneck. While our experiments used a manageable number of tasks, the GRAPE framework is flexible and can be adapted for scalability using two primary strategies:

Method 1: Task Reweighting with Partial Feedback

Instead of computing gradients for all $N$ tasks at every update, we can sample a small subset of $m << N$ tasks. This can be done in a principled manner using techniques from the multi-armed bandit literature (e.g., EXP3.M [2]), which balance exploiting known slow-learning tasks with exploring others. This would reduce the per-update cost from $O(N)$ to $O(m)$ without losing the core benefit of the algorithm.

Denote task weights as $z(t)$ as in the paper, define a sampling distribution: $p_n^{(t)} = (1-\epsilon) \cdot z_n^{(t)} + \frac{\epsilon}{N}$

to balance exploitation (focusing on tasks currently believed to be important) with exploration (ensuring all tasks remain observable). For each sampled task $n \in S_t$, we compute the alignment score to approximate ROI scores:

$A^{(t)} =$ a^{(t)}{1}, \hdots, a^{(t)}{N} $$, where $a^{(t)}_{n} = \langle \Delta_{\theta} l_n^{val}, d_t \rangle$

where $\Delta_\theta l_n^{val}$ is the normalized validation gradient for task $n$, and $d_t$ is the current model update (a weighted combination of training domain gradients). To convert this alignment signal into a bandit-style reward in [0,1], we rescale it as follows: $r_{n}^{(t)} = \frac{1}{2} (1 - a_{n}^{(t)}/C)$ where $C > 0$ is a scaling constant that upper-bounds the magnitude of alignment scores (e.g., a norm bound). This ensures that $r_{t,n} \in$ 0,1 $$, with larger rewards assigned to more misaligned tasks—consistent with GRAPE's principle of prioritizing underperforming tasks. We then update the task weights using the standard multiplicative-weights rule (as in EXP3.M [2]), with importance-weighted feedback: $z_{n}^{(t+1)} \leftarrow z_{n}^{(t)} · exp(\eta_{z} · r_{n}^{(t)}/p_{n}^{(t)} · 1_{n \in S_t})$

This procedure allows GRAPE to focus on the most corrective tasks using only a small number of gradient evaluations per step.

Hierarchical Task Reweighting

For corpora with thousands of fine-grained partitions, we can first group similar domains or tasks into semantically coherent clusters based on metadata or embeddings similarity. GRAPE would then operate on these high-level clusters, reweighting groups of tasks rather than individual ones. This dramatically reduces the number of entities to reweight, making the process highly scalable.

[1] Data Shapley in One Training Run. Wang et. al, 2025.

[2] Algorithms for adversarial bandit problems with multiple plays. T. Uchiya, A. Nakamura, and M. Kudo, 2010.

Thank you for your time! Your answers here are very interesting, and I hope they'll make it into a revision of the paper.

We are grateful to the reviewer for their careful assessment of our work and the insightful feedback. We agree with the suggested points and provide our responses below:

Q1: Hyperparameter tuning & Runtime:

We perform a hyperparameter sweep on the reweighting intervals of $\Delta T_{\alpha}, \Delta T_{z} = 50, 100, 200$ for both GRAPE and DoGE. We train 125M model for 50k steps and report the final task accuracies and run time as follows:

These results show that GRAPE consistently outperforms DoGE across all tested update frequencies. We also observe that the most frequent updates ($\Delta T=50$) do not yield the best performance. This is likely because the Rate-of-Improvement (RoI) metric becomes more reliable after the model has trained for a sufficient number of steps on the current data mixture, making overly frequent weight adjustments sub-optimal.

A side note on efficiency improvement: Regarding efficiency, while GRAPE has a theoretical FLOP overhead of ~ 28%, we note that the practical runtime is also largely influenced by the extra GPU memory occupation for storing extra gradients. A promising direction for future work is to incorporate a recently proposed "ghost gradient computation" technique [1], which could further improve the wall-clock efficiency for both GRAPE and DoGE.

We will add a discussion on the hyperparameter tuning in the updated version of our paper.

Q2: GRAPE with exact ROI scores

We intended to predict the Ratio-of-Improvements at the next step using a look-ahead strategy, where the first-order approximation stands. Also, the gradient alignment between tasks and training set can provide more information about the optimization directions on the landscape than the exact ROI score.

To validate this, we ran an experiment with the exact ROI scores, where we calculated ROI directly to update task weights at every reweighting step and then update domain weights according to the task weights. The 5-shot accuracy on the target tasks are presented as follows:

The results above show that the underlying GroupDRO framework in GRAPE is more effective than the directly computed ROI scores. We will add more discussions and analysis on this ablation in the revised version of the paper.

Q3: Apply on tasks with priority preference

GRAPE is primarily designed for robustness, where the goal is to achieve balanced performance by automatically identifying and lifting the lagged tasks. In practice, we can run GRAPE with task-weights initialized non-uniformly according to the pre-defined task priority. However, the principle of GRAPE that continuously optimizes the worst-improved task could conflict with the predefined priority-order. For instance, GRAPE might upweight a low-priority task if it falls significantly behind, which would counter the user's explicit preference.

Therefore, if a practitioner has pre-defined, static priorities for tasks, applying DoGE with a static non-uniform task-weights would be more appropriate.

[1] Data Shapley in One Training Run. Wang et. al, 2025.

We thank the reviewer for the thoughtful questions and constructive suggestions! We provide our response as follows:

Q1: Clarification on the baselines

We provide detailed implementation descriptions for all baseline methods in Appendix D.1 (line 620-674). We also provided the distribution of regmix-optimized domain weights in Figure 7,8,11,12,16,17. The actual domain weights from RegMix on the CLIMBLAB dataset with 6 target tasks is presented as follows:

We will add the tables with the actual domain weights from all baseline methods in our revised version for better readability and clarity!

Q2: Performance after SFT

To investigate if GRAPE's pretraining advantages remains after post-training, we conducted supervised fine-tuning experiments on the pretrained 125M models. We present the 5-shot accuracy scores before and after supervised fine-tuning on PIQA, SciQ, ARC-Easy, ARC-Challenge as follows:

The results show that GRAPE keeps the advantage after the SFT stage in most of the tasks. On ARC-Challenge, GRAPE outperforms DoGE and performs comparably as the uniform baseline. We did not conduct the RL here since we did not apply a relevant task as our target, but how to select the optimal data for RL is definitely a critical and challenging direction to explore.

Q3: Limited training length and baselines in large-scale experiments

Due to the limitation of computation budget, we did not manage to pretrain models larger than 0.7B on sufficient numbers of tokens. As the reviewer notes, we prioritized comparing against DoGE as it was the strongest baseline under the limited resources but did not include RegMix and CRISP baselines in 0.7B experiments. We will be sure to add the results for RegMix and CRISP to the 0.7B experiments in the final version of the paper to ensure a complete comparison.

We thank the reviewer for the effort to read through our work and the constructive feedback! These questions have helped us clarify key aspects of our work. We provide our responses as follows:

Q1: Hyperparameter tuning & Runtime:

We perform a hyperparameter search on the reweighting intervals $\Delta T_{\alpha}, \Delta T_{z} = 50, 100, 200$ for both GRAPE and DoGE. We train 125M model for 50k steps and report the final task accuracies and run time as follows:

These results show that GRAPE consistently outperforms DoGE across all tested update frequencies. We also observe that the most frequent updates ($\Delta T$=50) do not yield the best performance. This is likely because the Rate-of-Improvement (RoI) metric becomes more reliable after the model has trained for a sufficient number of steps on the current data mixture, making overly frequent weight adjustments sub-optimal.

A side note on efficiency improvement: Regarding efficiency, while GRAPE has a theoretical FLOP overhead of ~ 28%, we note that the practical runtime is also largely influenced by the extra GPU memory occupation for storing extra gradients. A promising direction for future work is to incorporate a recently proposed "ghost gradient computation" technique [1], which could further improve the wall-clock efficiency for both GRAPE and DoGE.

We will add a discussion on the hyperparameter tuning in the updated version of our paper.

Q2: Average versus worst-case accuracy

We did not report worst-case accuracy as it can be volatile and skewed by a single, exceptionally difficult task, which may not reflect the model's broader capabilities across a diverse task suite.

However, to directly evaluate the worst-case performance against our optimization objective (cross-entropy loss), we provide extensive results on worst-case log-perplexity in Table 5, Table 8, and Table 9. The results demonstrate that GRAPE consistently achieves superior (lower) worst-case perplexity compared to baselines across numerous task configurations.

Q3: Justification on the small learning rate assumption

The assumption of a small learning rate is a standard condition for applying the first-order Taylor expansion to approximate the one-step loss improvement, which is a common practice in the theoretical analysis of gradient-based algorithms. Empirically, it can be justified as the learning rates used in large-scale language model pretraining are typically considered small relative to the optimization landscape.

Q4: Assumptions on Theorems

• On Theorem 2.1: we agree with the reviewer that the general LLM training is a highly non-convex scenario. However, the convex assumption is commonly used for theoretical analysis for optimization to provide theoretical grounding and intuition for the algorithm's behavior in an idealized setting, such as the very first work on influence function [2]. Our empirical results show that even in the non-convex LLM pretraining practice, GRAPE with the underlying group-DRO algorithm can still provide valuable insight on the data mixture selection and brings significant empirical improvements.
• On Theorem 2.2: We thank the reviewer for this sharp observation. We agree that the μ-strongly convex condition is not necessary to prove monotonic variance reduction. We will remove this assumption in the revised version of our paper.

Q5: Scaling behaviors of GRAPE

We thank the reviewer for this insightful question! While GRAPE outperforms all baselines at both the 125M and 0.7B scales, the relative improvement over the uniform baseline is less pronounced for the larger model. We hypothesize a few reasons:

• (1) Limited Training Tokens: due to the limited computation budget, The 0.7B model was trained on 13B tokens, which may be insufficient to show the full benefits of a dynamic data curriculum at that scale;
• (2) Sub-optimal Hyperparameters: The optimal hyperparameters may differ across model scales, and our settings may perform better on 125M models than 0.7B models;
• (3) Shifting Task Dynamics: while the model is scaling up, the conflict pattern among various target tasks can be shifted. For example, larger models can be more sensitive to task conflicts or easier to generalize across various domains. Understanding the scaling effect of the multi-task learning behavior is an intriguing direction to explore. We will leave it to our future work.

Q6: Comparison to the oracle of data selection

While the true oracle of data mixture is hard to define in the context of LLM pretraining, we apply CRISP as a proxy for a static oracle baseline. CRISP identifies the domain sampling weights as the closest training data distribution to the target data distribution according to the sequence embedding similarity. Therefore, the CRISP mixture is determined using an auxiliary embedding model (e.g. sentence-bert) before the training begins.
Our results show that GRAPE's dynamic approach significantly outperforms this static oracle. In Table 1,5,6,7,8, GRAPE largely improves perplexity and reasoning accuracies compared to CRISP; according to Figure 19,20, in the multilingual setting, the static mixtures from CRISP can sabotage performance on most languages, whereas GRAPE delivers consistent gains on all the target languages, leveraging the dynamic data curriculum.

[1] Data Shapley in One Training Run. Wang et. al, 2025.

[2] Understanding Black-box Predictions via Influence Functions. Koh and Liang, 2018.