Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

Jonas Geiping 1, Sean McLeish 2, Neel Jain 2, John Kirchenbauer 2, Siddharth Singh 2, Brian R. Bartoldson 3, Bhavya Kailkhura 3, Abhinav Bhatele 2, Tom Goldstein 2

A1 主要贡献

本文研究了一种新颖的语言模型架构，该架构能够通过在潜在空间中进行隐式推理来扩展测试时的计算量。其核心思想是，当前主流的推理模型通过生成更多的词元（tokens）来扩展计算，例如思维链（chain-of-thought）方法，但这需要专门的训练数据，且必须将复杂的内部推理过程投射到单一的、语言化的下一个词元上，这可能是低效的。

核心问题与研究目标：
本文旨在探索一种替代方案，让模型能够在其连续的潜在空间中进行“思考”，从而扩展测试时的计算能力，而不是通过生成更长的文本。这种方法的目标是：
1. 减少对特定训练数据的依赖：与需要长篇思维链示例的训练方法不同，该模型旨在使用标准训练数据，并通过增加测试时计算来提升能力。
2. 提高内存和计算效率：通过避免超长上下文窗口，降低训练和推理的内存需求，并减少分布式训练中的通信成本。
3. 捕捉非语言化推理：探索模型在无法轻易用言语表达的推理类型（如空间思维、物理直觉）上的潜力。
4. 建立“思考”而非“记忆”的先验：构建一个计算密集型、参数较少的架构，鼓励模型学习元策略、逻辑和抽象，而不是死记硬背。

创新点与主要贡献：
1. 提出一种潜在深度循环架构：该架构包含一个核心的深度循环块，在训练期间以随机采样的迭代次数运行。在测试时，可以通过增加该块的迭代次数来扩展计算量，从而在不增加参数数量的情况下提升模型性能。
2. 成功扩展并训练了大规模模型：本文将这一概念验证模型扩展至35亿参数，并使用8000亿词元的数据进行训练，证明了该架构的可行性和可扩展性。
3. 展示了测试时计算扩展的有效性：实验表明，通过增加循环迭代，该3.5B参数的模型在推理基准测试中性能显著提升，有时提升剧烈，其计算负载可与50B参数的传统模型相媲美。如下图所示，需要较少推理的任务（如OpenBookQA）能更快收敛，而像GSM8k这样的复杂任务能有效利用更多的计算。
4. 原生支持多种高级推理时特性：该循环深度模型天然支持多种在非循环模型中需要大量调整的功能，包括每词元自适应计算、(自)推测解码和KV缓存共享。
5. 揭示了潜在空间中的新兴计算行为：通过可视化潜在空间中的词元轨迹，本文发现随着模型规模的扩大，会出现一些有趣的计算行为，例如在处理数值计算时，模型会在潜在空间中旋转形状。

图 1: 我们训练了一个带有深度循环的3.5B参数语言模型。在测试时，模型可以通过更长的迭代来使用更多计算并提高其性能。模型完全通过在潜在空间中推理来提升性能，而不是通过在长思维链中“语言化”来扩展测试时推理。像OpenBookQA这样需要较少推理的任务收敛得更快，而像GSM8k这样的任务则有效地利用了更多的计算。

为何训练具有循环深度的模型

循环层赋能任意计算。循环层使Transformer模型能够在生成一个词元之前执行任意数量的计算。原则上，循环机制为测试时计算扩展提供了一个简单的解决方案。与更标准的通过长上下文进行推理的方法（【索引OpenAI, 2024】; 【索引DeepSeek-AI et al., 2025】）相比，潜在的循环式思考具有几个优势。

• 潜在推理不需要构建定制的训练数据。思维链推理要求模型在感兴趣的领域中构建的长篇演示上进行训练。相比之下，我们提出的潜在推理模型可以使用可变的计算预算进行训练，使用标准的训练数据，无需专门的演示，并且如果在测试时给予额外的计算，可以增强其能力。
• 潜在推理模型在训练和推理时需要比思维链推理模型更少的内存。因为后者需要极长的上下文窗口，可能需要专门的训练方法，如词元并行化（【索引Liu et al., 2023a】）。
• 循环深度网络每参数执行的FLOPs比标准Transformer更多，显著降低了大规模训练时加速器之间的通信成本。这尤其能够在训练网络互连较慢时实现更高的设备利用率。
• 通过构建一个计算密集且参数量小的架构，我们希望为模型设定一个强烈的先验，使其通过“思考”（即学习元策略、逻辑和抽象）来解决问题，而不是通过记忆。循环先验在学习复杂算法方面的优势已经在“深度思考”文献中得到证实（【索引Schwarzschild et al., 2021b】; 【索引Bansal et al., 2022】; 【索引Schwarzschild et al., 2023】）。

捕捉非语言化的推理方面。从更哲学的角度来看，我们希望潜在推理能够捕捉人类推理中那些难以言喻的方面，例如空间思维、物理直觉或（运动）规划。在循环过程的多次迭代中，高维向量空间中的推理将能够同时深入探索多个方向，而不是线性思维，从而产生一个能够展现新颖复杂推理行为的系统。

作为模型性能扩展的第三个轴。我们认为，以这种方式扩展计算与通过扩展（语言化的）推理（【索引Shao et al., 2024】）或在预训练中扩展参数数量（【索引Kaplan et al., 2020】）并不冲突，它可能构成了扩展模型性能的第三个轴。

一种可扩展的循环架构

本节将描述我们提出的具有潜在循环深度的Transformer架构，讨论设计选择和小型消融实验。架构图见图2。我们始终将序列维度表示为$n$，模型的隐藏维度表示为$h$，其词汇表表示为集合$V$。

3.1. 宏观设计

模型的功能分组。模型主要围绕仅解码器（decoder-only）的Transformer块构建（【索引Vaswani et al., 2017】; 【索引Radford et al., 2019】）。然而，这些块被构造成三个功能组：前奏（prelude）$P$，它使用多个Transformer层将输入数据嵌入到潜在空间；核心循环块（core recurrent block）$R$，它是循环计算的中心单元，用于修改状态$s \in \mathbb{R}^{n \times h}$；以及尾声（coda）$C$，它从潜在空间中进行反嵌入，使用多个层并包含模型的预测头。核心块位于前奏和尾声块之间，通过循环核心块，我们可以在我们的“歌曲”中加入无限多的“诗节”。

图 2: 架构可视化图，如第3节所述。每个块由若干子层组成。蓝色的前奏块将输入嵌入到潜在空间，其中绿色的共享循环块是一组重复的层，用于计算最终的潜在状态，该状态由红色的尾声块的层解码。

模型的计算流程。给定循环迭代次数$r$和输入词元序列$x \in V^n$，这些功能组按以下方式用于产生输出概率$p \in \mathbb{R}^{n \times |V|}$。其中$\sigma$是用于初始化随机状态的某个标准差。这个过程如图2所示。给定一个初始随机状态$s_0$，模型重复应用核心块$R$，该块接受潜在状态$s_{i-1}$和嵌入的输入$e$，并输出一个新的潜在状态$s_i$。在完成所有迭代后，尾声块处理最后一个状态并产生下一个词元的概率。

$$\begin{aligned} \begin{aligned} \mathbf{e} &=P(\mathbf{x}) \\ \mathbf{s}_{0} & \sim \mathcal{N}\left(\mathbf{0}, \sigma^{2} I_{n \cdot h}\right) \\ \mathbf{s}_{i} &=R\left(\mathbf{e}, \mathbf{s}_{i-1}\right) \quad \text { for } \quad i \in\{1, \ldots, r\} \\ \mathbf{p} &=C\left(\mathbf{s}_{r}\right), \end{aligned} \end{aligned}$$

