跳转至

LoopRPT: 循环语言模型的强化预训练框架

会议: NeurIPS 2026
arXiv: 2603.19714
领域: 预训练
关键词: 循环语言模型, 强化预训练, 隐式推理, 动态早退, EMA教师

一句话总结

LoopRPT 针对循环语言模型将多步推理压缩在隐状态迭代而非外显 Token 的结构特性,构建基于 EMA 动态教师与高熵困难 Token 筛选的分步密集强化信号,通过加噪隐状态采样联合优化早退策略与前序表征,在降低平均循环步数的同时显著提升复杂推理精度。

研究背景与动机

大语言模型当前主要依赖思维链(Chain-of-Thought, CoT)等外显自回归生成方式实现复杂多步推理。这种范式本质上将推理计算递延至后训练阶段,且在生成过程中伴随显著的序列长度膨胀与端到端计算开销。相比之下,循环语言模型(Looped Language Models, LoopLMs,如 Ouro)采用参数共享的 Transformer 主干,在固定参数规模下通过隐空间内的循环迭代隐式精炼内部表征,将多步推理有效压缩至潜在转移过程中,展现出极高的计算参数比与内在推理潜力。

然而,如何通过强化学习有效激励此类循环架构面临着基础性结构矛盾。现有的可验证奖励强化学习(RLVR)与后训练范式绝大多数建立在外显输出 Token 的终态正确性上,所得奖励高度稀疏。对于依靠内部多步隐式迭代展开推理的 LoopLM,终态稀疏信号在跨越多次循环迭代与动态早退机制(Adaptive Early Exit)时会遭遇严重的信用分配(Credit Assignment)障碍,极易引发策略崩溃或浅层早退坍缩。与此同时,直接在此类模型上微调标准 NTP 无法激发隐式多步探索,而直接引入外显 CoT 甚至会导致预测精度断崖式下跌。

本文的切入角度是:将因果语言建模中的下一个 Token 预测重新形式化为「下一个 Token 推理」任务,利用预训练语料本身的自监督真值构建密集奖励。核心 idea:LoopRPT 提出面向循环语言模型的强化预训练框架,利用 EMA 动态教师与自适应时间惩罚为中间隐式推理步骤提供稠密步级奖励,并通过对循环隐状态注入高斯扰动的策略梯度与步权表征学习,联合优化早退策略与前序步骤表征。

方法详解

整体框架

LoopRPT 的核心目标是在自回归预训练阶段直接塑造循环模型的隐式思考轨迹,促使模型在更少的循环迭代内完成充分表征,同时维持困难位置的推理准确率。整个训练流程主要由三个协同阶段构成:首先,利用 EMA 教师模型在最后一步的前向输出计算序列熵,动态筛选出最需要深层推理的高熵困难 Token;其次,基于教师模型在早退判定阈值下的参考步与真值对数几率,构建结合精度增益与难度感知时间惩罚的分步奖励表(Step-wise Reward Table);最后,通过对隐状态注入高斯噪声展开多次探索性 Rollout,结合策略梯度(Policy Gradient)更新动态早退门控,并利用步级优势加权优化主干网络的前向预测表征。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列上下文 x<t"] --> B["高熵困难 Token 动态筛选<br/>基于 EMA 教师终步分布熵 Ht 选 top-ρ%"]
    B --> C["分步密集奖励构建<br/>EMA 教师动态参考基线 + 难度感知时间惩罚"]
    C --> D["双轨联合优化"]
    subgraph D["双轨联合优化"]
        direction TB
        D1["加噪隐状态 Rollout 策略梯度<br/>h(k)+ϵ 扰动采样与组归一化 PG"]
        D2["步权表征学习<br/>基于优势加权的确定性前向 NTP 优化"]
    end
    D --> E["收敛输出:早退门控校准 + 高质量前序隐状态表征"]

关键设计

