Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

AI summary

Overview

Research area: Reinforcement-learning post-training of large language models for reasoning, specifically multi-task variants of Group-Relative Policy Optimization (GRPO).

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods, GRPO's advantage estimation, KL regularization, and distributionally robust / minimax optimization.

Scope: The paper proposes and evaluates MT-GRPO, a multi-task post-training algorithm that adapts task weights toward underperforming and slowly improving tasks while a ratio-preserving sampler keeps realized gradient contributions aligned with those weights.

What This Paper Is About

RL post-training with GRPO reliably improves LLMs on single reasoning tasks, but deploying a general-purpose reasoner requires competence across many tasks at once. A naive multi-task version of GRPO optimizes average reward, so strong gains on some tasks mask stagnation on others, and tasks differ widely in how often prompts produce zero advantages (and therefore zero gradients), which further distorts each task's effective contribution. The goal is a post-training algorithm that raises worst-task accuracy without sacrificing average accuracy.

Key Contributions

1. A robustness-aware multi-task objective. A constrained formulation that maximizes average task reward subject to a bound on the disparity between any two tasks, with a tunable parameter controlling the trade-off between worst-task robustness and average performance.
2. A task reweighting framework (IWU). An improvement-aware weight update that combines each task's reward and its measured per-step improvement, rather than relying on the GRPO loss, which the authors argue is an unreliable signal for comparing tasks.
3. A ratio-preserving, acceptance-aware batch construction (RP sampler). A sampler that enforces target task proportions in the batch after zero-gradient prompts are filtered out, using oversampling and acceptance-aware resampling so that learned task weights translate into actual gradient contributions.
4. Empirical evaluation at two scales. Post-training 3B and 7B models on reasoning tasks spanning planning (Countdown, Zebra puzzles), inductive reasoning (ARC), natural language QA (SciKnowEval), and mathematics, in both a controlled 3-task setting and a 9-task setting, with comparisons against GRPO, SEC-GRPO, DAPO, and SEC-DAPO.

Main Findings

• Worst-task gains over baselines: MT-GRPO achieves 16–28% and 6% absolute improvement on worst-task performance over standard GRPO and DAPO respectively, while maintaining competitive average accuracy.
• Faster time to robustness: MT-GRPO requires 50% fewer training steps to reach 50% worst-task accuracy in the 3-task setting, and in several cases baselines fail to reach target thresholds within the training budget.
• λ controls the trade-off: In the 9-task experiment, increasing λ consistently improves worst-task accuracy (by 16% over GRPO and 6% over DAPO at λ = 1.2) but reduces average accuracy. λ = 0.1 yields the highest average per-task relative change.
• λ = 0.2 vs λ = 0.25 in the 3-task setting: λ = 0.25 yields higher worst-task accuracy than λ = 0.2, while λ = 0.2 achieves better average accuracy.
• Weight reallocation drives the gains: Early in training Zebra outperforms Countdown, but this reverses after approximately 50 steps; MT-GRPO reallocates weight toward Zebra and reduces emphasis on Countdown, whereas DAPO and SEC-DAPO continue to prioritize Countdown.
• Zero-gradient rates are task-dependent: ARC exhibits a much higher ratio of zero-gradient prompts than Zebra, causing ARC to be underrepresented in training batches relative to its intended weight; the RP sampler aligns realized batch proportions with the learned task weights.
• Strict worst-task optimization collapses: With ε = 0, task weights rapidly collapse to the current worst task and oscillate as the worst task shifts, leaving Countdown near-zero weighting after early steps.
• Scaling to a larger model and diverse tasks: With Qwen2.5-7B on MATH, ARC, and SciKnowEval, MT-GRPO (λ = 0.2) achieves a 12.6% worst-task accuracy gain over DAPO and reaches DAPO's final worst-task accuracy of 25.0% in approximately 60% less training time (5.5 vs 13.3 hours).
• Joint training beats sequential training: On the nine-task setting, MT-GRPO (λ = 1.2) reaches 46.3% worst-task accuracy versus 41.1% for the strongest sequential variant (240 steps per family, two orderings, GRPO and DAPO); sequential training is sensitive to task ordering and prone to forgetting earlier task families.
• Difficulty-wise reallocation: For smaller λ, MT-GRPO shows negative relative change on easy tasks and positive relative change on hard tasks, indicating optimization effort shifts toward harder tasks.

Methodology in Plain English

The authors frame multi-task post-training as maximizing the average KL-regularized reward across K tasks, where each task has its own dataset of prompts with verifiable answers and its own reward function that checks correctness and formatting. They note two failures of simply averaging GRPO losses: the average objective permits some tasks to be sacrificed for others, and tasks differ in how often a prompt's sampled responses all receive identical rewards, which produces zero advantage and therefore no gradient.

