RecurrentGPT：通过递归调制在Transformer中实现表达性深度

# RecurrentGPT：通过循环调制实现Transformer的深度表达能力  
来源：https://arxiv.org/html/2608.15062  

###### 摘要  
扩大Transformer语言模型规模时，表达能力与内存效率之间存在内在矛盾。虽然跨层使用独立权重可以保持从输入关联到抽象精炼的功能专化，但这会导致巨大的内存占用。相反，标准的深度共享机制强制应用均匀变换，会破坏表征多样性并降低建模质量。我们提出**RecurrentGPT**，一种循环深度Transformer架构，其核心由固定深度的前奏块和尾声块夹着一个共享核心，该核心被迭代**R**次。受门控循环神经网络启发，我们采用轻量级投影和逐元素更新门——该门基于隐藏状态、固定前奏输出以及每步重采样的噪声进行条件调制，以控制循环更新过程。这使得模型能够在多次循环中将输入特化到同一组少量层，而非需要大量不同层来实现功能多样性。在**等计算量**约束下，3层RecurrentGPT在相似训练与推理计算量下匹配12层GPT-2 Small基线的准确率，并在所有九种“规模-预算”组合中超越MoR和重尾深度采样方法；在中等和大规模设置下，其在标准令牌预算时接近稠密模型质量，在预算翻倍时甚至超越稠密模型。在**等参数量**约束下，更深的循环在匹配参数与数据预算时实现2.76的验证损失（非循环对应模型为2.84）。实验结果表明，自适应深度复用是以参数换质量的合理策略：在大规模设置下，参数量减少63%、峰值解码内存降低59%，但编译生成延迟增加10%。  

## 1 引言  
扩大Transformer语言模型规模带来了显著收益，但这必然需要同步增加参数量与训练计算量。在标准架构中，深度与参数被刚性耦合：每增加一层都引入全新权重集，造成内存瓶颈，限制了固定硬件预算下的有效深度。然而，除内存效率外，更深层的计算动机源于智能本身的本质——图灵的基础洞见指出：**有限**的状态与符号集合，通过**迭代**应用，足以计算任何可计算对象——这表明推理系统的力量不在于参数广度，而在于迭代深度。这一原则在测试时计算范式中得到惊人体现：模型在推理时分配更多计算以提升推理能力。虽然现有方法通过输出空间的思维链步骤实现这一计算（过程中消耗序列长度），但循环深度提供了互补的实现路径：通过在输入表征上迭代共享变换，模型能在隐藏状态中“深度思考”，无需生成额外令牌或存储额外参数。  

权重共享是应对内存-深度耦合的自然方案，允许在不增加参数量的情况下提升有效深度。先前研究利用循环深度作为效率工具在减少参数量时匹配质量，或作为性能工具在匹配参数量时提升准确性。然而，这些视角的统一处理尚缺失，且更深层的架构矛盾依然存在：标准权重共享对不断演变的表征施加完全相同的变换。由于第一次循环后的隐藏状态与第八次循环后存在根本差异，静态变换会破坏使深度有价值的功能多样性。这引发了一个问题：单一固定变换能否有意义地服务于如此多样的表征阶段？还是会不可避免地破坏功能多样性？这些问题促使了RecurrentGPT的设计。  

为解决此矛盾，RecurrentGPT引入了逐元素学习门控机制，该门控基于当前隐藏状态、原始输入的固定**前奏**表征以及随机噪声进行条件判断。这在每次循环时构建了独特的上下文关联输入，使单个权重张量能表现为多个特化层。门控初始化确保训练开始时残差流几乎无变化地通过所有循环。随着训练推进，门控逐渐显现，共享块学会选择性地精炼表征。因此，共享变换在权重上固定但在行为上动态——模型在每次迭代中思考方式不同，尽管复用相同参数。结合训练期间的深度采样，该架构支持推理时的连续计算-质量权衡，并促进在**等计算量**与**等参数量**两种模式下的评估。  