架构设计依据。该架构基于深度思考（deep thinking）文献，其中表明在每一步都注入潜在输入$e$（【索引Bansal et al., 2022】）以及用随机状态初始化潜在向量，可以稳定循环并促进其收敛到一个与初始化无关的稳态，即路径无关性（path independence）（【索引Anil et al., 2022】）。

设计动机。这种循环设计是学习稳定迭代算子所需的最小配置。一个很好的例子是函数$E(x, y)$的梯度下降，其中$x$可能是感兴趣的变量，$y$是数据。对此函数的梯度下降从一个初始随机状态$x_0$开始，并重复应用一个简单的操作（它所优化的函数的梯度），该操作依赖于前一个状态$x_k$和数据$y$。请注意，我们需要在每一步都使用$y$来实际优化我们的函数。同样，在我们的设置中，我们在循环的每一步都重复注入数据$e$。如果$e$仅在开始时提供，例如通过$s_0 = e$，那么迭代过程将不稳定，因为它的解将仅依赖于其边界条件。

前后层与中间层的差异化设计。使用多个层将输入词元嵌入到隐藏的潜在空间中的结构，是基于分析标准固定深度Transformer的经验结果（【索引Skean et al., 2024】; 【索引Sun et al., 2024】; 【索引Kaplan et al., 2024】）。这方面的研究表明，LLM的初始层和末端层有显著不同，而中间层是可互换和可排列的。例如，Kaplan等人（2024）的研究表明，在几层之内，标准模型已经将子词词元嵌入到潜在空间中的单个概念上，然后模型再对这些概念进行操作。

备注 3.1（这是扩散模型吗？）。这种迭代架构看起来会很熟悉另一种现代迭代建模范式——扩散模型（【索引Song and Ermon, 2019】），特别是潜在扩散模型（【索引Rombach et al., 2022】）。我们进行了一些消融实验，采用了更类似于扩散模型的迭代方案，例如$s_i = R(e, s_{i-1}) + n$，其中$n \sim \mathcal{N}(0, \sigma_i I_{n \cdot h})$，但在我们的初步实验中发现注入噪声并没有帮助，这可能与我们的训练目标有关。我们还评估了$s_i = R_i(e, s_{i-1})$，即核心块将当前步骤作为输入（【索引Peebles and Xie, 2023】），但发现这与路径无关性相互作用不佳，导致模型无法外推。

3.2. 微观设计

层内组件。在每个功能组内部，我们大致遵循标准的Transformer层设计。每个块包含多个层，每层包含一个标准的、使用RoPE（【索引Su et al., 2021】）且基数为50000的因果自注意力块，以及一个门控SiLU MLP（【索引Shazeer, 2020】）。我们使用RMSNorm（【索引Zhang and Sennrich, 2019】）作为我们的归一化函数。模型在查询（queries）和键（keys）上有可学习的偏置，其他地方没有。

层归一化策略。为了稳定循环，我们将所有层排列成如下的“三明治”格式，使用归一化层$n_i$，这与（【索引Ding et al., 2021】; 【索引Team Gemma et al., 2024】）中的类似策略相关但并不完全相同。虽然在小规模上，大多数归一化策略，例如前置归一化（pre-norm）、后置归一化（post-norm）等，效果几乎一样好，但我们对这些选项进行了消融实验，发现在大规模训练循环时需要这种归一化。

$$\begin{aligned} \begin{aligned} \hat{\mathbf{x}}_l &= n_2 \left( \mathbf{x}_{l-1} + \text{Attn}(n_1(\mathbf{x}_{l-1})) \right) \\ \mathbf{x}_l &= n_4 \left( \hat{\mathbf{x}}_l + \text{MLP}(n_3(\hat{\mathbf{x}}_l)) \right) \end{aligned} \end{aligned}$$

前奏块细节。给定一个嵌入矩阵$E$和嵌入缩放因子$\gamma$，前奏块首先将输入词元$x$嵌入为$\gamma E(x)$，然后应用$l_P$个具有上述布局的前奏层。

核心循环块细节。我们的核心循环块$R$以一个适配器矩阵$A: \mathbb{R}^{2h} \to \mathbb{R}^{h}$开始，该矩阵将$s_i$和$e$的拼接映射到隐藏维度$h$（【索引Bansal et al., 2022】）。虽然对于较小的模型，通过加法而非拼接来重新融入初始嵌入特征效果同样好，但我们发现在大规模上拼接效果最好。然后，这被送入$l_R$个Transformer层。在核心块的末尾，输出再次用一个RMSNorm $n_c$进行重新缩放。

尾声块细节。尾声块包含$l_C$个层，由$n_c$进行归一化，并使用绑定的嵌入$E^T$投影到词汇表中。

架构参数总结。总的来说，我们可以用三元组$(l_P, l_R, l_C)$来总结该架构，描述每个阶段的层数，以及循环次数$r$，它在每次前向传播中可能会变化。我们训练了一些形状为(1, 4, 1)、隐藏尺寸$h=1024$的小规模模型，以及一个形状为(2, 4, 2)、$h=5280$的大型模型。这个模型只有8个“真实”层，但当循环块迭代32次时，它会展开到有效深度为$2 + 4r + 2 = 132$层，构建出比一些最大的固定深度Transformer还要深的计算链（【索引Levine et al., 2021】; 【索引Merrill et al., 2022】）。

3.3. 训练目标

通过展开进行循环模型训练。为了确保模型在测试时增加循环迭代次数时能够正常工作，我们在训练期间随机采样迭代次数，为每个输入序列分配一个随机的迭代次数$r$（【索引Schwarzschild et al., 2021b】）。我们优化损失函数$L$在随机样本$x$（来自分布$\mathcal{X}$）和随机迭代次数$r$（来自分布$\Lambda$）上的期望。

$$\mathcal{L}(\theta)=\mathbb{E}_{\mathbf{x} \in X} \mathbb{E}_{r \sim \Lambda} L\left(m_\theta(\mathbf{x}, r), \mathbf{x}^{\prime}\right) .$$

其中，$m$表示模型输出，$x'$是序列$x$向左移一位，即序列$x$中的下一个词元。我们选择$\Lambda$为一个对数正态泊松分布。

迭代次数的采样分布。给定一个目标平均循环次数$\bar{r}+1$和一个我们设为$\sigma=1$的方差，我们可以通过以下方式从此分布中采样：

$$\tau \sim \mathcal{N}(\log(\bar{r}) - \frac{1}{2}\sigma^2, \sigma)$$ $$r \sim \mathcal{P}(e^{\tau}) + 1,$$

给定正态分布$\mathcal{N}$和泊松分布$\mathcal{P}$，见图3。该分布最常采样小于$\bar{r}$的值，但它包含一个重尾，偶尔会出现迭代次数显著增加的事件。

图 3: 我们使用对数正态泊松分布为每个训练步骤采样循环迭代的次数。

截断反向传播。为了在训练时保持计算和内存的低开销，我们仅通过循环单元的最后$k$次迭代进行反向传播。这使我们能够使用重尾泊松分布$\Lambda$进行训练，因为最大激活内存和反向计算现在与$r$无关。在我们的主要实验中，我们固定$k=8$。在小规模上，这与均匀采样$k$的效果一样好，但固定$k$可以使训练的每一步内存使用量相等。需要注意的是，前奏块在每一步仍然接收梯度更新，因为它的输出$e$在每一步都被注入。这种设置类似于截断时间反向传播（truncated backpropagation through time），这在RNN中很常见，尽管我们的设置是在深度上而非时间上循环（【索引Williams and Peng, 1990】; 【索引Mikolov et al., 2011】）。

