Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

Shyam Sundhar Ramesh; Xiaotong Ji; Matthieu Zimmer; Sangwoong Yoon; Zhiyong Wang; Haitham Bou Ammar; Aurelien Lucchi; Ilija Bogunovic

Source-linked AI summary

Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

Shyam Sundhar Ramesh, Xiaotong Ji, Matthieu Zimmer, Sangwoong Yoon, Zhiyong Wang, Haitham Bou Ammar, Aurelien Lucchi, Ilija Bogunovic

arXiv:2602.05547v3Feb 5, 2026cs.CLcs.AIcs.LG

TL;DR

Multi-task RL post-training must support reliable performance across diverse reasoning tasks, but average-focused objectives can leave some tasks stagnating and zero-gradient rates can distort optimization. MT-GRPO combines improvement-aware task reweighting with ratio-preserving sampling, improving worst-task performance while maintaining competitive average accuracy. Its scope is currently bounded by the training task mixture and verifiable-reward setting.

Problem
Average-focused multi-task GRPO can produce imbalanced outcomes, while uneven zero-gradient rates distort effective task contributions during optimization.
Method
MT-GRPO combines improvement-aware task reweighting with ratio-preserving, acceptance-aware batch construction to align effective gradients with target task weights.
Results
MT-GRPO consistently improves worst-task performance while maintaining competitive average accuracy across multi-task settings.
Takeaways & Limitations
Explicitly optimizing task-wise robustness is practical and beneficial for building general-purpose reasoning models.
Takeaways & Limitations
MT-GRPO is scoped to the training task mixture and relies on verifiable rewards; robustness outside that mixture and in non-verifiable domains remains open.

Abstract

from arXiv · show

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.

1. Introduction

The paper frames multi-task RL post-training as necessary for reliable general-purpose reasoning, because average-focused optimization can leave substantial task imbalances. MT-GRPO addresses this with improvement-aware reweighting and ratio-preserving sampling, improving worst-task performance while preserving competitive average accuracy.

General-purpose deployment requires reliable competence across diverse reasoning skills rather than specialization on isolated benchmarks.
Average-focused multi-task post-training can let strong gains on some tasks mask stagnation or interference on others.
MT-GRPO uses improvement-aware task reweighting to promote balanced progress and improve worst-task performance without sacrificing average performance.
MT-GRPO uses ratio-preserving, acceptance-aware batch construction so learned task weights translate into actual gradient signals.
MT-GRPO improves worst-task accuracy across settings while maintaining competitive average performance and reallocating effort toward weaker or slowly improving tasks.

2. Problem Formulation

The paper formulates joint post-training over K reasoning tasks with task-specific datasets, verifiable rewards, and KL-regularized performance objectives. It identifies average-reward optimization and uneven zero-gradient rates as obstacles to balanced multi-task learning.

The setting contains K reasoning tasks, each with a disjoint prompt dataset Dk and a task-specific reward function Rk evaluating correctness and formatting.
Each task’s performance metric Jk(θ) is its KL-regularized expected reward, with regularization limiting deviation from a reference policy.
Standard multi-task GRPO averages task losses, allowing strong gains on some tasks to compensate for substantial underperformance on others.
Different tasks have widely varying zero-gradient rates because prompts with identical rollout rewards produce zero advantages and no policy gradient.
Filtering zero-gradient prompts can bias post-filtered batches toward tasks with fewer such prompts, especially when weaker tasks are explicitly upweighted.

3. Multi-Task Post-Training Objective

The objective seeks high average rewards while bounding disparities across tasks, then uses a max–min surrogate for tractable optimization. The section identifies degenerate worst-task weighting and introduces improvement-aware reweighting to promote balanced progress.

