多任务GRPO：跨任务的可靠LLM推理

提出MT-GRPO算法，通过自适应任务加权和比例保持采样实现多任务推理的均衡优化

前置知识

GRPO（Group-Relative Policy Optimization）

GRPO是一种广泛用于LLM后训练的策略优化方法。其核心思想是：对每个prompt采样一组G个回复，通过组内归一化计算相对优势函数 $A(x, y_i) = \frac{R(x, y_i) - \text{mean}(\{R(x, y_j)\}_{j=1}^G)}{\text{std}(\{R(x, y_j)\}_{j=1}^G)}$。相比PPO等方法，GRPO无需价值函数，利用组内相对比较构建稳定的优势估计。当所有采样回复获得相同奖励时，优势为零，该prompt不产生梯度更新。

本文是GRPO的多任务扩展，理解GRPO的基本原理是理解MT-GRPO创新点的前提

分布鲁棒优化（Distributionally Robust Optimization, DRO）

DRO是一类优化框架，目标是在最坏情况下保证性能。在多任务学习中，DRO表现为最大化最差任务的性能（minimax目标），通过学习任务权重z，在权重空间和模型参数空间之间交替优化。传统DRO假设损失函数是可靠的标量信号，但在GRPO后训练场景中这一假设不成立。

MT-GRPO的理论框架基于DRO，但针对GRPO的特殊性做出了关键改进

零梯度问题（Zero-Gradient Problem）

在GRPO中，当某个prompt的所有采样回复获得相同奖励时，组内归一化后的优势全为零，导致该prompt对策略更新的梯度贡献为零。不同任务的零梯度prompt比例差异很大，这会导致实际训练批次中各任务的梯度贡献比例偏离预期的任务权重。

这是MT-GRPO需要解决的核心技术挑战之一，直接催生了ratio-preserving sampler的设计

KL散度正则化

在RL后训练中，KL正则化项 $\tau \text{KL}(\pi_\theta \| \pi_{\text{ref}})$ 惩罚后训练策略偏离参考策略的程度，防止模型在优化特定任务奖励时丧失通用能力。$\tau$ 控制正则化强度。

论文的任务级性能度量 $J_k(\theta)$ 包含KL正则化项，影响任务权重更新的信号

研究动机

当前基于GRPO的LLM后训练主要针对单一任务或基准进行优化，将每个任务视为孤立的优化目标。然而，真实世界部署需要模型在多种推理任务上都表现可靠。直接将GRPO应用于多任务场景会导致严重的性能不平衡：一些任务主导优化过程而其他任务停滞不前。论文实验中，标准GRPO在三任务设置下，简单任务Countdown获得高准确率，但困难任务ARC仅30.3%、Zebra仅46%，最差任务性能远低于平均性能。更严重的是，不同任务的零梯度prompt比例差异巨大（如ARC的零梯度比例远高于Zebra），这使得即使显式上调弱任务权重，实际梯度贡献仍然偏离预期，导致优化信号失真。

本文的目标是本文的目标是开发一种多任务后训练算法，使得单一模型能够在多种推理任务上获得均衡且可靠的性能。具体来说，算法需要同时优化两个目标：最大化平均任务性能，以及最小化任务间性能差距（即提升最差任务性能）。论文提出了一个可调参数 $\lambda$ 来控制这两个目标之间的权衡，使用户可以根据实际需求在稳健性和平均性能之间做出选择。

与已有工作不同的是，本文的独特切入角度在于识别并解决了GRPO在多任务场景中的两个结构性限制。首先，传统DRO方法假设损失函数是可靠的性能指标，但GRPO的损失在所有回复正确或全错时都为零，无法区分任务的优劣，论文因此提出使用真实任务奖励而非GRPO损失来更新任务权重。其次，现有方法忽略了零梯度率的跨任务差异，论文创新性地提出了ratio-preserving sampler来确保实际梯度贡献与目标任务权重一致。这两个问题在先前的鲁棒优化文献中从未出现，是GRPO后训练场景独有的挑战。

核心方法

MT-GRPO的整体思路可以概括为：在标准GRPO的基础上引入两个互补机制来解决多任务训练中的不均衡问题。直觉上，如果一个任务表现差或进步慢，就应该给它更多训练机会；但如果直接按任务权重采样，由于不同任务的零梯度率不同，实际梯度贡献会偏离预期。因此，MT-GRPO一方面动态调整任务权重（改进感知的任务重加权），另一方面确保实际训练批次中各任务的比例符合这些权重（比例保持采样器）。技术路线是：在每个训练步骤中，(1) 按当前任务权重采样并过滤零梯度prompt构建训练批次；(2) 用加权GRPO损失更新模型参数；(3) 基于任务奖励和改进信号更新任务权重。