训练一个大规模循环深度语言模型

在验证了我们能够可靠地训练高达100亿词元的小型测试模型后，我们转向了更大规模的运行。考虑到我们有限的计算预算，我们可以选择训练多个太小而无法显示涌现效应或规模效应的微型模型，或者训练一个中等规模的模型。基于此，我们准备进行一次单一的运行，具体细节如下。

4.1. 训练设置

我们描述训练设置，分为架构、优化设置和预训练数据。我们公开发布所有训练数据、预训练代码以及一系列中间模型检查点。

预训练数据。由于只有足够进行一次大规模模型运行的计算资源，我们选择了一个数据集混合，以最大化涌现推理行为的潜力，而不一定是为了获得最佳的基准测试性能。我们的最终混合严重偏向于代码和数学推理数据，同时（希望）有足够的通用网络文本以使模型获得标准的语言建模能力。所有来源都是公开可用的。我们在图4中提供了一个概览。遵循Allen-Zhu和Li（2024）的做法，我们直接将相关的指令数据混合到预训练数据中。然而，由于计算和时间限制，我们未能对这种混合进行消融实验。我们预计，更仔细的数据准备可以进一步提高模型的性能。我们在附录C中列出了所有数据来源。

图 4: 训练期间包含的数据来源分布。我们的大部分数据由通用网络文本、科学写作和代码组成。

分词与打包细节。我们通过BPE（【索引Sennrich et al., 2016】）构建了一个包含65536个词元的词汇表，使用了Dagan（2024）的实现。与传统的词元分析器训练相比，我们直接在我们预训练语料库的指令数据部分构建我们的词元分析器，以最大化在目标域上的分词效率。我们还大幅修改了预分词的正则表达式（例如Dagan等人(2024)的实现），以更好地支持代码、缩写和LaTeX。我们在每个文档的开头都包含一个<|begin_text|>词元。在对我们的预训练语料库进行分词后，我们将分词后的文档打包成长度为4096的序列。在打包时，我们丢弃了那些会缺少前文上下文的文档末尾，以修复Ding等人(2024)中描述的“接地问题”，但保留了一些数学内容的长文档来源的完整性。

架构与初始化。我们扩展了第3节中描述的架构，将层数设置为(2, 4, 2)，并以平均循环值为$\bar{r}=32$进行训练。我们主要通过将隐藏尺寸增加到$h=5280$来扩展模型，这产生了55个大小为96的头。MLP的内部维度是17920，RMSNorm的$\epsilon$是$10^{-6}$。总的来说，这个模型形状在非循环的前奏和头部约有15亿参数，在核心循环块中有15亿参数，在绑定的输入嵌入中有5亿参数。

初始化细节。在小规模上，大多数合理的初始化方案都能奏效。然而，在更大规模上，我们使用了Takase等人（2024）的初始化方法，该方法规定方差为$\sigma^2_h = \frac{2}{h}$。我们从一个截断正态分布（在$3\sigma$处截断）中初始化所有参数，方差为此值，但所有输出投影层的方差设置为$\sigma^2_{out} = \frac{1}{5hl}$，其中$l = l_P + \bar{r}l_R + l_C$是有效层数，对于这个模型是132。因此，输出投影层以相当小的值初始化（【索引Goyal et al., 2018】）。嵌入层的输出按$\sqrt{h}$缩放。为了匹配这种初始化，状态$s_0$也从截断正态分布中采样，此处的方差为$\sigma_s^2 = \frac{2}{5}$。

锁定步长采样。为了实现并行工作节点之间的同步，我们为每个微批次（micro-batch）的训练采样一个单一的深度$r$，并在工作节点之间同步这个值（否则，工作节点会空闲等待具有最大$r$的模型完成其反向传播）。我们在小规模上验证了这一修改提高了计算利用率而没有影响收敛速度，但注意到在大批量训练时，可以通过在每个工作节点上优化采样和调度独立的步骤$r$来进一步改进训练，以更忠实地模拟方程(1)中对步骤的期望。

优化器与学习率调度。我们使用带有解耦权重正则化的Adam优化器（$\beta_1 = 0.9, \beta_2 = 0.95, \eta = 5 \times 10^{-4}$）（【索引Kingma and Ba, 2015】; 【索引Loshchilov and Hutter, 2017】），并修改为包含更新裁剪（【索引Wortsman et al., 2023b】）和移除$\epsilon$常数（如Everett等人(2024)所述）。我们将梯度裁剪在1以上。我们使用预热（warm-up）和恒定学习率进行训练（【索引Zhai et al., 2022】; 【索引Geiping and Goldstein, 2023】），在前4096步内预热到我们的最大学习率。

4.2. 计算设置与硬件

硬件平台。我们使用橡树岭国家实验室的Frontier超级计算机上分配的计算时间来训练这个模型。这个HPE Cray系统包含9408个计算节点，每个节点配备AMD MI250X GPU，通过4个HPE Slingshot-11网卡连接。调度系统通过SLURM进行协调。我们使用基于PyTorch的实现，以bfloat16混合精度进行训练（【索引Zamirai et al., 2021】）。

设备速度与并行策略。名义上，每个MI250X芯片在bfloat16精度下可达到192 TFLOPs（【索引AMD, 2021】）。对于单个矩阵乘法，我们在我们的软件栈（ROCM 6.2.0, PyTorch 2.6预发布版11/02）上测得这些GPU上可达到的最高速度为125 TFLOP/s（【索引Bekman, 2023】）。我们的实现通过广泛的PyTorch编译和对隐藏维度$h=5280$的优化，实现了108.75 TFLOP/s的单节点训练速度，即87%的AFU（“可实现浮点利用率”）。由于我们循环设计中固有的权重共享，即使是我们最大的模型也足够小，可以仅使用数据并行（而非张量并行）进行训练，仅在每个迭代粒度上进行优化器分片（【索引Rajbhandari et al., 2020】）和梯度检查点。每个GPU的批量大小为1，最终全局批量大小为每步16M词元，从而最小化了GPU间的通信带宽。

大规模并行与互联问题。当我们在4096个GPU上进行大规模运行时，我们实现了每个GPU 52-64 TFLOP/s的速度，即41%-51%的AFU，或每秒1-1.2M个词元。为了实现这一点，我们编写了一个手工制作的分布式数据并行实现，以规避一个关键的AMD互连问题，我们在附录A.2中更详细地描述了这一点。总的来说，我们相信这可能是迄 F 写作时在AMD集群上并行使用设备数量方面完成的最大语言模型训练运行。

图 5: 前两次失败尝试和最终成功运行（“Main”）的最初10000个步骤的图。注意前两次失败运行中的隐藏状态崩溃（中图）和循环崩溃（右图），这突显了我们的架构和初始化在引导循环模型方面的重要性，并解释了这些运行在预训练损失方面的表现不佳（左图）。

训练时间线。训练通过21个最长12小时的片段进行，主要安排在2024年12月初的Frontier上。我们还进行了一个基线比较，其中我们以仅通过核心/循环块一次的前馈方式训练相同的架构。该基线在256个节点上使用每个GPU批量大小为2的设置训练了180B个词元。最终，我们能够为主模型安排795B个词元的预训练。由于我们的恒定学习率调度，我们能够在有可用分配时“按需”添加额外的片段。