Multi-task objective: The objective maximizes average rewards while bounding the reward difference between any two tasks.These constraints promote balanced competence rather than allowing gains on some tasks to offset marginal progress elsewhere.
Robustness constraint: ε = 0 enforces strict equality of task performance, while larger values relax the constraint toward average-reward optimization.
Surrogate objective: The max–min surrogate uses a learned task distribution z to optimize the robustness objective through a Lagrangian reformulation.
Degenerate dynamics: Strict worst-task optimization can concentrate training on one task, leaving non-worst tasks systematically under-optimized.For ε = 0, the inner optimization places all weight on the lowest-reward task; softmax updates can amplify this collapse.
Improvement-aware reweighting: Improvement-aware reweighting combines task reward with task-level improvement to prioritize tasks that are underperforming and under-improving.The update accounts for stagnation or deterioration rather than repeatedly upweighting every low-reward task.
Improvement-aware reweighting: The resulting updates prevent collapse onto a single task and promote balanced progress across tasks.

4. Algorithm

MT-GRPO jointly learns policy parameters and task-sampling weights, updating weights with reward and improvement signals. Its ratio-preserving sampler keeps realized post-filtered batch proportions aligned with those weights despite zero-gradient prompts.

Algorithm overview: MT-GRPO jointly learns policy parameters and a task distribution that governs prompt sampling during training.
Adaptive task reweighting: Adaptive task reweighting increases sampling for tasks that underperform or improve slowly while balancing worst-task robustness against average performance.The parameter λ controls this trade-off: larger values emphasize worst-task robustness, while smaller values favor average performance.
Ratio preservation: Uneven zero-gradient rates can make effective gradient contributions diverge from intended task weights, because filtered prompts contribute no update.The discrepancy is especially consequential when tasks have substantially different zero-gradient prevalences.
Ratio preservation: The Ratio-Preserving Sampler enforces target task proportions in the post-filtered batch using oversampling and acceptance-aware resampling.
Algorithm overview: Task weights, sampling, and subsequent reward and improvement updates form an adaptive loop for reliable performance across tasks.

5. Practical Findings and Solutions

Zero-gradient prompts distort multi-task training because their prevalence varies across tasks. MT-GRPO corrects this by enforcing target proportions after filtering and by anticipating filtering during sampling.

Uneven zero-gradient rates: Zero-gradient prompts arise when all sampled responses for a prompt receive identical rewards, causing the prompt’s gradient to vanish.
Uneven zero-gradient rates: Different zero-gradient prevalences cause tasks with more informative prompts to contribute disproportionately to parameter updates.Thus, equal assigned task weights do not guarantee equal effective training proportions.
Target-ratio enforcement: MT-GRPO enforces target task proportions in the post-filtered batch after zero-gradient prompts are removed.
Acceptance-aware sampling: Acceptance-aware sampling oversamples tasks with higher expected filtering and prioritizes deficient tasks during resampling.This reduces regeneration overhead while preserving the proportions induced by task-weight updates.

6. Experiments

Experiments across three-task, nine-task, and larger-model mixtures show that MT-GRPO improves worst-task performance while maintaining competitive average accuracy. Adaptive reweighting and ratio-preserving sampling redirect optimization toward weaker or slowly improving tasks, with λ controlling the robustness–average-performance trade-off.

Experimental settings: MT-GRPO is evaluated on three-task, nine-task, and heterogeneous larger-model mixtures spanning planning, inductive reasoning, mathematics, and natural-language question answering.The experiments use Qwen-2.5-3B and Qwen2.5-7B settings, including Countdown, Zebra, ARC, MATH, and SciKnowEval.
Experiment 1: MT-GRPO achieves higher worst-task accuracy than all baselines while improving or maintaining average accuracy in the controlled three-task experiment.It also attains the highest average per-task relative change, indicating more balanced improvements across tasks.
Optimization dynamics: MT-GRPO reallocates weight toward under-performing tasks, whereas baselines continue prioritizing high-performing Countdown and underrepresent ARC relative to its intended weight.The ratio-preserving sampler aligns realized batch proportions with learned task weights and supports higher ARC performance.
Experiment 1: 50% fewer training steps are required for MT-GRPO to reach specified worst-task accuracy thresholds than baselines, which sometimes fail to reach them within the training budget.The result indicates faster progress on the weakest task, not only a stronger final worst-task score.
Experiment 2: Increasing λ improves worst-task accuracy but reduces average accuracy, while smaller λ favors slower-improving and harder tasks and higher average relative change.In the nine-task setting, MT-GRPO surpasses the strongest sequential-training variant on worst-task accuracy, 46.3% versus 41.1%.
Experiment 3: MT-GRPO gains persist on Qwen2.5-7B and domain-diverse MATH, ARC, and SciKnowEval, achieving a 12.6% worst-task accuracy gain over DAPO and a ∼60% training-time reduction.It reaches DAPO’s final worst-task accuracy of 25% in ∼5.5 hours versus ∼13.3 hours.

