Scaling Latent Reasoning via Looped Language Models

文章标题: 通过循环语言模型扩展潜在推理
作者/机构: Rui-Jie Zhu, Zixuan Wang, Kai Hua, Tianyu Zhang, Ziniu Li, Haoran Que, Boyi Wei, Zixin Wen, Fan Yin, He Xing, Lu Li, Jiajun Shi, Kaijing Ma, Shanda Li, Taylor Kergan, Andrew Smith, Xingwei Qu, Mude Hui, Bohong Wu, Qiyang Min, Hongzhi Huang, Xun Zhou, Wei Ye, Jiaheng Liu, Jian Yang, Yunfeng Shi, Chenghua Lin, Enduo Zhao, Tianle Cai, Ge Zhang, Wenhao Huang, Yoshua Bengio, Jason Eshraghian; 隶属于 ByteDance Seed, UC Santa Cruz, Princeton University, Mila - Quebec AI Institute, University of Montreal, Peking University, Carnegie Mellon University, University of Pennsylvania, Conscium, University of Manchester, M-A-P 等多家机构。

A1 主要贡献

核心问题与研究目标:
大型语言模型（LLMs）的发展历史上主要依赖于模型规模的扩大，但这导致了部署基础设施要求高、延迟和成本增加等问题。因此，参数效率（在固定的参数预算内实现更好的模型能力）变得至关重要。传统的提升参数效率的途径（扩大训练语料库、利用思维链等推理时计算）面临数据稀缺和输出序列膨胀等限制。本文探索了第三种途径：通过架构创新在固定参数预算内实现动态计算。

解决方案与创新点:
本文提出并开源了一个名为Ouro的预训练循环语言模型（Looped Language Models, LoopLM）系列。该模型通过在预训练阶段引入潜在空间的迭代计算来内置推理能力，而不是像思维链（CoT）那样将推理推迟到训练后。其核心设计是递归地应用共享参数（一组权重绑定的层在前向传播中被迭代重用），从而实现动态计算。

主要贡献如下:
* 卓越的参数效率: 通过在7.7T tokens上进行预训练，证明了1.4B和2.6B参数的LoopLM在多数基准测试上能达到4B和8B标准Transformer的性能，实现了2-3倍的参数效率提升，这对于资源受限环境的部署至关重要。

• 熵正则化的自适应计算: 传统的自适应退出机制容易崩溃到浅层或过度使用深层循环。本文提出了一种新的训练目标，通过在均匀先验下进行熵正则化，鼓励模型在探索阶段无偏地探索计算深度。随后进入一个专注训练阶段，根据输入难度调整计算与性能的权衡，分配循环步数。

• 对循环机制的深入理解: 通过受“语言模型物理学”框架启发的对照实验，发现循环并未增加原始知识的存储容量（循环和非循环模型每个参数约存储2比特），但显著增强了需要事实组合和多跳推理等任务中的知识操纵能力。

• 提升安全性与忠实性: LoopLM在HEx-PHI安全基准上降低了有害性，且安全性随着循环步数的增加而提高（包括外推步数）。与CoT相比，其迭代的潜在更新产生的推理轨迹与最终输出更一致，表明其具有更高的因果忠实性，而非事后合理化。

本研究将循环深度确立为继模型大小和数据量之后的第三个扩展轴，并公开发布了Ouro模型家族（1.4B和2.6B参数），以展示LoopLM在大规模应用中的优势。

A2 相关工作

视角1：参数共享以提升模型效率

该视角将LoopLM视为一种参数共享技术：一个或多个Transformer块，甚至子模块（如注意力、FFN），在模型的深度上被重用，从而在不改变计算量的情况下减少参数。现代Transformer时代最著名的例子是ALBERT【23, Albert: A lite bert for self-supervised learning of language representations, 2019, arXiv】，它结合了参数重用和嵌入分解来大幅减少总参数量。在LLM广泛采用之前，参数共享在机器翻译领域得到了广泛探索【24, Recurrent stacking of layers for compact neural machine translation models, 2019, AAAI】；Takase等人【25, Lessons on parameter sharing across layers in transformers, 2021, arXiv】系统地研究了共享策略以平衡压缩和准确性。随着模型规模的增长，对参数重用的兴趣有所下降，但为了缩小LLM的内存占用而再度兴起。例如，Megrez2【26, Megrez2 technical report, 2025, arXiv】在标准的专家混合（MoE）模型中跨层重用专家，为内存有限的边缘LLM部署展示了一条可行的道路。

视角2：潜在推理与迭代求精

该视角将LoopLM的迭代过程视为潜在推理，其中每一步都是一个非语言的“思考”，用于提炼模型的内部表示。经验表明，增加循环步数可以提高复杂推理任务的性能【7, Reasoning with latent thoughts: On the power of looped transformers, 2025, arXiv】【17, Scaling up test-time compute with latent reasoning: A recurrent depth approach, 2025, arXiv】。一些模型通过将隐藏状态反馈到输入中来使这个过程显式化。例如，Coconut模型【27, Training large language models to reason in a continuous latent space, 2024, arXiv】插入一个“连续思考”令牌，该令牌源自前一步最后一层的隐藏状态，使模型能在连续的潜在空间中“思考”。CoTFormer【28, Cotformer: More tokens with attention make up for less depth, 2023, WANT@ NeurIPS】则将激活值交错回输入序列，然后将这个增强后的序列重新应用于共享层。这些显式反馈循环与隐式的LoopLM变体形成对比，后者的整个思考过程都包含在隐藏状态从一个循环步到下一个的演变中。因此，视角1（模型压缩）和视角2（潜在推理）都利用共享参数的迭代来提高参数效率，并正在被探索用于增强推理能力和高效的序列长度扩展（例如，PHD-Transformer【29, Efficient pretraining length scaling, 2025, arXiv】）。