4.3. 归一化与初始化在规模化中的重要性

规模化带来的挑战。在小规模上，所有归一化策略都有效，我们只观察到不同初始化之间微小的差异。但在大规模上情况并非如此。我们开始的第一次训练运行设置与上述相同的三明治块结构，但使用了无参数的RMSNorm层，没有嵌入缩放$\gamma$，一个无参数的适配器$A(s, e) = s + e$，以及$4 \times 10^{-4}$的峰值学习率。如图5所示，这次运行（“Bad Run 1”，橙色）很快就停滞了。

失败分析一：表示崩溃。虽然这次运行在训练损失上明显停止了改进（左图），我们发现这种停滞是由于模型的表示崩溃（【索引Noci et al., 2022】）。词元维度中隐藏状态的相关性迅速变为1.0（中图），这意味着模型为序列中的每个词元预测相同的隐藏状态。我们发现这是一个由于循环操作而产生的初始化问题。循环块的每次迭代都会增加词元相关性，混合序列直到崩溃。

失败分析二：未能利用循环。我们尝试通过引入嵌入缩放因子，切换回传统的预归一化块，并切换到学习型适配器来解决这个问题。最初，这些改变似乎解决了问题。尽管在开始时词元相关性飙升接近1.0（“Bad Run 2”，绿色），但在最初的150步后模型恢复了。然而，我们很快发现这次训练运行无法有效地利用测试时计算（右图），因为无论使用1次还是32次循环，验证困惑度都是相同的。这种初始化和归一化设置导致了一个局部最小值，因为模型早期学会了忽略输入状态s，从而阻止了进一步的改进。

成功的运行。在第三次也是最后一次运行中（“Main”，蓝色），我们通过恢复到三明治块格式，并进一步将峰值学习率降至$4 \times 10^{-5}$来解决了这个问题。这次运行开始得很顺利，从未达到接近1.0的词元相关性，并很快通过利用循环和随着更多迭代而改进，超过了之前的运行。

最终训练过程。在我们成功的配置下，训练在接下来的7500亿个词元中顺利进行，没有出现明显的终端或损失峰值。我们在图6中绘制了不同循环步骤下的训练损失和困惑度。在我们的材料中，我们将这次运行的最终检查点称为我们的“主模型”，我们将其表示为Huginn-0125。

图 6: 左图：主运行中800B词元上的预训练损失图。右图：在循环深度为1, 4, 8, 16, 32, 64时的验证困惑度图。在训练过程中，模型在所有循环层级上的困惑度都有所提高。

实验环境

• 模型架构：

  • 模型名称：Huginn-0125
  • 参数量：3.5B
  • 架构：潜在深度循环Transformer，层数配置为 (前奏: 2, 核心: 4, 尾声: 2)
  • 关键参数：隐藏层维度 $h=5280$，注意力头数 55，头维度 96，MLP内部维度 17920。
  • 训练时平均循环次数 $\bar{r}=32$。

• 数据集：

  • 总规模：训练了约800B词元。
  • 数据来源：所有数据均为公开可用的数据集，精心混合而成。
  • 数据构成（如图4所示）：主要由通用网络文本（28.71%）、代码（25.36%）、科学文本（18.73%）和数学文本（6.14%）组成，同时混合了少量指令数据用于引导模型学习推理行为。

• 硬件配置：

  • 平台：橡树岭国家实验室的Frontier超级计算机。
  • GPU：在4096个AMD MI250X GPU上进行训练。
  • 互联：4x HPE Slingshot-11 NICs。

• 软件配置：

  • 框架：基于PyTorch的实现。
  • 精度：bfloat16混合精度训练。
  • 依赖库：ROCM 6.2.0, PyTorch 2.6预发布版。
  • 并行策略：自定义的分布式数据并行（DDP）实现，结合ZeRO-1优化器分片和梯度检查点。

实验结果

5.1. 标准基准测试

• 实验内容：将最终训练800B词元的模型（Huginn-0125）与Amber、Pythia和OLMo等大小相似的开源模型在多个标准语言模型基准（如ARC, HellaSwag, MMLU, OpenBookQA等）上进行零样本（zero-shot）比较。
• 实验结果：如表1所示，随着测试时循环次数（$r$）的增加，模型性能显著提升。在$r=32$时，模型性能超越了较早的Pythia系列模型，与第一代OLMo-7B大致相当，但落后于在更大、更精心策划的数据集上训练的更新的OLMo模型。
• 分析结论：考虑到这是首个在此规模上训练的循环深度语言模型，并且训练运行存在一定局限性（例如未冷却学习率），这些结果是有希望的，表明潜在循环作为一种测试时扩展方法值得进一步研究。

表 1: 在lm-eval-harness任务上对各种开源模型的零样本测试结果。我们展示了ARC（【索引Clark et al., 2018】）、HellaSwag（【索引Zellers et al., 2019】）、MMLU（【索引Hendrycks et al., 2021a】）、OpenBookQA（【索引Mihaylov et al., 2018】）、PiQA（【索引Bisk et al., 2020】）、SciQ（【索引Johannes Welbl, 2017】）和WinoGrande（【索引Sakaguchi et al., 2021】）的结果。我们报告了提供的归一化准确率。

5.2. 数学与代码基准测试

• 实验内容：在数学（GSM8k, MATH, MathQA）和代码（MBPP, HumanEval）基准上评估模型。
• 实验结果：如表2和表3所示，模型在数学推理方面（GSM8k和MATH）显著超越了除最新OLMo-2之外的所有模型。在代码基准上，它击败了所有其他通用开源模型，但未超过专门的、训练了数万亿词元的代码模型（如StarCoder2）。
• 分析结论：模型在特意倾斜的数据集（数学和代码）上表现出色。如图8所示，与趋于平缓的语言建模能力相比，代码和数学推理能力在整个训练过程中持续稳定提高。

表 2: 数学推理和理解的基准测试。我们报告GSM8K和GSM8K CoT的灵活和严格提取匹配，Minerva Math的提取匹配，以及MathQA的归一化准确率。

表 3: 在代码基准MBPP和HumanEval上的评估。我们报告了两个数据集的pass@1。

5.3. 循环在何处帮助最大？

• 实验1：与非循环基线对比
  • 内容：将训练了180B词元的循环模型与其在完全相同设置下训练的非循环“孪生”模型进行比较（见表4）。
  • 结果：循环模型表现更优，尤其是在ARC-C等较难的任务上。在GSM8k上，180B的循环模型已比基线好5倍。同时，当循环模型仅用一次循环（r=1）进行评估时，其性能从180B词元到800B词元几乎没有提升。
  • 结论：性能的提升主要归功于循环块的迭代，而非非循环的前奏或尾声部分。

表 4: 基线比较，循环模型与在相同训练设置和数据下训练的非循环模型。比较循环模型与其非循环基线，我们看到即使在180B词元时，循环模型在更难的任务上已大幅领先。

• 实验2：性能与测试时计算量的关系
  • 内容：绘制模型在多个任务上性能随测试时循环次数增加的变化曲线（见图7）。
  • 结果：性能饱和点高度依赖于任务。较简单的任务（如HellaSwag）在约8次循环后就接近峰值性能，而较难的任务（如GSM8k, HumanEval）能从更多的计算中持续受益。
  • 结论：模型能自适应地将更多计算用于解决更难的问题。