7. Conclusion

The paper concludes that MT-GRPO combines improvement-aware task reweighting with ratio-preserving sampling to improve worst-task performance across tasks. It presents task-wise robustness optimization as practical for general-purpose reasoning models, while identifying out-of-mixture and non-verifiable domains as open boundaries.

Conclusion: MT-GRPO uses improvement-aware task reweighting and ratio-preserving sampling to promote balanced progress and align task weights with actual gradient contributions.The method is presented as a robustness-aware post-training algorithm for improving LLM reasoning across tasks.
Conclusion: Across experiments, MT-GRPO improves worst-task performance while maintaining competitive average accuracy.The conclusion frames task-wise robustness as practical and beneficial for building general-purpose reasoning models.
Limitations: MT-GRPO’s robustness objective is scoped to the training task mixture, leaving robustness to tasks outside that mixture for future work.The training pipeline also relies on verifiable rewards, and behavior in non-verifiable domains remains open.

Impact Statement

Broad competence across diverse reasoning skills is important for reliable LLM deployment, motivating multi-task post-training that balances task performance. MT-GRPO addresses this goal through robustness-aware task weighting and ratio-preserving optimization.

Broad competence across diverse reasoning skills is crucial for LLM deployment in healthcare, education, legal systems, and autonomous decision-making.
The paper develops MT-GRPO to promote balanced reasoning capabilities across tasks for reliable LLM deployment.
The paper states that it does not anticipate unethical applications because the method is designed to enhance AI reasoning capabilities.
Naive multi-task post-training can produce imbalanced outcomes, task interference, and negative transfer across heterogeneous tasks.
MT-GRPO incorporates task-wise robustness into the multi-task RL objective and introduces adaptive reweighting with ratio-preserving batch construction.

C.3. Derivation of Improvement-Aware Task Reweighting

The improvement-aware reweighting derivation chooses task weights and policy updates jointly, prioritizing tasks that underperform or improve slowly while controlling update magnitude.

Small-step policy improvement is approximated by the inner product between each task gradient and the shared update direction.
The method formulates a minimax objective that selects an update direction and task weights to prioritize underperforming or under-improving tasks.
The objective balances worst-case improvement and worst-task performance while regularizing update magnitude to reduce Taylor-approximation error.
A surrogate improvement signal avoids costly or noisy inner products between per-task gradients.
The practical update performs one gradient-descent step on task weights, using softmax parameterization to keep weights in the task simplex.

D.2. Ablation Study

The ablations show that adaptive reweighting improves worst-task accuracy but can distort effective gradient contributions, while ratio-preserving sampling corrects this mismatch. Combining both components gives the strongest balanced performance.

GRPO IWU 0.25 assigns ARC a mean weight of ∼0.80, but its realized post-filtered batch ratio is only ∼0.45.
GRPO RPS outperforms DAPO in Experiment 1 on worst-task accuracy, 54.6% versus 52.6%, and average accuracy, 68.9% versus 65.2%.
The nine-task ablation confirms that MT-GRPO with λ=1.2 achieves the best worst-task accuracy, reinforcing the necessity of both components.
MT-GRPO achieves the best worst-task accuracy in both experiments: 58.8% at λ=0.25 and 46.3% at λ=1.2, with competitive average accuracy.
Sequential training underperforms MT-GRPO on worst-task accuracy across both GRPO and DAPO variants.

D.4. Analysis of RP SAMPLER : Effect of RPS and Acceptance-Aware Sampling

