Pith. sign in

REVIEW 4 major objections 6 minor 1 cited by

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

T0 review · 4 major / 6 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read GRAPE steers pretraining data toward tasks that are learning slowest, lifting multiple benchmarks at once.

desk verdict GRAPE is a genuinely new multi-target data-mixing method with consistent empirical gains and a theory section that should be treated as motivation, not proof. read the letter →

arxiv 2505.20380 v1 pith:DKGLUWDE submitted 2025-05-26 cs.LG

Simin Fan , Maria Ios Glarou , Martin Jaggi This is my paper
classification cs.LG
keywords datamixtureoptimizationdomainreweightinggroupDROmulti-taskpretrainingRate-of-Improvementgradientalignmentlow-resourcemultilingualadaptivecurriculum
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

GRAPE addresses a gap in pretraining-data curation: existing reweighting methods optimize a mixture for one target or for average performance, which lets a dominant task or conflicting gradients distort the model. The paper's proposal is to run two coupled reweighting loops during pretraining. An inner step, based on group distributed-robust optimization (a worst-case weighting over tasks), upweights targets with the smallest normalized one-step improvement; an outer step then shifts the sampling weights of source data domains toward domains whose gradients align with the gradients of those lagging tasks. The paper reports that this feedback loop improves average 5-shot accuracy on six reasoning benchmarks and log-perplexity on eight low-resource languages, and that at 0.7B scale it reaches the uniform-mixture baseline's average accuracy with roughly 40 percent fewer tokens. If correct, this makes multi-target robustness a property of the data mixture itself rather than a post-hoc balancing act.

What carries the argument

The mechanism is the Rate-of-Improvement (RoI), defined as $r_t^n = (l_n(\theta_t) - l_n(\theta_{t+1})) / l_n(\theta_t)$, the scale-invariant relative drop in each target task's loss over one step. A first-order Taylor expansion rewrites RoI as a weighted sum of gradient alignments between each task's normalized validation gradient and each domain's training gradient. GRAPE's update rules are multiplicative mirror-descent steps: tasks with low alignment to the current update direction get larger weight, and domains whose gradients align well with the weighted validation gradients get larger weight. The minimax objective is the formal statement of this negative feedback loop, and the two exponentiated update rules are the machine that runs it.

What would settle it

Take a fixed pretraining corpus and a fixed target-task suite, and measure the rank correlation between the gradient alignments GRAPE computes at each update and the actual per-task loss reductions over the following block; a correlation near zero or negative would indicate the signal driving the reweighting does not predict learning. A second, cheaper check is a control run that reweights domains and tasks at the same frequency with random directions of the same magnitude; if random reweighting matches GRAPE's gains, the alignment signal is not the cause of the improvement.

Watch

Extended reading notes

Core claim

The central claim is that a minimax loop over task weights and domain weights calibrates a pretraining mixture for simultaneous strong performance on multiple targets. Concretely, the inner minimization finds the task distribution that concentrates on the tasks with the lowest Rate-of-Improvement under the current mixture, and the outer maximization moves domain sampling weight toward the domains whose gradients align with the normalized validation gradients of those tasks. GRAPE implements this as multiplicative mirror-descent updates with Bregman regularization, yielding a negative feedback loop: struggling tasks get more weight, domains that help them get more weight, and the loop repeats. The paper presents empirical evidence on six reasoning benchmarks and on multilingual pretraining from six source languages toward eight low-resource targets.

Load-bearing premise

The load-bearing premise is that the alignment between a task's gradient and a domain's gradient reveals which data will help that task; if the alignment is noisy or misleading, the whole reweighting loop chases spurious correlations.

Editorial extensions

