LoopRPT: Reinforcement Pre-Training for Looped Language Models¶
会议: NeurIPS 2026
arXiv: 2603.19714
OpenReview: https://openreview.net/forum?id=h8EZBjmDcp
代码: 待开源(评审阶段提供匿名仓库)
领域: LLM 预训练
关键词: 循环语言模型、强化预训练、潜空间推理、自适应早退、EMA 教师模型
一句话总结¶
针对循环语言模型中隐式多步推理与传统基于输出 Token 的稀疏强化学习信号结构性错配的问题,LoopRPT 提出了首个专用于循环架构的强化预训练框架,通过高熵 Token 筛选、基于动态 EMA 教师的步级密集奖励与高斯噪声隐空间采样,实现了潜空间表征质量提升与早退推理效率的帕累托最优。
研究背景与动机¶
现代大型语言模型(LLM)的推理能力主要依赖显式文本生成,即通过链式思维(Chain-of-Thought, CoT)在后训练阶段激发多步推理。这种范式不仅严重依赖大量冗长 Token 的逐步输出,造成推理延迟与显存开销剧增,而且未能充分挖掘海量预训练语料本身的深层结构化信息。与此相对,循环语言模型(Looped Language Models, LoopLMs,如 Ouro)采用参数共享的循环主干网络,将多步推理压缩在连续的潜空间(Latent Space)迭代更新中,无需吐出中间中间 Token 即可实现更深度的知识表征计算与自适应早退(Early Exit)。
然而,如何通过强化学习有效激励和优化循环语言模型的潜空间计算依然是一个开放挑战。现有的可验证奖励强化学习(RLVR)方法天然依赖最终输出 Token 的离散正确性反馈,这种稀疏奖励机制与循环模型内部展开的连续隐式推理存在根本性的结构错配。稀疏奖励无法为循环网络中各个中间循环步骤提供精确的信用分配(Credit Assignment),导致动态早退策略发生崩溃或退化;标准 GRPO 等策略梯度算法在循环架构上甚至会出现灾难性性能衰退。
因此,核心矛盾在于:循环架构需要密集、连续且稳定的步级监督信号来引导隐式状态的精细演化,但海量无监督语料仅提供下一个 Token 的真实标签。近期涌现的强化预训练(Reinforcement Pre-Training, RPT)将下一个 Token 预测重新形式化为推理任务,为利用自监督语料构建密集强化学习信号开辟了道路。核心 idea:将下一个 Token 预测重构为潜空间多步推理任务,利用动态 EMA 教师模型在潜空间循环各步构建结合精度增益与难度自适应时间惩罚的密集奖励,并通过注入高斯噪声的潜空间采样与步加权表示学习,联合优化自适应早退策略与内部循环表征。
方法详解¶
整体框架¶
LoopRPT 是一个针对循环语言模型深度推理与早退校准的强化预训练框架。整个训练管线分为两阶段协同推进:在无梯度前向阶段,算法首先借助动态指数移动平均(EMA)教师模型评估序列中各位置的预测不确定性,仅筛选出最具学习梯度的 Top-\(\rho\%\) 高熵困难 Token;随后在该子集上构建教师参考步与步级精度-时间权衡奖励表。在梯度优化阶段,一方面利用确定性前向传播计算步级优势并执行步加权表示学习,强化浅层潜状态的预测能力;另一方面在循环隐状态中注入各向同性高斯噪声进行多路潜空间展开(Latent Rollouts),在连续潜空间内探索并借助组归一化策略梯度更新早退门控策略。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入序列 Context x_{<t}"] --> B["高熵困难 Token 筛选<br/>教师模型末步预测分布计算香农熵 H_t,提取 Top-ρ 困难位置"]
B --> C["基于 EMA 教师的步级密集奖励<br/>确定教师参考步 t_{ref},计算精度增益 Δ_{acc} 与难度感知时间惩罚 C(k)"]
C --> D["噪声隐空间采样与策略梯度<br/>潜状态注入高斯噪声 ϵ~N(0,σ²I),组归一化奖励更新早退门控 π_θ"]
C --> E["步加权表示学习<br/>利用优势函数 ReLU(A) 与退出概率加权多步交叉熵,强化早步表征"]
D --> F["联合多目标损失优化<br/>L = α L_{PG} + β L_{rep} + γ L_{ent} + δ L_{KL} 并动量更新 EMA 教师"]
E --> F
F --> G["优化后循环模型<br/>自适应早退计算量显著降低,困难 Token 推理精度提升"]
关键设计¶
1. 基于熵的高难 Token 筛选:剔除平凡信息并放大梯度能量
对序列中的每一个 Token 均匀施加强化学习监督极不经济,因为绝大多数文本 Token 属于易预测的平庸转移,其梯度信号微弱且容易淹没关键推理信号。为了最大化训练收益,LoopRPT 在前向阶段使用当前 EMA 教师模型 \(\bar{\theta}\) 展开至最大循环步数 \(K\) 的输出分布计算预测香农熵:
$\(H_t = -\sum_{v \in \mathcal{V}} p_{\bar{\theta}}(v \mid x_{<t}) \log p_{\bar{\theta}}(v \mid x_{<t})\)$
并在每个训练样本的有效响应位置中计算 \((1-\rho)\) 分位数阈值 \(q_i\),仅保留前 \(\rho=20\%\) 的高熵位置进入强化学习目标。理论证明表明,在教师分布作为标签代理的假设下,交叉熵梯度的期望平方范数下界受限于 \(1 - \|\bar{p}\|_2^2 \ge 1 - e^{-H_t}\);高熵位置直接对应了更大的期望梯度能量,从而在数学上保证了训练信号能够高度聚焦于需要复杂推理的关键 Token。
2. 基于 EMA 教师的步级密集奖励:稳定信用分配与难度感知时间惩罚
由于循环模型缺乏外部显式步骤的监督,直接利用学生模型自身的变动输出作为基线会导致严重的“目标追逐”(Target Chasing)与自蒸馏崩溃。LoopRPT 维护一个动量参数为 \(\phi=0.995\) 的 EMA 教师模型 \(\bar{\theta}\),并在给定的累积退出概率阈值 \(\tau\) 下计算教师的置信退出步:
$\(t_{\text{ref}} = \min\Bigl\{k: \sum_{j \le k} \pi_{\bar{\theta}}(j) \ge \tau\Bigr\}\)$
将教师在 \(t_{\text{ref}}\) 处的真实 Token 对数概率作为基线 \(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}})\)。其中惩罚系数 \(\lambda_t = \lambda_{\text{base}}(1 + \lambda_{\text{scale}}(1 - d_t))\) 与归一化难度得分 \(d_t = \text{Clamp}(H_t^{\text{ref}}/\log|\mathcal{V}|, 0, 1)\) 负相关。这意味着对于低熵的简单 Token,时间惩罚权重大幅增加以惩罚冗余迭代;而对于高熵的极难 Token,惩罚显著放宽以允许充分计算。最终在第 \(k\) 步构建出密集标量奖励 \(R(k) = \Delta_{\text{acc}}(k) - C(k)\),并对其跨循环步做均值方差标准化获得步优势 \(\widehat{A}(k)\)。
3. 噪声隐空间采样与步加权表示联合优化:平滑策略探索与早步表征压缩
循环语言模型在单 Token 内的推理发生在不可微的门控动作与连续的隐藏状态交界处。为在不破坏自回归生成的同时引入策略梯度探索,LoopRPT 在循环迭代中向隐状态注入高斯微扰:
$\(\mathbf{h}^{(k)} \leftarrow \mathbf{h}^{(k)} + \boldsymbol{\epsilon}^{(k)}, \quad \boldsymbol{\epsilon}^{(k)} \sim \mathcal{N}(0, \sigma^2 \mathbf{I})\)$
通过为每个选中 Token 独立采样 \(G=8\) 条随机隐轨迹,平滑原本脆弱的确定性退出曲面。在采样步 \(t^{(g)} \sim \pi_\theta^{(g)}\) 处获取奖励 \(r^{(g)} = R(t^{(g)})\),经组内均值方差归一化得到优势 \(A^{(g)}\),构建无偏估计的策略梯度损失 \(\mathcal{L}_{\text{PG}}\)。与此同时,为从根本上提升早退步的预测表征质量,算法在无噪声主干上执行步加权表示学习,加权权重 \(w_k = \text{softmax}[\pi_\theta(k)(1 + \text{ReLU}(\widehat{A}(k)))]\) 既贴合模型退出概率分布,又倾向于强化具有正向强化优势的高收益步,强制模型在更少循环次数内完成特征解耦与模式收敛。
损失函数 / 训练策略¶
LoopRPT 的联合优化目标由四项损失加权构成: $\(\mathcal{L} = \alpha \mathcal{L}_{\text{PG}} + \beta \mathcal{L}_{\text{rep}} + \gamma \mathcal{L}_{\text{ent}} + \delta \mathcal{L}_{\text{KL}}\)$ 各项具体形式与作用如下: 1. 策略梯度损失 \(\mathcal{L}_{\text{PG}}\):\(\mathcal{L}_{\text{PG}} = -\mathbb{E}_g [A^{(g)} \log \pi_\theta^{(g)}(t^{(g)})]\),使用组归一化优势优化退出门控参数。 2. 步加权表征损失 \(\mathcal{L}_{\text{rep}}\):\(\mathcal{L}_{\text{rep}} = -\sum_{k=1}^K w_k \ell_\theta^{(k)}\),使用优势加权的跨步交叉熵更新底层参数共享主干网络。 3. 早退熵正则项 \(\mathcal{L}_{\text{ent}}\):\(\mathcal{L}_{\text{ent}} = -\mathbb{E}[\sum_k \pi_\theta(k) \log \pi_\theta(k)]\),防止早退门控概率分布过快塌缩为确定性一步退出,保持动态深度的探索多样性。 4. K3 形式 KL 散度约束 \(\mathcal{L}_{\text{KL}}\):基于学生与教师真实对数概率差值 \(\Delta^{(k)} = \ell_{\bar{\theta}}^{(k)} - \ell_\theta^{(k)}\),计算数值稳定的 K3 散度代理 \(\mathrm{K3}^{(k)} = \exp(\Delta^{(k)}) - \Delta^{(k)} - 1\),防止更新步幅过大偏离慢速演进的教师参照系。
超参数配置上,权重设为 \(\alpha=1.0, \beta=1.0, \gamma=0.01, \delta=0.0001\);噪声尺度 \(\sigma=0.1\);时间惩罚基底 \(\lambda_{\text{base}}=0.02, \lambda_{\text{scale}}=0.5\);选词比率 \(\rho=0.2\)。优化器采用 AdamW,在 8 张 NVIDIA A100 GPU 上混合精度训练 3 个 epoch,最长序列长度 4096,学习率 \(5 \times 10^{-6}\) 并搭配余弦学习率衰减。每次梯度步后以动量 \(\phi=0.995\) 更新 EMA 教师。
实验关键数据¶
主实验¶
论文在竞赛级数学推理数据集 Omni-Math(4,428 道题)及一系列下游通用/推理基准上评估 LoopRPT。评估模型为 Ouro-1.4B 与 Ouro-2.6B,并与参数匹配的标准密集 LLM(如 Qwen3-1.7B/4B、Llama3.2-1.2B/3B、Gemma3)及继续预训练基线进行严格对比。
首先是 Omni-Math 验证集上的下一 Token 推理(Next-token Reasoning, NTR)准确率与平均循环退出步数(表 1):
| 模型 | 方法模式 | 简单 Token 准确率 (↑) | 简单退出步 (↓) | 中等 Token 准确率 (↑) | 中等退出步 (↓) | 困难 Token 准确率 (↑) | 困难退出步 (↓) |
|---|---|---|---|---|---|---|---|
| Qwen3-1.7B | Vanilla (NTP) | 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, K=4) | 74.62 | 4.00 | 57.28 | 4.00 | 33.79 | 4.00 |
| Ouro-1.4B | + LoopRPT (Peak) | 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 (Adap.) | 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, K=4) | 74.33 | 4.00 | 57.19 | 4.00 | 34.52 | 4.00 |
| Ouro-2.6B | + LoopRPT (Peak) | 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 (Adap.) | 76.07 (+1.56) | 2.05 (-1.19) | 60.21 (+2.86) | 2.18 (-1.17) | 37.24 (+2.89) | 2.28 (-1.23) |
其次是在同算力匹配受控实验(Compute-Matched Controls)下的下游基准表现(表 3):所有对照组使用相同基础模型、数据集、更新步数与学习率调度,唯一变量为优化目标:
| 规模 | 方法变体 | NTR-Hard 准确率 | GSM8K (数学) | MATH500 (奥数) | BBH (复杂推理) | MBPP+ (代码) | ARC-C (科学常识) | HellaSwag |
|---|---|---|---|---|---|---|---|---|
| 1.4B | Ouro 原始基线 | 33.91 | 78.17 | 20.60 | 70.97 | 60.85 | 60.95 | 74.30 |
| 1.4B | ++ Continued-NTP | 34.29 | 77.63 | 29.40 | 71.10 | 62.17 | 60.58 | 74.51 |
| 1.4B | ++ GRPO / RLVR | 33.00 | 4.17 | 12.60 | 28.84 | 20.58 | 49.23 | 74.36 |
| 1.4B | ++ LoopRPT | 34.82 | 78.92 | 37.00 | 71.37 | 63.76 | 61.10 | 74.93 |
| 2.6B | Ouro 原始基线 | 34.35 | 81.76 | 43.00 | 77.98 | 64.81 | 66.12 | 79.28 |
| 2.6B | ++ Continued-NTP | 35.21 | 82.41 | 43.20 | 78.19 | 63.76 | 66.75 | 79.74 |
| 2.6B | ++ GRPO / RLVR | 33.00 | 6.63 | 18.00 | 29.36 | 57.58 | 49.89 | 78.07 |
| 2.6B | ++ LoopRPT | 37.24 | 85.36 | 43.50 | 78.24 | 65.08 | 66.89 | 80.03 |
消融实验¶
为验证各核心模块的不可或缺性,论文在 Ouro-1.4B 上针对高熵筛选、高斯噪声、时间惩罚及各项损失进行全面消融(表 5):
| 模型配置变体 | 困难集 Peak 准度 | 困难集 Adap 准度 | 困难集平均步数 | 全量平均 Peak 准度 | 全量平均 Adap 准度 | 全量平均退出步数 | 关键变化特征与结论 |
|---|---|---|---|---|---|---|---|
| 完整 LoopRPT-1.4B | 34.74 | 34.82 | 3.07 | 56.14 | 55.85 | 2.79 | 达到最佳精度与推理步数帕累托前沿 |
| w/o 高斯隐空间噪声 | 34.69 | 34.76 | 3.29 (↑0.22) | 55.86 | 55.57 | 3.17 (↑0.38) | 退出步数明显变保守,证明扰动有助于平滑探索 |
| w/o \(\mathcal{L}_{\text{PG}} + \mathcal{L}_{\text{KL}}\) | 34.59 | 34.17 | 3.28 (↑0.21) | 55.89 | 55.51 | 3.14 (↑0.35) | 去掉策略梯度后自适应退出的早退校准能力显著恶化 |
| w/o \(\mathcal{L}_{\text{rep}} + \mathcal{L}_{\text{ent}}\) | 34.02 | 33.75 | 3.46 (↑0.39) | 55.52 | 55.19 | 3.24 (↑0.45) | 掉点最剧烈;中间层表征退化导致精度与计算开销双输 |
| w/o 困难 Token 筛选 | 34.60 | 34.52 | 2.89 (↓0.18) | 55.76 | 55.41 | 2.49 (↓0.30) | 全量均匀更新过度倾向平庸浅层,导致困难集准确率下降 |
| w/o 动态时间惩罚 \(C(k)\) | 34.57 | 34.80 | 3.30 (↑0.23) | 55.94 | 55.78 | 3.19 (↑0.40) | 失去早退激励,模型退化为趋向使用更多循环深度 |
关键发现¶
- GRPO 稀疏奖励在循环模型上完全崩溃:直接套用标准的以输出 Token 为奖励的 GRPO 训练循环网络,会导致 GSM8K(81.76% 暴跌至 6.63%)和 BBH(77.98% 暴跌至 29.36%)发生灾难性崩溃。这证实了最终输出的稀疏奖励无法跨循环时间步完成 credit assignment,必须引入密集步级奖励。
- 计算-精度实现双向帕累托占优:在 Ouro-2.6B 困难 Token 上,LoopRPT 在将平均推理步数从 3.51 步压缩到 2.28 步(节约 35% 循环计算量)的同时,准确率反而从 34.35% 提升到 37.24%(+2.89%),在 A100 上使端到端生成物理延迟降低 39.5%。
- 中间层表征质量本质提升,而非盲目早退:对隐藏状态逐层探测(Fig. 7)发现,LoopRPT 在 Step-1 与 Step-2 阶段的预测准确率显著超过原始模型。这证明模型并非单纯学习投机取巧地提前退出,而是成功将需要多步计算的推理模式压缩内化到了更浅的表征中。
- No-CoT 极度鲁棒性:在剥除所有提示词脚手架的 No-CoT 0-shot 极端设置下(Table 10),LoopRPT 在 GSM8K 上的优势进一步放大(1.4B 提升 +13.65%,2.6B 从 41.32% 暴增至 67.32%,净增 +26.00%),证明其激发的是模型内在的潜空间深度推理能力,而非表层 CoT 模式匹配。
亮点与洞察¶
- 将预训练视作潜在强化的新范式:不同于依赖下游人工标注或特定验证器的传统 RLVR,LoopRPT 将无标签语料本身的下一个 Token 视作天然验证信号,利用 RPT 范式打通了预训练大规模自监督与强化学习潜空间探索的通道。
- EMA 教师充当动态防崩锚点:引入动量平滑的 EMA 教师既提供了随学生能力进化的动态基准,又在数学上保证了步级参考阈值的稳定性,彻底杜绝了模型陷入死循环或无节制投机早退的自蒸馏退化陷阱。
- 连续潜空间的高斯策略梯度:跳过离散 Token 解码开销,直接在隐状态中添加微小高斯扰动进行策略梯度探索,为循环神经网络与扩散类连续状态模型的强化学习提供了一个极其轻量可复用的 trick。
局限与展望¶
- 训练深度受限于预设最大循环步数:作者指出当前训练仍然受限于预设的静态展开上限(\(K=4\)),无法在训练阶段探索任意长度的渐进式推理链;未来需要探索支持动态训练展开深度的无限步强化算法。
- 数据域偏向数学与代码:当前核心预训练主要在 Omni-Math 上验证,尽管在 OpenThoughts-3 上验证了通用扩展性,但在千亿级参数模型与更大规模通用网络多语言海量语料下的 Scaling Law 仍需进一步验证。
- 分布外泛化时早退阈值校准:在 AIME 等远超训练难度的极高难场景下,早退门控策略可能面临校准漂移,如何动态保证测试时早退阈值对分布偏移的鲁棒性是关键演进方向。
相关工作与启发¶
- vs. Ouro (Zhu et al., 2025):Ouro 构建了循环语言模型基座并提出了基于静态熵阈值的启发式早退机制,但缺乏强化学习训练,在面临复杂推理时过早退出导致性能瓶颈;LoopRPT 为其注入了步级强化预训练目标,联合校准了早退策略与内部表征。
- vs. RPT (Dong et al., 2025; Hatamizadeh et al., 2025):原版 RPT 针对标准非循环 Transformer,把计算前向当作强化探索;LoopRPT 将这一思想无缝嫁接到循环参数共享架构中,拓展了时间轴上的潜空间密集强化。
- vs. RLTT (Williams & Tureci, 2026):RLTT 同样探索循环架构的潜空间奖励分配,但它属于基于最终答案正确性的下游后训练(RLVR)任务;LoopRPT 则是无需人工答案标注、可规模化扩展的大规模预训练目标。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将强化预训练引入循环语言模型,巧妙设计了潜空间高斯策略梯度与步级 EMA 密集奖励。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1.4B/2.6B 双模型尺度,对比了 NTP、GRPO 等算力匹配基线,包含了深度的逐层探测、No-CoT 验证与消融分析。
- 写作质量: ⭐⭐⭐⭐⭐ 动机清晰,数学推导与直觉解释结合紧密,实验设计层层推进且逻辑自洽。
- 价值: ⭐⭐⭐⭐⭐ 为解决推理模型计算冗余、实现端侧高效深度思考提供了一条兼具理论与实效的全新技术路径。