A3 方法细节

本节正式定义了基于因果Transformer的LoopLM架构，并介绍了用于自适应潜在推理的训练方案。图3展示了该架构在训练和推理时的概览。我们的目标是让模型能够为每个令牌和每个样本选择循环步数，从而在简单的输入上花费更少的计算，在困难的输入上花费更多的计算，同时在可用步数充足时不牺牲准确性。

3.1 LoopLM架构

非循环与循环语言模型的形式化定义。首先，我们定义一个标准的非循环语言模型（LM）。设$emb(\cdot) : R^{|V|} \rightarrow R^d$为令牌嵌入层，$T_{\theta}(\cdot) : R^{M \times d} \rightarrow R^{M \times d}$为一个由$\theta$参数化的因果Transformer层，其隐藏层大小为$d$，输入长度为$M$。$lmhead(\cdot) : R^d \rightarrow R^{|V|}$是词汇量为$V$的去嵌入层。一个非循环LM由L个层堆叠而成，其中$\circ$表示函数复合：

$$F(\cdot) := \text{lmhead} \circ \mathcal{M}^L \circ \text{emb}(\cdot), \qquad \mathcal{M}^L(\cdot) := \mathcal{T}_{\theta_L} \circ \cdots \circ \mathcal{T}_{\theta_1}(\cdot)$$

接着，我们定义循环模型。设$t \in \{1, \ldots, T_{max}\}$为循环步数（即循环深度）。循环模型$F^{(t)}$将同一个深度为L的层堆栈重用$t$次：

因此，当$t=1$时，模型即为非循环模型$F^{(1)} \equiv F$。

单步损失函数的定义。如图3（左）所示，在每个循环步骤$t$，模型都会产生一个语言建模头的输出。我们将单个步骤$t$的标准交叉熵损失定义为$L^{(t)}$：

$$\mathcal{L}^{(t)}=\mathbb{E}_{x_{1: M}}\left[\sum_{\ell=1}^{M-1}-\log p_\theta^{(t)}\left(x_{\ell+1} \mid x_{1: \ell}\right)\right],$$

其中，$p_{\theta}^{(t)}(\cdot | x_{1:\ell}) = \text{softmax}(\text{lmhead}(h_{\ell}^{(t)}))$，$x_{1:\ell}$表示输入的前$\ell$个令牌，$h_{\ell}^{(t)}$是在位置$\ell$经过$t$次循环后的隐藏状态。需要注意的是，这是单个循环步骤的独立损失。包含所有步骤的总训练目标将在后续部分定义。

引入自适应计算的动机。以往的研究【7, Reasoning with latent thoughts: On the power of looped transformers, 2025, arXiv】【11, A little depth goes a long way: The expressive power of log-depth transformers, 2025, arXiv】表明，增加循环步数$t$对推理任务有益。然而，这会增加计算量，并且并非所有令牌都需要很多步的计算【30, Physics of language models: Part 2.1, grade-school math and the hidden reasoning process, 2024, arXiv】【31, Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for llm reasoning, 2025, arXiv】。因此，将计算预算用在正确的令牌上至关重要。这通过下一节描述的门控机制来实现。

3.2 通过门控机制实现自适应计算

退出门的引入与概率定义。为了实现自适应计算，我们在每个步骤$t \leq T_{max}$处添加一个与LM头并行运行的退出门（见图3）。在每个循环$t$，该门会输出一个瞬时的（每一步的）退出概率：

$$\lambda_{t}(x)=\sigma\left(\text { Linear }_{\phi}\left(h^{(t)}\right)\right) \in(0,1)$$

其中$h^{(t)}$是步骤$t$的最后一层隐藏状态，$\phi$是门的参数。

生存概率与退出概率分布的构建。我们定义$S_t(x)$为“生存”概率，即在前$t$步没有退出的概率：

$$S_t(x)=\prod_{j=1}^t(1-\lambda_j(x)),\quad S_0(x)\equiv 1,$$

那么，首次在步骤$t$退出的未归一化概率为：

为了获得一个在退出步骤上的有效离散分布，我们将剩余的概率质量分配给最后一步：

带提前退出的推理机制。如图3（右）所示，我们从学习到的退出分布$\{p_{\phi}(t | x)\}_{t=1}^{T_{max}}$中推断出一个退出步骤，从而实现高效推理。到步骤$n$为止的累积退出概率为：

$$\operatorname{CDF}(n \mid x)=\sum_{t=1}^n p_\phi(t \mid x)=1-\prod_{j=1}^n\left(1-\lambda_j(x)\right), \quad n<T_{\max }, \quad \operatorname{CDF}\left(T_{\max } \mid x\right)=1 .$$ <p>
给定一个阈值$q \in [0, 1]$，我们在累积概率首次超过$q$的步骤终止： 阈值$q$控制了计算与准确性的权衡：较小的$q$倾向于提前退出（计算量少），而较大的$q$则允许更深的计算。在实践中，$q$可以全局选择、按任务校准或设置步数的上下限。这种基于分位数的确定性策略避免了采样，同时与学习到的分布保持一致。

门控参数的学习阶段。门控参数$\phi$（以及通过$\{\lambda_t\}$决定的$p_{\phi}$）分两个阶段学习：
* 阶段I：在预训练期间，通过优化一个熵正则化的目标函数（第3.3节），门与语言模型共同学习。
* 阶段II：我们冻结语言模型，并微调$\phi$以锐化$p_{\phi}$（即调整深度分配），而不改变令牌级别的预测。

完整的训练目标在下一节描述。

3.3 阶段I：学习熵正则化目标