If this is right

  • A single pretraining run can be aimed at a whole suite of tasks, so gains on one benchmark no longer have to come at the cost of another.
  • The slowest-improving tasks act as a dynamic curriculum, with training attention shifting from foundational reading comprehension to physical and commonsense reasoning in the later stages of training.
  • The two-loop update is cheap enough to run during pretraining: roughly 15 percent extra gradient computations on the SlimPajama setup, about 7 percent of that coming from the new task-reweighting step.
  • In multilingual pretraining, the same mechanism accelerates learning of low-resource languages by at least 60 percent in log-perplexity terms, according to the paper's experiments.
  • The paper's variance-reduction result predicts that task performances become more balanced over time, counteracting the tendency of average weighting to let one task dominate.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The task-weight trajectories the paper plots could be extracted once and reused as a fixed curriculum for a larger model or a longer run, avoiding the per-training reweighting cost; the paper does not test this reuse, but the trajectories' smooth stage-wise structure makes it plausible.
  • The multilingual weight trajectories hint that the alignment signal is capturing linguistic relatedness, for instance Slavic targets pulling up Slavic source languages. A controlled experiment that varies only phylogenetic distance between source and target languages would test whether the learned weights track that distance; the paper reports the trajectory correlations but does not run that test
  • The same Rate-of-Improvement logic could be pushed from task-level to cluster-level or sample-level weighting inside a task, which the paper itself names as future work. A natural first step would be to apply RoI to clusters of validation examples and see whether the curriculum sharpens further.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 6 minor

Summary. The paper introduces GRAPE, a domain- and task-reweighting algorithm for multi-target LLM pretraining. GRAPE maintains two distributions: task weights z over N target validation tasks and domain weights α over K pretraining domains. Using a minimax formulation motivated by group DRO, the inner step upweights tasks with the lowest Rate-of-Improvement, approximated by the alignment between normalized task gradients and the training update direction, while the outer step upweights domains whose gradients align with the task-weighted validation gradient. The authors report improved 5-shot reasoning accuracy over Uniform, DOGE, RegMix, CRISP, and related baselines on ClimbLab and SlimPajama, and improved log-perplexity across 8 low-resource target languages in Wiki-40B experiments. They also state convergence and variance-reduction theorems in Section 2.2 and appendices, and provide ablations over 12 task combinations and 3 progress metrics.

Significance. If the empirical results are taken at face value, GRAPE is a useful practical contribution: it is simple, adds modest computational overhead (about 15% in gradient computations in the main configuration), and the paper provides unusually broad validation, including 12 multi-task configurations, two model scales, two reasoning corpora, and a multilingual setting. The weight-trajectory analyses are also a constructive feature, connecting the learned curriculum to interpretable task and language relations. These strengths are substantial and go beyond a single benchmark pair. However, the paper's theoretical claims, as written, are not reliable, and the central gradient-alignment assumption is not independently validated. The empirical core is plausible and likely worth publishing after the theory is corrected or dropped and the evaluation protocol is clarified.

major comments (4)
  1. [2.2 / Appendix C.1] Theorem 2.1 is not established by its proof. The theorem statement assumes only L-smooth losses and bounded gradients, but the proof in Appendix C.1 invokes "µ-strong convexity" without stating or defining this assumption. The key step asserting that the reweighting strategy ensures Q_t > G^2/2 after sufficient iterations is asserted without proof, and Q_t is only upper-bounded, not lower-bounded, by the stated gradient bound. The displayed convergence bound mixes the strong-convexity term L∥θ_1−θ*∥^2/(2T) with the alignment term −(1/T)Σ Q_t/L without a valid derivation, and the theorem's conclusion about a "neighborhood of the Pareto optimal solution" is not linked to the proof, which concerns the worst-case loss max_n l_n. As stated, the O(1/T) claim is unsupported. The authors should either provide a correct proof under precise assumptions or remove or substantially weaken this theorem.
  2. [Appendix C.2 / Theorem 2.2] The variance-reduction theorem rests on Lemma 2, which is incorrect. Lemma 2 treats b_t − a_t, where a_t = min_n l_n(θ_t) and b_t = max_n l_n(θ_t), as a constant denominator in an affine transformation and concludes that the variance of normalized losses is Var(l_i)/(b_t−a_t)^2. But b_t−a_t is a random function of the very same losses, so the classical affine variance identity does not apply. Additionally, Lemma 1 is stated without proof. Theorem 2.2's conclusion of monotonic variance reduction therefore does not follow from the arguments given. This is a load-bearing issue for the theoretical narrative in Section 2.2, and the theorem should be repaired with a valid proof or removed from the claims.
  3. [2.1 / Algorithm 1] The algorithm's update rules do not match the theoretical derivation in a way that matters for the mechanism. Equations (3) and (4) are derived from the Taylor expansion of RoI at θ_t using a plain gradient descent update of the form θ_{t+1} = θ_t − γ_t d_t. Algorithm 1 lines 8 and 15, however, compute the gradient-alignment inner products at θ_{t+1}, after an AdamW step, using one minibatch per domain or task and with updates applied only every 100 steps. No evidence is provided that these single-batch post-AdamW alignments at θ_{t+1} are predictive of the actual per-task loss improvement over the subsequent update horizon. Since the z and α updates in Equations (3)–(4) are the whole mechanism that reweights the data mixture, this gap is central. I would like to see a concrete validation, for example a measured correlation between the Algorithm 1 alignment scores and the realized RoI on the same tasks over a window, or a comparison against a variant that updates weights using directly measured per-task loss differences.
  4. [3.1 / Table 1] The evaluation is partly in-sample with respect to the optimization signal. The target validation losses are used to compute task weights during training, and Table 1 then reports accuracy on the same target tasks; it is not stated whether these scores are on the validation sets used for reweighting or on separate test sets. If they are the same splits, the improvement over Uniform partly reflects in-sample optimization on the target benchmarks, and the claim of "robust performance across multiple target tasks" should be supported by additional held-out tasks or by a clear statement of which split is used for every reported number. The comparison is somewhat mitigated because DOGE, RegMix, and CRISP also use the target tasks, but the headline claim about generalization would be stronger with a held-out evaluation.