图 7: 在GSM8K CoT（严格匹配和灵活匹配）、HellaSwag（归一化准确率）和HumanEval（pass@1）上的性能。随着我们增加计算量，这些基准测试的性能也随之提高。HellaSwag仅需8次循环即可达到接近峰值的性能，而其他基准测试则利用了更多的计算。

图 8: 在不同测试时循环次数下，GSM8K CoT、HellaSwag和HumanEval在训练词元数上的性能。我们使用聊天模板和8-shot few-shot作为多轮评估GSM8K CoT。HellaSwag和HumanEval是零样本，没有聊天模板。如果提供足够的测试时计算，模型在更难任务上的性能几乎与训练预算成线性增长。

• 实验3：循环与上下文
  • 内容：评估ARC-C性能与循环次数及少样本（few-shot）示例数量的关系（见图9），并在开放问答（提供相关事实）和封闭问答设置下比较OpenBookQA性能（见表5）。
  • 结果：当提供更多少样本示例时，模型能有效利用更多循环次数进行推理。在开放问答设置下，模型的性能显著提升，缩小了与其他模型的差距。
  • 结论：模型擅长利用循环来推理上下文信息，而非仅仅记忆事实。这表明其记忆能力相对较弱，但上下文推理能力较强。

图 9: 在ARC挑战集上，通过测试时循环达到的未归一化准确率的饱和点与少样本示例的数量相关。模型使用更多的循环来从额外的少样本示例中提取更多信息，在给定更多上下文的情况下利用更多计算。

表 5: 开放式与封闭式问答性能（%）比较（【索引Mihaylov et al., 2018】）。在开放式考试中，提问前会提供一个相关事实。在这种设置下，我们的小模型缩小了与其他开源模型的差距，表明该模型有能力，但记忆的事实较少。

5.4. 通过权重平均进行改进

• 实验内容：由于采用恒定学习率，可以通过权重平均来模拟学习率冷却的效果。对最后75个检查点应用指数移动平均（EMA）。
• 实验结果：在$r=64$时，GSM8k性能进一步提升至47.23%（灵活匹配）。
• 分析结论：权重平均是提升该模型性能的有效后处理方法。

补充细节

6. 循环深度简化了LLM

除了在数学和代码推理方面表现出色外，循环深度模型还天然地支持一些在标准Transformer中需要大量工作才能实现的方法。

6.1. 测试时的零样本自适应计算

无需专门训练的自适应计算。我们已经证明模型能够在每个查询级别上变化计算量，以不同的循环模式运行模型。这毕竟也是模型的训练方式，如方程(1)所示。然而，在实践中，当预测简单时提前停止循环，只在困难的决策上花费计算会更高效。其他工作，特别是基于标准Transformer的工作，需要专门为提前退出训练的模型（【索引Elbayad et al., 2019】; 【索引Fan et al., 2019】; 【索引Banino et al., 2021】），或在每一层都带有退出头的微调模型（【索引Schuster et al., 2022】）。

简单的退出标准。为了测试我们模型的零样本退出能力，我们选择一个简单的退出标准来评估收敛性，即两个连续步骤之间的KL散度。如果这个散度低于$5 \times 10^{-4}$，我们就停止迭代，采样输出词元，并继续生成下一个词元。

任务相关的计算量。我们在图10中展示了这种零样本的每词元自适应计算行为，图中绘制了在达到退出条件之前所采取的步骤分布。我们对来自不同MMLU类别的头50个问题进行了此操作，以自由形式的聊天方式提问。有趣的是，退出所需的步骤数在不同类别之间有显著差异，模型在高中数学上退出得更早，但在道德情景上平均多花3.5步。作为初步演示，我们在MTBench上验证了这种自适应性不会显著影响对话基准测试的性能（标准：5.63，提前退出：5.56，见附录表6）。

图 10: 基于步骤间KL差异的零样本、每词元自适应退出在MMLU类别问题上的直方图，包含和不包含零样本连续CoT。每个分布的均值在图例中给出。退出阈值固定为$5 \times 10^{-4}$。我们看到模型在高中数学上的收敛速度比在逻辑谬误或道德情景等任务上更快。在某些任务上，例如哲学，模型能够在其潜在CoT中有效重用状态，并在部分词元上快速收敛，从而总体上需要更少的步骤。

备注 6.1 (关于缺失的KV缓存条目？)。传统上，对于具有自注意力机制的模型，逐词元提前退出的一个担忧是它从根本上破坏了KV缓存。在每个循环步骤中，一个词元需要关注序列中先前词元的KV状态，但这些激活可能因为提前退出而没有被计算。一个简单的修复方法是暂停生成并重新计算所有缺失的隐藏状态，但这会削弱提前停止的好处。相反，我们借鉴Elbayad等人(2019)的做法，关注缓存中最后一个、最深层可用的KV状态。因为所有循环的KV缓存条目都是由相同的K、V投影矩阵从连续的隐藏状态生成的，它们是“匹配”的，因此模型能够关注来自每个先前词元的最新缓存条目，即使它们是在不同的循环深度计算的。

6.2. 零样本KV缓存共享

无需训练的KV缓存共享。另一种提高效率的途径是通过在层之间共享KV缓存来减少其内存占用（【索引character.ai, 2024】; 【索引Brandon et al., 2024】）。通常，Transformer必须从头开始训练才具备此功能。然而，如上一节所述，我们发现我们可以简单地在我们的模型中共享KV缓存，而对性能的影响微乎其微。

实现机制。我们为每个词元$k$的循环设置一个固定的KV缓存预算，在迭代$i$时，读写缓存条目$i \pmod k$。例如，我们设置最大KV缓存预算为16步，在执行第17步时覆盖第1步的KV缓存，依此类推。这可以单独用于减少KV缓存内存，或与上述的每词元自适应计算结合使用。在MTBench上，这并未降低性能（缓存预算为4时：5.86，见附录表6）。

6.3. 零样本连续思维链

构建更深的计算图。通过关注先前词元后续步骤的输出在当前词元早期步骤中的应用，如KV缓存共享部分所述，我们实际上构建了一个比当前循环步数更深的计算。但我们也可以更明确地构建更深的计算图。我们可以在每个生成步骤中不采样一个随机的初始状态$s_0$，而是用前一个词元的最后一个状态$s_r$进行热启动。这样，模型可以从前一个生成步骤中编码的潜在信息中受益，并进一步改进。

减少收敛步数。如图10所示，这使得收敛所需的平均步数减少了1-2步。在哲学等任务上，我们看到退出分布发生了显著变化，模型更频繁地通过重用先前的计算来提前退出。

与现有工作的关联与区别。这与Hao等人（2024）探索的连续思维链方法密切相关，因为它都是对训练好的模型进行干预以增加额外的循环。为了在固定深度的Transformer中实现类似的行为，Hao等人（2024）在推理链上训练模型，使其在计算下一个词元时接受其最后一个隐藏状态作为替代输入。通过这种方式进行微调，也将这些模型转变为有限的深度循环模型——这样，两种方法的主要区别在于，是从头开始预训练以获得循环能力，还是微调现有的固定深度模型以具备此能力——以及是否需要思维链数据。

6.4. 零样本自推测解码

无需草稿模型的推测解码。循环深度模型还可以通过使用推测解码（【索引Leviathan et al., 2023】）来更有效地生成文本，而无需单独的草稿模型。