问题：朴素梯度下降导致概率崩溃。在对下一个词元预测损失进行朴素梯度下降时，更深的循环通常会降低单步损失$L^{(t)}$（见公式2），直到某个深度；超过该深度后，收益递减，梯度会将概率质量转移到后面的步骤。随着$p_{\phi}$集中在后面的步骤，这些步骤会接收到更多的训练信号，其损失会进一步下降，这反过来又将更多的概率质量拉向末端。这种自增强效应导致$p_{\phi}$崩溃到$t = T_{max}$。为了解决这个问题，引入一个熵项来惩罚向最深步骤的崩溃，从而在$p_{\phi}$中保持足够的分散性以反映输入难度。

熵正则化目标函数。给定单步损失$L^{(t)}$和来自公式3的退出步分布$p_{\phi}(t | x)$，我们的训练目标结合了下一个词元预测和熵正则化：

$$ \mathcal{L}=\underbrace{\sum_{t=1}^{T_{\max }} p_{\phi}(t \mid x) \mathcal{L}^{(t)}}_{\text {expected task loss }}-\underbrace{\beta H\left(p_{\phi}(\cdot \mid x)\right)}_{\text {entropy regularization }}, \quad H\left(p_{\phi}(\cdot \mid x)\right)=-\sum_{t=1}^{T_{\max }} p_{\phi}(t \mid x) \log p_{\phi}(t \mid x) . $$

直观上，期望任务损失通过在步骤$t$退出的概率来加权每个$L^{(t)}$。系数$\beta$控制了探索-利用的权衡：较大的$\beta$鼓励更高熵（更具探索性）的$p_{\phi}$，而较小的$\beta$则允许$p_{\phi}(t|x)$在模型对最优深度有信心时将大部分概率质量放在一个特定的步骤上。

变分推断视角：使用均匀先验。公式4中的目标可以看作是一个证据下界（ELBO）损失，其中退出步骤$z \in \{1, \ldots, T_{max}\}$是一个潜变量，其变分后验是学习到的退出分布$p_{\phi}(z=t|x)$，其先验是$\pi(t)$。负ELBO为：

$$\mathcal{L}_{\text{ELBO}} = \sum_{t=1}^{T_{\max}} p_\phi(t \mid x) \mathcal{L}^{(t)} + \beta \operatorname{KL}(p_\phi(\cdot \mid x) \| \pi(\cdot)).$$

使用均匀先验$\pi_t = 1/T_{max}$，KL散度项变为：

因此，最小化ELBO等价于（在常数$\log T_{max}$的意义上）公式4中的目标。这揭示了熵项实际上是一个KL正则化器，并阐明了期望损失是在所有退出步骤上进行边缘化。这也将其与PonderNet【32, Pondernet: Learning to ponder, 2021, arXiv】等自适应计算方法联系起来，后者也优化一个ELBO以实现动态停止。

为何选择均匀先验。不同的先验编码了不同的深度偏好。几何先验，如参考文献【32, Pondernet: Learning to ponder, 2021, arXiv】中所用，或泊松-对数正态先验【17, Scaling up test-time compute with latent reasoning: A recurrent depth approach, 2025, arXiv】，会温和地偏好较早的停止。而均匀先验是深度无偏的。我们采用均匀先验，是为了将由输入难度驱动的退出决策与任何全局计算偏好分离开；然后，熵项防止模型总是崩溃到使用$T_{max}$。与几何先验的经验比较见附录A节。

3.4 阶段II：专注的自适应门训练

训练目标。在此阶段，我们冻结语言模型的参数，只训练退出门，使其根据已实现的性能增益来做终止决策。我们使用一个贪心信号，该信号平衡了额外循环带来的边际改进与增加的计算成本。

量化性能增益。为确保门控不改变语言模型的表示，我们在每个$i$处计算一个分离的（detached）每步损失$L_{i,stop}^{(t)}$。

$$I_i^{(t)}=\max(0,\mathcal{L}_{i,\text{stop}}^{(t-1)}-\mathcal{L}_{i,\text{stop}}^{(t)})$$

其中，较大的$I_i^{(t)}$表示正在持续改进；较小的值表示增益已经停滞，LoopLM应选择提前退出。

理想继续概率作为标签。我们通过计算理想的继续概率来实现这一点，这是一个训练标签，指示是继续（接近1）还是退出（接近0）：

$$w_i^{(t)} = \sigma(k \cdot (I_i^{(t)} - \gamma))$$

其中斜率$k=50.0$，阈值$\gamma=0.005$，因此$w_i^{(t)} \approx 1$建议继续，$w_i^{(t)} \approx 0$建议退出循环。

自适应退出损失函数。在步骤$t$的自适应退出损失采用门预测的继续概率$1 - \lambda_i^{(t)}$与理想标签$w_i^{(t)}$之间的二元交叉熵，并在序列长度$M$上取平均：

$$\begin{aligned} \mathcal{L}_{\text {adaptive }}^{(t)}=-\frac{1}{M} \sum_{i=1}^{M}\left[w_{i}^{(t)} \log (\underbrace{1-\lambda_{i}^{(t)}}_{\begin{array}{c} \text { predicted } \\ \text { continuation } \end{array}})+\left(1-w_{i}^{(t)}\right) \log (\underbrace{\lambda_{i}^{(t)}}_{\begin{array}{c} \text { predicted } \\ \text { exit } \end{array}})\right]. \end{aligned}$$

总的自适应损失在所有循环步骤上取平均：