MT-GRPO的核心创新在于两点。第一，改进感知的任务重加权（IWU）：不同于传统DRO仅根据损失更新权重，MT-GRPO使用真实任务奖励 $J_k(\theta)$ 和任务级改进信号 $I_k^{(t)} = J_{\text{GRPO},k}(\theta_{t+1}) - J_{\text{GRPO},k}(\theta_t)$ 的加权组合来更新任务权重。这解决了GRPO损失的歧义性问题——一个任务可能因为所有回复都正确而损失为零（表现好），也可能因为所有回复都错而损失为零（表现差）。第二，比例保持采样器（RP Sampler）：通过过采样和接受感知重采样机制，确保过滤零梯度prompt后的批次中各任务比例与目标任务权重一致。这两个机制共同作用，使得任务权重能够真实反映为梯度贡献。

方法步骤详情

MT-GRPO的执行流程如下：输入包括K个任务的数据集 $\{D_1, \ldots, D_K\}$、批次大小B、每个prompt的rollout数N、初始策略参数 $\theta_0$、初始任务权重logits $\xi_0$。在每个训练步骤t：(1) 批次构建：使用RP Sampler根据当前任务权重 $z_t$ 采样B个prompt，采样时先根据Multinomial分布确定各任务的目标数量 $(n_1, \ldots, n_K) \sim \text{Multinomial}(B, z)$，然后过采样以补偿零梯度过滤，维护过滤率估计 $\rho_k$ 和膨胀因子 $m_k = \min\{1/(1-\rho_k), M_{\text{acc}}\}$，直到各任务满足目标数量或达到重采样预算；(2) 策略更新：用加权GRPO损失 $\theta_{t+1} = \theta_t + \gamma_t \sum_{k=1}^K z_k \nabla_\theta J_{\text{GRPO},k}(\theta_t)$ 更新模型参数；(3) 计算各任务改进信号 $I_k^{(t)}$；(4) 权重更新：使用IWU计算综合信号 $s_k^{(t)} = I_k^{(t)} + \lambda J_k(\theta_t)$，梯度 $g_k^{(t)} = z_{k,t}(s_k^{(t)} - \sum_j z_{j,t} s_j^{(t)})$，更新logits $\xi_{t+1} = \xi_t - \beta g_t$，得到新权重 $z_{t+1} = \text{Softmax}(\xi_{t+1})$。

技术新颖性

MT-GRPO的技术新颖性体现在三个方面。首先，识别并形式化了GRPO后训练场景中两个前所未有的问题：GRPO损失作为任务性能指标的歧义性，以及零梯度率的跨任务差异对任务权重实现的影响。其次，改进感知任务重加权的引入，通过将改进信号 $I_k^{(t)}$ 纳入权重更新，防止了权重退化为单任务优化（如论文Figure 2所示，仅使用奖励会导致权重坍缩到当前最差任务）。参数 $\lambda$ 提供了从最差任务鲁棒性到平均性能的连续调节能力。第三，ratio-preserving sampler的设计，通过膨胀采样和接受率跟踪机制，在保持计算效率的同时确保了任务权重与实际梯度贡献的一致性。这些创新共同填补了鲁棒优化理论与GRPO后训练实践之间的空白。

实验结果

论文在Qwen-2.5-3B模型上进行了两类实验。实验1（三任务设置）：在中等难度的Countdown、Zebra、ARC三个任务上，MT-GRPO（$\lambda=0.2$）最差任务准确率达45.8%，相比GRPO的30.3%提升15.5个百分点，相比DAPO的24.6%提升21.2个百分点；平均准确率58.8%，优于GRPO的52.6%和DAPO的51.0%。MT-GRPO达到50%最差任务准确率所需的训练步数比基线少约50%。实验2（九任务设置）：在easy/medium/hard三种难度的Countdown、Zebra、ARC共9个任务上，MT-GRPO（$\lambda=1.2$）最差任务准确率达46.3%，相比GRPO提升16个百分点，相比DAPO提升6个百分点。$\lambda$ 参数有效控制了鲁棒性与平均性能的权衡：较大的 $\lambda$ 提升最差任务性能但降低平均准确率，较小的 $\lambda$ 则促进更均衡的任务改进。

查看结构化数据

