Skip to Content
ArchitectureTransformerRecurrent Transformer2025Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

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

论文链接：https://arxiv.org/pdf/2502.05171 

代码链接：https://github.com/seal-rg/recurrent-pretraining 

摘要

我们研究了一种新型语言模型架构，该架构能够通过在潜在空间中进行隐式推理来扩展测试时的计算能力。我们的模型通过迭代循环块来实现，从而在测试时展开到任意深度。这与主流推理模型通过生成更多 token 来扩展计算能力的做法截然不同。与基于思维链的方法不同，我们的方法不需要任何专门的训练数据，可以处理较小的上下文窗口，并且可以捕捉难以用文字表示的推理类型。我们将概念验证模型扩展到了 3.5B 参数和 800B token。我们证明，该模型在推理基准测试中能够提升性能，有时甚至能显著提升，计算负载最高可达 50B 参数。

1.Scaling by Thinking in Continuous Space

人类解决某些问题时自然会比解决其他问题付出更多的脑力。虽然人类能够通过将中间结果口头表达并记录下来来进行长时间的思考，但在说出答案的第一个字之前，大脑中大量的思考是通过复杂且反复出现的神经活动模式进行的。

早期提升语言模型能力的尝试主要集中在扩展模型规模上，但这需要海量的数据和计算资源。近年来，研究人员探索了通过扩展测试时计算来增强模型推理能力的方法。主流方法是在后训练使用较长的思维链示例，以提升模型在其上下文窗口中表达中间计算过程的能力，从而实现思维的外化。

然而，必须将昂贵的内部推理结果始终投射到单个口语指令上的限制似乎是一种浪费；如果模型能够在连续的潜在空间中进行原生“思考”，它们的能力可能会更强。释放这一未开发的额外计算维度的一种方法是向模型中添加循环单元。该单元以循环方式运行，迭代地处理和更新其隐藏状态，从而使计算能够无限期地进行下去。虽然这并非目前的主流范式，但这一思想是机器学习的基础，并且每隔十年就会被重新发现，例如循环神经网络、扩散模型以及 universal 或 looped Transformer 模型。

本文证明，深度循环语言模型能够高效学习，训练效率高，并且在测试时计算规模扩大的情况下，性能显著提升。我们提出的 Transformer 架构基于一个潜在深度循环模块，该模块在训练过程中运行随机抽取的迭代次数。我们证明，这种范式可以扩展到数十亿个参数和超过 5000 万亿个预训练 token。在测试时，该模型可以通过在潜在空间中进行循环推理来提升性能，使其能够与那些受益于更多参数和训练数据的其他开源模型相媲美。此外，我们还证明，循环深度模型在推理时自然地支持许多在非循环模型中需要大量调优和研究才能实现的特性，例如逐 token 自适应计算、（自）推测解码和键值缓存共享。最后，我们通过追踪潜在空间中的 token 轨迹来结束这项研究，结果表明，随着规模的扩大，许多有趣的计算行为会自然而然地涌现出来，例如模型在潜在空间中旋转形状以进行数值计算。

2.Why Train Models with Recurrent Depth?

循环层使 Transformer 模型能够在发出 token 之前执行任意次数的计算。原则上，循环机制为测试时计算扩展提供了一种简单的解决方案。与更标准的长上下文推理方法相比，潜在循环思维具有以下几个优势。

  • 潜在推理无需构建定制的训练数据。而链式推理则要求模型使用在目标领域内构建的长上下文数据进行训练。相比之下，我们提出的潜在推理模型可以使用可变的计算预算，利用标准训练数据（无需专门的演示数据）进行训练，并在测试时通过增加计算资源来提升其性能。
  • 潜在推理模型在训练和推理过程中所需的内存比链式推理模型要少。由于链式推理模型需要极长的上下文窗口，因此可能需要诸如 token 并行化之类的专门训练方法。
  • 与标准 Transformer 模型相比，循环深度网络每个参数的浮点运算能力更强，从而显著降低了大规模加速器之间的通信成本。这尤其有利于在使用速度较慢的互连网络进行训练时提高设备利用率。
  • 通过构建计算密集型、参数数量少的架构，我们希望为那些通过“思考”（即学习元策略、逻辑和抽象）而非记忆来解决问题的模型设定一个强有力的先验。循环先验在学习复杂算法方面的强大作用已在“深度思考”文献中得到证实。

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

我们认为，以这种方式扩展计算能力与通过扩展（口语）推理扩展或在预训练中扩展参数数量并不矛盾，它可能构建第三个维度来扩展模型性能。

3. A scalable recurrent architecture

本节将描述我们提出的具有潜在循环深度的 Transformer 架构，并讨论设计选择和小规模消融。架构图如图 2 所示。我们始终将序列维度记为 nn，模型的隐藏维度记为 hh，词表记为集合 VV。