自适应损失的意义。公式6中的自适应损失训练门在步骤$t$使其预测与从实际性能改进中导出的理想行为相匹配：
* 预测概率：门生成$\lambda_i^{(t)}$（退出概率）和$1-\lambda_i^{(t)}$（继续概率）。
* 目标标签：理想行为被编码为$w_i^{(t)}$（目标继续概率）和$1-w_i^{(t)}$（目标退出概率）。
该公式同时惩罚了两种失败模式：
* 思考不足（Underthinking）：门在应该继续时退出（标签$w_i^{(t)}$大，但预测的退出概率$\lambda_i^{(t)}$也大）。
* 过度思考（Overthinking）：门在应该退出时继续（标签$w_i^{(t)}$小，但预测的退出概率$1-\lambda_i^{(t)}$小）。
优化公式6训练门选择一个贪婪的退出步骤，用额外的计算来换取可测量的改进。经验评估见第5.4.1节。

4. 训练循环语言模型

整体流程。我们的Ouro模型家族的端到端训练流程如图4所示。总共使用7.7T个token来训练基础模型Ouro-1.4B和Ouro-2.6B。最后的推理监督微调（Reasoning SFT）产生了Ouro-1.4B-Thinking和Ouro-2.6B-Thinking变体。本节详细介绍了每个训练阶段使用的架构、数据构成和具体配置。前四个阶段的训练方案概览见表1。

表1 Ouro 1.4B和2.6B的训练方案。

4.1 Transformer架构

核心组件。Ouro模型使用标准的仅解码器Transformer【33, Attention is all you need, 2017, NeurIPS】，优先考虑循环计算机制的简洁实现，没有额外的修改。核心架构由一个循环应用的Transformer块堆栈组成。每个块使用带旋转位置嵌入（RoPE）【34, Roformer: Enhanced transformer with rotary position embedding, 2023】的多头注意力（MHA）。每个块中的前馈网络（FFN）使用SwiGLU激活函数【35, Glu variants improve transformer, 2020, arXiv】。为了增强训练稳定性（这对于深度循环计算尤其关键），我们采用了三明治归一化结构，在注意力和FFN子层之前都放置一个RMSNorm层【17, Scaling up test-time compute with latent reasoning: A recurrent depth approach, 2025, arXiv】。两个模型都使用来自SmolLM2模型【36, Smollm2: When smol goes big–data-centric training of a small language model, 2025, arXiv】的49,152个token的词汇表。该分词器针对代码和拉丁字母语言进行了优化。架构细节总结在表2中。

表2 Ouro模型架构配置。两个模型共享相同的词汇表和核心组件类型，仅在参数数量和层深度上有所不同。

4.2 数据

数据构成概述。数据设定了基础模型的能力边界。我们的语料库跨越多个阶段，涵盖了网络文本、数学、代码和长上下文文档，旨在构建核心语言理解能力，同时加强推理、编码和长上下文处理技能。除了标准的网络爬取数据，我们还加入了针对数学推理和代码生成的目标数据集，以提高解决复杂问题的能力。表3总结了每个训练阶段的数据构成和规模。

表3 训练语料库统计。由于数据在预训练期间是随机抽样的，数据集大小不直接对应于看到的token总数。

训练数据分阶段策略。为确保可复现性，我们的训练语料库完全由开源数据集组成，数据统计如表4所示。我们将数据分为四个阶段，每个阶段的构建策略与现代预训练中常用的Warmup-Stable-Decay（WSD）学习率调度器【37, Understanding warmup-stabledecay learning rates: A river valley loss landscape perspective, 2024, arXiv】对齐。

表4 阶段1（稳定训练I & II）的数据构成。总数据集大小：6T tokens。

阶段1：预训练。此阶段支持训练的预热和稳定阶段。语料库主要由Web CommonCrawl（CC）数据组成。由于我们计划在超过2T tokens上训练模型，许多流行的开放语料库规模太小（例如，Fineweb-Edu为1.3T tokens【38, The fineweb datasets: Decanting the web for the finest text data at scale, 2024, NeurIPS】，DCLM为2.6T tokens【39, Datacomp-lm: In search of the next generation of training sets for language models, 2024, NeurIPS】）。因此，我们使用Nemotron-CC【40, Nemotron-cc: Transforming common crawl into a refined long-horizon pretraining dataset, 2024, arXiv】（6.3T tokens）作为稳定阶段的主要数据集。为了让模型具备基本的中文能力，我们加入了Ultra-FineWeb-zh【41, Ultra-fineweb: Efficient data filtering and verification for high-quality llm training data, 2025, arXiv】和MAP-CC【42, Chinese tiny llm: Pretraining a chinese-centric large language model, 2024, arXiv】。然而，由于分词器中没有中文字汇，汉字会被碎片化为多个字节级子词元，因此我们从阶段2开始移除了中文数据。为了增强编码和数学能力，我们加入了OpenCoder【43, Opencoder: The open cookbook for top-tier code large language models, 2024】和MegaMath【44, Megamath: Pushing the limits of open math corpora, 2025, arXiv】。数据集比例详情见表4。

阶段2：持续训练（CT）退火。CT退火阶段引入更高质量的数据，在退火学习率下增强模型。词元序列长度扩展到16K，超过了大多数样本的长度以最小化截断。我们从Nemotron-CC的高质量子集构建语料库，并增加了HQ MegaMath、Nemotron-CC-Math-v1【45, Nemotron-cc-math: A 133 billion-token-scale high quality math pretraining dataset, 2025】【46, Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model, 2025】、OpenCoder-Annealing【43, Opencoder: The open cookbook for top-tier code large language models, 2024】、Nemotron-pre-training-Code-v1【46, Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model, 2025】和Nemotron-pre-training-SFT-v1【46, Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model, 2025】。数据构成见表5。

表5 阶段2（CT退火）的数据构成。总数据集大小：1.4T tokens。

阶段3：长上下文训练（LongCT）。LongCT阶段扩展了模型的长上下文能力。我们采用了ProLong【47, How to train long-context language models (effectively), 2024, arXiv】的64K长度子集，包含20B个token，用于训练模型处理更长的序列，并提高其处理长上下文的能力。

