跳转至

Credit Assignment with Resets in Language Model Reasoning

会议: NeurIPS 2026
arXiv: 2605.25507
领域: Reasoning
关键词: 强化学习, 信用分配, 语言模型推理, 状态重置, 保守策略迭代

一句话总结

本文提出基于重置的强化学习后训练机制(RRPO 与 SRPO),通过模型自定位失败轨迹的首个错误推理步并重置采样反事实后缀,从理论和实证上解决了可验证奖励下全轨迹均匀分配信用的根本缺陷。

研究背景与动机

在大语言模型的多步复杂推理后训练中,带可验证奖励的强化学习(RLVR,如 GRPO 等)已成为提升模型数学与逻辑推理能力的核心范式。然而,现行方法普遍采用稀疏的终末结果奖励,并将该单一标量奖励均匀广播分摊到整条推理轨迹的每一个 token 上。在长链条推理场景下,这种均匀分配完全抹杀了不同推理步骤对成败的真实因果贡献:无论前期构思多么精妙、或是后半段哪一步推导突然出现逻辑跳跃,所有步骤都被赋予完全相同的梯度权重,导致策略梯度信号严重稀释甚至产生有害更新。

尽管直觉上可以通过过程奖励模型(PRM)或逐步人类标注来提供细粒度反馈,但这些方案依赖高昂的人工标注成本或复杂的辅助模型训练流水线,且容易受到奖励欺骗与分布偏移的干扰。人类和生物认知系统在面对失败决策时,往往能够心理回溯至关键分歧点,并通过模拟“如果当时采取另一举动”的反事实推演来归因责任。这启发我们思考:能否仅利用大语言模型自身的能力,在不引入外部监督的前提下,将状态重置机制引入后训练框架,从而实现高效、精准的步骤级信用分配?

本文的切入角度是将强化学习经典理论中的状态重置(Resets)重构为语言模型推理任务的信用分配原语。核心 idea:通过让模型自我定位错误思维步骤并重置采样替代后缀,将全轨迹粗粒度更新转化为针对可改进前缀的反事实后缀信用分配,并在理论上证明其相比随机重置能实现样本复杂度与迭代改进速度的大幅提升。

方法详解

整体框架

本文构建了面向思维级推理的 Thought MDP,将离散的自然语言推理切分为语义自洽的思维步骤。整个训练管线围绕构建双分组混合经验池展开:对每个提示词,系统首先采样基础组(Base Group)以维持策略的全面探索覆盖;当遇到错误轨迹时,利用随机重置或自定位机制截取可改进前缀,由此分支出共享前缀组(Shared-Prefix Group)并采样多条替代后缀。通过组内相对优势归一化与前缀梯度掩码,策略更新被精确聚焦于分歧后缀。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入提示词与 Thought MDP"] --> B["阶段 1:基础组多路 Rollout 与失败种子检测"]
    B --> C["阶段 2:重置状态定位<br/>随机重置 (RRPO) vs 自定位重置 (SRPO)"]
    C --> D["阶段 3:共享前缀反事实后缀重采样"]
    D --> E["阶段 4:组内相对优势估计与前缀梯度掩码优化"]

关键设计

1. Thought MDP 建模与自定位重置机制:将粗粒度序列拆解为语义自洽的决策基元

传统的 token 级切分过于细碎,无法承载稀疏奖励下的归因语义;而依赖外部模型打断或固定 token 窗口又会强行割裂语言模型的完整语义表达。本文在输入提示词中引入特定的思维分隔标识,将生成过程形式化为 Thought MDP:初始状态为提示词 \(x_0\),在第 \(h\) 步模型生成一段以特定停止符结尾的思维块 \(\tilde{y}_h \sim \pi_h(\cdot \mid x_0, \tilde{y}_{1:h-1})\)。在这种结构化表达下,RRPO(Random-Reset Policy Optimization)均匀随机选择截断点 \(h^* \sim \mathrm{Unif}\{1, \ldots, H_s\}\);而 SRPO(Self-Reset Policy Optimization)则要求模型自身对失败轨迹进行因果诊断,直接输出首个错误思维步骤的索引 \(h^*\),将前置未出错的前缀 \(x^* = (x_0, \tilde{y}_{1:h^*-1})\) 作为重置状态。这消除了外部逐步标注的需求,使自监督反事实采样成为可能。

2. CPI-CARO 理论框架与信用分配预言机:可改进状态采样的严格理论优势