3.1 Macroscopic Design

该模型主要围绕仅解码器（decoder-only）的 Transformer 块构建。然而，这些块被组织为三个功能组：（1）前奏（prelude）PP，它使用多个 Transformer 层将输入数据嵌入到潜在空间中；（2）随后是核心循环块（core recurrent block）RR，它是循环计算的中心单元，用于修改状态 s∈Rn×h\mathbf{s} \in \mathbb{R}^{n \times h}；（3）最后是尾声（coda）C)C)，它使用若干层从潜在空间中进行解嵌入，并且还包含模型的预测头。核心块被置于前奏块与尾声块之间，通过循环执行核心块，我们便可以在我们的歌曲中加入无限数量的诗节。

给定循环迭代次数 rr 和输入 token 序列 x∈Vn\mathbf{x} ∈ V^n，这些组按以下方式使用以生成输出概率 p∈Rn×∣V∣\mathbf{p} ∈ \mathbb R^{n×|V|}。

e=P(x)s0∼N(0,σ2In⋅h)si=R(e,ei−1)fori∈{1,...,r}p=C(sr),\begin{array}{cc} \mathbf{e}=P(\mathbf{x})\\ \mathbf{s}_0\sim\mathcal N(\mathbf{0},σ^2I_{n·h})\\ \mathbf{s}_i=R(\mathbf{e},\mathbf{e}_{i-1})\quad for\quad i\in\{1,...,r\}\\ \mathbf{p}=C(\mathbf{s}_r), \end{array}

其中 σσ 是用于初始化随机状态的标准差。该过程如图 2 所示。给定初始随机状态 s0\mathbf s_0，模型重复应用核心模块 RR，该模块接收潜在状态 si−1\mathbf s_{i−1} 和嵌入输入 e\mathbf e，并输出新的潜在状态 si\mathbf s_i。完成所有迭代后，尾部模块处理最后一个状态，并生成下一个 token 的概率。

该架构基于深度思考文献，其中表明，在每一步中注入潜在输入 e\mathbf e 并用随机状态初始化潜在向量可以稳定递归过程，并促进收敛到与初始化无关的稳态，即路径独立性。

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

使用多层结构将输入 token 嵌入到隐藏的潜在空间中，是基于对标准固定深度 Transformer 模型的实证分析结果。该系列研究表明，LLM 的初始层和末端层存在显著差异，而中间层则可以互换和排列。例如，Kaplan et al. (2024) 指出，标准模型只需几层就能将子词 token 嵌入到潜在空间中的单个概念中，模型随后便基于这些概念进行运算。

Remark 3.1 (Is this a Diffusion Model?)。这种迭代架构与另一种现代迭代建模范式——扩散模型（尤其是潜在扩散模型）——非常相似。我们运行了几个与扩散模型更为相似的迭代方案的消融实验，例如 si=R(e,si−1)+n\mathbf s_i = R(\mathbf e, \mathbf s_{i−1}) + \mathbf n，其中 n∼N(0,σiIn⋅h)\mathbf n ∼ \mathcal N (0, σ_iI_{n·h})，但发现注入噪声在我们的初步实验中并没有帮助，这可能与我们的训练目标有关。我们也评估了 si=Ri(e,si−1)\mathbf s_i = R_i(\mathbf e, \mathbf s_{i−1})，即以当前步骤作为输入的核心模块，但发现它与路径独立性存在不良交互，导致模型无法进行外推。

3.2 Microscopic Design

在每个组内，我们大致遵循标准的 Transformer 层设计。每个模块包含多个层，每个层包含一个标准的因果自注意力模块（使用基数为 50000 的 RoPE）和一个门控 SiLU MLP。我们使用 RMSNorm 作为归一化函数。该模型在 query 和 key 上具有可学习的偏差，其他位置则没有。为了稳定递归，我们按照以下“三明治”格式排列所有层，并使用归一化层 nin_i，这与(Ding et al., 2021; Team Gemma et al., 2024) 中的类似策略相关但不完全相同：

