Imbalanced Gradients in RL Post-Training of Multi-Task LLMs

📄 arXiv: 2510.19178v2 📥 PDF

作者: Runzhe Wu, Ankur Samanta, Ayush Jain, Scott Fujimoto, Jeongyeol Kwon, Ben Kretzu, Youliang Yu, Kaveh Hassani, Boris Vidolov, Yonathan Efroni

分类: cs.LG, cs.AI

发布日期: 2025-10-22 (更新: 2025-10-26)

💡 一句话要点

揭示RL后训练中多任务LLM梯度不平衡问题，并指出其对模型优化的负面影响。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多任务学习 强化学习 大型语言模型 梯度不平衡 后训练

📋 核心要点

  1. 现有LLM多任务后训练方法假设各任务梯度大小相似，忽略了任务间差异可能导致的优化偏差。
  2. 该研究揭示了RL后训练中梯度不平衡现象，发现大梯度任务并不一定带来更大的性能提升。
  3. 分析表明，梯度不平衡并非由训练奖励或优势等因素导致，而是源于任务本身的固有属性。

📝 摘要（中文）

大型语言模型（LLM）的多任务后训练通常通过混合来自不同任务的数据集并联合优化它们来实现。这种方法隐含地假设所有任务都贡献了相似大小的梯度；当这个假设不成立时，优化会偏向于大梯度任务。然而，本文表明，这种假设在强化学习（RL）后训练中是不成立的：某些任务产生明显更大的梯度，从而使更新偏向于这些任务。只有当更大的梯度意味着任务上更大的学习收益（即，更大的性能提升）时，这种梯度不平衡才是合理的——但我们发现事实并非如此。大梯度任务可以实现与小梯度任务相似甚至更低的学习收益。进一步的分析表明，这些梯度不平衡不能用典型的训练统计数据（如训练奖励或优势）来解释，这表明它们源于任务之间的内在差异。这告诫人们不要进行简单的混合数据集，并呼吁未来对LLM进行有原则的梯度级别校正。

🔬 方法详解

问题定义：论文旨在解决在多任务强化学习（RL）后训练大型语言模型（LLM）时，由于不同任务梯度大小不平衡而导致的模型优化偏差问题。现有方法通常简单地混合不同任务的数据集进行联合训练，忽略了任务间梯度差异可能导致模型过度拟合大梯度任务，而忽略小梯度任务。这种做法的痛点在于无法保证所有任务都能得到充分学习，最终影响模型的整体性能。

核心思路：论文的核心思路是深入分析RL后训练中不同任务的梯度大小，并验证梯度大小与任务学习收益之间的关系。通过实验发现，梯度大的任务并不一定带来更大的性能提升，这表明简单的梯度大小并不能作为任务重要性的有效指标。因此，需要对梯度进行校正，以避免模型优化过程中的偏差。

技术框架：论文没有提出一个全新的技术框架，而是侧重于对现有多任务RL后训练流程的分析和诊断。其研究方法主要包括：1) 实验验证梯度不平衡现象的存在；2) 分析梯度不平衡的原因，排除训练奖励和优势等因素的影响；3) 评估梯度大小与任务学习收益之间的相关性。

关键创新：论文最重要的技术创新点在于发现了RL后训练中梯度不平衡现象，并证明了这种不平衡与任务学习收益之间不存在正相关关系。这一发现挑战了现有方法的隐含假设，为未来的研究方向提供了新的思路，即需要设计更有效的梯度校正方法，以解决多任务学习中的优化偏差问题。

关键设计：论文的关键设计在于实验分析，包括：1) 选择合适的RL任务和LLM模型；2) 设计实验来测量不同任务的梯度大小；3) 设计实验来评估不同任务的学习收益；4) 分析梯度大小与学习收益之间的相关性。论文没有涉及具体的参数设置、损失函数或网络结构的设计，而是侧重于对现有训练流程的分析和诊断。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在RL后训练中，某些任务的梯度显著大于其他任务，但大梯度任务的学习收益并不一定高于小梯度任务。例如，某些大梯度任务的性能提升甚至低于小梯度任务。这表明简单的混合数据集训练方法会导致优化偏差，需要进行梯度级别的校正。

🎯 应用场景

该研究成果可应用于提升多任务强化学习的训练效果，尤其是在涉及复杂任务和大型语言模型的场景下。通过对梯度进行校正，可以避免模型过度拟合某些任务，从而提高模型在所有任务上的泛化能力。这对于开发通用型AI系统具有重要意义，例如智能助手、机器人控制等。

📄 摘要（原文）

Multi-task post-training of large language models (LLMs) is typically performed by mixing datasets from different tasks and optimizing them jointly. This approach implicitly assumes that all tasks contribute gradients of similar magnitudes; when this assumption fails, optimization becomes biased toward large-gradient tasks. In this paper, however, we show that this assumption fails in RL post-training: certain tasks produce significantly larger gradients, thus biasing updates toward those tasks. Such gradient imbalance would be justified only if larger gradients implied larger learning gains on the tasks (i.e., larger performance improvements) -- but we find this is not true. Large-gradient tasks can achieve similar or even much lower learning gains than small-gradient ones. Further analyses reveal that these gradient imbalances cannot be explained by typical training statistics such as training rewards or advantages, suggesting that they arise from the inherent differences between tasks. This cautions against naive dataset mixing and calls for future work on principled gradient-level corrections for LLMs.