为了从理论层面严谨刻画自定位重置优于随机重置的数学根基,本文在保守策略迭代(Conservative Policy Iteration, CPI)框架下引入信用分配重置预言机(Credit-Assignment Reset Oracle, CARO)。在阈值 \(\tau > 0\) 下定义 \(\tau\)-可改进状态集 \(\mathcal{G}_{h,\tau} \coloneqq \{x \in \mathcal{X} : \max_{y \in \mathcal{Y}} A_h^\pi(x, y) \geq \tau\}\),其策略覆盖率为 \(p_\pi\)。传统的 CPI-RR 均匀抽取重置状态,导致优势估计信号被稀释在无提升空间的正确状态上;而 CPI-CARO 仅在 \(\mathcal{G}_{h,\tau}\) 上采样并实施策略更新,推导出带信用感知的模拟引理: $$ \sum_{h=1}^H |d_{\pi_\alpha}^h - d_\pi^h|{\mathrm{TV}} \leq \alpha H^2 p\pi $$ 该引理证明,限制在可改进集更新不仅将单步提升下界提升了 \(1/p_\pi\) 倍,更将样本复杂度从 \(\tilde{O}\left(\frac{|\mathcal{Y}| H^2 R_{\max}^2}{\tau^2 p_\pi^2}\right)\) 降至 \(\tilde{O}\left(\frac{|\mathcal{Y}| H^2 R_{\max}^2}{\tau^2}\right)\),在理论上节省了 \(1/p_\pi^2\) 级别的样本预算。

3. 共享前缀组归一化与前缀梯度掩码:将策略梯度精准锚定在分歧决策上

