基于群组分布鲁棒优化的LLM推理强化学习 Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning
提出多对手GDRO框架，通过动态难度分组和自适应资源分配提升LLM推理性能
前置知识
Group Relative Policy Optimization (GRPO)
GRPO是PPO的一种计算高效替代方案，它通过在同一个提示的多个采样响应之间构建组内基线来消除学习值函数的需求。具体来说，对于每个提示，GRPO生成n个响应，计算每个响应的奖励，然后通过组内标准化得到优势估计。这种方法避免了学习单独的值网络，简化了训练流程，同时保持了策略梯度的有效性。GRPO通过裁剪代理目标函数（类似于PPO）来更新策略，确保训练的稳定性。
本文提出的框架是在GRPO基础上构建的，理解GRPO的机制是理解本文如何通过动态调整提示采样分布和rollout分配来改进GRPO的前提。
Group Distributionally Robust Optimization (GDRO)
GDRO是一种优化框架，旨在最小化在最坏情况下的子群体风险。它假设数据可以划分为多个组（例如不同的难度级别），然后优化目标是在所有组上表现最差的那个组的风险最小化。GDRO通过引入一个对抗性分布q，在概率单纯形上最大化期望损失，从而迫使学习者改善在困难组上的性能。这种方法在处理数据异质性和隐藏分层问题时特别有效。
本文的核心思想就是将GDRO引入LLM推理的后训练阶段，通过定义动态难度组并优化最坏情况性能来提升模型的鲁棒性和推理能力。
Online Difficulty Classifier (在线难度分类器)
本文提出的在线难度分类器基于模型的实时通过率(pass@k)来动态地将提示分配到不同的难度组。它维护一个滑动窗口来估计每个提示的通过率，然后根据通过率将提示分配到离散的准确率区间（例如0-0.1，0.1-0.2等）。为了稳定性，它采用了滞后机制：只有当移动平均准确率跨越区间边界超过一个阈值δ时，才会重新分配提示。这种动态分组使得框架能够适应模型能力的变化，而无需依赖静态的数据集元数据。
这是本文方法的基础组件，它使得GDRO框架能够工作在动态变化的难度组上，从而实现自适应的课程学习。
EXP3P算法 (Exponential-weight algorithm for Exploration and Exploitation with Penalizations)
EXP3P是一种bandit算法，用于在对抗性环境中进行决策。它维护一个关于各个臂（选项）的权重分布，并根据观察到的损失来更新这些权重。更新规则是指数权重的：表现差的臂的权重会指数级减少，而表现好的臂的权重会相对增加。EXP3P通过混合一个均匀分布来确保探索，并且具有无遗憾保证。在本文中，GDRO-EXP3P是专门为GDRO问题设计的EXP3变体。
本文使用GDRO-EXP3P来实现两个对抗性控制器（Prompt-GDRO和Rollout-GDRO），理解这个算法是理解本文如何在线更新对抗性分布的关键。
Shadow Price (影子价格)
在经济学和优化理论中，影子价格是与约束条件相关联的拉格朗日乘子，它表示放松约束一单位所能带来的目标函数改善量。在本文的Rollout-GDRO中，影子价格μ对应于平均rollout预算约束的拉格朗日乘子。它用于权衡方差减少和计算成本：更高的μ意味着计算资源更昂贵，因此控制器会更保守地分配rollouts。通过双重上升法更新μ，确保总计算预算保持不变。
理解影子价格的概念有助于理解Rollout-GDRO如何将资源分配问题形式化为一个经济控制问题，以及如何在不增加总计算成本的情况下优化资源分配。
研究动机
当前主流的大语言模型推理强化学习范式，如Group Relative Policy Optimization (GRPO)，依赖于静态均匀性假设：从训练分布中均匀采样提示，并为每个提示分配固定的计算预算（例如固定数量的rollouts）。然而，推理数据本质上是异构且重尾的，包含不同难度的子领域（如基础代数与奥林匹克数论）。在均匀采样下，优化被最频繁、通常也更容易的模式主导，学习信号集中在“简单核心”上，而困难长尾中的错误持续存在。同时，静态的rollout预算分配无法捕捉到不同提示对探索需求的差异性：已解决的提示产生低方差梯度，而高熵的“前沿”提示需要大量探索来减少梯度方差。这种刚性导致计算资源被浪费在冗余验证上，同时关键失败模式得不到充分探索。
本文的目标是本文旨在提出一个优化优先的框架，通过动态调整训练分布来超越静态均匀性模型。具体目标包括：1) 设计一个在线难度分类器，基于模型的实时表现动态地将提示划分为不同的难度组；2) 实现一个提示重加权对抗器（Prompt-GDRO），通过EMA去偏的乘性权重bandit采样器，针对密集难度边际进行目标采样，提升对持久困难组的权重，而无频率偏差；3) 实现一个rollout分配对抗器（Rollout-GDRO），使用影子价格控制器在固定平均预算下重新分配rollouts，最大化对困难任务的梯度方差减少，实现计算中性；4) 在DAPO 14.1k数据集上使用Qwen3-Base模型（1.7B、4B、8B）验证框架的有效性，目标是在pass@8准确率上相对于GRPO基线实现显著提升。
与已有工作不同的是，与已有的课程学习、自适应采样或资源分配工作相比，本文的独特切入角度在于将分布鲁棒优化（GDRO）的视角引入LLM推理的后训练阶段，并将其形式化为一个包含两个独立对抗器的多对手零和游戏。已有工作通常只关注数据选择或计算分配中的一个方面，而本文同时优化这两个控制旋钮。更重要的是，本文提出的框架是数据无关的：难度分组完全基于模型的实时表现（在线pass@k），而非依赖于静态的数据集元数据或人工定义的分组。这使得框架能够自动发现并提升对当前策略仍有学习价值的提示子集，而不是重复训练已解决的、低价值的样本。此外，本文提供了理论保证，证明了Prompt-GDRO对应于一个熵正则化的GDRO代理目标，并具有无遗憾游戏解释。
核心方法
本文的方法可以类比为一个智能的“教练”系统，它有两个独立的“助手”来优化训练过程。第一个助手（Prompt-GDRO）负责决定“练什么题”，它动态地识别出学生（模型）仍然薄弱的题目类型（难度组），并增加这些题目的练习频率。第二个助手（Rollout-GDRO）负责决定“每道题练多少遍”，它将有限的计算资源（rollout次数）从已经掌握良好的题目重新分配到需要更多探索的难题上。这两个助手都基于同一个实时“仪表盘”（在线难度分类器）的信息，该仪表盘根据学生的实时表现（pass@k）对题目进行动态分类。整个系统被形式化为一个零和游戏：两个助手（对抗器）试图最大化学生的损失（暴露弱点），而学生（策略）则试图最小化损失（提升能力）。通过这种对抗性互动，系统自动生成一个适应性的课程，引导训练资源流向当前最具学习价值的“推理前沿”。
本文的核心创新点在于提出了一个多对手GDRO框架，它同时优化两个独立的控制旋钮：数据分布（提示采样）和计算分配（rollout预算）。与已有方法最本质的区别在于，本文将这两个旋钮都形式化为GDRO框架内的对抗性游戏，并通过在线难度分组将它们耦合起来。具体来说，Prompt-GDRO通过EMA去偏的指数权重更新来实现，它针对密集难度边际（平均损失）而非累积损失，从而避免了频率偏差，确保稀有但困难的组能够获得足够的权重。Rollout-GDRO则通过影子价格控制器实现，它将资源分配问题形式化为一个约束优化问题，在固定平均rollout预算下，通过双重上升法更新拉格朗日乘子μ，动态地将更多rollouts分配给高方差的组。这种设计使得两个对抗器都能在不增加总计算成本的情况下，显著提升梯度信号的质量和针对性。
方法步骤详情
本文方法分为三个主要步骤。第一步是动态分组：在线难度分类器为每个提示维护一个唯一的标识符（UID）和基于滑动窗口的pass@k统计量。在训练步骤t，对每个提示采样k个rollouts，计算任意一个正确的指示函数，然后通过长度为H的滑动窗口估计pass@k。根据估计的pass@k将提示分配到离散的准确率区间（例如[0,0.1)为accbin 0，[0.1,0.2)为accbin 1等），形成动态分组gt。为了稳定性，引入滞后机制：只有当移动平均准确率跨越区间边界超过阈值δ时才重新分配。第二步是Prompt-GDRO：对于每个难度组b，维护一个EMA难度分数St(b)，它通过衰减系数β对平均提示级损失进行平滑。然后通过指数化裁剪后的分数得到未归一化的权重ωt(b)，并通过混合均匀分布γ得到最终的采样概率qt(b)。在实际实现中，并不物理重采样数据，而是通过缩放每个提示的rollout优势来实现重加权。第三步是Rollout-GDRO：将rollout数量nb视为离散变量，通过拉格朗日松弛法求解约束优化问题。为每个组维护一个独立的EXP3P实例，其中“臂”是离散的rollout计数。在每个步骤，选择最大化联合概率且严格满足全局预算等式约束的配置{nb}。双重变量μ通过双重上升法更新，确保方法在计算上保持中性。
技术新颖性
本文的技术新颖性体现在多个方面。首先，这是首次将分布鲁棒优化（GDRO）系统地引入LLM推理的后训练阶段，并将其形式化为包含两个独立对抗器的多对手游戏。已有工作大多只关注数据选择或计算分配中的一个方面，而本文同时优化两者。其次，本文提出了EMA去偏的EXP3P算法变体，用于实现Prompt-GDRO。该算法针对密集难度边际（平均损失）进行更新，而非累积损失，从而避免了频率偏差，确保稀有但困难的组能够获得足够的权重。这与标准的GDRO实现有本质区别。第三，本文将rollout分配问题形式化为一个影子价格控制的资源分配问题，并证明了方差最优分配遵循平方根定律。这种经济学视角为计算资源的自适应分配提供了新的理论框架。第四，本文提出的数据无关的在线难度分类器，使得整个框架无需依赖静态的数据集元数据，能够自动适应模型能力的变化。最后，本文提供了理论分析，证明了Prompt-GDRO对应于熵正则化的GDRO代理目标，并具有无遗憾游戏解释，为方法的有效性提供了理论保障。
实验结果
本文在DAPO 14.1k数学推理数据集上使用Qwen3-Base模型（1.7B、4B、8B）进行了全面实验，结果表明提出的两个机制均能显著提升推理性能。Prompt-GDRO在pass@8准确率上平均提升了10.6%，具体来说，在1.7B模型上提升了9.74%（从50.74%到55.68%），在4B模型上提升了13.13%（从56.31%到63.70%），在8B模型上提升了8.96%（从62.04%到67.60%）。Rollout-GDRO同样表现出色，平均提升了10.1%，在1.7B模型上提升了10.64%（到56.14%），在4B模型上提升了10.59%（到62.27%），在8B模型上提升了9.20%（到67.75%）。值得注意的是，这些提升是在计算中性的前提下实现的：Prompt-GDRO保持rollout预算不变，Rollout-GDRO保持平均rollout预算不变。除了pass@8，论文还报告了在MATH 500、AIME、AMC、MINERVA、OLYMPIAD和GPQA等基准上的mean@8准确率。例如，在4B模型上，Prompt-GDRO在MATH 500上达到75.78%（基线72.05%），在AIME上达到12.92%（基线11.25%），在AMC上达到64.06%（基线60.94%），在MINERVA上达到26.72%（基线17.79%），在OLYMPIAD上达到40.98%（基线30.48%），在GPQA上达到40.88%（基线35.54%）。这些结果表明，框架在多种数学推理任务上都带来了显著且一致的改进。定性分析揭示了一个涌现的课程：对抗器将资源转移到不断演变的推理前沿，增强了推理模型的性能。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理综合 (pass@8) | pass@8准确率 | Prompt-GDRO: 1.7B: 55.68%, 4B: 63.70%, 8B: 67.60%; Rollout-GDRO: 1.7B: 56.14%, 4B: 62.27%, 8B: 67.75% | GRPO基线: 1.7B: 50.74%, 4B: 56.31%, 8B: 62.04% | Prompt-GDRO: +9.74%, +13.13%, +8.96%; Rollout-GDRO: +10.64%, +10.59%, +9.20% |
| MATH 500 | mean@8准确率 | Prompt-GDRO: 1.7B: 63.20%, 4B: 75.78%, 8B: 76.18%; Rollout-GDRO: 1.7B: 63.98%, 4B: 75.20%, 8B: 77.88% | GRPO: 1.7B: 50.62%, 4B: 72.05%, 8B: 73.45% | 显著提升，例如4B模型Prompt-GDRO提升3.73个百分点 |
| AIME (2024 & 2025平均) | mean@8准确率 | Prompt-GDRO: 1.7B: 6.88%, 4B: 12.92%, 8B: 16.04%; Rollout-GDRO: 1.7B: 7.50%, 4B: 13.96%, 8B: 15.63% | GRPO: 1.7B: 5.42%, 4B: 11.25%, 8B: 14.38% | 提升明显，例如4B模型Prompt-GDRO提升1.67个百分点 |
| MINERVA | mean@8准确率 | Prompt-GDRO: 1.7B: 14.61%, 4B: 26.72%, 8B: 32.17%; Rollout-GDRO: 1.7B: 17.28%, 4B: 26.47%, 8B: 29.55% | GRPO: 1.7B: 14.56%, 4B: 17.79%, 8B: 28.17% | 4B模型Prompt-GDRO提升8.93个百分点，提升尤为显著 |
局限与改进
本文的局限性主要体现在以下几个方面。首先，实验范围相对有限：所有实验都在DAPO 14.1k数据集和Qwen3-Base模型上进行，尚未在更广泛的数据集、模型架构或任务领域（如代码生成、创意写作）上验证框架的通用性。其次，论文没有进行全因子消融实验来全面评估各个组件（如分组方式、提示重加权、rollout分配、在线分组选择）的独立贡献和交互作用。第三，系统开销是一个实际限制：与标准GRPO相比，Prompt-GDRO和Rollout-GDRO引入了额外的在线机制（如维护难度分类器、更新对抗性分布、在Rollout-GDRO中执行动态规划选择），这增加了训练时间的开销。论文报告了在Qwen3-4B上，GRPO的优势阶段需要0.043秒/步，Prompt-GDRO需要0.355秒/步，Rollout-GDRO需要0.446秒/步。第四，框架对在线分组和奖励噪声敏感：基于pass@k的在线分组估计在训练早期可能具有高方差，而分组边界的不连续重分配可能会引入偏差。最后，本文将两个对抗器作为独立模块进行分析和评估，没有研究它们的联合训练动态，这可能限制了性能的进一步提升。
独立分析的弱点
尽管本文取得了显著成果，但仍存在几个值得改进的弱点。首先，系统开销较大，限制了其在大规模训练中的实用性。建议探索异步和流式变体，例如延迟分组更新、批量对抗器步骤或部分卸载的记账机制，以在保持核心目标的同时最小化训练减速。其次，在线分组估计的稳定性有待提高。在训练早期，pass@k估计可能具有高方差，导致分组不稳定。建议引入基于贝叶斯估计器或置信区间分配规则的、具有不确定性感知的分组机制。第三，框架目前只使用了二元的正确/错误奖励信号，没有利用更丰富的过程级监督（如步骤级奖励或验证器信号）。整合过程级监督信号可以稳定难度估计，并提供更细粒度的学习信号。第四，两个对抗器的独立评估可能忽略了它们的协同效应。建议进行系统性的联合训练研究，并探索阶段化课程策略，例如先使用Prompt-GDRO建立基础课程，再引入Rollout-GDRO进行精细资源分配。最后，框架的超参数（如分组数量、平滑窗口H、滞后阈值δ、学习率ηq等）需要更广泛的搜索。初步实验表明性能在约6个分组时达到峰值，但这需要更鲁棒的验证。
未来方向
基于本文的成果，未来研究方向可以从多个维度展开。首先，最直接的方向是研究两个对抗器的联合训练动态，探索它们之间的耦合效应，并设计多时间尺度的优化策略。其次，可以扩展框架的应用范围，将其应用于更广泛的推理任务（如代码生成、逻辑谜题）、更大的模型规模（如70B+）以及更复杂的数据混合。第三，可以探索更丰富的歧义集和解决方案，例如f-散度DRO、Wasserstein DRO或χ²-不确定性集，这些可能更适合未来的规模化需求。第四，可以将对抗性分布塑造与经验生成和持续后训练相结合，形成一个闭环系统：模型提出新问题或扰动，评估自身失败，然后使用Prompt-GDRO/Rollout-GDRO对结果前沿进行优先级排序。这自然地连接到自训练和自改进范式，如STaR式引导、Quiet-STaR式隐式“思考”训练、自奖励/基于评判的优化以及RLAIF式可扩展反馈。最后，可以设计针对特定推理目标（如安全性、风险规避、约束满足）的对抗器，超越当前的准确性目标，连接到依赖于规则或AI生成反馈的对齐框架。
复现评估
本文的复现性评估如下。数据方面，论文使用了公开可用的DAPO 14.1k数据集（HuggingFace链接已提供），这是一个经过处理的英文数学推理数据集，便于获取。模型方面，使用了Qwen3-Base模型（1.7B、4B、8B），这些是公开可用的预训练模型。代码方面，论文没有明确说明是否开源，但方法描述相对详细，包含算法伪代码和实现细节，理论上可以根据论文复现。算力需求方面，实验在多个模型规模上进行，最小的1.7B模型可能需要中等规模的GPU资源，而8B模型则需要更强大的计算资源。复现难度中等：核心算法（在线分组、EMA更新、EXP3P）相对直接，但实现细节（如动态规划选择rollout配置、双重上升更新）需要仔细处理。系统开销可能是一个挑战，尤其是在大规模训练时。总体而言，对于有RL和LLM训练经验的团队，复现本文结果应该是可行的，但需要一定的工程努力来优化效率和稳定性。
论文图表