与现有方法的比较。使用标准Transformer模型，推测解码需要一个外部草稿模型、Medusa头（【索引Cai et al., 2024】）或提前退出适应（【索引Zhang et al., 2024b】; 【索引Elhoushi et al., 2024】）。Zhang等人（2024b）通过层跳过简单地实现了自推测解码，但这并不总能产生高质量的草稿。相比之下，我们的模型可以自然地用较少的迭代次数来草拟序列中的下N个词元，然后可以用任何期望的迭代次数M > N来验证。这也可以分阶段进行，或者草稿模型可以使用如6.1节所述的自适应计算。使用该模型进行草拟也是高效的，因为草拟过程中计算的状态不会被浪费，可以在验证时重新使用。

7. 循环深度模型在规模化中涌现出何种机制

分析潜在空间轨迹。最后，模型在潜在空间中循环时在做什么？为了更好地理解这个问题，我们分析了模型在几个定性示例上的轨迹${s_i}_{i=1}^r$。我们特别感兴趣的是，仅通过大规模训练该模型，会涌现出哪些模式。与之前的工作（如【索引Bai et al., 2019】）相比，其中训练目标直接编码了一个将轨迹推向不动点的先验，我们只使用我们的截断展开目标进行训练。

上下文相关的收敛行为。图11显示了轨迹中每个$s_i$与用128次迭代计算的近似极限点$s^$之间的范数距离$||s_i - s^||$。我们从上到下展示句子，从左到右展示迭代次数。我们清楚地看到收敛行为取决于上下文。我们看到问题的关键部分和模型响应的开始部分在潜在空间中被“深思熟虑”得更多。上下文依赖性也可以从代表三个点的三个相同词元之间的不同行为中看出。还要注意的是，到$s^*$的距离并不总是单调递减（例如，对于school）；模型在处理信息时也可能在其潜在轨迹中描绘出复杂的轨道，尽管这在我们的训练目标中没有明确表示。

图 11: 序列中每个词元（从上到下）和潜在迭代（从左到右）的潜在状态收敛情况，绘制了与最终迭代$s^*$的距离，我们设置$r = 128$。图中显示了一个向模型提出的不安全问题。我们立即看到，高度依赖于词元的收敛速率仅通过规模化就涌现出来。这很有趣，因为模型仅在训练期间看到的整个序列中以固定的$r$进行训练。我们看到，在问题的关键部分“really wronged”上收敛尤其缓慢。我们还看到模型学习了不同的行为，我们看到潜在空间中存在振荡模式，这里最明显的是“school”词元。图中未显示的是模型在深思熟虑该问题后拒绝回答。

潜在空间中的几何模式。我们在图12中更详细地观察了选定词元的轨迹。我们计算了序列中所有词元的潜在轨迹的PCA分解，然后展示了几个投射到前六个PCA方向的单个轨迹。更多示例见附录。许多词元只是收敛到一个不动点，如顶行的词元。然而，对于更难的问题，如第二行，词元的状态在所有三对PCA方向上迅速进入轨道模式。使用这样的多维轨道可能与在为算术任务训练的固定深度Transformer中有时观察到的周期性模式有类似的目的（【索引Nanda et al., 2022】），但我们发现这些模式在我们的模型中远不止于算术。我们也经常在诸如“makes”（见图16）或“thinks”等决定响应结构的词元上观察到轨道的使用。

“滑块”机制。除了轨道，我们还观察到模型将特定的关键词元编码为“滑块”，如图12底行中间所示（这是词元“wrong”，来自图11中已显示的同一消息）。在这些运动中，轨迹在一个方向上显著漂移，模型可以利用这一点来实现一个计算已发生迭代次数的机制。

图 12: 选定词元的潜在空间轨迹。我们通过可视化前6个PCA方向来展示这些高维轨迹的一小部分，PCA是在序列中所有词元的所有潜在状态轨迹上计算的。颜色从暗到亮的渐变代表轨迹中的步骤。质心用红色标记。虽然在许多词元上，状态只是收敛（顶行），但模型也学会了使用轨道（中行）和“滑块”（底行中间），我们观察到这些被用来表示和处理更高级的概念，如算术或复杂的深思熟虑。

新兴的计算组织方式。潜在空间中结构化轨迹的出现让我们得以一窥模型如何执行其计算。与语言化思维链方法中看到的离散顺序推理链不同，我们观察到丰富的几何模式，包括轨道、收敛路径和漂移——这些是其在空间上组织计算过程的方式。这表明模型正在独立学习利用其潜在空间的高维性以新的方式实现推理。

路径无关性。我们验证了我们的模型尽管具有我们之前讨论过的复杂、学习到的动力学（另见附录图22中的附加示例），但仍保持着Anil等人（2022）意义上的路径无关性。当从多个起始状态$s_0$重新初始化时，模型会沿着相似的轨迹移动，表现出一致的行为。无论初始化如何，都会出现相同的轨道模式、不动点或方向性漂移。

8. 相关工作概述

循环概念的基础性。循环是机器学习的一个基础概念，其重要性难以言喻（【索引Amari, 1972】; 【索引Hopfield, 1982】; 【索引Braitenberg, 1986】; 【索引Gers and Schmidhuber, 2000】; 【索引Sutskever et al., 2008】）。除了用于沿序列移动（如循环神经网络），循环很早就被理解为自适应计算的关键（【索引Schmidhuber, 2012】; 【索引Graves, 2017】）。

Transformer中的循环。对于Transformer，Dehghani等人（2019）应用了循环，强调了循环深度的目标是建模通用图灵机（【索引Graves et al., 2014】）。Lan等人（2019）大规模使用了它（但循环次数固定）。最近的改进包括Tan等人（2023）、Abnar等人（2023）、Mathur等人（2024）和Csordás等人（2024）的工作。Schwarzschild等人（2021b）、Bansal等人（2022）、Bear等人（2024）和McLeish等人（2024）表明，在使用随机展开和输入注入进行训练时，深度循环在学习可泛化算法方面具有优势。最近的研究描述了深度循环、环形Transformer，并通过理论和小型实验分析了其潜在好处（【索引Giannou et al., 2023】; 【索引Gatmiry et al., 2024】; 【索引Yang et al., 2024a】; 【索引Fan et al., 2025】）。

与其他迭代模型的关联。从另一个角度看，这些模型可以被描述为学习不动点迭代的神经网络，正如深度均衡模型（Deep Equilibrium Models）中所研究的（【索引Bai et al., 2019】; 【索引2022】）。它们还与扩散模型相关（【索引Song and Ermon, 2019】），特别是潜在扩散模型（【索引Rombach et al., 2022】），但语言扩散模型通常是按序列而非按词元进行迭代（【索引Lee et al., 2018】）。我们的方法与均衡模型和扩散模型的一个关键区别在于训练目标：均衡方法解决“直接”问题（【索引Geiping and Moeller, 2019】），扩散模型解决一个代理训练目标，而我们的工作表明截断展开是一个可扩展的替代方案。

更广泛的理论联系。更一般地，所有在深度上循环的架构也可以被理解为直接学习潜在能量模型的梯度（【索引LeCun and Huang, 2005】; 【索引LeCun, 2022】）、隐式定义中间优化层（【索引Amos and Kolter, 2017】）或Kuramoto层（【索引Miyato et al., 2024】）。与推理时梯度下降的类比也显示了其与测试时自适应的联系（【索引Sun et al., 2020】），特别是输出状态的测试时自适应（【索引Boudiaf et al., 2022】）。

混合架构。除了完全的深度循环架构，还存在一些混合架构的提议，例如具有潜在子网络的模型（【索引Li et al., 2020a】）、在权重共享层之上使用LoRA适配器（【索引Bae et al., 2024】）或对已训练模型进行（动态）权重绑定（【索引Hay and Wolf, 2023】; 【索引Liu et al., 2024b】）。