1. 高熵困难 Token 动态筛选:消除无效监督并聚焦高学习价值梯度 在标准因果语料中,绝大多数 Token 属于常见搭配或语法虚词,其预测难度极低,盲目在所有位置铺设复杂多步强化学习不仅计算开销巨大,还会引入低信噪比梯度。LoopRPT 采用教师模型的最终循环步骤 \(K\) 对应的隐状态输出分布计算香农熵 \(H_t = -\sum_{v \in \mathcal{V}} p_{\bar{\theta}}(v \mid x_{<t}) \log p_{\bar{\theta}}(v \mid x_{<t})\)。理论分析证明,高熵位置在局部一致性假设下严格对应更大的交叉熵期望梯度范数下界(\(\mathbb{E}[\|\nabla_z \mathcal{L}\|_2^2] \ge 1 - e^{-H_t}\)),代表该位置潜藏最丰富的推理监督信号。框架按序列分位数动态截取前 \(\rho\%\)(如 20%)的高熵 Token 施加后续强化损失,使模型算力高度集中在具备复杂逻辑依赖的预测节点上。

2. 基于 EMA 教师的分步密集奖励与难度感知时间惩罚:解决隐式推理信用分配与早退折冲 为了消除终态稀疏奖励对循环多步迭代的信用分配困难,LoopRPT 维护一个以动量更新的指数移动平均(EMA)教师模型 \(\bar{\theta}\)。针对每个被选中的 Token,教师模型根据累积早退概率超过阈值 \(\tau\) 的最小步数定义动态参考步 \(t_{\text{ref}}\),并将教师在该步对应的真值对数几率作为动态基线 \(b_{\text{ref}} = \ell_{\bar{\theta}}^{(t_{\text{ref}})}\)。第 \(k\) 步隐式状态获得的精度增益定义为 \(\Delta_{\text{acc}}(k) = \ell_{\theta}^{(k)} - b_{\text{ref}}\)。为了促使模型尽可能在较早步数完成计算而不诱发盲目早退,设计了难度感知的时间惩罚项: $\(C(k) = \lambda_t (k - t_{\text{ref}}), \quad \lambda_t = \lambda_{\text{base}} \left(1 + \lambda_{\text{scale}} (1 - d_t)\right)\)$ 其中归一化难度得分 \(d_t = \text{Clamp}(H_t^{\text{ref}} / \log |\mathcal{V}|, 0, 1)\)。当预测位置较为简单(熵低)时,时间惩罚权重 \(\lambda_t\) 增大,强制压制冗余步数;而在极难位置(熵高)时,惩罚被自适应削弱,为复杂隐式计算预留充足预算。总步级奖励为 \(R(k) = \Delta_{\text{acc}}(k) - C(k)\)。

3. 加噪隐状态 Rollout 策略梯度:构建平滑连续动作空间并校准早退门控 循环模型的推理过程在连续隐空间内闭环演化,缺乏离散 Token 采样所天然具备的多样性。LoopRPT 在循环转移过程中显式注入各向同性高斯噪声 \(\mathbf{h}^{(k)} \leftarrow \mathbf{h}^{(k)} + \boldsymbol{\epsilon}^{(k)}\)(\(\boldsymbol{\epsilon}^{(k)} \sim \mathcal{N}(0, \sigma^2 \mathbf{I})\)),独立展开 \(G\) 次加噪轨迹。这在理论上等价于对原本非平滑的策略目标进行高斯平滑化(Gaussian Smoothing),赋予探索鲁棒性界。在每个 Rollout 轨迹中依据门控分布采样退出步 \(t^{(g)} \sim \pi_{\theta}^{(g)}\) 并获取对应步级奖励 \(r^{(g)} = R(t^{(g)})\),经组内归一化计算优势 \(A^{(g)}\) 后执行策略梯度反向传播 \(\mathcal{L}_{\text{PG}} = -\mathbb{E}_g [A^{(g)} \log \pi_{\theta}^{(g)}(t^{(g)})]\),从而精细校准早退门控在不确定性扰动下的决策边界。