阶段4：中期训练。此阶段使用一组多样化且质量极高的数据，包括⟨问题, 答案⟩和⟨问题, 思维链, 答案⟩样本，以进一步发展高级能力。我们整合了20多个开源SFT数据集以增加数据广度，并进行了彻底的去污染，以避免与主流评估基准重叠。所有样本都转换为ChatML格式，以减少后续后训练阶段的对齐税。处理后，我们得到182B个token，从中随机抽取90B个token。为了稳定训练分布，我们重放了30B来自阶段1的token和180B来自阶段2的token，有效数据量达到300B。因此，此阶段在多样化的监督信号下巩固和扩展了预训练期间获得的能力。

4.3 训练稳定性与自适应配置

训练框架与稳定性优先策略。我们使用flame【48, Flame: Flash language modeling made easy, 2025】框架进行预训练，该框架基于torchtitan【49, Torchtitan: One-stop pytorch native solution for production ready LLM pretraining, 2025, ICLR】。在训练过程中，我们优先考虑稳定性而非激进的扩展，并根据对训练动态的经验观察做出了一些关键调整。这些决策对于实现循环架构的稳定收敛至关重要，因为循环架构表现出与标准Transformer不同的优化特性。

为稳定性减少循环步数。我们在阶段1a（稳定训练I）中最初使用8个循环步骤的实验导致了损失尖峰和梯度振荡。我们推测这是由于梯度流经多个循环迭代时被放大，从而放大了微小的扰动。因此，我们在阶段1b（图4中的稳定训练II）中将循环步数从8减少到4，这在计算深度和训练稳定性之间取得了平衡。

批处理大小扩展。为了进一步增强稳定性，我们将批处理大小从4M逐步增加到8M个token。更大的批处理大小提供了更稳定的梯度估计，这对于循环架构尤其重要，因为梯度在多次迭代中的流动会引入额外的方差。

KL散度系数的降低。我们策略性地将公式4中的$\beta$从阶段1a的0.1降低到后续阶段的0.05。这种降低有两个目的：（1）它减少了任务损失和KL惩罚之间的冲突梯度，从而实现更稳定的优化；（2）它减少了来自均匀先验的“拉力”，使模型在探索有益的深度模式时有更大的自由度，而不会受到人为限制。这一调整使模型能够在没有过度约束的情况下学习有用的深度模式。

优化配置。在所有阶段，我们都使用AdamW优化器，权重衰减设置为0.1，$\beta_1 = 0.9$，$\beta_2 = 0.95$，梯度裁剪值为1.0。这些保守的设置是专门为维持循环架构的稳定性而选择的。

学习率考量。我们经验性地发现，循环架构比参数匹配的Transformer需要更小的学习率。由于计算限制，我们没有进行详尽的学习率扫描，而是采用了优先考虑稳定收敛而非可能更快但风险更高的学习率方案。

序列长度递增。序列长度在不同阶段逐步增加：两个预训练阶段为4K，CT退火阶段为16K，长上下文训练为64K，中期训练为32K。这种递增方式在扩展上下文能力的同时稳定了优化过程和训练吞吐量。

4.3.1 分阶段训练细节

• 阶段1a：预训练阶段I（探索）。我们以8个循环步骤初始化训练。学习率遵循Warmup-Stable调度，峰值为$3 \times 10^{-4}$。序列长度为4K，初始批大小为4M tokens，为保证稳定性逐渐增加到8M。在此阶段，我们观察到训练不稳定性，这促使了后续的架构调整。

• 阶段1b：基于稳定性的预训练阶段II与升级。在阶段1a发现稳定性问题后，我们将循环步数从8减少到4。为了在提高稳定性的同时保持计算效率，我们将方法分为两个变体：

  • 1.4B Ouro：使用原始的24个预训练层，循环4次。
  • 2.6B Ouro：通过层复制将24层升级到48层，循环4次。
    我们架构的循环特性使得这种升级过程特别平滑，因为跨迭代共享的权重自然地促进了层复制，而没有标准Transformer升级中常见的典型不稳定性。

• 阶段2：CT退火。学习率退火至$3 \times 10^{-5}$，同时让模型接触高质量的训练数据。循环步数保持为4，这已被证明是稳定-性能权衡的最佳选择。数据构成如表5所示进行了仔细平衡。

• 阶段3：LongCT。批大小保持在8M tokens。降低的KL系数（$\beta = 0.05$）即使在64K长度的序列下也继续提供稳定的训练动态。
• 阶段4：中期训练。学习率进一步降低到$1 \times 10^{-5}$，并采用余弦调度器，以帮助模型更好地吸收这个多样化、高质量的数据集。

4.4 监督微调

数据构成。我们在一个包含约830万个样本的多样化语料库上进行SFT，这些样本来自高质量的公共数据集。如表6所示，我们的训练混合数据侧重于数学推理（350万样本）和代码生成（320万样本），同时也包含了科学推理（80.8万样本）和对话能力（76.7万样本）。

数据来源。对于数学推理，我们结合了OpenThoughts3【50, Openthoughts: Data recipes for reasoning models, 2025, arXiv】和AceReason-1.1-SFT【51, Acereason-nemotron 1.1: Advancing math and code reasoning through sft and rl synergy, 2025, arXiv】，以提供对问题解决策略的全面覆盖。我们的代码训练数据聚合了多个来源，包括AceReason-1.1-SFT、OpenCodeReasoning【52, Opencodereasoning: Advancing data distillation for competitive coding, 2025, arXiv】、Llama-Nemotron-Post-Training-Dataset【53, Llama-nemotron: Efficient reasoning models, 2025, arXiv】和OpenThoughts3，确保了对不同编程范式和推理模式的广泛接触。科学推理能力通过OpenThoughts3和Llama-Nemotron-Post-Training-Dataset来培养，而对话熟练度则通过OO1-Chat-747K和DeepWriting-20K【54, Reverse-engineered reasoning for open-ended generation, 2025, arXiv】数据集来增强。