minor comments (6)
  1. [3.2] There is a typo: "Ukrian" should be "Ukrainian."
  2. [3.1] The text referring to "Figure 3.1" should be "Figure 1"; the current cross-reference appears to be broken.
  3. [2.1] Equation (2) uses both ∆_K and ∆_k for the simplex; the notation should be made consistent across the paper.
  4. [4] The hyperparameters µ_α, µ_z, ∆T_α, ∆T_z, and β are fixed throughout, and the paper acknowledges that no sweep was performed. A small sensitivity analysis over at least one of these parameters would make the reported results more convincing.
  5. [Appendix D / Table 3] Table 3 includes a 1M model configuration but the experiments report 125M and 0.7B models; the 1M row should either be removed or explicitly referenced.
  6. [4 / Equation (5)] The notation for the EMA loss l_t^{ema,n} is introduced after the equation that uses it; moving the definition before the display would improve readability.

Circularity Check

2 steps flagged · score 6.0 of 10

Partial circularity: the benchmark headline is evaluated on the same target validation sets that drive the reweighting, and Theorem C.2's variance-reduction proof assumes its own conclusion in an unproved Lemma 1.

  1. fitted input called prediction [Section 3.1 ('Setup'), Algorithm 1 lines 6-9, Table 1]
    "For each target task Tn, we use its standard validation set to compute the task loss ln(θt) and the Rate-of-Improvement r(t)n needed for GRAPE's updates during training. ... 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."

    The target validation sets are inputs to the algorithm (Equation 1, Equation 3, Algorithm 1 lines 6-9): the same six benchmarks' losses set the task weights z and the domain weights α, and then Table 1 reports 5-shot accuracy on exactly those tasks. The headline 'reasoning performance across 6 benchmarks' is therefore an in-sample measure of the objective being optimized, not an independent prediction on held-out target tasks. The reported improvement is statistically aligned with the fitted validation-loss signal by construction, fitting the 'fitted input called prediction' pattern.

  2. other [Appendix C.2, Lemma 1 and Theorem C.2]
    "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). ... Theorem C.2 ... the variance decreases monotonically until reaching a minimal basin, i.e., σ2t+1 ≤ σ2t for all t ≥ T0 for some finite T0."

    Lemma 1 is asserted without proof, and its content is exactly the catch-up property needed to conclude variance reduction: if initially slower-improving tasks are guaranteed to improve faster, then with Lemma 2's affine variance identity, σ2t+1 ≤ σ2t follows directly. The theorem therefore 'derives' its own conclusion by postulating it in an unproved lemma; the derivation chain reduces to the assumption. Theorem 2.1's proof contains the same pattern, asserting without derivation that 'Qt becomes increasingly positive' and that 'we eventually reach a point where Qt > G2/2, ensuring consistent progress.'