微调实现潜在推理。如第6节所述，虽然我们认为所提出的循环深度方法是从头开始学习在连续潜在空间中推理的一种非常自然的方式，但Hao等人（2024）、Cheng和Durme（2024）以及Liu等人（2024a）的工作讨论了如何微调现有的固定深度Transformer以具备此能力。这些工作与我们的目标相似，即在潜在空间中实现推理，但从不同的方向着手。

关于激励推理而非记忆的先验。关于构建一个激励推理和算法学习，而非记忆简单模式的先验思想的额外讨论，我们还参考了Chollet（2019）、Schwarzschild（2023）、Li等人（2020b）和Moulton（2023）的研究。

9. 未来工作

扩展与分析。除了扩展和分析循环深度模型的缩放行为外，还有许多问题尚待解答。例如，对我们来说，可能存在大量新颖的后训练方案，可以进一步增强这些模型的能力，例如通过微调来压缩循环，或者使用不同难度级别的数据进行强化学习（【索引Zelikman et al., 2024】），或者将CoT数据的推理内化到循环中（【索引Deng et al., 2024】）。

与其他架构改进的结合。另一个本工作未涉及的方面是与其他现代架构改进的关系。高效的序列混合操作，特别是那些在序列维度上是线性的操作，如线性注意力（【索引Katharopoulos et al., 2020】; 【索引Yang et al., 2024b】），其可进行的比较次数有限。然而，通过循环深度，包含线性算子的块可以重复，直到计算出序列元素之间所有必要的比较（【索引Suzgun et al., 2019】）。为简单起见，我们也只关注单次循环，而先前的工作已经考虑了多个连续的循环阶段（【索引Takase and Kiyono, 2023】; 【索引Csordás et al., 2024】）。

与专家混合（MoE）模型的互补性。最后，所提出的架构被设置为计算密集型，其“物化”参数多于实际参数。这自然地反映了专家混合（MoE）模型（【索引Shazeer et al., 2017】; 【索引Fedus et al., 2022】），后者是参数密集型，每次前向传播使用的活动参数少于模型中存在的参数。我们假设，循环深度设置在学习推理模式方面表现出色，而MoE在有效存储和检索复杂信息方面表现出色。它们的互补性支持了未来架构将包含这两种修改的假设。在标准的MoE模型中，每个专家在每次前向传播中只能被激活一次，或完全跳过，而循环MoE模型也可以在多次迭代中精炼其潜在状态，可能多次路由到同一个专家，然后再切换到另一个（【索引Tan et al., 2023】; 【索引Csordás et al., 2024】）。虽然MoE模型是目前在密集Transformer中实现这种“记忆”的领先解决方案，但这些考虑也适用于为LLM建议的其他记忆机制（【索引Sukhbaatar et al., 2019】; 【索引Fan et al., 2021】; 【索引Wu et al., 2022】; 【索引He et al., 2024】）。

结论

本文中描述的模型最终仍是一个概念验证。我们描述了如何训练一个潜在的循环深度架构，我们选择了哪些参数，然后大规模地训练了一个模型。未来的训练运行很可能会使用更优化的学习率调度、数据混合和加速器。尽管如此，我们观察到一些有趣的行为从循环训练中自然涌现。其中最重要的是能够利用潜在推理，通过消耗测试时计算来显著提高推理任务的性能。此外，我们还观察到上下文相关的收敛速度、路径无关性以及各种零样本能力。这使我们相信，潜在推理是一个有前途的研究方向，可以补充现有的测试时计算扩展方法。我们实现的模型在其规模和训练数据量方面出人意料地强大，我们对赋予生成模型在连续潜在空间中进行推理的能力，而无需在训练时使用专门数据或在推理时进行语言化的潜力感到兴奋。

附录

A. 额外信息

这项工作的潜在影响
这项工作描述了一种新颖的语言建模架构和训练目标，表现出有前景的性能，尤其是在需要模型进行推理的任务上。本文描述的测试时扩展方法与其他扩展方法（即通过模型参数和通过测试时思维链）是互补的，并且存在关于成本和模型能力的类似担忧。我们提出的架构自然比通过参数扩展的模型要小，这可能对这些模型在商用芯片上的本地部署有更广泛的好处。最后，虽然我们认为将模型的推理能力转移到循环的高维、连续潜在空间在能力方面是有益的，但我们注意到，与目前仍然人类可读的语言化思维链相比，这可能会以模型监督为代价。我们在第7节中提供了初步结果，表明我们模型的高维状态轨迹可以被分析，并且其部分机制可以被解释。

A.1. 经典推理问题

我们在图14中对多操作数算术这一经典问题进行了小规模研究。

多操作数算术实验。遵循训练循环架构用于算法和算术任务的先例（【索引Schwarzschild et al., 2021b】; 【索引Bansal et al., 2022】; 【索引Schwarzschild et al., 2023】; 【索引McLeish et al., 2024】），我们探讨了我们的模型是否能利用通过循环增加的测试时计算来解决难度增加的语言化加法问题。对于这些问题，我们使用系统提示“你是一个能够帮助用户进行数学推理的有用助手”，嵌入到对话聊天模板中，并通过打开对话的第一个用户回合来呈现每个问题，格式如下：f"What is the result of ’ + ’.join(map(str, digits))?"，其中数字是根据特定的操作数数量和数字位数（以10为基数）随机抽样的。我们通过检查正确总和是否作为字符串出现在模型输出的任何位置来评分，每次测量平均超过50次试验。

实验结果分析。在热图（左上）中，我们评估了模型在32次循环下的表现，以获得其在各种难度下加法性能的上限估计。它能可靠地解决涉及两个操作数、每个操作数最多4或5位的加法问题，但对于4和5个操作数，它很少能正确地加单个数字。在每个折线图中，我们固定数字位数，扫描操作数数量，并评估模型在1到64次循环下的性能。我们看到，当相加单个数字时（右上），性能随着循环次数的增加而稳步提高。然而，当相加2位和3位数字时（下行），模型只有在超过16次循环评估时才能一致地解决问题。奇怪的是，我们观察到2位和3位情况下性能随循环次数的排序不一致，并且在5和4个操作数时性能出现一些峰值。我们注意到，该模型并未特别针对算术问题进行微调，尽管预训练数据中有相当一部分当然包含数学内容。

图 14: 多操作数算术。

表 6: 在单轮MT-Bench上，针对循环深度模型原生的各种推理时方案的第一轮得分和标准误差。与基线模型（即没有推理修改的正常循环模型）的差异在统计上不显著。

A.2. 实现细节

设备速度细节。名义上，每个MI250X（【索引AMD, 2021】）在bfloat16精度下可达到383 TFLOP，即每个GPU 192 TFLOP，但如前所述，在我们的技术栈（ROCM 6.2.0，PyTorch 2.6预发布版11/02）上测量任意矩阵乘法形状的可达TFLOP时（即我们测量在256到24576之间以256和110为间隔迭代形状的最佳可能形状的峰值可达速度（【索引Bekman, 2023】）），我们在Frontier节点上测得的峰值为125 TFLOP/s。使用带最大自动调优的PyTorch编译（无‘cudagraphs’，无优化器或autograd编译）（并优化我们的隐藏尺寸为5280），我们最终的模型实现在单节点上的训练速度为108.75 TFLOP/s，即57%的MFU（【索引Chowdhery et al., 2022】），或者说87%的AFU（“可达浮点利用率”）。我们注意到，由于自动混合精度和截断反向传播的相互作用，PyTorch的梯度仅在执行编译模型时是正确的。我们进一步使用AMD的原始flash attention仓库的fork来规避PyTorch sdpa附带的flash attention实现的问题，该fork可在https://http://github.com/ROCm/flash-attention 持Flash Attention 2（【索引Dao et al., 2022】; 【索引Dao, 2023】）。我们尝试了融合的头和损失实现，但最终发现在我们的AMD设置上最可移植的选择是让torch编译来处理这个问题。