训练配置。我们使用LlamaFactory代码库【55, Llamafactory: Unified efficient fine-tuning of 100+ language models, 2024, arXiv】训练2个epoch，最大序列长度为32K。我们采用Adam优化器，学习率为$2 \times 10^{-5}$，$\beta = (0.9, 0.95)$，并应用余弦衰减调度以实现稳定收敛。

表6 监督微调数据构成。训练语料库包含四个关键能力领域的830万个样本。

4.5 强化学习尝试

探索性实验。在SFT阶段之后，我们使用DAPO【56, Dapo: An open-source llm reinforcement learning system at scale, 2025, arXiv】和GRPO【57, Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024, arXiv】在DAPO-17K数据集上进行了探索性的RLVR（带可验证奖励的强化学习）对齐实验。这些尝试并未比最终的SFT检查点带来显著的性能提升。主要问题源于模型的动态提前退出机制。vLLM/SGLang通过固定的执行路径提供快速的rollouts，但这在LoopLM的可变深度计算下会失效。

失败的尝试。我们尝试了两种方法，均未成功：
1. 离策略rollouts：我们在vLLM中生成完整的四步rollouts，为每个token产生四个logit候选项。然后，我们选择第一个超过终止阈值的token来模拟提前退出。对于更新，我们使用截至该步骤的累积损失，丢弃后续的token和损失。这种离策略不匹配，即token是用最终深度生成的，而损失是在较早的深度计算的，并未提高性能。
2. 固定的4轮RL：为了避免离策略问题，我们在固定的四个循环步骤上进行rollouts和更新。训练正常进行，但性能没有超过SFT检查点。一个可能的原因是规模：在已经进行了广泛的SFT之后，这些较小的模型可能没有太多RL增益的空间。有趣的是，尽管是在四轮上训练的，模型在推理时仍然在有益的情况下使用更少的轮数。这种泛化背后的机制尚不清楚。

未来工作。我们将继续探索该架构的RL对齐，因为我们正在开发能够完全支持LoopLM动态计算的基础设施。

A4 实验环境与结果

实验环境

• 数据集:

  • 预训练: 使用了总计7.7T tokens的数据，主要包括Nemotron-CC (网页数据)、MAP-CC (网页数据)、Ultra-FineWeb-zh (网页数据)、OpenCoder-pretrain、MegaMath-web、ProLong-64K等。
  • 监督微调 (SFT): 使用了约830万个样本，涵盖数学(OpenThoughts3, AceReason)、代码(AceReason, OpenCodeReasoning)、科学和聊天等领域。
  • 评测基准: 涵盖通用能力(MMLU, MMLU-Pro, BBH, ARC-C, HellaSwag, Winogrande)、数学与代码(GSM8K, MATH500, HumanEval, MBPP)以及高级推理(AIME, OlympiadBench, GPQA, SuperGPQA, BeyondAIME, HLE)和安全性(HEx-PHI)。

• 模型架构:

  • Ouro-1.4B: 14亿参数，24层Transformer块，隐藏层维度2048。
  • Ouro-2.6B: 26亿参数，48层Transformer块，隐藏层维度2048。
  • 共性: 均为仅解码器架构，使用多头注意力(MHA)、SwiGLU激活函数、旋转位置编码(RoPE)和三明治RMSNorm归一化。词汇表大小为49,152。

• 硬件配置: 论文未明确提及具体的硬件配置，但从训练7.7T tokens的规模来看，应使用了大规模的GPU集群。

• 软件配置:
  • 预训练框架: flame【48】，基于torchtitan【49】。
  • SFT框架: LlamaFactory【55】。
  • 评测框架: lm-eval-harness【58】和evalplus【59】。

实验结果

5.1 基础模型评估

• 实验内容: 将Ouro 1.4B和2.6B基础模型与主流的开源模型（如Qwen, Gemma, Llama3系列）在通用知识、推理、数学、科学、编码和多语言能力上进行全面比较。
• 实验结果与结论:
  • Ouro 1.4B（4个循环步）的性能与4B参数的Qwen3-Base相当，在BBH、GSM8K和MATH500等挑战性推理任务上甚至更优（见表7）。
  • Ouro 2.6B在推理密集型基准上超越了高达8B参数的密集模型，如在MMLU-Pro、BBH和MATH500上均超过8B的Qwen3-Base（见表8）。
  • 结果验证了循环架构通过迭代计算能显著增强模型的推理和知识操纵能力。

表7 1.4B LoopLM模型与1-4B参数基线的比较。最佳分数加粗，次佳分数加下划线。LoopLM的列已用灰色突出显示。

表8 2.6B LoopLM模型与3-12B参数基线的比较。最佳分数加粗，次佳分数加下划线。LoopLM的列已用灰色突出显示。

5.2 推理模型评估

• 实验内容: 评估经过SFT的Ouro-Thinking模型（4个循环步）在AIME、OlympiadBench、GPQA等高难度数学和科学基准上的表现。
• 实验结果与结论: 迭代推理带来了持续的性能增益。Ouro-1.4B-Thinking的性能达到了4B基线的水平，而Ouro-2.6B-Thinking则与8B基线相当或更优（见表9）。

表9 不同基准的性能比较。对于AIME24和AIME25，我们报告了pass@1/pass@10指标。最佳分数加粗，次佳分数加下划线。

5.3 不同循环深度的性能与外推能力