To fix the first problem, they add constraints bounding the reward gap between any two tasks, with a parameter ε controlling how strict the constraint is. Using a Lagrangian reformulation, they convert this into a min-max problem: a learned distribution z over tasks is optimized inside the objective, with a regularizer penalizing deviation from uniform weighting. Setting ε = 0 gives a strict worst-task objective, but the authors show that in the GRPO setting the loss cannot distinguish a task the model fails completely from one it solves perfectly, since both can yield zero GRPO objective values. Their fix is to decouple reweighting from the GRPO loss and update task weights using true task-level rewards, parameterizing weights as a softmax over logits and updating the logits by gradient descent so that tasks below the weighted average reward gain weight.

Because strict worst-task optimization collapses all weight onto one task, they add an improvement signal: the per-step change in each task's GRPO objective. Weights are updated using a λ-weighted combination of task reward and task improvement, so tasks that are stagnating are prioritized rather than only those with low absolute reward; large λ approaches strict worst-task maximization and small λ promotes balanced improvement.

For the second problem, the sampler draws target post-filtered counts from a multinomial over the task weights, generates samples, filters prompts that produce zero gradients, and resamples until each task's count is met or a regeneration budget is exhausted. To avoid excessive resampling when a task has a high filtering rate, the sampler maintains a running estimate of each task's filtering rate and inflates that task's sampling probability accordingly, capped by a maximum inflation factor.

The resulting algorithm alternates: build a batch with the ratio-preserving sampler, take a standard GRPO update, measure each task's improvement, then update the task weights. Baselines are uniform-sampling GRPO, SEC-GRPO, DAPO with dynamic sampling, and SEC-DAPO. Evaluation centers on worst-task accuracy (minimum accuracy across tasks), plus average accuracy and average per-task relative change normalized against the DAPO baseline. Datasets come from Chen et al. (2025b), generated via Stojanovski et al. (2025), with three difficulty levels per task family and 10k training and 200 evaluation instances per level.

Why This Matters

Impact on research. The paper identifies a structural property of GRPO-style objectives that prior robust-optimization work does not face: the GRPO loss is ambiguous as a task-difficulty signal, and zero-gradient prompt rates differ across tasks, so upweighting a weak task does not automatically translate into more gradient signal. It also challenges the default practice of treating benchmarks as isolated optimization targets, and shows joint multi-task training can outperform sequential training on worst-task accuracy.

Real-world applications:

• General-purpose reasoning assistants that must handle math, logic puzzles, and domain QA in one deployed model rather than as separate specialists.
• Domain QA and scientific question answering, where a model that is excellent at one subject area but weak in another is unreliable in practice.
• Planning and constraint-satisfaction workloads such as Countdown and Zebra-style puzzle solving used in agentic settings.
• Training pipelines where compute is the binding constraint, since MT-GRPO reaches target worst-task accuracy in fewer steps.

Industry relevance. The author list spans UCL, Huawei Noah's Ark Lab, UNIST, University of Edinburgh, and University of Basel, and the methods target the standard GRPO post-training pipeline used in current LLM development, so the technique is applicable wherever teams fine-tune reasoning models across multiple verifiable-reward tasks.

Future Directions

• Automating the choice of λ, since the experiments show it governs a real worst-task versus average-accuracy trade-off and the best value differs between the 3-task and 9-task settings (λ = 0.25/0.2 versus λ = 1.2).
• Scaling further in task count and model size, beyond the 3-task and 9-task settings on 3B and 7B models reported here.
• Reducing the resampling overhead of the ratio-preserving sampler in regimes with very high zero-gradient rates, where the regeneration budget is exhausted.
• Addressing forgetting and ordering sensitivity observed in sequential training, and extending the approach to task families beyond those with verifiable rewards.
• Clarifying the role of the l2 regularization alternative to improvement-aware reweighting, which the paper analyzes in Appendix F rather than in the main text.

Target Audience

Researchers and engineers working on RL-based LLM post-training, multi-task learning, and robustness-aware optimization, particularly those already familiar with GRPO and looking for principled ways to prevent task imbalance. It also suits practitioners building multi-domain reasoning models who care more about worst-case reliability than average benchmark scores.

Authors’ abstract

RL-based post-training with GRPO is widely used to improve large language models on individual reasoning tasks. However, real-world deployment requires reliable performance across diverse tasks. A straightforward multi-task adaptation of GRPO often leads to imbalanced outcomes, with some tasks dominating optimization while others stagnate. Moreover, tasks can vary widely in how frequently prompts yield zero advantages (and thus zero gradients), which further distorts their effective contribution to the optimization signal. To address these issues, we propose a novel Multi-Task GRPO (MT-GRPO) algorithm that (i) dynamically adapts task weights to explicitly optimize worst-task performance and promote balanced progress across tasks, and (ii) introduces a ratio-preserving sampler to ensure task-wise policy gradients reflect the adapted weights. Experiments on both 3-task and 9-task settings show that MT-GRPO consistently outperforms baselines in worst-task accuracy. In particular, MT-GRPO achieves 16-28% and 6% absolute improvement on worst-task performance over standard GRPO and DAPO, respectively, while maintaining competitive average accuracy. Moreover, MT-GRPO requires 50% fewer training steps to reach 50% worst-task accuracy in the 3-task setting, demonstrating substantially improved efficiency in achieving reliable performance across tasks.