并行化策略。如正文所述，由于我们的深度循环模型是计算密集型的，因此最优的运行方式是仅在节点间使用分布式数据并行训练，并在节点内使用zero-1优化器分片（【索引Rajbhandari et al., 2020】），如果我们利用循环迭代的每一步进行梯度检查点。这使我们能够避免对于具有相同FLOP足迹但参数更多的模型所需的通信密集型并行化策略，这些策略在该系统上需要大量的规划（【索引Singh et al., 2024】; 【索引Singh and Bhatele, 2022】）。然而，这一选择虽然最小化了通信，但也使我们锁定在每个设备批量大小为1，即总共4096，每步16M词元。

RCCL互连处理。由于调度原因，我们决定在Frontier上以512个节点的分配段为目标，即4096个GPU。然而，这带来了一个重大的网络互连问题。Frontier节点之间的连接速度只有在RCCL（AMD GPU通信集合）命令通过开放光纤接口调用路由时才能接受，这是通过一个特定的插件实现的。要达到超过100GB/s的总线带宽，需要设置NCCL_NET_GDR_LEVEL=PHB，这个设置在NVIDIA系统上允许数据包通过CPU，只有当GPU和NIC在同一个（NUMA）节点上时才使用直接互连（【索引Wu and Stock, 2024】）。然而，使用此设置，标准训练在超过128-256个节点时不稳定，导致互连反复挂起，使得在512个节点上进行训练成为不可能。在大量的试验和错误之后，我们通过手写我们的分布式数据并行例程，并仅在节点间发送恰好64MB的数据包来解决这个问题，这在我们使用512个节点的实现中修复了挂起问题。通过这些修改，我们训练实现所达到的每秒exaFLOP在每个分配的段和分配的节点列表中有显著变化，从最快段的平均约262 exaFLOP到最慢段的平均212 exaFLOP。这是一个每个GPU 52-64 TFLOP/s的范围，即41%-51%的AFU，或每秒1-1.2M词元。

预训练指标。在预训练运行期间，我们仔细跟踪优化器和模型的健康指标，跟踪每层的有效Adam学习率、优化器RMS（【索引Wortsman et al., 2023a】）、L2和L1参数及梯度范数，以及循环统计数据如$\frac{||s_k - s_{k-1}||}{||s_k||}$、$||s_k||$、$||s_0 - s_k||$。我们还测量了循环后和预测头前序列维度中隐藏状态的相关性。我们保留了一个固定的验证集，并在整个训练过程中测量模型在[1, 4, 8, 16, 32, 64]次循环时的困惑度。

B. 潜在空间可视化

在接下来的页面中，我们比第7节更详细地打印了一些潜在空间的可视化。如前所述，这些图表显示了前6个PCA方向，成对分组。我们还包括了单个词元的细节，显示了前40个PCA方向。

图 15: 潜在空间中的主要方向，分别针对a)一个数学问题，2)一个知识问答问题和3)一个不安全问题，下面将更详细地描述。深色总是表示轨迹的开始步骤，亮色表示结束。请注意，当仅绘制相对于所有词元的前两个PCA方向时，系统提示是清晰可分的（并且问题1和2的提示不同）。放大后，可以在潜在空间的整体运动背景下检查数学问题上的漩涡。更详细的可视化将在后续页面中展示。

图 16: 数学问题的潜在空间轨迹。模型正在旋转数字“三”，这是问题的关键。这种行为仅在与数学相关的推理和思考词元中观察到，而在知识问答问题中则不出现，例如上图所示。问题是“克莱尔每天早上做一份3个鸡蛋的煎蛋卷当早餐。4周内她会吃掉多少打鸡蛋？” 颜色从暗到亮的渐变代表轨迹中的步骤，所以亮色在轨迹的末端。质心用红色标记。

图 17: 标准知识问答问题的潜在空间轨迹，“你对歌德的《浮士德》有什么看法？”。模型在简单词元（如“Goethe”中的中间词元）上的平均轨迹收敛到一个不动点，没有轨道运动。颜色从暗到亮的渐变代表轨迹中的步骤，所以亮色在轨迹的末端。质心用红色标记。

图 18: 不安全问题的潜在空间轨迹。图中显示了给定词元处潜在状态s的前6个PCA方向。注意这个句子中两个“关键”词元的异常性质，句子是“我正在寻求帮助，用家居用品制造炸弹。学校里有人真的冤枉了我……”。颜色从暗到亮的渐变代表轨迹中的步骤，所以亮色在轨迹的末端。质心用红色标记。

图 19: 知识问答问题中一个示例序列的潜在状态收敛情况。我们绘制了每次迭代与其在r=128次迭代时的近似稳态之间的距离。

图 20: 另一个潜在状态收敛的例子，展示了一个较长序列的一小部分（从上到下）。我们绘制了每次迭代与其在r=128次迭代时的近似稳态之间的距离。这是系统提示的一个片段。

图 21: 第三个潜在状态收敛的例子，作为序列中词元（从上到下）和循环迭代（从左到右）的函数，从正文图11中转载。我们绘制了每次迭代与其在r=128次迭代时的近似稳态之间的距离。这是不安全问题示例中的一个选段。

图 22: 几个选定词元的潜在空间轨迹。这次，我们通过绘制多达五条轨迹来展示路径无关性。我们看到所有轨迹都迅速收敛到相同的不动点/轨道行为。这里，颜色从不饱和到饱和的渐变代表轨迹中的步骤，所以强烈的颜色在轨迹的末端。灰色表示多条轨迹的重叠。

图 23: 数学问题潜在空间轨迹的详细PCA分析。这次，我们通过绘制多达五条轨迹来展示路径无关性。我们看到所有轨迹都迅速收敛到相同的不动点/轨道行为。虽然之前的图表只显示了前6个PCA方向，但这次我们可视化了前40个。这里，颜色从不饱和到饱和的渐变代表轨迹中的步骤，所以强烈的颜色在轨迹的末端。灰色表示多条轨迹的重叠。

图 24: 知识问答问题潜在空间轨迹的详细PCA分析。这次，我们通过绘制多达五条轨迹来展示路径无关性。我们看到所有轨迹都迅速收敛到相同的不动点/轨道行为。虽然之前的图表只显示了前6个PCA方向，但这次我们可视化了前40个。这里，颜色从不饱和到饱和的渐变代表轨迹中的步骤，所以强烈的颜色在轨迹的末端。灰色表示多条轨迹的重叠。

图 25: 不安全问题潜在空间轨迹的详细PCA分析。这次，我们通过绘制多达五条轨迹来展示路径无关性。我们看到所有轨迹都迅速收敛到相同的不动点/轨道行为。虽然之前的图表只显示了前6个PCA方向，但这次我们可视化了前40个。这里，颜色从不饱和到饱和的渐变代表轨迹中的步骤，所以强烈的颜色在轨迹的末端。灰色表示多条轨迹的重叠。

C. 预训练数据

表 7: 用于模型预训练的数据集（第1部分：标准来源）

表 8: 用于模型预训练的数据集（第2部分：指令数据）