我们总结贡献如下：  
- **架构**：提出RecurrentGPT，采用逐元素门控与随机扰动防止表征坍缩，使共享核心在循环步骤中表现出差异性。  
- **统一评估**：在三种规模下评估RecurrentGPT的**等计算量**与**等参数量**模式。在**等计算量**模式中，其在所有九种“规模-预算”组合中领先MoR和重尾深度采样，在标准预算下匹配或领先RRT和Ouro，并在标准预算下小规模设置中使用64%更少参数匹配稠密GPT-2基线，在中大规模设置中随令牌预算增长逐渐接近稠密模型。在**等参数量**模式中，其在匹配参数量时提升验证损失，但推理计算成本更高。  
- **涌现式早退**：无需辅助损失，RecurrentGPT自然支持推理时早退，仅运行半数循环步骤时仍保持约92%准确率。  

图1：RecurrentGPT概览。(a) 门控循环单元在每个序列位置重用相同权重矩阵——在*序列*维度上循环。(b) 标准Transformer堆叠**L**个独立层，每层拥有独立权重，对整个序列应用一次。(c) RecurrentGPT引入在*深度*维度上的循环，对整个序列应用每次循环  

## 2 相关工作  
跨重复计算步骤共享权重的概念可追溯至Rumelhart、Hinton和Williams的1986年反向传播开创性论文，他们简略提出了同步迭代网络：一个每次迭代对应绑定权重层的网络。后续工作通过时间反向传播将其扩展至时序序列建模，启发了沿序列维度应用共享权重的循环架构——RNN、LSTM和GRU，逐时间步处理输入。相反，本工作在整体输入上应用循环深度，沿网络深度迭代共享权重变换以逐步精炼输入表征。  

在Transformer时代，最早关于循环深度的论文之一是ALBERT，其提出所有Transformer层共享相同权重的BERT模型。通用Transformer进一步增加了逐令牌自适应停止机制。然而，两种方法共享**整个**层栈，迫使每层无论输入表征阶段如何都施加相同变换。相关研究通过为前奏层、共享核心层和尾声层分配不同角色来放宽此约束。类似地，有研究将预训练大语言模型的早期、中期和晚期层分别识别为编码、推理和解码阶段，提出仅循环中间层而保持其余层固定。我们的RecurrentGPT遵循选择性循环中间层的原则。先前研究显示前馈网络层主要存储事实知识，而计算深度驱动推理能力，这促使通过循环提升推理能力。  

我们沿两种评估模式组织循环深度Transformer的先前工作：在**等参数量**模式中——通过循环层增加计算量同时保持参数量固定，以计算换精度；在**等计算量**模式中——通过循环层减少参数量同时匹配训练与推理计算量，以内存换效率。我们的工作涵盖两种模式，**等计算量**为主要贡献，互补的**等参数量**结果见第4节。  

相关工作系动态调整计算量而非统一固定：自适应计算时间根据学习到的终止信号按位置停止计算；混合深度采用路由方法动态分配令牌到不同层子集；最近的混合递归将此思想扩展至循环设置，将令牌子集路由到不同递归步数。虽然三者都在令牌级别通过停止、层路由或递归路由自适应计算，但都需要专用终止信号或路由机制。相反，RecurrentGPT使用更简单的门控机制调制每次循环步骤中共享块更新的吸收量，无需路由或停止逻辑。  

权重绑定迭代趋向固定点是深度均衡模型的特征，其直接求解均衡并通过隐式梯度微分，避免了存储展开轨迹的内存成本。RecurrentGPT共享权重绑定更新但不采用均衡目标：我们展开固定数量离散步骤并反向传播，深度采样训练每个步骤作为出口而非趋向单一极限。  

缩放定律已确立模型规模与数据如何共同决定性能，Pythia提供了跨规模的控制比较。值得注意的是，直接评估参数共享Transformer时发现：循环模型在匹配参数量时表现更好，但在匹配计算量时表现更差。深度与宽度的相互作用受到专门关注：理论分析表明增加宽度可补偿算法推理任务中深度的减少，而实证证明缩放定律对深度-宽度比例敏感。在我们的RecurrentGPT中，我们提出新型循环深度架构，旨在改善损失与参数量、损失与训练计算量的缩放定律，在模型规模与训练数据增长时在**等计算量**与**等参数量**模式下均展示持续收益。  