4. 步权表征学习与稳定约束:强化前序隐状态预测质量并防止策略坍缩 除门控策略外,主干特征提取能力必须同步演进。在无噪声确定性前向传播中,LoopRPT 根据门控概率与步级标准化优势 \(\widehat{A}(k)\) 计算复合权重 \(w_k = \text{softmax}[\pi_{\theta}(k)(1 + \text{ReLU}(\widehat{A}(k)))]\),用于加权最小化各循环步的负对数似然 \(\mathcal{L}_{\text{rep}} = -\sum_{k=1}^K w_k \ell_{\theta}^{(k)}\)。该项促使主干网络在获得高综合回报的早期循环步中即建立起高辨识度的语义特征。此外,算法引入门控分布熵正则项 \(\mathcal{L}_{\text{ent}}\) 防止探索过早退化,并基于 K3 代理损失引入针对 EMA 教师的信任域约束 \(\mathcal{L}_{\text{KL}}\),构成四位一体的联合优化目标 \(\mathcal{L} = \alpha \mathcal{L}_{\text{PG}} + \beta \mathcal{L}_{\text{rep}} + \gamma \mathcal{L}_{\text{ent}} + \delta \mathcal{L}_{\text{KL}}\)。

损失函数 / 训练策略

训练总目标函数为: $\(\mathcal{L} = \alpha \mathcal{L}_{\text{PG}} + \beta \mathcal{L}_{\text{rep}} + \gamma \mathcal{L}_{\text{ent}} + \delta \mathcal{L}_{\text{KL}}\)$ 超参数设置严格对齐论文配置:策略梯度权重 \(\alpha=1.0\),表征损失权重 \(\beta=1.0\),熵正则权重 \(\gamma=0.01\),KL 惩罚权重 \(\delta=0.0001\)。时间惩罚基准系数 \(\lambda_{\text{base}}=0.02\),缩放系数 \(\lambda_{\text{scale}}=0.5\);高熵筛选比例 \(\rho=0.2\);Rollout 采样组大小 \(G=8\),隐状态扰动方差 \(\sigma=0.1\)。EMA 教师动量参数设为 \(\phi=0.995\)。优化器采用 AdamW,初始学习率 \(5 \times 10^{-6}\),余弦退火配合 \(3\%\) 步数预热,序列截断长度 4096。

实验关键数据

主实验

论文在 Omni-Math 验证集(划分 Easy、Medium、Hard 三档熵难度)以及 11 个通用基准(涵盖常识逻辑、高阶数学竞赛与代码合成)上全面评估了 LoopRPT 对 Ouro-1.4B 与 Ouro-2.6B 的增强效果。

下表展示了下一个 Token 推理(Next-Token Reasoning, NTR)主实验数据(Peak 为打满 \(K=4\) 循环步,Adap. 为自适应早退):

模型 方法 Easy 准确率 (%) Easy 平均步数 Medium 准确率 (%) Medium 平均步数 Hard 准确率 (%) Hard 平均步数
Qwen3-1.7B Vanilla 47.49 – 32.18 – 19.19 –
Qwen3-1.7B +CoT 10.70 (-36.79) – 8.60 (-23.58) – 7.44 (-11.75) –
Ouro-1.4B Peak 74.62 4.00 57.28 4.00 33.79 4.00
Ouro-1.4B +LoopRPT 75.38 (+0.76) 4.00 58.29 (+1.01) 4.00 34.74 (+0.95) 4.00
Ouro-1.4B Adap. 74.40 3.34 57.20 3.53 33.91 3.75
Ouro-1.4B +LoopRPT 75.00 (+0.60) 2.50 (-0.84) 57.72 (+0.52) 2.81 (-0.72) 34.82 (+0.91) 3.07 (-0.68)
Ouro-2.6B Peak 74.33 4.00 57.19 4.00 34.52 4.00
Ouro-2.6B +LoopRPT 76.89 (+2.56) 4.00 61.15 (+3.96) 4.00 38.10 (+3.58) 4.00
Ouro-2.6B Adap. 74.51 3.24 57.35 3.35 34.35 3.51
Ouro-2.6B +LoopRPT 76.07 (+1.56) 2.05 (-1.19) 60.21 (+2.86) 2.18 (-1.17) 37.24 (+2.89) 2.28 (-1.23)