full rationale

GRAPE's update rules are derived from a first-order expansion of Rate-of-Improvement and a regularized minimax objective; that algebraic derivation is self-contained and not circular. The gradient-alignment proxy's empirical validity is a correctness risk, but it is not a definitional equivalence, so it does not by itself raise the circularity score. The circularity issues are two. First, the headline empirical claim is evaluated on the same target benchmarks whose validation losses drive the reweighting (Eq. 1, Eq. 3, Algorithm 1), so the reported 'reasoning performance across 6 benchmarks' is an in-sample measure of the optimized objective rather than an independent prediction. Second, the variance-reduction theorem's proof postulates as unproved Lemma 1 the exact catch-up mechanism the theorem is supposed to establish; combined with Lemma 2's affine variance identity, Theorem C.2 follows from its own assumption. Theorem 2.1's proof similarly assumes the key alignment-growth property without derivation. The self-citations to the authors' DOGE/DGA gradient-alignment work are not separately load-bearing here, because GRAPE re-derives the alignment mechanism from the RoI Taylor expansion; however, the multilingual experiments use held-out test log-perplexity and provide partial independent evidence. Overall, there is partial circularity in the in-sample evaluation and in the convergence/variance proofs, but the core minimax derivation and the breadth of ablations retain independent content.

Assumptions & free parameters 4 free parameters · 4 assumptions · 0 invented entities

The method's only fitted quantities are its reweighting hyperparameters (µα, µz, update frequencies), chosen by hand and not swept. The core steering signal is the gradient inner product between task validation gradients and domain training gradients, an assumption inherited from the authors' DOGE and DGA papers. The theory adds L-smoothness and bounded-gradient assumptions that are not satisfied by LLM training, and the proofs invoke strong convexity despite the theorem statements omitting it. No invented entities are introduced.

free parameters (4)
  • µα (domain weight regularization coefficient) = 1e-4
    Controls the step size of domain weight updates; chosen by hand, no sweep (Appendix D.1).
  • µz (task weight regularization coefficient) = 1.5e-5
    Controls the step size of task weight updates; chosen by hand, no sweep (Appendix D.1).
  • Update frequencies ∆Tα and ∆Tz = 100 steps
    Period for reweighting updates; chosen for efficiency, effect on performance not analyzed.
  • β (EMA decay for ROI-ema) = 0.7
    Only used in the ROI-ema ablation variant; chosen by hand.
assumptions (4)
  • domain assumption Target validation losses are reliable progress signals
    RoI in Equation 1 is computed from target validation losses during training; noisy single-batch estimates are used in Algorithm 1.
  • domain assumption Gradient inner product predicts transfer
    Equations 3-4 assume that alignment <∇ log l_n, g_k> indicates which domains help which tasks; this is the core steering signal, inherited from DOGE.
  • standard math L-smoothness and bounded gradients
    Assumed in Theorems 2.1 and 2.2; LLM losses are non-convex, and the theorems' statements omit strong convexity even though the proofs use it.
  • ad hoc to paper Tasks that improve slowly deserve more weight
    Heuristic design principle stated in Section 1; motivates the DRO formulation but is not derived from an external objective.

how reviews work

0 comments
Cite this review

Pith. "Pith review of GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining." pith.science (2026). https://pith.science/paper/DKGLUWDE

@misc{pith2026250520380,
  author       = {Pith},
  title        = {Pith review of: GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/DKGLUWDE}},
  note         = {Machine review of arXiv:2505.20380}
}
read the original 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 Group Robust Multi-target Adaptive PrEtraining (GRAPE), a novel multi-source-multi-target domain reweighting framework designed to calibrate pretraining data mixtures for robust performance across multiple target tasks simultaneously. GRAPE dynamically adjusts sampling weights across source domains (domain weights) while concurrently modulating task weights that quantify the relative importance of each individual target task. This adaptive process prioritizes tasks based on their learning difficulty throughout training. We formulate this interleaved reweighting mechanism as a minimax optimization problem: The inner maximization adjusts task weights leveraging group distributed-robust-optimization (DRO), where those tasks demonstrating the least improvement under the current data mixture are prioritized with higher weights; The outer minimization then optimizes domain weights to maximize loss reduction on the prioritized tasks. Experiments on ClimbLab and SlimPajama datasets demonstrate that GRAPE consistently outperforms baseline methods in terms of reasoning performance across 6 benchmarks. Furthermore, when applied to multilingual targets, GRAPE effectively identifies optimal training mixtures from mainstream languages, achieving superior language modeling capabilities across 8 low-resource target languages.