任务	指标	本文	基线	提升
3-task setting (Countdown/Zebra/ARC medium)	Worst-Task Accuracy	45.8% (λ=0.2)	GRPO: 30.3%, DAPO: 24.6%	+15.5pp over GRPO, +21.2pp over DAPO
3-task setting (Countdown/Zebra/ARC medium)	Avg Accuracy	58.8% (λ=0.2)	GRPO: 52.6%, DAPO: 51.0%	+6.2pp over GRPO, +7.8pp over DAPO
9-task setting (easy/medium/hard variants)	Worst-Task Accuracy	46.3% (λ=1.2)	GRPO: ~30%, DAPO: ~40%	+16% over GRPO, +6% over DAPO
3-task setting	Steps to 50% worst-task accuracy	~200 steps	GRPO/DAPO: ~400+ steps	50% fewer training steps

局限与改进

论文存在以下局限性：首先，实验规模有限，仅在Qwen-2.5-3B（3B参数）上验证，未在更大规模模型（如7B、70B）上测试，模型规模对MT-GRPO效果的影响尚不清楚。其次，任务类型局限于规划任务（Countdown、Zebra）和归纳推理任务（ARC），未涉及代码生成、数学推理等更广泛的任务类型。第三，论文采用ReasoningGym框架生成的数据集，这些合成任务与真实世界推理任务存在差异。第四，MT-GRPO引入了额外超参数 $\lambda$ 和任务权重更新的学习率 $\beta$，虽然论文展示了 $\lambda$ 的调节效果，但最优值需要根据任务组合进行调优。第五，ratio-preserving sampler的过采样机制在任务零梯度率极高时可能导致采样效率下降。此外，论文未讨论MT-GRPO在任务数量更多（如几十个任务）时的扩展性和计算开销。

独立分析的弱点

论文存在几个值得关注的弱点。(1) 单一模型规模验证：仅在3B参数模型上实验，可能无法代表更大规模模型的行为，建议在7B-70B规模上验证。(2) 合成任务局限：ReasoningGym生成的任务相对标准化，真实世界的推理任务可能具有更复杂的难度分布和任务交互，需要在自然任务集合上验证。(3) 超参数敏感性：$\lambda$ 和 $\beta$ 的选择对性能影响显著，但论文缺乏自动调节这些超参数的机制，可以探索基于元学习的自适应调节。(4) 缺乏理论收敛保证：虽然论文提供了直觉分析，但MT-GRPO在非凸设置下的收敛性质缺乏严格理论分析。(5) 任务权重的可解释性：Softmax参数化使得权重的语义不够直观，可以考虑提供任务难度和改进趋势的可视化工具。

未来方向

基于本文成果可以延伸以下研究方向。(1) 自适应 $\lambda$ 调节：当前 $\lambda$ 是固定超参数，可以开发基于任务改进动态自动调节 $\lambda$ 的机制，实现更智能的鲁棒性-性能权衡。(2) 大规模任务扩展：将MT-GRPO应用于几十甚至上百个任务的场景，研究权重学习在高维任务空间中的行为。(3) 与课程学习结合：MT-GRPO的自适应权重与课程学习思想互补，可以探索在任务内部引入难度递进。(4) 跨模态任务：将方法扩展到多模态推理任务（如视觉推理、代码生成）。(5) 理论分析：建立MT-GRPO在GRPO后训练场景下的收敛保证，特别是零梯度过滤对收敛速度的影响。(6) 在线学习率调节：将权重更新学习率 $\beta$ 也做成自适应的，避免对固定步长的依赖。

复现评估

论文具有良好的可复现性。代码已在GitHub开源（https://github.com/rsshyam/MT-GRPO），使用Qwen-2.5-3B作为基础模型，这是公开可获取的模型。数据集来自ReasoningGym框架，也是公开资源。训练在标准硬件上进行，论文提供了详细的训练配置（包括批次大小、rollout数量、学习率等）。算法伪代码完整清晰（Algorithm 1和2），关键实现细节（如RP Sampler的过采样机制、IWU的梯度计算）都有明确描述。主要复现难点在于：(1) 需要GPU资源进行多任务GRPO训练；(2) 超参数 $\lambda$ 需要针对不同任务组合调优；(3) 零梯度率的估计需要在训练过程中动态更新。总体来说，按照论文描述和开源代码，中等计算资源即可复现核心结果。

多任务GRPO：跨任务的可靠LLM推理 Multi-Task GRPO: Reliable LLM Reasoning Across Tasks