• 实验内容: 分析Ouro模型在不同循环计算深度（T=1到8）下的性能。模型训练时的最大深度为T=4，因此T>4属于外推情况。
• 实验结果与结论:
  • 基础模型: 在标准基准上，性能通常在训练深度T=4达到峰值，外推时性能有中度下降（见表10、11）。然而，模型的安全性即使在外推区域（T>4）也持续提升。
  • 推理模型: 在T=1时性能很低，说明迭代对于复杂任务至关重要。性能峰值出现在训练深度附近（T=3或T=4），外推时性能下降（见表12、13）。

表10 Ouro 1.4B基础模型在不同循环步数下的性能（C-QA是CommonsenseQA【66】）。步骤5-8代表外推，因为模型训练时最多4步。性能在训练深度（T=4）达到峰值，然后下降。

表11 Ouro 2.6B基础模型在不同循环步数下的性能（C-QA是CommonsenseQA【66】）。步骤5-8代表外推，因为模型训练时最多4步。性能在训练深度（T=4）附近最强，外推时表现出不同的下降模式。

表12 Ouro-1.4B-Thinking模型按循环步数的性能。模型在T=4时训练。性能在T=4或T=5附近达到峰值。所有分数均为百分比（0-100）。

表13 Ouro-2.6B-Thinking模型按循环步数的性能。模型在T=4时训练。性能在T=3或T=4达到峰值。所有分数均为百分比（0-100）。

5.4 提前退出与自适应计算效率

• 实验内容: 对比了三种提前退出策略（静态退出、隐藏状态差异阈值、学习门控）在MMLU上的准确率-效率权衡曲线，并研究了KV缓存共享对推理效率的影响。
• 实验结果与结论:
  • 退出策略: 经过专门自适应训练的学习门控在所有计算预算下都达到了最佳准确率，证明了基于损失改进的训练信号的有效性。即使未经专门训练，学习门控也优于静态基线（见图5）。
  • KV缓存共享: 推理的解码阶段，复用最后一步的KV缓存可以在几乎不损失性能的情况下，将内存开销降低4倍（见表14）。

表14 解码期间的KV缓存共享策略。最后一步和平均策略都实现了最小的性能损失，同时将内存减少了4倍。

6. 从参数化知识视角理解LoopLM的优越性

• 实验内容: 通过一系列完全可控的合成任务，探究LoopLM的性能优势来源：是增加了知识容量，还是提升了知识操纵能力。
• 实验结果与结论:
  • 知识容量: 在Capo任务（记忆合成传记）上的实验表明，LoopLM并未增加知识容量，循环和非循环模型都维持在约2比特/参数的水平（见图6左）。
  • 知识操纵: 在Mano任务（模块化算术）和多跳问答任务上，LoopLM在相同参数量下显著优于非循环模型，甚至优于计算量相同的模型。这表明LoopLM具有更好的知识操纵归纳偏置，学习组合和操纵知识的样本效率更高（见图6右，图7）。

7. 安全性、忠实性与一致性

• 实验内容:

  • 安全性: 使用HEx-PHI数据集评估模型安全性，并对良性与有害提示的隐藏表示进行PCA分析。
  • 忠实性: 使用Quora Question Pairs数据集，通过探测器分析中间隐藏状态的可预测性以及不同循环步骤间答案的一致性，来评估推理过程的忠实性。

• 实验结果与结论:

  • 安全性: 随着循环步数的增加，模型的安全性（有害性得分降低）得到提升，即使在外推情况下也是如此。PCA分析表明，更多的循环步骤有助于模型更好地区分良性和有害提示（见图8a, 8b）。
  • 忠实性: 与标准模型（如Qwen3）答案在早期就已确定不同，LoopLM在不同循环步骤间的答案存在系统性分歧，表明模型在迭代过程中会真实地更新其决策，这是一种更忠实的推理过程（见图9）。

A5 结论

本文介绍了Ouro，一个通过在7.7T tokens的预训练中直接集成迭代计算和自适应深度，从而展现出卓越参数效率的循环语言模型家族。我们的1.4B和2.6B模型持续地达到或超过了4B和8B标准Transformer的性能，展示了2-3倍的效率增益。我们证明了这种优势并非源于知识存储的增加，而是根本上更强的知识操纵能力，这一点得到了合成实验和理论分析的支持。我们还提出了一种使用带有均匀先验的熵正则化的实用训练目标来学习自适应深度，并验证了高效的KV缓存共享策略，使得LoopLM在现实世界部署中变得可行。

除了性能之外，LoopLM架构还表现出独特的属性：其迭代求精过程提供了一个因果忠实的推理轨迹，缓解了标准CoT中常见的事后合理化问题；其安全对齐随着循环步数的增加而独特地改善，即使在进行外推时也是如此。这项工作将迭代潜在计算确立为超越参数和数据的第三个关键扩展轴。未来的研究应侧重于增强在更深层次上的性能外推，并探索更复杂的循环机制，从而在一个数据受限的时代，巩固这种参数高效方法作为必要的发展方向。

A6 附录

B 从LoopLM的物理学角度理解

B.4 案例研究：MMLU不同类别的改进。为了在广泛的真实世界基准上验证我们从合成任务中得出的发现，我们对MMLU所有57个子类别的性能增益进行了细粒度分析。我们的假设是，如果LoopLM主要增强知识操纵和推理能力，那么最大的性能增益应该出现在程序性、重推理的任务中，而重知识、基于检索的科目改进应该较小。我们通过比较单个循环步骤（Loop 1）与完全训练深度（Loop 4）的准确性来衡量相对改进。分析强烈支持了我们的假设。改进最显著的类别是那些需要逻辑、数学或程序性推理的类别；相反，那些更依赖于检索特定、记忆化事实或细微世界知识的类别增益最微小。例如，初等数学提升了155.6%，而全球事实仅提升了8.3%。这一鲜明对比表明，迭代计算并非简单地增加模型可访问的知识，而是在主动执行复杂主题（如逻辑和数学）所需的多步符号操纵。

