Title: GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM pretrain

data selection
Abstract: 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 $\textbf{G}$roup $\textbf{R}$obust Multi-target $\textbf{A}$daptive $\textbf{P}$r$\textbf{E}$training (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 ($\textit{domain weights}$) while concurrently modulating $\textit{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 $\texttt{ClimbLab}$ and $\texttt{SlimPajama}$ datasets demonstrate that GRAPE consistently outperforms baseline methods in terms of reasoning accuracies 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.

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

Abstract

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 G roup R obust Multi-target A daptive P r E training (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.

1 Introduction

The composition of pretraining data, typically encompassing diverse sources, topics, and languages [Gao et al., 2020, Together Computer, 2023], is a critical determinant of large language model (LLM) capabilities. While various data sources are conventionally combined using fixed, heuristically determined mixing ratios, recent research indicates that adaptively adjusting these ratios ( domain weights ) via domain reweighting during pretraining can substantially enhance downstream task performance [Xie et al., 2023, Liu et al., 2025, Fan et al., 2024a]. However, prevailing domain reweighting methods are largely confined to the multi-source-single-target paradigm, focusing exclusively on source-target relevance while failing to address complex interactions across multiple target tasks. In particular, task-adaptive domain reweighting approaches [Grangier et al., 2025, Fan et al., 2024b] curate data mixtures based on performance on one specific task, while often leading to overfitting with compromised generalization ability. General domain reweighting techniques such as Liu et al. [2025] and Fan et al. [2024a] adjust training data mixtures for average performance across various target tasks, assessed via mean stochastic task losses or gradients. Yet, due to the heterogeneous dynamics of multi-task learning, average assessment proves suboptimal when confronted with dominant gradients - where loss or gradient magnitude on certain tasks overshadow others, or conflicting gradients - where improving one task hinders another. Conversely, traditional multi-task learning (MTL) strategies like gradient surgery [Yu et al., 2020, Liu et al., 2023, Bohn et al.,

2024] require substantial data from each target task and often involve complex gradient manipulations that scale poorly to foundation models, rendering them infeasible for LLM pretraining.

To facilitate effective multi-task adaptation during LLM pretraining, we propose G roup R obust Multitarget A daptive P r E training (GRAPE), the first domain reweighting algorithm explicitly designed for the multi-source-multi-target setting. Instead of optimizing for average task performance, GRAPE calibrates domain weights to maximize improvement on the most challenging tasks at each stage of training, guided by the key principle:

Tasks that are improving slowly deserve more attention.

Specifically, GRAPE employs group distributed-robust-optimization (DRO) [Duchi and Namkoong, 2020, Sagawa et al., 2020] to adaptively assign higher task weights to target tasks that are currently learning slowest. We introduce the Rate-of-Improvement ( RoI , Equation 1), a normalized, step-wise metric assessing task-specific learning progress. Unlike prior metrics like excess loss [Xie et al., 2023], RoI avoids reliance on auxiliary models and inherently accounts for varying task difficulties by normalizing improvement against current performance.

We formulate this bifold reweighting process as a minimax optimization problem (Equation 2), jointly optimizing: (1) domain weights ( α ), the sampling distribution over source domains; and (2) task weights ( z ), a distribution over target tasks reflecting their priority during the evolution of training. The inner maximization phase uses DRO to identify the task distribution z that emphasizes the least-improved tasks (minimal RoI). Tasks demonstrating minimal RoI under the current domain mixture are thus prioritized in subsequent optimization steps; The outer minimization phase then optimizes the domain weights α to maximize the expected improvement according to this dynamically determined worst-case task distribution. This interleaved mechanism enables GRAPE to balance domain utilization against multi-task performance adaptively via a dynamic curriculum, preventing dominance by any single task and fostering robust generalization.

Contributions. Our primary contributions are threefold:

We introduce GRAPE, the first principled multi-target domain reweighting algorithm for LLM pretraining. We provide a rigorous derivation grounded in optimization theory, including a theoretical analysis of its properties (§ 2).

We empirically validate GRAPE's efficacy, demonstrating superior performance over baseline methods on reasoning benchmarks (§ 3.1) and low-resource multilingual language modeling (§ 3.2).

Through comprehensive ablation studies and comparisons of different progress assessment metrics used in DRO, we offer insights on the training dynamics of multi-target domain reweighting, shed light on the future research in data-centric LLM pretraining and curriculum learning (§ 4).

2 Methodology: Multi-Target Domain Reweighting via GRAPE

This section formalizes the problem of domain reweighting for multi-target adaptation and presents GRAPE, our proposed algorithm that optimizes pretraining data composition for robust performance across multiple target tasks.

Problem Formulation. Let D train ≜ {D 1 , . . . , D K } represent a pretraining corpus partitioned into k source domains based on meta-attributes (e.g., source, topic, etc.). During training, data is sampled from D train according to domain weights α ∈ ∆ K , where ∆ K is the probability simplex in R K . Let D tgt ≜ {T 1 , . . . , T N } be a set of N target tasks, which can be presented as different benchmarks. Our objective is to find the optimal domain weights α that yield strong performance across all tasks T n ∈ D tgt. We introduce task weights z ∈ ∆ N ⊂ R N to represent the relative importance assigned to each target task during training. The core idea is to dynamically update z to reflect task learning progress and subsequently adjust α to optimize the training trajectory based on z .

Let θ t denote the model parameters at step t , and let l n ( θ t ) be the stochastic loss evaluated on target task T n for n ∈ [ N ] . We assume optimization proceeds via gradient descent: θ t +1 = θ t -γ t d t , where γ t is the learning rate and d t is the update direction. Let l k ( θ t ) , g k ( θ t ) be the stochastic loss and gradient computed on data from training domain D k for k ∈ [ K ] . The overall parameter update direction is a weighted combination of stochastic domain gradients: d t := ∑ K k =1 α t k g k ( θ t ) , where the domain weights α t are used as the dynamic importance sampling weights at step t .

2.1 GRAPE: Prioritizing the Lagged Tasks by Group DRO

Algorithm. To achieve balanced improvement across target tasks, GRAPE employs Group Distributed Robust Optimization (Group DRO) to jointly adapt task weights ( z ) and domain weights ( α ). We assess task improvement using the Rate-of-Improvement (RoI) , inspired by Liu et al. [2023]:

Here, l n ( θ t ) -l n ( θ t +1 ) approximates the one-step improvement in loss for task T n . Normalizing by l n ( θ t ) provides a scale-invariant measure of relative progress, preventing bias towards tasks with intrinsically higher loss magnitudes. For sufficiently small learning rates γ t , a first-order Taylor expansion yields:

where ∇ θ log l n ( θ t ) = ∇ θ l n ( θ t ) l n ( θ t ) . Following the principles of Group DRO [Sagawa et al., 2020, Qi et al., 2021], we formulate a minimax objective to optimize the domain weights ( α ) for the worst-case (minimum) RoI score across tasks:

Here, the inner minimization finds the task weights z that concentrate on the tasks with the lowest expected RoI; The outer maximization finds the domain weights α that maximize this minimum expected RoI. Following Qi et al. [2021], Namkoong and Duchi [2016], we introduce Bregman divergence regularization terms h ( α ) and h ( z ) relative to the previous weights α ( t -1) and z ( t -1) to stabilize the updates and manage stochastic variance:

where Ψ( b ) = ∑ i b i log( b i ) is the negative entropy. Solving this regularized minimax problem (details in Appendix B) yields multiplicative update rules resembling online mirror descent:

In practice, the expectations are replaced by stochastic estimates using mini-batches. The inner product ⟨∇ θ log l n ( θ t ) , g k ( θ t ) ⟩ captures the alignment between the gradient direction of task n ∈ [ N ] (normalized by loss) and the gradient direction of domain k ∈ [ K ] . ∇ θ ℓ ( θ t ; x ) denotes the (training loss) gradient from a training batch x sampled according to α t -1 , and ∇ θ log ℓ ( θ t ; y ) denotes the task-weighted normalized validation gradient from a validation batch y sampled according to z t -1 . The pseudocode is presented in Algorithm 1.

Efficiency analysis. Updating α and z at every step (Equation 3, Equation 4) requires N target gradients and K domain gradients computations, which is computationally expensive. We mitigate this by performing updates periodically, where domain weights ( α ) and task weights ( z ) are updated every ∆ T z and ∆ T α steps. For our experiments on SlimPajama with K = 7 and N = 6 , using ∆ T z = 100 , ∆ T α = 100 , the computational overhead is approximately ( N +1)∆ T -1 z +( K +1)∆ T -1 α = 15% in terms of gradient computations relative to standard training. Compared to prior domain reweighting algorithm DOGE [Fan et al., 2024a], our task reweighting step only introduce ( N +1)∆ T -1 z = 7% overhead in computations. The memory overhead stems primarily from storing gradients for each domain/task during the update step. If model parameters require m storage, AdamW optimizer stores

Algorithm 1 G roup R obust Multi-target A daptive P r E training (GRAPE)

≈ 4 × m (model parameters, gradients, and two EMA states). The peak memory usage at reweighting steps stores one additional gradient ( m ), leading to ≈ 25% memory overhead.

Interpreting the Reweighting Mechanism. The update rules (Equation 3, Equation 4) reveal an elegant negative feedback regulation at the heart of GRAPE's reweighting mechanism: Task weights z n increase for tasks where the normalized gradient ∇ log ℓ ( θ , y ) has low alignment with the overall update direction d t , indicating slow progress under the current regime (Equation 3). Subsequently, domain weights α j increase for domains whose gradients g k align well with the normalized gradients of these prioritized, struggling tasks (represented by ∇ log ℓ ( θ , y ) where y ∼ mix( z ) ) (Equation 4). This mechanism continuously directs training focus towards underperforming tasks by upweighting relevant source domains. The full derivation is provided in Appendix B.

2.2 Theoretical Properties of GRAPE

The GRAPE algorithm, formulated as a regularized minimax optimization problem solved via multiplicative updates, possesses desirable theoretical properties under standard assumptions. The update rules (Equation 3, Equation 4) are instances of online mirror descent applied to the variables ( α , z ) of the minimax game. Theorem 2.1 suggests that under strong convexity and smoothness assumptions, the algorithm converges towards the set of Pareto optimal solutions at an O (1 /T ) rate, where T is the number of training iterations. However, we acknowledge the limitation imposed by the strong convexity assumption. In practical non-convex scenarios of LLM pretraining, convergence guarantees are typically weaker. The proof of Theorem 2.1 is presented in Appendix C.1.

Theorem 2.1 (Convergence of GRAPE) . Let the loss functions l n ( θ ) be L -smooth for all n ∈ [ N ] and the norm of stochastic gradients be upper-bounded by G . If the learning rate γ t satisfies γ t ≤ 1 L and the regularization parameters µ α , µ z are chosen such that µ α > 0 and µ z > 0 , then the GRAPE algorithm with update rules given by the above equations converges to a neighborhood of the Pareto optimal solution at a rate of O (1 /T ) , where T is the number of training iterations. Specifically, for any ε > 0 , there exists T 0 such that for all T > T 0 :

Furthermore, Theorem 2.2 posits that under similar assumptions of smoothness and bounded gradients, the variance of task performances, σ 2 t = Var n ∈ [ N ] ( l n ( θ t )) , tends to be monotonically decrease after an initial phase σ 2 t +1 ≤ σ 2 t for t ≥ T 0 . It indicates that the algorithm actively counteracts the divergence of the task performances, which promotes more uniform progress across the task suite compared to static weighting or simple averaging. The proof is presented in Appendix C.2.

Theorem 2.2 (Monotonic Variance Reduction of Task Performance) . Let σ 2 t = Var n ∈ [ N ] ( l n ( θ t )) denote the variance of task performances at iteration t . Let the loss functions l n ( θ ) be L -smooth for all n ∈ [ N ] and the norm of stochastic gradients be upper-bounded by G , and assuming the task losses are L -smooth and µ -strongly convex, the variance decreases monotonically until reaching a minimal basin, i.e., σ 2 t +1 ≤ σ 2 t for all t ≥ T 0 for some finite T 0 .

3 Experiments

To evaluate the efficacy of GRAPE, we conduct experiments in two distinct multi-target pretraining scenarios: (1) optimizing a generic English language model for diverse reasoning tasks, and (2) optimizing the data mixture from mainstream languages for low-resource language modeling. For all experiments, we train various scales of decoder-only transformers following Vaswani et al. [2023] from scratch. More details on architecture and hyperparameters are provided in Appendix D.1.

3.1 Domain Reweighting for Multi-task Reasoning

Setup. We consider two pretraining corpora, ClimbLab [Diao et al., 2025] with K = 20 source domains clustered by topics; and SlimPajama with K = 7 domains classified by collection sources. We first apply GRAPE to standard English language model pretraining, where the objective is to dynamically adapt the data mixture to maximize performance across a suite of reasoning benchmarks simultaneously. We select N = 6 diverse reasoning tasks spanning scientific, logical, physical, and commonsense reasoning: ARC-Easy (ARC-E), ARC-Challenge (ARC-C) [Clark et al., 2018],

SciQ [Welbl et al., 2017], PIQA [Bisk et al., 2019], LogiQA [Liu et al., 2020], and HellaSwag [Zellers et al., 2019]. For each target task T n , we use its standard validation set to compute the task loss l n ( θ t ) and the Rate-ofImprovement r ( t ) n needed for GRAPE's updates during training. We update task weights z every ∆ T z = 100 steps and domain weights α every ∆ T α = 100 steps. Initial weights α 0 and z 0 are set to uniform distributions. We use the AdamW optimizer with standard hyperparameters for LLM pretraining. Experiments were conducted using 4 × H100 80GB GPUs. Implementation details and hyperparameter settings are provided in Appendix D. We report the results on ClimbLab in the following sections and present the results on SlimPajama in subsection D.3.

Baselines. On experiments with 0 . 7 B scale models, we compare GRAPE to two baseline methods: (1) Uniform : training data are sampled uniformly from each source domains; (2) DOGE [Fan et al., 2024a]: domain weights are dynamically adjusted by gradient alignment with uniform task weights across N target tasks. On smallscale ( 125 M) models, we employ a comprehensive suite of baseline approaches: (3) DoGE-PCGrad : a multi-task learning extension of DOGE, where the target gradient are dynamically updated with a gradient surgery algorithm PCGrad [Yu et al., 2020]; (4) RegMix [Liu et al., 2025]: predict the domain weights using regression towards the optimal average loss across all target tasks; (5) CRISP [Grangier et al., 2025]: determine the domain

0.32

0.64

0.96

1.28

1.6

0.36

0.38

0.40

0.42

0.44

5-shot Accuracy

Uniform

DoGE

ClimbMix

RegMix

CRISP

DoGE-PCGrad

GRAPE

GRAPE

-climbmix

(a) 125M-ClimbLab

1.3

2.6

3.9

5.2

6.5

7.8

9.1 10.411.7 13

0.42

0.44

0.46

0.48

0.50

5-shot Accuracy

Uniform

DoGE

GRAPE

(b) 0.7B-ClimbLab

weight distribution by importance sampling, i.e. clustering all data samples from target tasks onto source domains; (6) ClimbMix [Diao et al., 2025]: the optimized domain weights curated by iteratively regression and bootstrapping; (7) GRAPE-CLIMBMIX : apply GRAPE with domain weights α initialized from ClimbMix weights.

GRAPE improves multi-task reasoning capability. According to Table 1 and Figure 1, GRAPE demonstrates consistent improvements on multi-task reasoning capabilities across different model scales, achieving superior average 5-shot reasoning accuracy and learning efficiency compared to various baseline methods. According to Figure 3.1, GRAPE and GRAPE-climbmix, significantly outperform a suite of baselines including DOGE, DOGE-PCGRAD, ClimbMix, RegMix, CRISP, and uniform domain sampling when training small-scale ( 125 M) models. On 0 . 7 B models (Figure 3.1), both DOGE and GRAPE substantially outperform the uniform sampling baseline, demonstrating the general efficacy of adaptive domain reweighting approaches. Notably, GRAPE achieves comparable average 5-shot accuracies to the uniform baseline while utilizing approximately 40% fewer training tokens. Furthermore, compared to DOGE, GRAPE not only exhibits an approximate 25% acceleration in reaching similar accuracy thresholds, but also obtains more well-rounded improvements across all tasks above uniform baseline, as benefit from its distinctive task reweighting mechanism.

Task weights evolution reflects learning curriculum. The evolution of task weights ( z t ) on 0.7B model, as depicted in Figure 2, reveals the dynamic learning curriculum adopted by GRAPE.

In the early stage, the reading comprehension tasks like ARC-E and ARC-C are mostly up-weighted. As training progresses, in the late stage, physical and commonsense reasoning tasks like PIQA and Hellaswag are steadily prioritized, demonstrating the emergence of a skill-wise learning curriculum that moves from foundational to more complex reasoning abilities. Concurrently, tasks like SciQ and LogiQA consistently receive lower weights. This suggest they may either benefit sufficiently

0

25k

50k

75k

100k

Optimization Steps

0.0

0.2

0.4

0.6

0.8

Weights

ARC-E

ARC-C

SciQ

PIQA

LogiQA

from cross-task generalization driven by the prioritized tasks or can be learnt more efficiently relative to other tasks. This adaptive curriculum, by directing resources to evolving bottlenecks, ensures that difficult tasks are not neglected, fostering robust and balanced multi-task reasoning capabilities.

C1

C2

C4

C5

C6

C7

C8

C9

C10

C11

C12

C13

C14

C15

C16

C17

C18

C19

C20

DoGE

RegMix

GRAPE ClimbMix CRISP

0.15

0.13

0.00

0.16

0.01

0.39

0.00

0.00

0.00

0.04

0.01

0.05

0.00

0.00

0.00

0.01

0.05

0.00

0.00

0.00

0.00

0.01

0.00

0.19

0.01

0.03

0.02

0.02

0.01

0.00

0.01

0.02

0.00

0.00

0.02

0.01

0.62

0.02

0.00

0.00

0.03

0.02

0.00

0.09

0.02

0.15

0.03

0.03

0.09

0.02

0.25

0.08

0.00

0.00

0.02

0.05

0.02

0.01

0.07

0.01

0.01

0.01

0.01

0.03

0.02

0.20

0.16

0.01

0.01

0.07

0.01

0.28

0.01

0.00

0.00

0.07

0.06

0.02

0.01

0.00

0.25

0.02

0.00

0.15

0.00

0.03

0.03

0.00

0.00

0.04

0.00

0.30

0.01

0.00

0.09

0.00

0.07

0.00

0.00

0.00

0.00

0.25

0.50

Cluster

Which data domains and topics are critical for general reasoning? Figure 3 presents the domain weight distributions across different methods across 20 data clusters within the ClimbLab corpus, whose main topics are detailed in Table 4. GRAPE notably assigns high importance on Cluster 1,4,10,12, which primarily focus on broad science, mathematics and education contents. It also

explains the constant low-priority on SciQ task indicated in Figure 2 since the training data mixture contains substaintial scientific-QA related contents. In contrast, other methods exhibit different topical biases: DOGE shows a strong preference on Cluster 6, featuring healthcare and genetic contents; REGMIX heavily utilizes Cluster 17 which with topics on health and medical research; CRISP mostly favors Cluster 11, indicating a great emphasis on software and programming. Notably, CRISP achieves significant improvement on LogiQA task, which indicates training on code data is critical to improve the logical reasoning ability of language models.

0.13

0.26

0.39

0.52

0.65

8.75

9.00

9.25

9.50

Log Perplexity

tr

0.13

0.26

0.39

0.52

0.65

7.5

8.0

8.5

uk

0.13

0.26

0.39

0.52

0.65

6

8

ro

0.13

0.26

0.39

0.52

0.65

8.5

9.0

9.5

pl

0.13

0.26

0.39

0.52

0.65

7.5

8.0

8.5

9.0

nl

0.13

0.26

0.39

0.52

0.65

6.5

7.0

7.5

8.0

pt

0.13

0.26

0.39

0.52

0.65

6.5

7.0

7.5

ca

0.13

0.26

0.39

0.52

0.65

8.0

8.5

9.0

da

Uniform

DoGE

GRAPE

3.2 Language Mixture for Multi-lingual Learning

Setup. In this scenario, we investigate GRAPE's ability to optimize the language composition of the pretraining corpus from mainstream languages to enhance language modeling capabilities across multiple low-resource target languages. The source corpus consists of data from K = 6 languages selected from the wiki40b dataset [Guo et al., 2020], including high-resource languages English ( en ), French ( fr ), German ( de ), Spanish ( es ), Russian ( ru ) and Italian ( it ). Each language constitutes a domain D i , and GRAPE adapts the sampling weights α ∈ ∆ k over these source languages. We select N = 8 low-resource languages as target tasks, distinct from the source languages: Turkish ( tr ), Danish ( da ), Catalan ( ca ), Polish ( po ), Romanian ( ro ),

Portuguese ( pt ), Dutch ( nl ), and Ukrian ( uk ). Performance is measured by the language modeling loss, i.e. the log-perplexity (log-PPL), on held-out test sets for each target language. Lower log-PPL indicates better performance.

GRAPE Enhances Multilingual Adaptation.

Figure 4 illustrates the test log-perplexity on each target low-resource language. GRAPE consistently outperforms both uniform sampling and DOGE baselines across all languages, effectively accelerating the low-resource language modeling. Specifically, GRAPE accelerates the low-resource language modeling by no less than 60% in terms of log-perplexity scores, achieving significantly lower final perplexity within the same training budget. While Fan et al. [2024a] showed DOGE greatly improve single-target language learning, its efficacy diminishes when confronted with multiple target languages simultaneously, where it struggles to deliver competitive performance with uniform task weights. This outcome underscores the sub-optimality of average task weighting strate-

4k

8k

12k

16k

20k

Steps

0.14

0.15

0.16

0.17

0.18

0.19

Weights

en

fr

de

es

it

ru

(a) Train Domain Weights

4k

8k

12k

16k

20k

Steps

0.075

0.100

0.125

0.150

0.175

0.200

Weights

tr

uk

ro

pl

nl

pt

ca

da

(b) Task Weights

gies for multi-task learning, especially when tasks exhibit conflicting characteristics. Notably, according to Figure 20, the offline reweighting algorithm such as RegMix and CRISP exhibit a strongly biased performance across various target languages: while RegMix slightly facilitate the learning on most of targets languages including Catalan, Portuguese, Ukrainian and Polish, it sacrifices the performance on Dutch and Danish; the purely embedding-based method CRISP only accelerates the learning on Turkish while sabotaging all the other languages.

Weight trajectory reveal linguistic relations. The dynamic interplay between task weights ( z t ) for target languages (Figure 5, a) and domain weights ( α t ) for source languages (Figure 5, b) offers valuable insights into GRAPE's adaptive learning strategy and its implicit discovery of inter-lingual relationships. For instance, the initial prioritization of target languages Ukrainian (uk; East Slavic) and Romanian (ro; Romance) corresponds with an increased weighting of linguistically related source languages, namely Russian (ru; East Slavic) and Italian (it; Romance). Subsequently, the sustained high task weight for Polish (pl; Balto-Slavic) aligns with the continued high domain weight of Russian, from the same broad linguistic family. For Turkish (tr), a Turkic language without close phylogenetic relatives within the predominantly Indo-European source domains, its persistent learning challenge is evident from its high task weight. Nevertheless, GRAPE facilitates its learning by strategically up-weighting source languages like German (de) and French (fr), presumably due to their rich linguistic feature sets or similarity in syntax structures.

4 Discussion and Limitations

Progress measurements for Group DRO. We evaluate the impact of different step-wise progress metrics on the performance of Group DRO within the GRAPE framework. In addition to the primary Rate-of-Improvement (ROI) metric (Equation 1), we investigate two alternatives: Gap-ofImprovement ( GOI ) and EMA-Rate-of-Improvement ( ROI-ema ). The step-wise improvmenet of task T n at step t are assessed as:

The exponential moving average loss l t ema,n is updated as: l t ema,n = β · l t -1 ema,n +(1 -β ) · l n ( θ t ) , where the hyperparameter β is set to 0 . 7 in our experiments. Substituting these alternative progress metrics for r n in the minimax objective Equation 2 yields modified update rules according to Equation 8 and 9. The core optimization principle for adjusting task weights ( z ) and domain weights ( α ) remains, but the specific gradient terms within the exponents change. More details are provided in Appendix E.1.

Results on the six reasoning tasks (Figure 6) indicate that using GOI (GRAPE-gap) prioritizes intrinsically easier tasks. This is attributed to the fact that tasks with larger loss values can exhibit larger absolute improvements ( l n ( θ t ) -l n ( θ t +1 ) ) while their relative progress is slow. The ROI metric, by normalizing improvement with the current loss l i ( θ t ) , effectively mitigates this bias, offering a fair as-

4k

8k

12k

16k

20k

Steps

0.25

0.50

0.75

Weights

arc_easy

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

arc_challenge

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

sciq

4k

8k

12k

16k

20k

Steps

0.2

0.4

0.6

piqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

logiqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

sessment of task learning progress less biased to easy ones. In comparison, the ROI -ema metric, which normalizes the absolute improvement by an exponential moving average of past losses, introduces a historical dependency. While this can smooth out noisy step-wise improvements, our experiments suggest it may also lead to slower adaptation of task weights since the smoothed historical loss might not accurately reflect the current learning state. Comprehensive results across all 12 evaluated task combinations are provided in Appendix E, further substantiating these observations.

GRAPE outperforms baselines on diverse sets of target tasks. We further assess the robustness and adaptability of GRAPE on 12 diverse sets of target tasks with varying degrees of diversity and complexity. The task configurations and detailed results are presented in Appendix E. According to Figure 27, to support both code comprehension (Kodcode) and mathematical reasoning (GSM8K) alongside commonsense reasoning task (Hellaswag), GRAPE effectively balance code-specific sources (e.g. GitHub, Stackexchange) and technical web text (e.g. CommonCrawl and C4) for mathematical reasoning with broader corpora. With high task diversity on 8 tasks Figure 28 where commonsense reasoning tasks are predominant, GRAPE largely upweights of general knowledge sources like the Book domain to cater to the majority need, while maintaining a smaller yet significant allocation from domain-specific sources (e.g., GitHub, StackExchange). The comprehensive results in Table 2 and Appendix E demonstrate that GRAPE maintains strong average performance on

diverse task combinations with various complexities, outperforming baselines that struggle to cater to disparate learning objectives with fixed data mixtures and uniform task priorities.

Hyperparameter tuning. The performance and efficiency of GRAPE is dependent on its hyperparameters, including the regularization coefficients µ α and µ z , as well as the update frequencies ∆ T α and ∆ T z . Due to the high resource demands associated with LLM pretraining, an exhaustive hyperparameter sweep was beyond the scope of the current work. We posit that the performance of GRAPE can be further improved through extensive hyperparameter tuning, tailored to specific datasets and task configurations. How to determine a theoretical near-optimal step-sizes of domain weights and task weights updates is a direction for future work.

Impact of training data quality and domain granularity. Our investigations reveal that the quality of the pretraining corpus and the granularity of its domain partitions critically influence the efficacy of domain reweighting algorithms, including GRAPE, DOGE, and offline methods such as REGMIX and CRISP. On ClimbLab with higher-quality data and fine-grained semantic-based clustering, the domain weights exhibit more precise and impactful adjustments in response to the task curriculum compared to the coarser, source-defined domains of SlimPajama, where the benefits of reweighting can be less distinct. Notably, the offline reweighting algorithm such as RegMix and CRISP are prone to over-concentrated on some particular domains, which leads to biased performance in multi-task learning. These findings underscore the importance of meticulously partitioned data domains for the benefits of adaptive data mixing strategies, suggesting a potential direction for future research.

Towards better understandings of scaling effects on learning curriculum. The efficacy of domain reweighting are observed to vary significantly with model scale, besides the domain granularity of the training corpus. The 0.7B model tended to develop a highly concentrated focus on a single challenging task in later training stages, whereas the 125M model maintained a more distributed prioritization across several tasks (10). This suggests that the larger-scale model might be more sensitive to the conflicts between targets. Understanding these scaling effects, how the ideal curriculum changes with model size, and how its efficacy is bounded by data quality and partitioning, is crucial to developing more robust and universally effective adaptive pretraining strategies.

Towards a sample-level DRO framework for finer-grained reasoning. The current formulation of GRAPE applies Group Distributed Robust Optimization (DRO) at the task level, prioritizing data from entire target tasks that exhibit slower learning progress. A promising avenue for future research is the extension of this framework to sample-level DRO, where individual data instances (e.g., specific questions or examples) would be dynamically identified and assigned higher importance. By directing attention to the hardest samples, it allows the model to focus intensely on correcting specific errors in the mid- or post-training stages, or during interactions with evolving environment.

5 Conclusion

This paper introduced GRAPE, a novel multi-target domain reweighting algorithm designed to calibrate pretraining data mixture for robust performance across diverse target tasks. By dynamically adjusting domain and task weights through a Group DRO minimax framework, GRAPE effectively prioritizes tasks demonstrating the slowest improvement. Empirical results on challenging reasoning tasks and multilingual modeling showcase GRAPE's significant advantages over existing baselines, highlighting the efficacy of its adaptive curriculum for multi-target learning.

References

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020.

Together Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data .

Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining, 2023. URL https://arxiv.org/abs/2305.10429 .

Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training, 2025. URL https://arxiv.org/abs/2407.01492 .

Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation, 2024a. URL https://arxiv.org/abs/2310.15393 .

David Grangier, Simin Fan, Skyler Seto, and Pierre Ablin. Task-adaptive pretrained language models via clustered-importance sampling, 2025. URL https://arxiv.org/abs/2410.03735 .

Simin Fan, David Grangier, and Pierre Ablin. Dynamic gradient alignment for online data mixing, 2024b. URL https://arxiv.org/abs/2410.02498 .

Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning, 2020. URL https://arxiv.org/abs/2001.06782 .

Bo Liu, Yihao Feng, Peter Stone, and Qiang Liu. Famo: Fast adaptive multitask optimization, 2023. URL https://arxiv.org/abs/2306.03792 .

Christian Bohn, Ido Freeman, Hasan Tercan, and Tobias Meisen. Task weighting through gradient projection for multitask learning, 2024. URL https://arxiv.org/abs/2409.01793 .

John Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization, 2020. URL https://arxiv.org/abs/1810.08750 .

Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization, 2020. URL https://arxiv.org/abs/1911.08731 .

Qi Qi, Zhishuai Guo, Yi Xu, Rong Jin, and Tianbao Yang. An online method for a class of distributionally robust optimization with non-convex objectives, 2021. URL https://arxiv. org/abs/2006.10138 .

Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016. URL https://proceedings.neurips.cc/paper_files/paper/2016/ file/4588e674d3f0faf985047d4c3f13ed0d-Paper.pdf .

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL https://arxiv.org/abs/ 1706.03762 .

Shizhe Diao, Yu Yang, Yonggan Fu, Xin Dong, Dan Su, Markus Kliegl, Zijia Chen, Peter Belcak, Yoshi Suhara, Hongxu Yin, Mostofa Patwary, Yingyan, Lin, Jan Kautz, and Pavlo Molchanov. Climb: Clustering-based iterative data mixture bootstrapping for language model pre-training, 2025. URL https://arxiv.org/abs/2504.13161 .

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL https://arxiv.org/abs/1803.05457 .

Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions, 2017. URL https://arxiv.org/abs/1707.06209 .

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. URL https://arxiv.org/abs/1911.11641 .

Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning, 2020. URL https: //arxiv.org/abs/2007.08124 .

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. URL https://arxiv.org/abs/1905.07830 .

Mandy Guo, Zihang Dai, Denny Vrandeˇ ci´ c, and Rami Al-Rfou. Wiki-40B: Multilingual language model dataset. In Nicoletta Calzolari, Frédéric Béchet, Philippe Blache, Khalid Choukri, Christopher Cieri, Thierry Declerck, Sara Goggi, Hitoshi Isahara, Bente Maegaard, Joseph Mariani, Hélène Mazo, Asuncion Moreno, Jan Odijk, and Stelios Piperidis, editors, Proceedings of the Twelfth Language Resources and Evaluation Conference , pages 2440-2452, Marseille, France, May 2020. European Language Resources Association. ISBN 979-10-95546-34-4. URL https://aclanthology.org/2020.lrec-1.297/ .

Mayee F. Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-it! a data-driven skills framework for understanding and training language models, 2023. URL https://arxiv.org/abs/2307.14430 .

Hang Wang, Minghao Xu, Bingbing Ni, and Wenjun Zhang. Learning to combine: Knowledge aggregation for multi-source domain adaptation, 2020. URL https://arxiv.org/abs/2007. 08801 .

Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training, 2023. URL https://arxiv.org/abs/2312.02406 .

Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning, 2024. URL https://arxiv.org/abs/2110.14048 .

Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks, 2018. URL https:// arxiv.org/abs/1711.02257 .

Aviv Navon, Aviv Shamsian, Idan Achituve, Haggai Maron, Kenji Kawaguchi, Gal Chechik, and Ethan Fetaya. Multi-task learning as a bargaining game, 2022. URL https://arxiv.org/abs/ 2202.01017 .

Jean-Antoine Désidéri. Multiple-gradient descent algorithm (MGDA) for multiobjective optimization. Comptes Rendus Mathematique , 350(5-6):313-318, 2012.

Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies, 2024. URL https://arxiv.org/abs/2404.06395 .

Contents

A Related Work

Data Reweighting for LLM Pretraining Pretraining aims to create general-purpose models, requiring training on massive datasets. Data selection methods can be used to determine the optimal dataset composition according to various objective functions. They operate at different granularities with domain-level data selection methods allowing for less fine-grained control than token-level and sample-level methods but offering more tractable optimization that scales better to the massive datasets typical for LLMs. Learned approaches for domain reweighting either determine fixed domain weights [Xie et al., 2023, Fan et al., 2024a] or dynamically adjust the weights during training of the final model [Fan et al., 2024b, Chen et al., 2023, Wang et al., 2020, Albalak et al., 2023]. Fixedweight approaches rely on proxy models which require significant upfront compute and assume transferability of domain weights across model and training data scales. Notably, RegMix [Liu et al., 2025] trains hundreds of proxy models under tight computational budgets to fit a regressor predicting the optimal mixture for a held-out validation domain, then employs the top-ranked mixture during large-scale training; it relies on a single static surrogate objective (Pile-CC validation loss) that may not generalize to varied downstream tasks. CLIMB [Diao et al., 2025] similarly uses regression but first partitions the entire pretraining corpus into 20 semantically coherent 'domains' via unsupervised clustering of document embeddings and then runs a more efficient iterative mixture search to optimize their weights. DoGE [Fan et al., 2024a] uses a small proxy model to learn static domain weights based on bilevel optimization and gradient alignment, upweighting domains whose gradients most closely align with the chosen target domain. DGA [Fan et al., 2024b] extends this approach to the online scenario by periodically updating weights based on gradient alignment on the large model.

Domain reweighting is in general a two-step process, where we first define an objective that captures model performance, and secondly optimize domain weights according to that objective. Existing methods consider only single-objective scenarios that optimize validation loss of a single target as a proxy for downstream task performance. They lack mechanisms to balance improvements across multiple target tasks with competing objectives. Our approach provides a method to balance multiple tasks during pretraining, building on the gradient alignment score idea and using validation loss of downstream tasks as a surrogate to downstream task performance.

Gradient Surgery for Multi-task Learning. The multi-task learning (MTL) problem has been intensively explored in literature [Yu et al., 2020, Liu et al., 2024, Chen et al., 2018, Navon et al., 2022], while most of the proposed methods requires sufficient amount of training data directly from the target tasks. Given N ≥ 2 different tasks associated with loss function l i ( θ ) , the goal is to optimize model parameters θ that perform well across all target tasks. A standard objective for MTL is minimizing the average loss over all tasks. However, optimizing this averaged loss objective often lead to suboptimal and biased multi-tasking performance because of tasks with conflicting gradients or dominant gradients with large magnitudes [Yu et al., 2020, Liu et al., 2024]. Existing methods aim to avoid gradient conflict by gradient surgery, where the target gradients are linearly combined with adaptive weights Chen et al. [2018], Yu et al. [2020]. For instance, PCGrad [Yu et al., 2020] heuristically removes inter-task gradient conflicts via projections at cost O ( | θ | N 2 ) . CAGrad [Liu et al., 2024] strikes a balance between worst-task improvement and average loss reduction by solving a constrained optimization problem that requires O ( N 3 ) operations. MGDA [Désidéri, 2012] frames MTL as a multi-objective optimization and solves a quadratic programming problem, guaranteeing Pareto-optimal updates. NASH-MTL [Navon et al., 2022] offers equivalent Pareto guarantees, by treating gradient conflicts as a bargaining game that maximizes the product of task improvements (sum of log utility functions). All these methods share a common drawback : substantial memory requirements to store task gradients, O ( | θ | N ) , and at least O ( | θ | N 2 ) for the optimization step. These computational and memory costs become prohibitive for large-scale models and numerous tasks. FAMO [Liu et al., 2023] attempts to address these efficiency concerns. It dynamically weights task losses using O (1) space and time per iteration by trying to enforce equal loss decrease rate among tasks. We apply this idea in the domain reweighting setting to develop a multi-task-adaptive domain reweighting method.

B Derivation of GRAPE

We start with the minimax problem from Equation (2):

With the regularization terms defined as:

Where D Ψ is the Bregman divergence with Ψ( b ) = ∑ i b i log( b i ) :

where the last simplification follows because ∑ i a i = ∑ i b i = 1 for distributions. For the inner minimization over z , we form the Lagrangian:

Taking the partial derivative with respect to z i and setting to zero:

Solving for z i :

Let d t = ∑ k j =1 α ( t ) j g j ( θ t ) be the update direction. Then:

Using the constraint ∑ N i =1 z i = 1 to solve for λ :

Let Z z = ∑ N i =1 z ( t -1) i exp ( -γ t µ z ⟨∇ θ log l i ( θ t ) , d t ⟩ ) , then exp ( -λ + µ z µ z ) = 1 Z z .

Substituting back, we get the update rule for task weights:

Now, for the outer maximization over α , we substitute the optimal z ( t ) back into the original objective function. Let f ij = ⟨∇ θ log l i ( θ t ) , g j ( θ t ) ⟩ .

The resulting optimization problem becomes:

We form the Lagrangian for this maximization problem:

Taking the partial derivative with respect to α j and setting to zero:

Solving for α j :

Using the constraint ∑ k j =1 α j = 1 to solve for ν :

Let Z α = ∑ k j =1 α ( t -1) j exp ( γ t µ α ∑ N i =1 z ( t ) i f ij ) , then exp ( -ν + µ α µ α ) = 1 Z α .

Substituting back, we get the update rule for gradient weights:

Expanding f ij back to its full form:

C Theoretic Properties of GRAPE

C.1 Proof of the Convergence Theorem

Theorem C.1 (Convergence of GRAPE) . Let the loss functions l n ( θ ) be L -smooth for all n ∈ [ N ] and the norm of stochastic gradients be upper-bounded by G . If the learning rate γ t satisfies γ t ≤ 1 L and the regularization parameters µ α , µ z are chosen such that µ α > 0 and µ z > 0 , then the GRAPE algorithm with update rules given by the above equations converges to a neighborhood of the Pareto optimal solution at a rate of O (1 /T ) , where T is the number of training iterations. Specifically, for any ε > 0 , there exists T 0 such that for all T > T 0 :

Proof. We establish the convergence of GRAPE by analyzing the dynamics of the minimax optimization problem and demonstrating that the algorithm makes consistent progress toward Pareto optimal solutions while balancing performance across tasks. First, we analyze the task weight updates. Given that z is updated according to:

The task weights adjust to increase focus on tasks with slower improvement rates. For any task i where ⟨ d t , ∇ θ log l n ( θ t ) ⟩ < 0 (indicating poor alignment between the current update direction and task gradient), z t n increases proportionally. This ensures more attention to tasks making less progress.

For each iteration t , let L ( θ t ) = max n ∈ [ N ] l n ( θ t ) represent our worst-case objective. Let n t = arg max n ∈ [ N ] l n ( θ t ) be the index of the task with highest loss at iteration t . Due to the exponential update rule, as t increases, z t n t approaches 1, while other weights approach 0, directing optimization toward the current worst-case task. Considering the domain weight updates:

These updates increase weights for domains whose gradients align well with the task gradients weighted by z . Since z concentrates on the worst-performing tasks, α increasingly favors domains beneficial to these tasks. Given the L -smoothness of the loss functions, we can upper-bound the progress on the worst-case objective:

Let Q t = ⟨∇ θ l n t ( θ t ) , d t ⟩ . Due to our update rules for z and α , as t increases, Q t becomes increasingly positive (as domain weights shift toward domains beneficial for the worst-performing task). With the condition γ t ≤ 1 L , we have:

Since the domain weights are probability distributions and gradients are bounded (due to L -smoothness), ∥ d t ∥ is bounded by some constant G . Therefore:

As optimization progresses and Q t increases due to our reweighting strategy, we eventually reach a point where Q t > G 2 2 , ensuring consistent progress on the worst-case objective. Summing over iterations and applying the µ -strong convexity, we obtain:

This implies:

Given our reweighting strategy ensures Q t > G 2 2 after sufficient iterations, this difference becomes negative, establishing convergence.

For the convergence rate, with constant step size γ t = 1 L , we can show:

Where L ∗ = min θ max n ∈ [ N ] l n ( θ ) and θ ∗ is the corresponding optimal parameter. As T →∞ , this bound approaches G 2 2 L -¯ Q L , where ¯ Q is the average alignment. Therefore, GRAPE converges to a neighborhood of the Pareto optimal solution at a rate of O (1 /T )

C.2 Proof of the Variance Reduction Theorem

Theorem C.2 (Monotonic Variance Reduction of Task Performance) . Let σ 2 t = Var n ∈ [ N ] ( l n ( θ t )) denote the variance of task performances at iteration t . Let the loss functions l n ( θ ) be L -smooth for all n ∈ [ N ] and the norm of stochastic gradients be upper-bounded by G , and assuming the task losses are L -smooth and µ -strongly convex, the variance decreases monotonically until reaching a minimal basin, i.e., σ 2 t +1 ≤ σ 2 t for all t ≥ T 0 for some finite T 0 .

Proof. According to the definition of Rate-of-Improvement (RoI) for task i at iteration t :

For any task i , by the L -smoothness assumption, we can bound the change in loss:

Let ∆ l i ( t ) = l i ( θ t +1 ) -l i ( θ t ) represent the change in loss for task i .

First, we establish a direct relationship between task weights and improvement rates. At each iteration, GRAPE updates task weights according to:

it reflects the principle that tasks with lower RoI receive higher weights in the next iteration.

When domain weights are subsequently updated, domains that yield better improvements for these higher-weighted (struggling) tasks are upweighted:

This creates a feedback mechanism that specifically targets and improves tasks with lower RoI values. This indicates that tasks with lower RoI experience proportionally greater increases in their improvement rates. We can formalize this in Lemma 1 :

Lemma 1 : There exists a constant β > 0 such that if r ( t ) i < r ( t ) j for tasks i and j , then after GRAPE's reweighting mechanism, the expected improvement in the next iteration satisfies: E [ r ( t +1) i ] -E [ r ( t ) i ] > E [ r ( t +1) j ] -E [ r ( t ) j ] + β ( r ( t ) j -r ( t ) i )

We further define the normalized task losses:

We show that the variance of normalized loss ˆ l i ( t ) is proportional to the variance of losses:

Lemma 2 (Variance of Normalized Losses) : The variance of normalized losses ˆ σ 2 t = Var i ∈ [ N ] ( ˆ l i ( t )) is directly proportional to the variance of original losses σ 2 t = Var i ∈ [ N ] ( l i ( θ t )) .

Proof. We begin with the definition of normalized losses:

Let's denote a t = min n l n ( θ t ) and b t = max n l n ( θ t ) for brevity. The normalized loss can be written as:

This is an affine transformation of the original losses. By the properties of variance for affine transformations, for any random variable X and constants c and d , we have:

In our case, viewing the task losses as samples from a distribution, we have c = 1 b t -a t and d = -a t b t -a t . Thus, we have,

Therefore:

Which establishes a direct proportionality between ˆ σ 2 t and σ 2 t , with the proportionality constant 1 ( b t -a t ) 2 .

Given the relationship between task weights and RoI established in Lemma 1 , tasks with higher normalized losses tend to have lower RoI values and thus receive higher weights in GRAPE. This leads to a correlation between ˆ l i ( t ) and the subsequent change in normalized loss ∆ˆ l i ( t ) = ˆ l i ( t +1) -ˆ l i ( t ) .

Specifically, for any pair of tasks i and j where ˆ l i ( t ) > ˆ l j ( t ) , GRAPE's reweighting mechanism ensures:

Therefore, we can express the change in variance:

Given that the tasks with higher normalized loss will experience greater improvements on normalized loss under GRAPE, this correlation term is negative:

for some constant κ > 0 , representing the strength of GRAPE's balancing effect.

The second term, 1 N ∑ N i =1 (∆ ˆ l i ( t ) -∆ ¯ ˆ l t ) 2 , represents the variance of the changes in normalized losses, which is bounded by O ( γ 2 t ) under the smoothness assumption. Thus:

For sufficiently small γ t , this difference becomes negative, establishing that ˆ σ 2 t +1 < ˆ σ 2 t . Following Lemma 2 , it indicates that σ 2 t +1 < σ 2 t .

D Domain Reweighting for Multi-task Reasoning

D.1 Implementation Details

On small-scale experiments with 125 M models, we employ cosine learning-rate scheduler with maximum lr = 1 . 5 e -4 . On large-scale runs with 0 . 7 B models, we adopt the WSD scheduler [Hu et al., 2024] where the learning rate remains constant ( lr max = 1 . 5 e -4 ) during training while linearly decaying to lr min = 1 . 5 e -5 at the last 20% of total iterations. The model architecture and hyperparameter settings are detailed in Table 3. The full implementation of GRAPE is open-sourced in https://github.com/Olivia-fsm/GRAPE_data_mixture_with_multi_target .

GRAPE. For GRAPE-specific parameters, we generally set the regularization coefficients µ α = 1 e -4 and µ z = 1 . 5 e -5 , corresponding to the step-size of γ/µ α = 1 . 5 , γ/µ z = 10 . Both domain weights ( α ) and task weights ( z ) were periodically updated every ∆ T α = ∆ T z = 100 steps. Initial domain weights ( α 0 ) and task weights ( z 0 ) were set to uniform distribution, unless otherwise specified. Gradients required for the GRAPE weight updates (e.g., ∇ log l i ( θ t ) on target tasks, g i ( θ t ) on source domains) were estimated stochastically using dedicated mini-batches processed within the update step.

DOGE. We implement DOGE following the out-of-domain generalization setting in [Fan et al., 2024a]. Domain weights ( α ) are updated every ∆ T α = 100 steps during training according to the alignment of gradients between each training data domain and the average gradient across all target tasks, which is equivalent to applying a uniform task weights ( z ) constantly. We apply the same regularization hyperparameters as GRAPE.

CLIMBMIX. We set the domain weights following Diao et al. [2025]. Specifically, we utilize the final iteration weights derived by iterative regression and bootstrapping, which is claimed as the optimal mixture in the original paper.

GRAPE-CLIMBMIX. While the GRAPE algorithm initializes domain weights ( α 0 ) as the uniform distribution, we implement a variant of GRAPE with α 0 initialized with the optimized CLIMBMIX weights [Diao et al., 2025]. The other hyperparameters for reweighting and optimization are kept consistent as GRAPE.

DOGE-PCGRAD. We tailor the DOGE [Fan et al., 2024a] method for multi-task learning by dynamically tune the target gradient with a gradient surgery method PCGrad [Yu et al., 2020], where the gradients from each target task are combined to avoid gradient conflicts. PCGrad eliminates conflicting gradient components by projecting each task gradient onto the normal plane of any conflicting task gradient. Conflict is indicated by a negative dot product between target gradients. We process gradients iteratively, taking each task gradient in a randomized sequence and adjusting it when conflicts are detected, continuing until all gradients have been properly modified. The random ordering of gradient processing helps prevent systematic bias toward specific tasks. The final target gradient, g PCGrad , is calculated as the average of the processed target gradients. At each reweighting step to calculate a conflict-free target gradient, g PCGrad . Domain weights ( α ) are updated every ∆ T α = 100 steps during training according to the alignment of gradients between each training data domain and g PCGrad . We apply the same regularization hyperparameters as GRAPE. With N target tasks, each reweighting step requires storing N gradients into memory, which incures N ×| θ | storage costs. Thus, the traditional gradient surgery algorithm for multi-task learning can lead to OOMissue due to the significant GPU memory load, which prevent it from scaling up to large models and datasets.

REGMIX. We implement REGMIX [Liu et al., 2025] to predict the optimal domain weights ( α ) by fitting a regression model. Following Liu et al. [2025], we train a total of 768 TinyLlama-1M proxy models with with various data mixture configurations: 512 of them are used to fit the regression model, while the other 256 of them are saved for evaluating the performance of the regressor. For each experiment run, we record the average validation loss across all target tasks; we then train the LightGBM regressor to predict the input domain weights which achieves the lowest average loss. The other hyperparameters for proxy model training are set to the default values in [Liu et al., 2025].

CRISP We compute embedding-based domain weights using CRISP [Grangier et al., 2025]. For each train and target domain we calculate the sequence embeddings using SBERT MiniLM-L6-v2 . To avoid overfitting, we select 5 representative centroids for each training domains D k , k ∈ [ K ] respectively using k-means. We then obtain the domain weights by applying a nearest-neighbor ( k = 5 ) classification on each piece of the samples from all target sets, i.e. mapping every individual target embedding to centroids within each training domains. We finally compute the distribution of target samples located in each training domain to get the final domain weights.

D.2 Results on ClimbLab

D.2.1 Description of Dataset

ClimbLab is a dataset introduced by Diao et al. [2025], which is partitioned into 20 clusters according to semantic-based sequence embeddings. The related topic on each cluster (i.e. train domain) are detailed in Table 4.

D.2.2 Learning Curriculum from Various Reweighting Approaches

Domain Weights from RegMix. We present the domain weights from RegMix in Figure 7 and Figure 8, with different sets of target tasks. The RegMix regressor ends up concentrating almost all of its weight to just two clusters. In both settings, Cluster 17 (advanced biomedical and clinical research content) receives the highest weight. Cluster 15 (film and cultural-arts) gets picked up for the general reasoning benchmarks showing that its narrative and conceptual text aligns with the story-and-commonsense inference style those tasks need. Including MedQA/MathQA, increases the weight of Cluster 9 (space science, urban planning), showing that its physics-heavy, formula-rich science content align well with the formal quantitative reasoning that MathQA requires.

0.6

Dataset Prior

RegMix

0.5

0.4

Weight

0.3

0.2

0.1

0.0

1234567891011121314151617181920

Domain (Cluster)

(a) Original vs. RegMix weights

6.50

Corr:0.98

6.25

LoSS

6.00

True

5.75

5.50

5.25

5.00

5.00

5.25

5.505.75

Predicted Loss

6.00

6.25

6.50

(b) Spearman correlation between predicted v.s. true loss

Domain Weights from CRISP. We present the domain weights from CRISP in Figure 9, with different sets of target tasks. In both settings CRISP assigns similar weights to the dataset clusters. Main focus is given on code data (Cluster 11), while Cluster 6 with technical life-science and automation content also receives significant attention. Including MedQA and MathQA, slightly

(a) Original vs. RegMix weights

(b) Correlation between Predicted

0.6

DatasetPrior

RegMix

0.5

6.50

0.4

6.25

Weight

0.3

SS07

6.00

True

5.75

0.2

5.50

0.1

5.25

5.00

0.0

Domain (Cluster)

5.00

5.25

5.50

Predicted Loss

5.75

6.00

6.25

6.50

v.s. True loss

boosts the weight of Cluster 10 (health and wellness content), reflecting MedQA's embeddings lie closest to clinical-style text.

(a) 6 Target Tasks

(b) 8 Target Tasks

0.25

0.20

0.20

0.15

0.15

0.10

0.10

0.05

0.05

0.00

1011

Clusters

121314151617181920

0.00

8

1011

Clusters

121314151617181920

Cluster 11: Software Development

Cluster 9: Astronomy & Space

Cluster 11: Software Development

Cluster 4: Education & Science

Cluster 6: Biotechnology & Healthcare

Cluster 12: Technology & Legal Content

Cluster 6: Biotechnology & Healthcare

Cluster 12: Technology & Legal Content

Cluster 4: Education & Science

Other clusters

Cluster 9: Astronomy & Space

Other clusters

Domain Weights Evolution from GRAPE and Scaling Effects. We observed that the task priority differs notably between model sizes. According to Figure 10, the 0 . 7 B parameter model tended to develop a highly concentrated focus on a single challenging task in later training stages, whereas the 125 Mmodel maintained a more distributed and stable prioritization across several tasks. This suggests that the larger-scale model might be more sensitive to the conflicts between targets. Additionally, in both cases, the task weight trajectories demonstrate a clear, stage-wise curriculum, which indicates a dynamic, adaptive curriculum is crucial to yield a better multi-task learning performance.

0

10k

20k

30k

40k

50k

Optimization Steps

0.0

0.2

0.4

0.6

0.8

Weights

(a) 125M-ClimbLab

0

25k

50k

75k

100k

Optimization Steps

0.0

0.2

0.4

0.6

0.8

Weights

ARC-E

ARC-C

SciQ

PIQA

LogiQA

(b) 0.7B-ClimbLab

D.3 Results on SlimPajama

We present the experimental results on SlimPajama in this section. The hyperparameter and baseline settings are constant with the experiments on ClimbLab.

D.3.1 Learning Curriculum from Various Reweighting Approaches

Domain Weights from REGMIX. Wepresent the domain weights derived from RegMix in Figure 11 and Figure 12, on various sets of target tasks. In both cases, the RegMix regressor learns a very spiky training mixture, where a large percentage ( ∼ 80 %) are allocated to web-crawled dataset such as C4, while nearly eliminates the domains other than C4, CC, and Book. Compared to 6 common-sense and logical related reasoning tasks, including domain-specific QA tasks - MathQA and MedQA - further increases the weights on C4, which can be attributed to its broad coverage on various knowledge sources.

(a) Original vs. RegMix weights

(b) Predicted vs. true loss

1.0

Dataset Prior

RegMix

7.25

Corr:0.98

0.8

7.00

Weight

0.6

6.75

6.50

Loss

0.4

6.25

0.2

6.00

0.0

5.75

5.50

5.25

5.4

5.6

5.8

6.0

6.2

6.4

Domain

Predicted Loss

(a) Original vs. RegMix weights

(b) Predicted vs. true loss

1.0

Dataset Prior

■

RegMix

7.25

Corr:0.98

0.8

7.00

Weight

0.6

6.75

True Loss

6.50

0.4

6.25

0.2

6.00

0.0

5.75

5.50

5.4

5.6

5.8

6.0

6.2

6.4

Domain

Predicted Loss

Domain Weights from CRISP. We present the domain weights derived from CRISP in Figure 13 on various sets of target tasks. CRISP assigns most of the weight to Github, followed by StackExchange. All other domains contribute minimally. A plausible rationale behind the CRISP weight distribution is that short Q&A tasks (ARC-E/C, SciQ, PIQA) embed closest to the concise question-answer format of StackExchange, while multi-step reasoning puzzles (LogiQA) and formula-heavy problems (MathQA) align more with GitHub's structured code snippets, and MedQA's longer clinical passages slightly with Book and arXiv. Including MedQA and MathQA brings clinical and formulaic embeddings that shift a bit of weight back toward Book and arXiv.

(a) 6 Target Tasks

(b) 8 Target Tasks

Top domains by weight

0.5

Top domains by weight

0.5

Other domains

Other domains

0.4

0.4

Weights

Weights

0.3

0.3

0.2

0.2

0.1

0.1

0.0

0.0

arxiv

book

github

book

github

wikipedia

wikipedia

stackexchange

stackexchange

D.3.2 Results on Various Task Configurations

Reweighting with 6 Target Tasks: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA. We present the full results on multi-task reasoning experiment, where the data mixture from SlimPajama dataset are adapted towards the optimal performance on 6 target tasks: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA. The Log-perplexity and 5-shot accuracy scores on each target reasoning tasks are presented in Figure 14, Table 5 and Table 6, where GRAPE consistently achieves better (lower) perplexity compared to the other 5 baselines. In contrast, CRISP, as an purely embedding-based reweighting algorithm, leads to significantly worse results on all benchmarks, in terms of perplexity and 5-shot accuracy scores.

arc_challenge

arc_easy

Log Perplexity

5.4

5.4

5.0

5.2

5.2

4.8

+DoGE

RegMix

5.0

5.0

CRISP

4.6

DoGE-PCGrad

4.8

4.8 -

GRAPE

4.6

4.6

4.4

4.4

4.4

4.2

4.2

4.2

4.0

4.0

4.0

X8

12k

16k

20k

X8

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Steps

logiqa

piqa

sciq

Log Perplexity

5.00

4.4

5.4

4.75

5.2

4.2

4.50

5.0

4.25

4.0

4.8 -

4.00

3.8

4.6

3.75

4.4

3.50

3.6

4.2

3.25

3.4

4k

8k

12k

16k

20k

4k

78

T

12k

16k

20k

4k

8k

T

12k

16k

20k

Steps

Steps

Steps

Reweighting with 8 Target Tasks: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA. We present the full results on multi-task reasoning experiment, where the data mixture from SlimPajama dataset are adapted towards the optimal performance on 8 target tasks: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA, covering a board topics and knowledge domains. The Log-perplexity and 5-shot accuracy scores on each target reasoning tasks are presented in Figure 15, Table 8 and Table 7, where GRAPE consistently outperforms the other 5 baselines on most of benchmarks (6 out of 8), while achieving comparable performance as RegMix and DoGEPCGrad on Hellaswag and PIQA. Notably, GRAPE improves the perplexity scores on MedQA by a large margin above all the other baselines, indicating its distinct adaptability to unique target tasks with domain-specific skill and knowledge. In contrast, CRISP demonstrates marginal improvements compared to uniform sampling baseline, which demonstrates its weakness in the multi-task learning setting.

arc_challenge

arc_easy

logiqa

Log Perplexity

5.4

5.4

5.0

4.4

Uniform

5.2

5.2

4.8

4.2

CRISP

DoGE

RegMix

5.0

5.0

4.6

DoGE-PCGrad

4.8 -

4.8

4.4

4.0

GRAPE

4.6

3.8 -

4.6

4.2

4.4

4.4

3.6

4.2

4.0

4k

X8

12k

16k

20k

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Steps

Steps

mathqa

medqa

piqa

sciq

Log Perplexity

5.4

5.0

5.4

5.0

4.5

5.2

4.8

5.2

5.0

4.6

5.0

4.0

4.8

4.4

4.8

3.5

4.6

4.2

4.6

3.0

4.4

4.0

4.2

3.8

4.4

2.5

3.6

4k

X8

12k

16k

20k

4k

X8

12k

16k

20k

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Steps

Steps

D.4 Results on Wiki-40b

We present detailed results on the multilingual pretraining experiments on Wiki-40B dataset in the following sections. Specifically, we experiment with two various target configurations: (1) target at 4 low-resource languages: Catalan, Danish, Romanian, Ukrainian; (2) target at 8 low-resource languages: Catalan, Danish, Romanian, Ukrainian, Polish, Portuguese, Turkish, Dutch, i.e. the setting we present in the main paper (§ 3.2).

D.4.1 Language Mixture from Various Reweighting Algorithms

Domain Weights from RegMix. We present the optimized domain weights distributions from RegMix in Figure 16 and 17. Specifically, with 4 target languages, RegMix upweights German (de) and Spanish (es) while reducing the weights assigned on French (fr) and Italian (it). With 8 target languages, RegMix significantly increases the domain weights on Russian (ru), while further decreases the proportion of English (en).

(a) Original vs. RegMix weights

0.6

Dataset Prior

RegMix

0.5

0.4

Weight

8'0

0.2

0.1

0.0

en

fr

de

es

it

r

Domain

10.0

True Loss

9.5

9.0

8.5

8.50

8.75

9.00

9.25

9.50

9.75

10.0010.25

Predicted Loss

(b) Spearman correlation between predicted v.s. true loss

(a) Original vs. RegMix weights

(b) Spearman correlation between predicted

0.6

Dataset Prior

RegMix

10.0

0.5

9.8

0.4

True Loss

9.6

Weight

0.3

9.4

0.2

9.2

0.1

9.0

8.8

0.0

en

fr

de

es

it

ru

9.0

9.2

9.4

9.6

9.8

Domain

Predicted Loss

v.s. true loss

Domain Weights from CRISP. We present the optimized domain weights distributions from CRISP in Figure 18, with various task configurations. Notably, CRISP yields very similar domain weights across 6 languages, where English (en) and Italian (it) are mostly upweighted, while they are proved not helpful for multi-language learning according to Figure 19 and 20.

(a) 4 target languages

(b) 8 target languages

0.30

Highest weighted domain

0.30

Highest weighted domain

Other domains

Other domains

0.25

0.25

Weights

0.20

Weights

0.20

0.15

0.15

0.10

0.10

0.05

0.05

0.00

0.00

de

es

de

D.4.2 Results on Various Task Configurations

Reweighting with 4 target languages: Catalan, Danish, Romanian, Ukrainian Figure 19 present the full results on the multilingual pretraining experiments, where the training langauge mixture are adapted to 4 target low-resource languages: Catalan, Danish, Romanian, Ukrainian. GRAPE significantly outperforms all the other baseline methods across all target languages. Notably, none of the other reweighting approaches can effectively facilitate the learning on Romanian, while only GRAPE unleashes the learning capability on it.

Catalan

Danish

Log Perplexity

6.9

8.8

6.8 -

8.6

6.7

6.6 -

8.4

6.5

8.2

6.4

8.0

6.3

6.2

7.8 +

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Romanian

Ukrainian

Perplexity

8.4

Uniform

8.2

DoGE

RegMix

7

8.0

CRISP

DoGE-PCGrad

7.8

GRAPE

7.6

5

4+

7.2

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Reweighting with 8 target languages: Catalan, Danish, Romanian, Ukrainian, Polish, Portuguese, Turkish, Dutch. Figure 20 present the full results on the multilingual pretraining experiments, where the training langauge mixture are adapted to 8 target low-resource languages: Catalan, Danish, Romanian, Ukrainian, Polish, Portuguese, Turkish, Dutch. GRAPE significantly outperforms all the other baseline methods across all 8 target languages. In contrast, the offline reweighting algorithm such as RegMix and CRISP exhibit a very biased performance across various target languages: RegMix accelerate the learning on Catalan, Portuguese, Ukrainian and Polish, while sacrificing the performance on Dutch and Danish; CRISP only accelerates the learning on Turkish while sabotaging all the other languages.

Catalan

Danish

Dutch

Polish

Log Perplexity

6.9

8.8

8.4

6.8

9.4

Uniform

8.6

8.2

9.2

DoGE

RegMix

8.0

9.0

CRISP

6.6

8.4

7.8

8.8

DoGE-PCGrad

GRAPE

6.5

8.2

98

7.6

6.4

8.4

6.3

8.0

7.4

8.2

6.2

7.8 +

7.2 +

8.0+

4k

X8

12k

16k

20k

4k

8k

12k

16k

20k

4k

X8

12k

16k

20k

4k

8k

12k

16k

20k

Steps

Steps

Steps

Steps

Portuhese

Romanian

Turkish

Ukrainian

Log Perplexity

7.1

9

9.6

8.4

8-

8.2

6.9

9.4

8.0

6.8 -

6

9.2

7.8

6.7

5

9.0

7.6

6.6

4

6.5

3

8.8

7.4

6.4

4k

8k

12k

16k

20k

4k

8k

12k

16k

20k

4k

X8

12k

16k

20k

7.2 +

4k

X8

T

12k

16k

20k

Steps

Steps

Steps

Steps

E Ablation Study

In this section, we present the full results on our ablation experiments, with three various progress measurement metrics for Group DRO, and 12 different target task configurations.

E.1 Progress Assessment Metrics for Group DRO

We evaluate the impact of different step-wise progress metrics on the performance of Group DRO within the GRAPE framework. In addition to the primary Rate-of-Improvement (ROI) metric (Equation 1), we investigate two alternatives: Gap-of-Improvement ( GOI ) and EMA-Rate-of-Improvement ( ROI-ema ). The step-wise improvmenet of task T n at step t are assessed as:

The exponential moving average loss l t ema,n is updated as: l t ema,n = β · l t -1 ema,n +(1 -β ) · l n ( θ t ) , where the hyperparameter β is set to 0 . 7 in our experiments. Substituting these alternative progress metrics for r n in the minimax objective Equation 2 yields modified update rules according to Equation 8 and 9. The core optimization principle for adjusting task weights ( z ) and domain weights ( α ) remains, but the specific gradient terms within the exponents change.

GOI :

ROI-ema :

Task Configurations. We evaluate the efficacy of GRAPE on 12 various target task combinations as follows:

T1: GSM8K, ARC-C, ARC-E

T2: GSM8K, Hellaswag

T3: GSM8K, PIQA

T4: GSM8K, LogiQA

T5: GSM8K, SciQ

T6: GSM8K, ARC-E, ARC-C, Kodcode

T7: GSM8K, Kodcode, Hellaswag

T8: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, Kodcode, GSM8K

T9: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA

T10: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA

T11: ARC-E, ARC-C, MathQA, MedQA

T12: LogiQA, Hellaswag, MathQA, MedQA

Results. We present the average and worst log-perplexity scores on each of 12 groups of ablation experiments in Table 9. Specifically, GRAPE outperforms uniform and DoGE across all task sets; on 7 out of 12 tasks, GRAPE with ROI metric outperforms the other two DRO variants with GOI and ROI -ema , respectively. We present the log-perplexity scores associated with task weights evolution trajectories for all 12 groups of ablations in Figure 21-32.

E.2 Task Combination 1: GSM8K, ARC-C, ARC-E

8k

12k

16k

20k

Steps

3.5

3.6

3.7

3.8

3.9

Log Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

5.0

5.5

6.0

arc_easy

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

5.0

5.5

arc_challenge

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

0.4

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.3

0.4

0.5

0.6

0.7

0.8

arc_easy

4k

8k

12k

16k

20k

Steps

0.1

0.2

0.3

0.4

0.5

DoGE

arc_challenge

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.10

0.15

0.20

0.25

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.3 Task Combination 2: GSM8K, Hellaswag

8k

12k

16k

20k

Steps

3.6

3.8

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

3.75

4.00

4.25

4.50

4.75

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.6

0.8

1.0

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.1

0.2

0.3

0.4

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.4 Task Combination 3: GSM8K, PIQA

Weights

8k

12k

16k

20k

Steps

3.6

3.7

3.8

3.9

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

4.00

4.25

4.50

4.75

5.00

piqa

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.4

0.6

0.8

1.0

DoGE

piqa

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.1

0.2

0.3

0.4

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.5 Task Combination 4: GSM8K, LogiQA

Weights

8k

12k

16k

20k

Steps

3.6

3.8

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

5.0

logiqa

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.4

0.6

0.8

1.0

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

DoGE

logiqa

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.1

0.2

0.3

0.4

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.6 Task Combination 5: GSM8K, SciQ

Weights

8k

12k

16k

20k

Steps

3.6

3.8

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

sciq

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.4

0.6

0.8

1.0

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

DoGE

sciq

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.05

0.10

0.15

0.20

0.25

0.30

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.7 Task Combination 6: GSM8K, ARC-E, ARC-C, Kodcode

8k

12k

16k

20k

Steps

3.5

3.6

3.7

3.8

3.9

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

5.0

arc_easy

4k

8k

12k

16k

20k

Steps

3.5

4.0

4.5

5.0

arc_challenge

4k

8k

12k

16k

20k

Steps

1.6

1.8

2.0

2.2

2.4

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.1

0.2

0.3

0.4

0.5

0.6

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.2

0.3

0.4

0.5

0.6

0.7

arc_easy

4k

8k

12k

16k

20k

Steps

0.2

0.3

0.4

arc_challenge

4k

8k

12k

16k

20k

Steps

0.00

0.05

0.10

0.15

0.20

0.25

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.05

0.10

0.15

0.20

0.25

0.30

0.35

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.8 Task Combination 7: GSM8K, Kodcode, Hellaswag

4k

8k

12k

16k

20k

Steps

3.6

3.8

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

1.50

1.75

2.00

2.25

4k

8k

12k

16k

20k

Steps

3.75

4.00

4.25

4.50

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

4k

8k

12k

16k

20k

Steps

0.4

0.6

0.8

1.0

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.05

0.10

0.15

0.20

0.25

0.30

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.9 Task Combination 8: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, Kodcode, GSM8K

4k

8k

12k

16k

20k

Steps

2.0

2.5

3.0

3.5

4.0

Perplexity

gsm8k

4k

8k

12k

16k

20k

Steps

4.4

4.6

4.8

arc_easy

8k

12k

16k

20k

Steps

4.2

4.4

4.6

4.8

arc_challenge

4k

8k

12k

16k

20k

Steps

1.6

1.8

2.0

2.2

2.4

8k

12k

16k

20k

Steps

3.8

4.0

4.2

4.4

4.6

sciq

8k

12k

16k

20k

Steps

4.2

4.4

4.6

4.8

piqa

8k

12k

16k

20k

Steps

3.6

3.7

3.8

3.9

4.0

4.1

logiqa

8k

12k

16k

20k

Steps

4.0

4.1

4.2

4.3

4.4

4.5

4.6

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.000

0.025

0.050

0.075

0.100

0.125

0.150

Weights

gsm8k

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

0.8

1.0

arc_easy

4k

8k

12k

16k

20k

Steps

0.00

0.05

0.10

0.15

0.20

arc_challenge

4k

8k

12k

16k

20k

Steps

0.00

0.02

0.04

0.06

0.08

0.10

0.12

4k

8k

12k

16k

20k

Steps

0.00

0.05

0.10

0.15

0.20

sciq

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

piqa

4k

8k

12k

16k

20k

Steps

0.000

0.025

0.050

0.075

0.100

0.125

0.150

logiqa

4k

8k

12k

16k

20k

Steps

0.00

0.05

0.10

0.15

0.20

0.25

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.10 Task Combination 9: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA

4k

8k

12k

16k

20k

Steps

4.25

4.50

4.75

Perplexity

arc_easy

4k

8k

12k

16k

20k

Steps

4.00

4.25

4.50

4.75

arc_challenge

4k

8k

12k

16k

20k

Steps

3.5

4.0

sciq

4k

8k

12k

16k

20k

Steps

4.2

4.4

4.6

4.8

piqa

8k

12k

16k

20k

Steps

3.6

3.8

logiqa

8k

12k

16k

20k

Steps

4.2

4.4

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.25

0.50

0.75

Weights

arc_easy

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

arc_challenge

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

sciq

4k

8k

12k

16k

20k

Steps

0.2

0.4

0.6

piqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

logiqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.05

0.10

0.15

0.20

0.25

0.30

0.35

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.11 Task Combination 10: ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

Perplexity

arc_easy

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

arc_challenge

4k

8k

12k

16k

20k

Steps

4.0

4.5

5.0

sciq

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

piqa

4k

8k

12k

16k

20k

Steps

4

5

logiqa

4k

8k

12k

16k

20k

Steps

4.0

4.5

5.0

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

mathqa

4k

8k

12k

16k

20k

Steps

4

6

medqa

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

Weights

arc_easy

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

arc_challenge

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

sciq

4k

8k

12k

16k

20k

Steps

0.2

0.4

piqa

4k

8k

12k

16k

20k

Steps

0.00

0.05

0.10

0.15

logiqa

4k

8k

12k

16k

20k

Steps

0.1

0.2

4k

8k

12k

16k

20k

Steps

0.25

0.50

0.75

mathqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

DoGE

medqa

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.10

0.15

0.20

0.25

0.30

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.12 Task Combination 11: ARC-E, ARC-C, MathQA, MedQA

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

Perplexity

arc_easy

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

arc_challenge

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

mathqa

4k

8k

12k

16k

20k

Steps

3

4

5

6

medqa

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.2

0.4

0.6

0.8

1.0

Weights

arc_easy

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

arc_challenge

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

mathqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

DoGE

medqa

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

Steps

0.10

0.15

0.20

0.25

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

(c) Domain Weights

E.13 Task Combination 12: LogiQA, Hellaswag, MathQA, MedQA

4k

8k

12k

16k

20k

Steps

3.50

3.75

4.00

4.25

4.50

Perplexity

logiqa

4k

8k

12k

16k

20k

Steps

4.0

4.5

5.0

4k

8k

12k

16k

20k

Steps

4.5

5.0

5.5

mathqa

4k

8k

12k

16k

20k

Steps

3

4

5

6

medqa

Uniform

DoGE

GRAPE

GRAPE-ema

GRAPE-gap

(a) Log-Perplexity Scores

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

Weights

logiqa

4k

8k

12k

16k

20k

Steps

0.2

0.4

0.6

0.8

1.0

4k

8k

12k

16k

20k

Steps

0.0

0.2

0.4

0.6

0.8

mathqa

4k

8k

12k

16k

20k

Steps

0.0

0.1

0.2

0.3

0.4

DoGE

medqa

GRAPE

GRAPE-ema

GRAPE-gap

(b) Task Weights

4k

8k

12k

16k

20k

0.10

0.15

0.20

Weights

arxiv

book

cc

c4

github

stackexchange

wikipedia

Steps

(c) Domain Weights

F Hyperparameter Tuning

We perform a hyperparameter search on the reweighting intervals ∆ T α , ∆ T α = 50 , 100 , 200 for both GRAPE and DoGE. We train 125 Mmodel for 50 k steps and report the final task accuracies and run time as in Table 10. These results show that GRAPE consistently outperforms DoGE across all tested update frequencies. We also observe that the most frequent updates ( ∆ 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitation of the work is discussed in section 4.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The proofs of Theorem 2.1 and Theorem 2.2 are provided in subsection C.1 and subsection C.2.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide the implementation details for GRAPE and baselines in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the anonymous link to the codebase. All data used in the experiments are open-sourced.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All implementation details are included in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: error bars are not reported because the pretraining experiments would be too computationally expensive

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We used 4 × H100 80GB GPUs for large-scale experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We confirm that we follow the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [No]

Justification: The work does not have significant societal impact since it mostly focuses on the data selection aspect for LLM pretraining.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer:[No]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: we properly cite the original paper that produced the code package or dataset.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: the paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: that the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: the LLM is used for editing and formatting purpose for this paper. The core method development in this work does not involve LLMs as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
comments: SUMMARY:
This paper addresses the challenge of balancing training across multiple source and target domains in language model pretraining. It optimizes the Distributed Robust Optimization (DRO) algorithm by using the rate of improvement on test sets as a metric, rather than just training set loss. This allows the method to dynamically adjust task weights to prioritize more challenging tasks, and subsequently adjust domain weights to better achieve this learning objectives. The efficacy of this approach is demonstrated through experiments in both monolingual and multilingual environments.

STRENGTHS AND WEAKNESSES:
Strengths:
1. Simplicity and Theoretical Soundness: The proposed method is straightforward and direct, backed by theoretical guarantees that support its overall effectiveness.
2. Clarity of Writing: The paper is well-written, making the methodology and results easy to understand.

Weaknesses:
1. Code Availability: The authors do not provide code for their implementation.
2. Limited Test Set Scope: The experiments utilize a relatively small number of test sets.

QUESTIONS:
1. Generalization to Unseen Tasks: The primary objective of pretraining LLMs extends beyond achieving superior performance on a predefined set of tasks (six in the main experiments of this paper). LLMs are also expected to exhibit strong generalization capabilities. How does this training method perform on tasks not included in the training process, particularly those requiring different types of abilities or knowledge?
2. Applicability to Broader Scenarios: Can this method be effectively extended to broader practical applications, such as instruction tuning or continued pretraining?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper points out that pretraining has two goals: to reduce the loss on domains in the pretraining corpus, and to improve on downstream tasks. The authors introduce GRAPE, an algorithm that balances these two goals. GRAPE works similar to DOGE, using the DOGE update rule for both pretraining domain weights and downstream task weights.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The core points of the paper---that pretraining aims to improve performance on many downstream tasks, and that existing algorithms don't explicitly do this---are well made. The authors also make thoughtful evaluations on the curriculum and multilingual training. The results are strong and GRAPE seems to work. I think GRAPE is a well-motivated and natural evolution of DoReMi and DoGE.\n\nWeaknesses: In using DoGE update rule twice, the GRAPE algorithm naturally more hypeparameters than DoGE. This can potentially be a challenge for hyperparameter turning.\n\nQuality: The paper is impressively thorough.\nClarity: I found the paper easy to follow. I clearly understood both GRAPE's connection to and clear separation from prior work.\nSignificance: To my knowledge, this is the first algorithm that adjusts both domain and downstream task weights.\nOriginality: Ideas build on prior work but are applied to a new setting.\n\nQUESTIONS:\n1. Do you have feedback for the reader on tuning hyperparameters? Hyperparameter tuning seems critical here because there are many of them.\n2. What do you think of other data mixing methods like Data Mixing Laws, Aioli, and Adaptive Data Optimization? These methods aren't derived from the DRO framework. I'd be interested in seeing a little commentary what the benefit of the minimax DRO framework is, over just a minimization framework seen in the methods above and RegMix.\n3. Some information about wall-clock time would be helpful. Although GRAPE is an online method, it will take longer than normal pretraining because you're taking lots of gradients on validation sets.\n4. Relatedly, how would GRAPE interact with pretraining corpora that have many partitions, like DCLM-Baseline or WebOrganizer? In these cases it seems like you would have many validation gradients, in which case you would be decreasing your throughput because you have to allocate more of each batch to computing validation gradients.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel method for dynamically rebalancing task weights in a bilevel optimization for data weighting where multiple tasks are simultaneously being optimized for. This approach significantly outperforms previous methods and baselines on a variety of experiments. The authors present interesting experiment setups which allow for deeper analysis of the functioning of the method. The paper concludes by presenting many compelling directions for extending this work.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper is clearly written, the algorithm presented is a novel approach to the problem. The results are very compelling. The presented analyses are very thorough and provide fascinating insight into the function of the method. The result is very likely to be used by the research community in future work.\n\nWeaknesses:\n\nmagical hyperparams:\n(This is bit of a nit, and does not detract terribly from the overall high quality of this work)\nIt makes sense that ∆T_z and ∆T_alpha should not be 1, for both computational cost and diminishing retuns on frequency of update. It would be very useful however to have some sense of what this tradeoff looks like in practice rather than simply punting the hyper-param tuning (lines 328-334). It does not need to be exhaustive, but it would be nice to include *some* analysis of this tradeoff. Did the authors ever run any exp with ∆T == 50? == 200? A simple sweep on a tiny model would suffice.\n\nQUESTIONS:\nThe authors present an approach where a first-order taylor expansion is used to approximate RoI in order to provide the setup for applying GroupDRO. Did the authors impelment the simpler version of just calculating RoI directly per-step and updating the task-weights via a simple heuristic (ie using the softmax of the inverse RoIs)? This seems a simple and efficient baseline that would be useful to compare against before getting more complicated.\n\nIn section 4, much space is spent on comparing GRAPE to GRAPE-ema and GRAPE-gap (l 285-314, fig 6, table 2). This is fine to establish that GRAPE is the better of the three methods but not terribly interesting beyond that. Relatively little space is spent on the next section (l 315-326) discussing the relative performance of GRAPE under different task distribution settings, where the results are referred to in the appendix. Is this not critical information to understand both how the method is functioning and in what settings it might be useful to be applied? Consider showing less of the GRAPE-ema comparison and more of the GRAPE-in-different-task-settings. \n\nDid the authors run any experiments where the relative improvements between tasks are not uniformly weighted? In practical terms, practitioners may prioritize some tasks over others in their downstream evaluations, so it would be very useful and interesting to know if the strong performance of GRAPE extended to these settings. Do the authors have any comments / intuitions on this?\n\nNits:\nline 288:  typo on \"improvement\"\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper docuses on optimizing data mixtures in pretraining language models with the goal of improving performance of the models on multiple target tasks at the same time. This is an extension of past research which set this up mostly with the goal of optimizing perormance on a single target task. \n\nThe paper introduces a method to adjust dynamically both sample weights from source domains, as well as task weights. Adjusting task weights dynamically allows the model to prioritize different tasks at different times in training. The learning problem is formalzied as a minimax optimization problem, where the task weights are selected using group distributed-robest-optimization and the domain weights are maximizing loss reduction on the weighted tasks.\n\nThe experiments are conducted on two setups:\n1. learning from ClimbLab and SlimPajama and testing on a set of 6 reasoning dataset using 5-shot in-context learning\n2. learning from Wikipedia data from 6 languages and evaluating using perplexity on 8 other languages\n\nSTRENGTHS AND WEAKNESSES:\nStrenghts:\n- new problem/experimental setup (multi-target testing)\n- theoretical justification of the algorithm \n- comparison with many baselines for the first experiment, including all natural ones from past work\n- experimentation on 2 different setups\n- feature analysis and interpretation of results\n\nWeaknesses:\n- evaluation for the multi-lingual setup is only done using perplexity, which may not be related to actual task performance. One could have evaluated instead on downstream tasks with ICL (e.g. NER or classification) or machine translation. These would make these results stronger\n\nQUESTIONS:\n-  baselines would require a bit more description for example it is not clear what the domain weights are trained and their frequency in regmix\n\n- line 244:  Ukrian > Ukrainian\n\nLIMITATIONS:\nHighlighting these limitations of the research, but I acknowledge the space and scope limitations of the current paper:\n\n- There is still to be studies what is happening after running SFT/RL on top of these trained models\n- Models trained are only up to 0.75 on limited training steps\n- Comprehensive baselines, but only for the first experiment; this is probably fine given that DoGE is the best non-GRAPE method in that experiment. But for completeness, it would be good to have more baselines for the other experiments", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of choosing the best pretraining data mixture that result in good accuracy on a set of downstream target tasks. Specifically, the authors propose optimizing over the relative proportions/weights of a set of pretraining domain to minimize the worst case task (GRAPE). Group distributionally robust optimization (DRO) is applied to choose domain weights that result in as much (normalized) loss reduction as possible. Experiments on training language models for reasoning tasks and multilingual capabilities show improvements over baselines.\n\nSTRENGTHS AND WEAKNESSES:\n- Overall, the problem (of finding good pretraining data for downstream tasks) is well-motivated. The paper proposes a new solution that leverages principles from robust optimization (building off of FAMO, which updates parameters to reduce task loss) and appears to work quite well.\n - The paper tries to upweight tasks that are improving more slowly, but I wonder if this is the right principle; if you have a task that has high error and cannot be improved, it seems like this will be maximally upweighted, which seems not ideal. Group DRO methods sometimes use a reference model that tries to take this into account; it would be nice to discuss this more.\n - Appendix D.1 says that DoGE used the same hyperparameters as GRAPE, but is there any reason to believe these are optimal? In general, for fair comparison, the hyperparameters need to be tuned carefully and separately for each method.\n\nQUESTIONS:\n- line 100: why is taking a small learning rate justified (in order to take the Taylor expansion)?\n - line 288: typo \"improvmenet\"\n - Theorem 2.1: this theorem seems to prove convergence to a global optimum, but this shouldn't be true for general (non-convex) loss functions?\n - Theorem 2.2: assuming that the task losses are strongly convex is a pretty strong/unrealistic assumption, but is this necessary, given that monotonicity if more of a local notion?\n - Given that the method is about minimizing the worst case loss, why is average accuracy what is mainly reported?\n - From 125M to 0.7B parameters, it seems like the improvement of GRAPE over the baseline is reduced; what are the prospects of GRAPE still outperforming the baseline when scaled up to 7B or beyond?\n - Can you comment on the runtime required of the different methods?\n - How well does GRAPE do compared to the oracle of selecting the best data for a given task? I'd like to understand how much of the value is being able to select data period versus being able to juggle multiple tasks.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-10"}