在构建得到共享前缀 \(x^*\) 后,模型并行生成 \(G\) 条后缀补全序列 \(y_{i,1:T_i}\)。如果像常规强化学习那样对整条拼接序列反向传播,已验证正确的前缀 token 仍会被动承担后缀成败带来的梯度震荡。为此,SRPO 采用双重设计:首先在共享前缀组内部独立执行组内相对优势归一化 \(\hat{A}_i = (r_i - \mu)/\sigma\),隔离不同分支的全局难度基线;其次对共享前缀实施梯度掩码(Prefix Masking),仅对后缀 token 施加策略梯度: $$ \mathcal{L}{\mathrm{SP}}(\theta) = -\frac{1}{G} \sum}^G \frac{1}{T_i} \sum_{t=1}^{T_i} \hat{Ai \log \pi\theta(y_{i,t} \mid x^*, y_{i,<t}) $$ 这种前缀掩码直接缩短了有效有效序列长度 \(T_i\),大幅提高了分歧后缀 token 处反向传播的单 token 梯度信噪比,实现了精准的局部策略塑形。

损失函数 / 训练策略

训练采用 LoRA 微调(rank=64, \(\alpha=64\)),使用 AdamW 优化器,学习率固定为 \(5 \times 10^{-5}\)。每个批次包含 32 个提示词,每个提示词在 \(1 \times 4\) 预算配置下分配 8 个 rollout(4 个基础组样本 + 4 个共享前缀重置样本)。总损失由全量基础组损失与共享前缀后缀掩码损失联合构成。值得注意的是,由于前缀掩码与组内优势归一化已具备极高的更新稳定性,消融证明直接弃用 PPO 的截断代理项(Clipped Surrogate)反而能在绝大多数任务上取得更优性能。

实验关键数据

主实验

在统一训练于 NuminaMath-Olympiads(400 题,2 epochs)的设定下,评估涵盖数学、科学推理、常识推理与策略推理共 10 个基准测试集。以下汇总 Qwen2.5-14B-Instruct 和 OLMo-3-7B-Instruct 的代表性主结果(准确率 %,3 次随机种子平均):

骨干模型 策略方法 Olympiads (oly) HMMT Nov MATH Level-5 StrategyQA CSQA Physics
Qwen2.5-14B Base 26.0 6.7 51.0 71.4 65.8 26.6
Qwen2.5-14B GRPO 25.0 ± 1.5 5.6 ± 1.6 52.1 ± 1.2 69.5 ± 2.8 66.1 ± 16.7 26.9 ± 6.8
Qwen2.5-14B SCoRe 25.3 ± 1.1 5.6 ± 1.6 54.1 ± 0.5 72.0 ± 0.5 78.0 ± 1.7 29.5 ± 4.1
Qwen2.5-14B Critique-GRPO 24.9 ± 0.7 7.8 ± 1.6 51.1 ± 0.2 72.1 ± 0.4 66.8 ± 0.7 24.9 ± 1.2
Qwen2.5-14B RRPO 25.7 ± 0.9 7.8 ± 1.6 52.7 ± 1.2 69.6 ± 3.9 77.7 ± 2.3 39.8 ± 13.0
Qwen2.5-14B SRPO (本文) 25.5 ± 1.2 6.7 ± 2.7 55.2 ± 0.7 74.9 ± 0.2 80.6 ± 0.9 45.5 ± 11.8
OLMo-3-7B Base 23.2 10.0 56.8 55.0 72.2 45.4
OLMo-3-7B GRPO 26.6 ± 0.2 11.1 ± 5.7 67.1 ± 1.2 62.0 ± 4.6 72.3 ± 0.6 42.3 ± 5.4
OLMo-3-7B RRPO 26.1 ± 0.9 10.0 ± 4.7 66.2 ± 6.0 68.1 ± 0.1 56.9 ± 25.7 33.9 ± 12.1
OLMo-3-7B SRPO (本文) 24.8 ± 1.9 15.6 ± 4.2 67.5 ± 3.5 66.6 ± 2.5 74.8 ± 1.0 54.5 ± 2.0

消融实验

在固定总计算预算(8 个 rollout/prompt)下,对重置采样策略分配方案以及 PPO 截断代理机制进行消融分析:

配置 核心机制设计 MATH Level-5 StrategyQA CSQA Physics
SRPO (\(1 \times 4\)) 4 基础 + 1 前缀 \(\times\) 4 后缀(平衡覆盖与深度) 67.5 ± 3.5 66.6 ± 2.5 74.8 ± 1.0 54.5 ± 2.0
SRPO (\(2 \times 4\)) 0 基础 + 2 前缀 \(\times\) 4 后缀(前缀多样性优先) 66.0 ± 3.1 65.9 ± 2.2 65.9 ± 3.3 43.2 ± 5.7
SRPO (\(1 \times 8\)) 0 基础 + 1 前缀 \(\times\) 8 后缀(极端后缀深度) 62.3 ± 2.5 64.1 ± 3.4 59.5 ± 14.9 43.4 ± 8.6
SRPO w/o Clip 纯前缀掩码组归一化,无 PPO 截断 68.2 69.6 73.4 52.0
SRPO w/ Clip 施加标准 PPO 裁剪限制 66.2 66.6 70.8 48.4

关键发现

  • 采样分配机制中,\(1 \times 4\) 组合展现出最优的泛化稳健性。完全舍弃基础组(如 \(2 \times 4\) 和 \(1 \times 8\))会导致从 prompt 开始的整体策略分布缺乏直接强化,引发探索漂移。
  • 在 LiveCodeBench 编码测试中,SRPO 收敛速度显著优于 GRPO 和 RRPO,达到同等通过率的速度快 2–3 倍。
  • 错误自定位的质量直接主导修复成功率:通过将模型自身定位与 Claude Opus 4.5 预言机标注对比,纯净前缀(Self \(\leq\) Oracle)的 Pass@4 纠错成功率达到 28.7%,而越界落入错误区域的前缀成功率暴跌至 16.3%,几乎呈现两倍的断崖差距。

亮点与洞察

  • 状态重置作为信用分配原语:巧妙地将强化学习中原本用于探索覆盖的 Reset 机制转换为反事实因果推断工具,通过共享前缀差异消除非关键上下文噪声。
  • 纯自监督因果定位:完全不需要人工或高阶 PRM 逐步标注,仅利用大模型内部自省定位错误步骤,零成本构建出高质量的 credit assignment 锚点。
  • 严密自洽的理论与实证对应:将 CPI 理论中的样本复杂度收益严格映射到语言模型的前缀梯度掩码与后缀组优势归一化上,闭环解释了经验增益的来源。

局限与展望

  • 自定位精度瓶颈:模型在长链推理中约有一半概率会越过真实错误节点导致前缀污染,使得纠错率显著下降;未来可探索针对自定位能力的联合训练。
  • 依赖可验证奖励机制:当前算法基于确定性的终末校验结果(数学/代码判题),在非可验证开放生成场景下的拓展仍有待探索。
  • 序列结构假设:方法高度依赖推理步骤清晰的 Thought MDP 结构,在非多步逻辑的弱结构任务中应用受限。

相关工作与启发

  • vs GRPO: GRPO 依赖全序列广播标量奖励,在长逻辑推理中产生严重信用稀释;SRPO 引入重置机制与前缀掩码,将更新专注在分歧后缀。
  • vs SCoRe: SCoRe 需要条件化拼接历史失败尝试进行二次完整重答,引入长上下文冗余;SRPO 直接回退到错误断点重开分支,算力利用率更高。
  • vs SPO-Tree: SPO-Tree 依赖启发式 token 计数或熵值突变作为切分点;SRPO 依赖语义完整的思维块并由模型自主语义定位,具备更强的因果解释性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [将状态重置重新诠释为强化学习信用分配原语,并成功与 CPI 理论及自定位能力结合]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 10 大推理基准、代码生成任务、严格的采样预算消融与定位质量深层因果审计]
  • 写作质量: ⭐⭐⭐⭐⭐ [理论推导严密清晰,算法设计自然流畅,图表与实验分析高度契合]
  • 价值: ⭐⭐⭐⭐⭐ [为推理模型的强化学习后训练提供了兼具理论保障与高效工程落地可能性的全新范式]