B.5 理论：使用LoopLM进行潜在思考。本节证明了LoopLM可以在$O(\log n)$步内解决图可达性问题（部分图知识在参数中学习）。这个结果与Transformer和带填充的循环Transformer的表达能力密切相关【11, A little depth goes a long way: The expressive power of log-depth transformers, 2025, arXiv】【12, Exact expressive power of transformers with padding, 2025, arXiv】【78, Understanding transformer reasoning capabilities via graph algorithms, 2024, NeurIPS】【79, Transformers, parallel computation, and logarithmic depth, 2024, arXiv】。我们正式定义了图可达性任务，即给定一个隐藏的知识图$G$和输入的上下文图$G_{ctx}$，判断查询节点对$(s,t)$是否在组合图$G+G_{ctx}$中连通。
定理1（非正式）。固定$n$为组合知识图$G$的最大尺寸。给定上下文图$G_{ctx}$的邻接矩阵和一个查询对$(s, t)$，存在一个与$G_{ctx}$无关的单层Transformer，通过$O(\log_2 D)$次循环，可以检查在组合知识图$(G + G_{ctx})$中是否存在从$s$到$t$的路径，其中$D$是$(G + G_{ctx})$的直径。
证明思路是通过在所有节点对的连通性上进行并行搜索，每次循环将可达距离加倍。这使得潜在推理比传统的CoT（$O(n^2)$）或连续CoT（$O(D)$）的循环观点在顺序计算步骤上呈指数级高效，从$O(n^2)$减少到$O(\log D)$。这种潜在的高效推理能力可能是LoopLM在知识操纵方面优越性的原因。

D LoopLM的缩放定律

RQ1：标准模型与LoopLM的性能差距如何？ 实验表明，标准模型的性能在相同条件下优于LoopLM，但这个差距随着模型规模的增大而减小（见图13）。这表明随着模型变大，LoopLM的性能越来越接近标准模型。

RQ2：循环步数如何影响LoopLM的总损失和分步损失？ 我们发现总损失$L_t$和分步损失$L_s$都可以用关于模型大小$N$、数据大小$D$和最大循环步数$T_m$（或当前步数$T$）的幂律公式很好地拟合，展示了LoopLM的可预测性。拟合公式如下：
$L_t = E + \frac{A}{(N+t_1)^\alpha} + \frac{B}{(D+t_2)^\beta} + \frac{C}{(T_m+t_3)^\gamma}$
$L_s = E + \frac{A}{(N+t_1)^\alpha} + \frac{B}{(D+t_2)^\beta} + \frac{C}{(T+t_3)^\gamma}$
拟合结果与真实损失曲线高度一致（见图15、16、17、18），证明了缩放定律的有效性。

RQ3：总损失和分步损失之间有何内在联系？ 总损失是各分步损失的加权和，权重由学习到的门控机制$q_{\phi}$决定。我们观察到一个现象：当模型尺寸不足时，为了降低总损失，模型会牺牲浅层分步损失（使其增加），以便将更多权重分配给性能更好的深层步骤。这表明，对于LoopLM，更大的模型尺寸可能更有效。此外，我们证明了可以通过分步损失的缩放定律和$q_{\phi}$的分布来推导出总损失的缩放定律，这揭示了两者之间的内在联系。

A7 补充细节

7.3 更多讨论

为安全关键部署提供实践优势。在安全关键的部署中，一个实际的障碍是模型的阐述性推理与其最终答案可能不一致。LoopLM架构通过暴露一系列与最终预测器强对齐的中间预测器来缩小这一差距，这些中间预测器可用于加速和预防性控制。我们总结了三个部署优势。

内置草稿模型用于推测性解码。LoopLM的架构天然适用于推测性解码。我们可以将较浅的循环步骤$t$的输出$Text(R_t)$作为“提议（proposal）”，而将最终步骤$T$的输出$Text(R_T)$作为“验证（verifier）”。这种“提议-验证”分解无需训练外部草稿模型，并且由于共享参数，可以重用激活和KV状态，从而降低了验证开销。

$$(\underbrace{\text{Text}(R_s)}_{\text{proposal}}, \underbrace{\text{Text}(R_T)}_{\text{verifier}}), \quad 1 \leq s < T.$$

加速与预防性安全的结合。利用相同的“提议-验证”结构，可以在推测性解码中交错进行安全检查，而无需额外模型。在步骤$s$：1）用$Text(R_t)$生成草稿token并在$Text(R_T)$下验证；2）在向用户展示任何token之前，对草稿分布或样本运行安全筛选；3）如果检测到违规，则在流式输出前停止或重路由响应。由于中间预测与最终答案分布良好对齐，早期步骤的输出成为最终输出的可靠代理，可用于早期筛选。

随时生成与单调求精。我们的训练目标优化了跨步骤的期望任务损失，同时保持了“越深越好”的特性。因此，对于下一个token的预测损失，有$E[L^{(t+1)}] \leq E[L^{(t)}]$。这意味着每个额外的循环都会将分布精炼为更高质量的预测。这产生了一种“随时可用（anytime）”算法：解码可以从任何中间步骤$s$开始，并继续流式传输，而后续步骤则继续验证或修正。与通常需要完成推理前缀才能输出答案的思维链流程不同，LoopLM在每一步都提供了一个单一的预测接口，使得在有延迟约束时能够立即回退到较小的计算预算。

$$\mathbb{E}[\mathcal{L}^{(t+1)}] \leq \mathbb{E}[\mathcal{L}^{(t)}], \quad 1 \leq t < T,$$