x^l=n2(x^l−1+Attn(n1(x^l−1)))‘‘‘‘‘mathxl=n4(x^l+MLP(n3(x^l)))\mathbf{\hat x}_{l}=n_2(\mathbf{\hat x}_{l-1}+Attn(n_1(\mathbf{\hat x}_{l-1}))) `` ```math \mathbf{x}_{l}=n_4(\mathbf{\hat x}_{l}+MLP(n_3(\mathbf{\hat x}_{l})))

虽然在小规模下，大多数归一化策略（例如pre-norm、post-norm等）的效果几乎一样好，但我们排除了这些选项，发现这种归一化是训练大规模递归所必需的。

  • 给定嵌入矩阵 EE 和嵌入缩放 γγ，前奏块首先将输入 token x\mathbf x 嵌入为 γE(x)γE(\mathbf x)，然后应用 lPl_P 个具有上述布局的前奏层。
  • 我们的核心循环模块 RR 以适配器矩阵 A:R2h→RhA : \mathbb R^{2h} → \mathbb R^h 开始，该矩阵将 si\mathbf s_i 和 e\mathbf e 的拼接映射到隐藏维度 hh。虽然对于较小的模型，通过加法而非拼接重新融合初始嵌入特征的效果同样出色，但我们发现拼接在大规模模型中效果最佳。然后，该模型被输入到 lRl_R 个 Transformer 层。在核心模块的末尾，输出再次使用 RMSNorm ncn_c 进行重新缩放。
  • 尾部包含 lCl_C 层，由 ncn_c 归一化以及使用绑定嵌入 ETE^T 投影到词表中。

总之，我们可以用三元组 (lP,lR,lC)(l_P, l_R, l_C) 来概括该架构，其表示每个阶段的层数，而 rr 表示循环次数，其在每次前向传播中可能有所不同。我们训练了多个形状为 (1,4,1)(1, 4, 1) 且隐藏层大小 h=1024h = 1024 的小规模模型，以及一个形状为 (2,4,2)(2, 4, 2) 且隐藏层大小 h=5280h = 5280 的大型模型。该大型模型只有 8 个“实际”层，但当循环块迭代（例如 32 次）时，其有效深度可达 2+4r+2=1322 + 4r + 2 = 132 层，构建的计算链深度甚至可以超过最大的固定深度 Transformer 模型。

3.3 Training Objective

Training Recurrent Models through Unrolling。为了确保模型在测试阶段增加循环迭代次数时仍能正常工作，我们在训练期间随机抽取迭代次数，为每个输入序列分配一个随机迭代次数 rr。我们优化损失函数 LL 对来自分布 XX 的随机样本 xx 和来自分布 ΛΛ 的随机迭代次数 rr 的期望值。

L(θ)=Ex∈XEr∼ΛL(mθ(x,r),x′).\mathcal L(\theta)=\mathbb E_{\mathbf x\in X}\mathbb E_{r\sim Λ}L(m_{\theta}(\mathbf x, r), \mathbf x').

这里，mm 代表模型输出，x′\mathbf x′ 是序列 x\mathbf x 左移后的下一个 token。我们选择 ΛΛ 为对数正态泊松分布。给定目标均值递增 rˉ+1\bar r + 1 和方差 σ=12σ = \frac{1}{2}，我们可以通过以下方式从该分布中采样：

τ∼N(log(rˉ)−12σ2,σ)(1)\tau\sim\mathcal N(log(\bar r)-\frac{1}{2}σ^2,σ)\tag{1} r\sim\mathcal P(e^{\tau})+1,\tag{2} ```math 给定正态分布 $\mathcal N$ 和泊松分布 $\mathcal P$，参见图 3。该分布最常采样小于 $\bar r$ 的值，但它包含一个厚尾，其中偶尔会出现需要进行更多迭代的事件。 **Truncated Backpropagation**。为了在训练阶段保持较低的计算量和内存占用，我们仅对循环单元的最后 $k$ 次迭代进行反向传播。这使得我们可以使用重尾泊松分布 $Λ$ 进行训练，因为最大激活内存和反向计算量现在与 $r$ 无关。在我们的主要实验中，我们将 $k$ 固定为 $8$。在小规模情况下，这种方法与均匀采样 $k$ 的效果一样好，但当 $k$ 固定时，每次训练步骤的总内存使用量相同。需要注意的是，由于 prelude 模块的输出 $\mathbf e$ 在每一步都会被注入，因此它仍然会在每一步中接收梯度更新。这种设置类似于 RNN 中常用的**截断时间反向传播**，尽管我们的设置是深度上的循环，而不是时间上的循环。 ## 4.Training a large-scale recurrent-depth Language Model 在验证我们能够可靠地训练包含高达 10B token 的小型测试模型后，我们开始进行更大规模的运行。鉴于我们有限的计算资源，我们可以选择训练多个规模过小的模型（这些模型不足以展现涌现效应或扩展性），或者训练一个中等规模的模型。基于此，我们准备了一次运行，详情如下。 ### 4.1 Training Setup <img src="https://i-blog.csdnimg.cn/direct/e482de3202e24c37a97ed91bb6e30aff.png" alt="" referrerpolicy="no-referrer" style="max-width: 100%; height: auto;" /> 我们详细描述了训练设置，包括架构、优化设置和预训练数据。我们公开了所有训练数据、预训练代码以及部分中间模型检查点。 **Pretraining Data**。由于计算资源仅够运行一次大规模模型，我们选择了一种数据集混合方案，旨在最大程度地提高模型涌现推理行为的潜力，而非追求最佳基准性能。最终的混合数据集侧重于代码和数学推理数据，并（希望）包含适量的通用网络文本，以使模型能够掌握标准的语言建模能力。所有数据源均公开可用。图 4 给出了概览。我们参考 Allen-Zhu and Li (2024) 的方法，将相关的指令数据直接混合到预训练数据中。然而，由于计算和时间的限制，我们无法对该混合数据集进行消融。我们预期，更细致的数据准备可以进一步提升模型的性能。所有数据源列于附录 C。 **Tokenization and Packing Details**。我们使用 Dagan (2024) 的实现，通过 BPE 构建了一个包含 65536 个 token 的词表。与传统的 tokenzier 训练方法相比，我们直接在预训练语料库的指令数据划分上构建 tokenzier，以最大限度地提高目标领域的 token 化效率。我们还大幅修改了预 tokenzier 的正则表达式（例如 Dagan 等人 (2024) 的正则表达式），以更好地支持代码、缩写和 LaTeX。我们在每个文档的开头添加了一个 `<|begin_text|>` token。在对预训练语料库进行 tokenizer 后，我们将 tokenizer 的文档打包成长度为 4096 的序列。打包时，我们丢弃了缺少上下文信息的文档末尾部分，以解决 Ding et al. (2024) 中描述的“grounding 问题”。此外，我们还完整保留了一些包含数学内容的长文档。 **Architecture and Initialization**。我们扩展了第 3 节中描述的架构，将层数设置为 $(2, 4, 2)$，并使用平均递归值 $\bar r = 32$ 进行训练。我们主要通过将隐藏层大小增加到 $h = 5280$ 来扩展模型，这产生了 55 个大小为 96 的头部。MLP 的内部维度为 17920，RMSNorm $ε$ 为 $10^{−6}$。总体而言，该模型在非递归前奏和头部中约有 1.5B 参数，在核心递归模块中约有 1.5B 参数，在绑定输入嵌入中约有 0.5B 参数。 在小规模下，大多数合理的初始化方案都能奏效。然而，在大规模下，我们采用 Takase et al. (2024) 的初始化方法，该方法规定方差为 $σ^2_h = \frac{2}{5h}$。除所有输出投影层外，我们使用截断正态分布（截断至 $3σ$）初始化所有参数，其中输出投影层的方差设置为 $σ^2_{out} = \frac{1}{5hl}$，$l = l_P + \bar rl_R + l_C$ 为有效层数，本模型中有效层数为 132。因此，输出投影层的初始化值相当小。嵌入层的输出乘以 $\sqrt{h}$ 进行缩放。为了与此初始化相匹配，状态 $s_0$ 也从截断正态分布中采样，其方差为 $σ^2_s = \frac{2}{5}$。 **Locked-Step Sampling**。为了实现并行工作节点间的同步，我们为每个 micro-batch 的训练数据采样一个深度 $r$，并在所有工作节点间同步该深度（否则，工作节点会因为等待深度最大的模型完成反向传播而处于空闲状态）。我们在小规模测试中验证了这种改进能够在不影响收敛速度的情况下提高计算利用率。但需要注意的是，在大 batch-size 规模下，可以通过在每个工作节点上对独立的步骤 $r$ 进行最优采样和调度，从而更准确地模拟公式 (1) 中步骤间的期望值，进一步提升训练效率。 **Optimizer and Learning Rate Schedule**。我们使用带有解耦权重正则化（$β_1 = 0.9, β_2 = 0.95, η = 5 × 10^{−4}$）的 Adam 优化器进行训练，并根据 Everett et al. (2024) 的方法进行了修改，包括更新裁剪和移除 $ε$ 常数。我们将大于 1 的梯度裁剪掉。我们使用预热阶段和恒定的学习率进行训练，在前 4096 步内逐渐达到最大学习率。 ### 4.2 Compute Setup and Hardware 我们使用橡树岭国家实验室 Frontier 超级计算机上的计算时间来训练此模型。这台 HPE Cray 系统包含 9408 个计算节点，每个节点配备 AMD MI250X GPU，并通过 4 个 HPE Slingshot11 网卡连接。调度系统由 SLURM 管理。我们使用基于 PyTorch 的实现，在 bfloat16 混合精度下进行训练。 ### 4.3 Importance of Norms and Initializations at Scale <img src="https://i-blog.csdnimg.cn/direct/c5475d7bc2414fb4a5a9fabb4fdf5f73.png" alt="" referrerpolicy="no-referrer" style="max-width: 100%; height: auto;" /> ## 5.Benchmark Results
Last updated on