## 3 用于语言建模的循环深度复用  
### 3.1 预备知识  
设$\mathbf{X}=(x_{1},\ldots,x_{T})$为词汇表$\mathcal{V}$中长度为$T$的令牌序列。标准自回归Transformer包含**L**个块，定义残差更新链：$\mathbf{h}^{(\ell)}=\mathbf{h}^{(\ell-1)}+\mathrm{Block}_{\ell}(\mathbf{h}^{(\ell-1)})$（$\ell=1,\ldots,L$），其中$d$为嵌入维度，$\mathbf{h}^{(0)}\in\mathbb{R}^{T\times d}$为初始令牌嵌入。每个块包含预归一化多头自注意力和逐令牌MLP。目标是最小化负对数似然：$\mathcal{L}=-\sum_{t}\log p(x_{t}\mid x_{<t})$  

### 3.2 循环Transformer架构  
RecurrentGPT将Transformer深度分解为三个功能组件：固定前奏块、共享核心块和固定尾声块。给定输入序列$\mathbf{X}$，前奏块首先处理输入：$\mathbf{h}_{\text{pre}}=\text{Prelude}(\mathbf{X})$。然后核心块被迭代**R**次，每次循环应用相同权重：$\mathbf{h}^{(r)}=\text{Core}(\mathbf{h}^{(r-1)})$（$r=1,\ldots,R$）。最后尾声块处理最终循环输出：$\mathbf{Y}=\text{Coda}(\mathbf{h}^{(R)})$。为防止表征坍缩并允许不同循环步骤的行为差异，我们引入门控机制：在每次循环时，更新后的隐藏状态通过基于当前隐藏状态、前奏输出和随机噪声的门控与残差流结合。  

### 3.3 门控循环更新  
核心思想是：$\mathbf{h}^{(r)}=(1-\mathbf{g}^{(r)})\odot\mathbf{h}^{(r-1)}+\mathbf{g}^{(r)}\odot\text{Core}(\mathbf{h}^{(r-1)})$，其中$\mathbf{g}^{(r)}\in(0,1)^{T\times d}$为逐元素门控值。门控计算为：$\mathbf{g}^{(r)}=\sigma\left(\frac{f_{\text{gate}}(\mathrm{LN}(\mathbf{h}^{(r-1)}),\mathrm{LN}(\mathbf{h}_{\text{pre}}))+\boldsymbol{\epsilon}_{g}}{\tau}\right)$，其中$\sigma$为sigmoid函数，$\tau$为温度，$\boldsymbol{\epsilon}_{g}\sim\mathcal{N}(0,\sigma_{g}^{2})$为高斯噪声。$f_{\text{gate}}$为两层MLP。训练时，循环深度从$\{1,\ldots,R\}$均匀采样。  

### 3.4 推理时深度自适应  
循环深度$R$为推理时计算量的上界，但门控机制允许实际计算深度自适应：若门控接近0，隐藏状态几乎不变，后续循环可提前终止。这实现了无需专用退出机制的自适应早退。  

### 3.5 伪代码  
```python
# 简化前向传播伪代码
h = prelude_blocks(x)          # [S, d]
h_pre = h.clone()              # 保存前奏表示
for r in range(R):             # 循环R次
    o = core_block(h)          # 核心变换
    h_tilde = project_to_do = B_shared(h_tilde) # 相同权重每步
    eps_g = sample_gaussian(0, sigma_g**2)
    g = sigmoid(f_gate(LN(h), LN(h_pre)) / tau + eps_g) # 门控在[0,1]
    h = g * h + (1 - g) * o   # 门控残差更新
for block in model.coda_blocks:
    h = h + block(h)
return lm_head(h)              # [S, |V|]
```  
RecurrentGPT前向传播伪代码  

### A.1 超参数表  
表5列出主论文中所有运行的完整训练配置。表5：训练配置。pp：前奏块；bb：共享块；RR：循环步数；cc：尾声块；dd：嵌入维度；hh：注意力头。批令牌数 = batch\_size×梯度累积×序列长度×GPU数量。所有运行使用AdamW（β1=0.9，β2=0.95），权重衰减0.1，梯度裁剪1.0，bfloat16混合精度，余弦学习率调度。所有循环运行使用门控偏置初始化+4（初始时g≈0.98）、状态噪声σx=0.1、门控噪声σg=0.1和门控温度τ=1.0。训练期间，循环深度从{1,...,R}均匀采样。  

## 附录B 额外消融实验  
### B.1 跨种子运行方差  
表6报告了小型**等计算量**配置及其稠密对应物的三个独立训练运行（种子1337、42、123），均采用完整20,000步预算。RecurrentGPT平均损失为3.145±0.004，稠密基线为3.188±0.056。