下表展示了 1.4B 级模型在下游跨领域通用任务与逻辑代码任务上的全面对比:

评测任务 / 数据集 指标 Ouro-1.4B Qwen3-1.7B LoopRPT-1.4B (本文) 相对基线提升
MMLU 5-shot Acc (%) 67.35 62.46 67.62 +0.27
MMLU-Pro 5-shot CoT Acc (%) 48.64 37.27 49.21 +0.57
BBH 3-shot CoT Acc (%) 70.97 53.51 71.37 +0.40
ARC-C 25-shot Acc (%) 60.95 55.72 61.10 +0.15
GSM8K 3-shot CoT Acc (%) 78.17 70.28 78.92 +0.75
MATH500 Greedy Acc (%) 20.60 – 37.00 +16.40
MATH (Full) Greedy Acc (%) 71.74 – 76.20 +4.46
MBPP pass@1 (%) 73.02 68.00 75.40 +2.38
MBPP+ pass@1 (%) 60.85 58.50 63.76 +2.91
HumanEval+ pass@1 (%) 67.01 59.80 67.07 +0.06

消融实验

论文在 Ouro-1.4B 上针对核心机制开展了详细消融(在 Hard 难度与整体平均 Overall 上的 Peak / Adap. 准确率及平均早退步数):

配置 Hard Adap. Acc (%) Hard 平均步数 Overall Adap. Acc (%) Overall 平均步数 说明
完整模型 (LoopRPT) 34.82 3.07 55.85 2.79 完整设计,兼具最优精度与早退效率
w/o 高斯隐状态噪声 34.76 3.29 (↑) 55.57 3.17 (↑) 缺乏隐空间连续扰动,探索受阻且早退收敛迟缓
w/o 策略梯度与KL (\(\mathcal{L}_{\text{PG}}+\mathcal{L}_{\text{KL}}\)) 34.17 3.28 (↑) 55.51 3.14 (↑) 移除门控 RL 优化,早退决策校准失效且精度下降
w/o 表征学习与熵正则 (\(\mathcal{L}_{\text{rep}}+\mathcal{L}_{\text{ent}}\)) 33.75 3.46 (↑) 55.19 3.24 (↑) 掉点最剧烈,表明对前序步特征的主干塑型是核心支柱
w/o 困难 Token 筛选 34.52 2.89 (↓) 55.41 2.49 (↓) 虽平均步数微降,但在 Hard 分割上显著掉点,出现虚假早退
w/o 难度感知时间惩罚 34.80 3.30 (↑) 55.78 3.19 (↑) 失去早退推进动力,平均推理步数显著回升

关键发现

  • 表征学习与策略梯度的协同性是核心支柱:消融显示去掉 \(\mathcal{L}_{\text{rep}} + \mathcal{L}_{\text{ent}}\) 会导致 Hard 分割上的精度发生最大幅度下滑(从 34.82% 骤降至 33.75%),且平均步数大幅回升至 3.46 步。这证明单纯依靠门控的强化探索无法使模型具备推理能力,必须依靠优势加权的表征损失直接打磨前序隐状态的特征表达质量。
  • 算力对齐对比彻底证伪了单纯数据累加与终态 RLVR:在同数据、同算力预算严格对齐的对照实验中,Continued-NTP 几乎没有带来通用任务收益;而直接套用输出 Token 级 GRPO 发生严重崩溃(GSM8K 准确率从 78.17% 暴跌至 4.17%,BBH 从 70.97% 骤降至 28.84%),直接证实了跨多步隐状态迭代时输出级稀疏奖励无法实现有效的信用分配,而分步稠密奖励是唯一稳定可行路径。
  • 物理推理延迟实现实质性压缩:在 A100 GPU 上的实测显示,自适应早退机制学到了极其锐利的置信度判定。在 4096 上下文长度下,Ouro-2.6B 的端到端单 Token 真实推理延迟降低了 39.5%,彻底打破了早退门控带来的额外寻址开销抵消算力节约的固有瓶颈。