RP SAMPLER preserves intended task proportions after zero-gradient filtering, improving the effective representation of weaker tasks. Although it adds per-step overhead, MT-GRPO reaches worst-task targets faster in wall-clock time.

MT-GRPO initially assigns Zebra a lower weight and increases it as relative task performance evolves, whereas baselines steadily reduce Zebra’s weight.
MT-GRPO reaches all specified worst-task accuracy thresholds in fewer steps than baselines in the nine-task setting.
Without ratio preservation, higher assigned ARC weights still produce a lower effective ARC representation than with RPS enabled.
Removing Acceptance-Aware Sampling increases average resampling rounds, showing that AAS improves sampling efficiency.
MT-GRPO averages 727 s/step versus DAPO’s 659 s/step, a +67 s or +10.2% overhead caused primarily by additional resampling.
Despite the per-step overhead, MT-GRPO achieves higher worst-task performance and reaches the 40% and 50% thresholds faster within an 80-hour budget.

E. Experimental Details and Reproducibility

The experiments use shared datasets, training infrastructure, reward settings, and sampling procedures to evaluate MT-GRPO reproducibly across multitask settings.

Training setup: All experiments use the Qwen-2.5-3B base model, the verl library, two NVIDIA H200 GPUs, and 720 training steps.The same versions of verl and relevant dependencies are used across methods.
Datasets: The datasets contain easy, medium, and hard variants of Countdown, Zebra, and ARC, with 1000 training and 200 test instances per variant.The datasets were generated using the ReasoningGym framework.
Rewards: Rewards are 1.0 for correct answers, 0.1 for incorrectly answered but correctly formatted responses, and 0 otherwise.This reward protocol follows Chen et al. (2025b).
Optimization: Advantages use 72 rollouts per prompt in Experiment 1 and 8 in Experiment 2, with temperature 1.0 and maximum lengths of 1024 prompt and 4096 response tokens.The KL-divergence coefficient is set to 0.
Optimization: Experiment 1 uses global and minibatch sizes of 32 and 8, while Experiment 2 uses 256 and 64; each global step performs four policy updates.Both experiments use AdamW with learning rate 1e −6 and betas (0.9, 0.99).
Sampling: The ratio-preserving sampler draws desired task counts from a multinomial distribution, filters zero-gradient samples, resamples deficiencies, and caps retained samples per task.It uses oversampling and acceptance-aware probabilities to preserve task-wise sampling ratios after filtering.

F. Stabilizing Task-weight Updates through Regularization

Regularization stabilizes task-weight dynamics but remains reward-only, whereas improvement-aware updates better prioritize tasks that are underperforming or improving slowly.

Regularized update: Regularization adds an ℓ2 shrinkage term to the task-weight logit update, damping logit growth and mitigating weight collapse.The regularized update is applied to z = Softmax(ξ).
Motivation: Absolute reward cannot distinguish a task improving rapidly from one with similar reward that has stagnated during training.This motivates tracking how each task’s loss evolves over time.
Experiment 1: In Experiment 1, weak regularization η = 1e −5 causes task-weight collapse onto the current worst-performing task and low worst-task and average accuracy.Countdown is largely ignored for most of training.
Experiment 1: With η = 1e −2 in Experiment 1, regularized updates achieve similar worst-task accuracy but higher average accuracy and average relative change than weak regularization.The gains are driven primarily by improved performance on the easiest task, Countdown.
Experiment 2: In Experiment 2, η = 5e −4 improves worst-task accuracy over baselines but remains below improvement-aware updates on worst-task accuracy and hard-task relative change.With η = 1e −2, worst-task accuracy falls below DAPO and average accuracy remains below DAPO.
Improvement-aware updates: Smaller improvement-aware trade-off parameters prioritize average accuracy and hard-task progress, while larger values focus primarily on worst-task performance.This exposes a controllable trade-off absent from reward-only reweighting.
Conclusion: Overall, regularization can prevent weight collapse but is insufficient by itself to reliably prioritize under-optimized tasks.Improvement signals provide a more effective basis for allocating emphasis to tasks that are underperforming or improving slowly.