LoopRPT：面向循环语言模型的强化预训练框架 LoopRPT: Reinforcement Pre-Training for Looped Language Models

LoopRPT 的训练分为两个阶段，交替进行。第一阶段是密集逐步奖励表构建：首先用教师模型（EMA）进行 K 步前向传播，计算最后一个推理步骤的 token 分布熵 $H_t = -\sum_v \bar{p}_\theta(v|x_{<t}) \log \bar{p}_\theta(v|x_{<t})$，选取熵最高的 top-$\rho$（论文中 $\rho=0.2$）比例位置作为困难 token。然后计算教师的退出分布 $\bar{\pi}_\theta(k)$ 及其 CDF，确定参考步骤 $t_{ref}$。接着用学生模型（无噪声）进行 K 步前向传播，对每一步计算准确率增益 $\Delta_{acc}^{(k)} = \ell_\theta^{(k)} - b_{ref}$ 和困难度感知的时间惩罚 $C^{(k)} = \lambda_t(k - t_{ref})$，其中 $\lambda_t$ 根据教师不确定性自适应调整（简单 token 惩罚更大，困难 token 惩罚更小）。逐步奖励为 $R^{(k)} = \Delta_{acc}^{(k)} - C^{(k)}$。第二阶段是联合优化：首先进行无噪声学生前向，计算步加权表示学习目标 $L_{rep} = -\sum_k w_k \ell_\theta^{(k)}$，其中权重 $w_k = \pi_\theta(k)(1 + \text{ReLU}(\hat{A}^{(k)))$。然后进行 G=8 次带噪声的隐空间 rollout（注入 $\sigma=0.1$ 的高斯噪声），对每个 rollout 重新计算奖励并采样退出步骤，通过 GRPO 风格的组归一化计算优势 $A^{(g)}$，优化策略梯度损失 $L_{PG}$。最后加入熵正则化 $L_{ent}$ 和 KL 正则化 $L_{KL}$ 防止退化，总损失为 $L = \alpha L_{PG} + \beta L_{rep} + \gamma L_{ent} + \delta L_{KL}$。

LoopRPT 的技术新颖性体现在几个与已有工作的本质区别上。与传统 RLHF/RLVR 相比，LoopRPT 不需要额外的奖励模型或人工标注，而是利用预训练语料本身的 next-token 预测正确性作为可验证奖励，并通过 EMA 教师将这种奖励细化到每个隐空间步骤。与原始 RPT（Dong et al., 2025）相比，LoopRPT 解决了 RPT 直接应用于循环架构时的两大困难：token 级反馈的稀疏性和隐空间不可采样性。具体地，原始 RPT 在标准 Transformer 上通过采样输出 token 来探索，但循环模型的推理在隐空间中展开，没有显式中间 token 可供采样——LoopRPT 通过向隐状态注入高斯噪声来创造性地解决这个问题，从理论上证明这优化了一个平滑的在策略目标 $J_\sigma(\theta)$（Proposition A.5）。与 Ouro 原始的训练方法相比，LoopRPT 引入了困难度感知的时间惩罚机制 $\lambda_t = \lambda_{base}(1 + \lambda_{scale}(1-d_t))$，避免了在困难 token 上强制提前退出，同时鼓励简单 token 更早退出。

从独立分析的角度看，LoopRPT 存在以下几个值得改进的弱点。第一，困难 token 选择依赖教师模型最后一轮的熵，这意味着选择质量受限于教师本身的校准程度——如果教师在某些 token 上过度自信（低熵但实际困难），这些 token 会被遗漏。改进方向可以考虑引入课程学习策略，随训练进程逐步调整选择比例 $\rho$，或结合多个推理步骤的熵变化来识别真正困难的 token。第二，噪声 rollout 的标准差 $\sigma=0.1$ 是全局固定的，没有根据 token 难度或推理步骤进行自适应调整——早期步骤可能需要更大探索噪声，后期步骤则需要更小噪声以保持精炼精度。第三，退出策略和表征优化虽然联合训练，但损失权重 $\alpha, \beta, \gamma, \delta$ 是手动调优的超参数（论文使用 1.0, 1.0, 0.01, 0.0001），缺乏自适应机制。可以考虑用元学习或不确定性加权来自适应调整这些权重。第四，EMA 教师的动量 $\phi=0.995$ 是固定的，但在训练早期（学生变化剧烈）和训练后期（学生趋于稳定）可能需要不同的动量值。

论文作者提出未来将扩展到更大规模和更广泛的数据混合，以及进一步改善早退出在分布偏移下的鲁棒性。基于 LoopRPT 的成果，还可以延伸以下几个方向：（1）将 LoopRPT 应用到其他循环架构（如 PonderLM、Parallel Loop Transformer），验证框架的架构通用性；（2）探索 LoopRPT 与显式 CoT 的协同——循环模型做隐式推理后再生成 CoT 可能获得两者的叠加收益；（3）将困难 token 选择从固定比例改为基于学习的动态选择，让模型自动决定哪些 token 需要更多推理；（4）探索 LoopRPT 在多模态循环模型中的应用，将隐空间推理优化扩展到视觉-语言任务；（5）研究 LoopRPT 训练后的模型在分布偏移（如不同领域数据、不同难度分布）下的早退出行为，建立更鲁棒的退出校准机制。

复现评估方面，LoopRPT 的复现条件较为友好。论文基于 Ouro 架构构建，而 Ouro 的代码和模型权重已开源（HuggingFace 上可获取）。训练数据 OMNI-MATH 同样是公开数据集。算力需求方面，1.4B 模型在 8×A100-80GB 上约需 2 小时，2.6B 约需 4 小时，这对学术实验室来说是可承受的。论文在附录中提供了完整的超参数表（Table 5）、训练算法伪代码（Algorithm 1-2）、评估设置（Table 6），以及详细的实现细节（使用 HuggingFace Accelerate、DDP、混合精度、梯度检查点）。然而，论文目前未提供 LoopRPT 训练代码的开源链接，核心算法实现（特别是 noisy latent rollout 和逐步奖励计算）需要读者自行实现。此外，一些关键实现细节分散在附录的不同位置（如 KL 代理在 Appendix D，hard-token 选择在 Appendix C），需要仔细阅读才能完整复现。总体复现难度为中等偏高——框架设计清晰但实现细节较多。