arXiv:2604.23318cs.CLcs.LG2026-04

用隐藏状态距离自动定位推理错误点,提升强化学习更新精度。

Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

论文配图:Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance
图 1 · 摘自论文原文
  • 通过跨度级隐藏状态分布的Wasserstein距离捕捉推理分歧点
  • 在数学与代码生成任务上超越标准GRPO,接近监督奖励模型效果
  • 无需额外标注或训练奖励模型,仅改写训练流程即可

组相对策略优化(GRPO)在具有可验证奖励的强化学习(RLVR)中采用粗粒度信用分配,为单个轨迹中的所有标记分配相同优势。过程奖励模型可提供更细粒度监督,但需步骤级标注或额外奖励建模。我们发现,隐藏状态分布包含可用于局部推理质量判断的有用信号,仅需RLVR中可用的结果级正确性标签即可提取。具体而言,在每个GRPO组内,正确与错误轨迹在推理分歧区域的跨度级隐藏状态分布间Wasserstein距离增大。该关联在不同样本间及单个轨迹内部均成立,表明隐藏状态分布差异可作为细粒度信用分配的自监督信号。我们通过分离定理形式化此观察:在温和结构假设下,当总体分布差距超过有限样本噪声时,分歧后跨度的Wasserstein距离大于分歧前跨度。受此启发,提出跨度级隐藏状态增强优势重加权(SHEAR),通过跨度级Wasserstein距离缩放令牌级优势,放大与对立组显著分离的标记的更新。该方法无需额外模型,仅需对训练流程进行最小修改。在五个数学推理基准和五个代码生成基准上的实验表明,SHEAR优于标准GRPO,且表现接近监督过程奖励模型,同时无需额外标注或奖励模型训练。

原文摘要 · Abstract (English)

Group Relative Policy Optimization (GRPO) performs coarse-grained credit assignment in reinforcement learning with verifiable rewards (RLVR) by assigning the same advantage to all tokens in a rollout. Process reward models can provide finer-grained supervision, but they require step-level annotation or additional reward modeling. We show that hidden-state distributions contain a useful signal for local reasoning quality that can be extracted using only outcome-level correctness labels available in RLVR. Specifically, within each GRPO group, the Wasserstein distance between span-level hidden state distributions of correct and incorrect rollouts increases around regions where their local reasoning quality diverges. This association holds both across examples and within individual trajectories, suggesting that hidden-state distributional divergence can serve as a self-supervision signal for fine-grained credit assignment. We formalize this observation with a separation theorem showing that, under mild structural assumptions, post-divergence spans have larger Wasserstein distances than pre-divergence spans whenever the population-level distributional gap exceeds finite-sample noise. Motivated by this result, we propose \textbf{S}pan-level \textbf{H}idden state \textbf{E}nabled \textbf{A}dvantage \textbf{R}eweighting (SHEAR), which modifies GRPO by using span-level Wasserstein distances to scale token-level advantages, amplifying updates on tokens whose hidden states are more separated from the opposing group. The method requires no additional model and only minimal changes to the training pipeline. Experiments on five mathematical reasoning benchmarks and five code generation benchmarks show improvements over standard GRPO and strong performance relative to supervised process reward models, while requiring no additional annotation or reward model training.

强化学习信用分配隐藏状态推理分析

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。