亮点与洞察

  • 将预训练自监督信号转化为密集步级强化奖励:巧借预训练语料本身的真实后继 Token 作为可验证真值,规避了传统 RLVR 严重依赖狭窄数学/代码验证器的局限,使得多步隐式思考能够享受海量无标注语料的强化缩放红利。
  • 动态 EMA 教师打破自蒸馏天花板并稳固参考锚点:采用动量系数为 0.995 的 EMA 教师充当动态基线与时间惩罚锚点,既避免了静态教师导致的探索上限受限,又消除了无教师时模型因投机取巧而发生早退坍缩的数值不稳定性。
  • 隐式推理与显式思维链的互补迁移价值:当移除下游测试的思维链脚手架(No-CoT 零样本评测)时,基线模型性能发生剧烈腰斩,而 LoopRPT 在 1.4B 和 2.6B 上的零样本直接回答表现分别反超 13.65 与 26.00 个百分点,表明隐空间多步迭代真正增强了模型的内生认知与鲁棒表征。

局限与展望

  • 训练期展开深度固定受限:目前训练阶段强制设定固定的最大循环截断深度(\(K=4\)),使得模型无法自主探索更深层次的动态计算地平线。未来可探索支持无限地平线或全动态训练步数的递归强化机制。
  • 高阶竞赛数学的分布外泛化仍有瓶颈:在距离预训练分布最远的 AIME 2024 等极难奥林匹克竞赛中,提升幅度趋于平缓(30 题中仍停留于 6.67%),表明在极小参数量(1.4B/2.6B)下仅凭隐空间迭代仍难以彻底解决长程超难符号规划。

相关工作与启发

  • vs Ouro (Zhu et al., 2025): Ouro 构建了循环语言模型与动态早退的基础架构,但仅依赖纯 NTP 监督训练,缺乏对隐状态中间步的高效引导,早退策略在测试时极易受置信度偏差影响。LoopRPT 首次为 LoopLM 引入强化预训练范式,构建了端到端对齐的步级奖励体系。
  • vs RPT / RLP (Dong et al., 2025; Hatamizadeh et al., 2025): 先前强化预训练工作聚焦于标准单向前馈 Transformer,利用预测增益激励外显搜索或表征更新。LoopRPT 拓展了 RPT 范式,首次将其与循环权重共享架构深度结合,将强化信号精准注入连续潜变量迭代过程中。
  • vs RLTT (Williams & Tureci, 2026): RLTT 同期探索了循环架构的隐式轨迹奖励分配,但其严格受限于任务特定的后训练(RLVR)阶段且依赖终态答案正确性。LoopRPT 则扎根于大规模无监督预训练阶段,借助自回归真值提供了泛化性更广的密集步级监督。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次将强化预训练引入循环语言模型,提出针对隐空间迭代的分步密集奖励与加噪探索机制。]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 NTR 三档难度、11 个下游基准、算力对齐控制组、细粒度消融以及真实物理延迟测试。]
  • 写作质量: ⭐⭐⭐⭐⭐ [动机叙述逻辑紧密,理论推导严密自洽,框架与算法设计高度结构化且易于复现。]
  • 价值: ⭐⭐⭐⭐⭐ [为后 Transformer 时代参数复用与测试时计算(TTC)在隐式空间的内生缩放指明了极具前景的研究路径。]