Figures

Figures reproduced from arXiv: 2505.20380 by the authors.

Figure 1
Figure 1. GRAPE facilitates multi-task reasoning. For 125M models, GRAPE and GRAPE-climbmix greatly outperform five baselines; For larger 0.7B models, GRAPE achieves comparable scores as uniform base￾line with 40% fewer tokens. 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 Tn, we use its standard vali￾dation set to compute the task lo… view at source ↗
Figure 2
Figure 2. Task weight evolution of GRAPE. In the early stage, the reading compre￾hension tasks like ARC-E and ARC-C are mostly up-weighted. As training pro￾gresses, in the late stage, physical and com￾monsense reasoning tasks like PIQA and Hellaswag are steadily prioritized, demon￾strating the emergence of a skill-wise learn￾ing curriculum that moves from founda￾tional to more complex reasoning abilities. Concurrently, tasks … view at source ↗
Figure 3
Figure 3. Domain weights attributions across 20 clusters in the ClimbLab dataset. Which data domains and topics are critical for general reasoning? [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗
Figures from the paper (29 more)
Figure 4
Figure 4. Figure 4: Low-resource language learning progress by Log-Perplexity. GRAPE significantly outperforms DOGE and Uniform sampling across all target languages. 3.2 Language Mixture for Multi-lingual Learning Setup. In this scenario, we investigate GRAPE’s ability to optimize the lan…
Figure 5
Figure 5. Figure 5: Weights evolution on multilingual pre￾training for low-resource language modeling. Portuguese (pt), Dutch (nl), and Ukrian (uk). Per￾formance is measured by the language modeling loss, i.e. the log-perplexity (log-PPL), on held-out test sets for each target language. L…
Figure 6
Figure 6. Figure 6: Task weights evolution on 6 reasoning tasks. Results on the six reasoning tasks ( [PITH_FULL_IMAGE:figures/full_fig_p008_6.png]
Figure 7
Figure 7. Figure 7: RegMix results on ClimbLab with LightGBM regression on 6 reasoning target tasks [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA] Domain Weights from CRISP. We present the domain weights from CRISP in [PITH_FULL_IMAGE:figures/full_fig_p022_7.png]
Figure 8
Figure 8. Figure 8: RegMix results on ClimbLab with LightGBM regression on 8 target tasks [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA]. boosts the weight of Cluster 10 (health and wellness content), reflecting MedQA’s embeddings lie closest to clinical-style text. (a) 6 Ta…
Figure 9
Figure 9. Figure 9: CRISP Domain weights across 7 data domains in ClimbLab. Domain Weights Evolution from GRAPE and Scaling Effects. We observed that the task priority differs notably between model sizes. According to [PITH_FULL_IMAGE:figures/full_fig_p023_9.png]
Figure 10
Figure 10. Figure 10: Task weights trajectory from GRAPE on ClimbLab from 125M/0.7B models. 23 [PITH_FULL_IMAGE:figures/full_fig_p023_10.png]
Figure 11
Figure 11. Figure 11: RegMix results on SlimPajama with 6 target tasks [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA] (a) Original vs. RegMix weights (b) Predicted vs. true loss [PITH_FULL_IMAGE:figures/full_fig_p024_11.png]
Figure 12
Figure 12. Figure 12: RegMix results on SlimPajama with 8 target tasks [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA] Domain Weights from CRISP. We present the domain weights derived from CRISP in [PITH_FULL_IMAGE:figures/full_fig_p024_12.png]
Figure 13
Figure 13. Figure 13: CRISP Domain weights across 7 data domains in SlimPajama. 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 …
Figure 14
Figure 14. Figure 14: Log-Perplexities on 6 target reasoning tasks [PITH_FULL_IMAGE:figures/full_fig_p025_14.png]
Figure 15
Figure 15. Figure 15: Log-Perplexities on 8 target datasets 26 [PITH_FULL_IMAGE:figures/full_fig_p026_15.png]
Figure 16
Figure 16. Figure 16: RegMix results on wiki-40b with LightGBM regression on 4 target languages [Catalan, Danish, Romanian, Ukrainian] 27 [PITH_FULL_IMAGE:figures/full_fig_p027_16.png]
Figure 17
Figure 17. Figure 17: RegMix results on wiki-40b with LightGBM regression on 8 target languages [Catalan, Danish, Romanian, Ukrainian, Polish, Portuguese, Turkish, Dutch] Domain Weights from CRISP. We present the optimized domain weights distributions from CRISP in [PITH_FULL_IMAGE:figure…
Figure 18
Figure 18. Figure 18: CRISP Domain Weigths across 6 high-resource languages in wiki-40b 28 [PITH_FULL_IMAGE:figures/full_fig_p028_18.png]
Figure 19
Figure 19. Figure 19: Log-perplexities on 4 target languages [Catalan, Danish, Romanian, Ukrainian] 29 [PITH_FULL_IMAGE:figures/full_fig_p029_19.png]
Figure 20
Figure 20. Figure 20: Log-perplexities on 8 target languages [Catalan, Danish, Romanian, Ukrainian, Polish, Portuguese, Turkish, Dutch] 30 [PITH_FULL_IMAGE:figures/full_fig_p030_20.png]
Figure 21
Figure 21. Figure 21: Ablation on task combination [GSM8K, ARC-C, ARC-E]. 33 [PITH_FULL_IMAGE:figures/full_fig_p033_21.png]
Figure 22
Figure 22. Figure 22: Ablation on task combination [GSM8K, Hellaswag]. 34 [PITH_FULL_IMAGE:figures/full_fig_p034_22.png]
Figure 23
Figure 23. Figure 23: Ablation on task combination [GSM8K, PIQA]. 35 [PITH_FULL_IMAGE:figures/full_fig_p035_23.png]
Figure 24
Figure 24. Figure 24: Ablation on task combination [GSM8K, LogiQA]. 36 [PITH_FULL_IMAGE:figures/full_fig_p036_24.png]
Figure 25
Figure 25. Figure 25: Ablation on task combination [GSM8K, SciQ]. 37 [PITH_FULL_IMAGE:figures/full_fig_p037_25.png]
Figure 26
Figure 26. Figure 26: Ablation on task combination [GSM8K, ARC-E, ARC-C, Kodcode]. 38 [PITH_FULL_IMAGE:figures/full_fig_p038_26.png]
Figure 27
Figure 27. Figure 27: Ablation on task combination [GSM8K, Kodcode, Hellaswag]. 39 [PITH_FULL_IMAGE:figures/full_fig_p039_27.png]
Figure 28
Figure 28. Figure 28: Ablation on task combination [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, Kodcode, GSM8K]. 40 [PITH_FULL_IMAGE:figures/full_fig_p040_28.png]
Figure 29
Figure 29. Figure 29: Ablation on task combination [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA]. 41 [PITH_FULL_IMAGE:figures/full_fig_p041_29.png]
Figure 30
Figure 30. Figure 30: Ablation on task combination [ARC-E, ARC-C, Hellaswag, SciQ, PIQA, LogiQA, MathQA, MedQA]. 42 [PITH_FULL_IMAGE:figures/full_fig_p042_30.png]
Figure 31
Figure 31. Figure 31: Ablation on task combination [ARC-E, ARC-C, MathQA, MedQA]. 43 [PITH_FULL_IMAGE:figures/full_fig_p043_31.png]
Figure 32
Figure 32. Figure 32: Ablation on task combination [LogiQA, Hellaswag, MathQA, MedQA]. 44 [PITH_FULL_IMAGE:figures/full_fig_p044_32.png]

Discussion (0). Sign in to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

    cs.CL 2026-02 conditional novelty 6.0 of 10

    MT-GRPO reweights tasks by reward and improvement and enforces those weights after zero-gradient filtering, improving worst-task accuracy by 6–28% over GRPO/DAPO baselines on 3- and 9-task setups.

Reference graph

Works this paper leans on

30 extracted references · 7 canonical work pages · cited by 1 Pith paper

  1. [1]

    The pile: An 800gb dataset of diverse text for language modeling, 2020

    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

    2020

  2. [2]

    Redpajama: An open source recipe to reproduce llama training dataset, 2023

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

    work page 2023

  3. [3]

    Le, Tengyu Ma, and Adams Wei Yu

    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

    arXiv 2023

  4. [4]

    Regmix: Data mixture as regression for language model pre-training, 2025

    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

    arXiv 2025

  5. [5]

    Doge: Domain reweighting with generalization estimation, 2024 a

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

    arXiv 2024

  6. [6]

    Task-adaptive pretrained language models via clustered-importance sampling, 2025

    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

    arXiv 2025

  7. [7]

    Dynamic gradient alignment for online data mixing, 2024 b

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

    arXiv 2024

  8. [8]

    Gradient surgery for multi-task learning, 2020

    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

    arXiv 2020

Show all 30 references
  1. [9]

    Famo: Fast adaptive multitask optimization, 2023

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

    2023 arXiv

  2. [10]

    Task weighting through gradient projection for multitask learning, 2024

    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

    2024 arXiv

  3. [11]

    Learning models with uniform performance via distributionally robust optimization, 2020

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

    2020 arXiv

  4. [12]

    Hashimoto, and Percy Liang

    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

    2020 arXiv

  5. [13]

    An online method for a class of distributionally robust optimization with non-convex objectives, 2021

    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

    2021 arXiv

  6. [14]

    Stochastic gradient methods for distributionally robust optimization with f-divergences

    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, ...

    2016

  7. [15]

    Gomez, Lukasz Kaiser, and Illia Polosukhin

    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

    2023 arXiv

  8. [16]

    Climb: Clustering-based iterative data mixture bootstrapping for language model pre-training, 2025

    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, 2...

    2025 arXiv

  9. [17]

    Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018

    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

    2018 arXiv

  10. [18]

    Liu, and Matt Gardner

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

    2017 arXiv

  11. [19]

    Piqa: Reasoning about physical commonsense in natural language, 2019

    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

    2019 arXiv

  12. [20]

    Logiqa: A challenge dataset for machine reading comprehension with logical reasoning, 2020

    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

    2020 arXiv

  13. [21]

    Hellaswag: Can a machine really finish your sentence?, 2019

    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

    2019 arXiv

  14. [22]

    W iki-40 B : Multilingual language model dataset

    Mandy Guo, Zihang Dai, Denny Vrande c i \'c , and Rami Al-Rfou. W iki-40 B : Multilingual language model dataset. In Nicoletta Calzolari, Fr \'e d \'e ric B \'e chet, Philippe Blache, Khalid Choukri, Christopher Cieri, Thierry Declerck, Sara Goggi, Hitoshi Isahara, Bente Maega...

    2020

  15. [23]

    Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré

    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

    2023 arXiv

  16. [24]

    Learning to combine: Knowledge aggregation for multi-source domain adaptation, 2020

    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

    2020 arXiv

  17. [25]

    Efficient online data mixing for language model pre-training, 2023

    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

    2023 arXiv

  18. [26]

    Conflict-averse gradient descent for multi-task learning, 2024

    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

    2024 arXiv

  19. [27]

    Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks, 2018

    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

    2018 arXiv

  20. [28]

    Multi-task learning as a bargaining game, 2022

    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

    2022 arXiv

  21. [29]

    Multiple-gradient descent algorithm ( MGDA ) for multiobjective optimization

    Jean-Antoine D\'esid\'eri. Multiple-gradient descent algorithm ( MGDA ) for multiobjective optimization. Comptes Rendus Mathematique, 350 0 (5-6): 0 313--318, 2012

    2012

  22. [30]

    Minicpm: Unveiling the potential of small language models with scalable training strategies, 2024

    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 L...

    2024 arXiv

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.