DeepLoop：循环Transformer的深度缩放

# DeepLoop: 循环Transformer的深度缩放
来源: https://arxiv.org/html/2607.13491
\\setlist

\[itemize\]leftmargin=\*,topsep=1pt,itemsep=1pt\\setlist\[enumerate\]leftmargin=\*,topsep=1pt,itemsep=1pt

Shuzhen Li1Yifan Zhang1,†Jiacheng Guo1 Quanquan Gu2Mengdi Wang1,† 1普林斯顿大学 2加州大学

###### 摘要

循环Transformer通过将一小组物理块多次应用来扩展序列计算，在不增加存储参数的情况下增加展开深度。这种复用改变了残差缩放问题：在非绑定的Transformer中，每个残差分支接收并应用自己的参数更新；而在循环Transformer中，一个共享的更新会聚合来自多次访问的梯度，并在下一次线性化前向传播中被相同的访问回读。我们通过一个由访问对齐系数κR\\kappa\_\{R\}控制的一阶扰动界来形式化这种绑定深度效应。当访问解相关时，该界恢复DeepNorm指数；但在保守的对齐模式下，当固定物理深度下循环次数增加时，要求指数从1/41/4增加到1/21/2。由此产生的方法DeepLoop保留了Post-LN DeepNorm架构，并为展开深度NN设置α=\(2N\)1/2\\alpha=\(2N\)^\{1/2\}和β=\(8N\)−1/2\\beta=\(8N\)^\{\-1/2\}。在GPT-2小规模和GPT-2中等规模的GPT风格循环语言模型上，当没有物理块被回访时，DeepLoop表现中性；一旦循环深度被激活，它能改善验证损失和下游准确率。这些结果表明，稳定的循环深度需要考虑到参数访问次数的残差缩放规则，而不仅仅是名义上的层数。

$\\dagger$$\\dagger$footnotetext:通讯作者: [email protected], [email protected]。项目页面: \\urlhttps://github.com/lszshu/DeepLoop

\(a\) 物理块 — K=2K\{=\}2个共享块，只存储一次 块1 \(φ1\\phi\_\{1\}\)attnffn 块2 \(φ2\\phi\_\{2\}\)attnffn loop×R\\times\\,R轮次 \(b\) 展开执行 — 相同的两个块在每一轮中被回访 x0\\mathbf\{x\}\_\{0\} xN\\mathbf\{x\}\_\{N\} attnffn φ1\\phi\_\{1\} attnffn φ2\\phi\_\{2\} attnffn φ1\\phi\_\{1\} attnffn φ2\\phi\_\{2\} attnffn φ1\\phi\_\{1\} attnffn φ2\\phi\_\{2\} 轮次r=1r\{=\}1 轮次r=2r\{=\}2 轮次r=3r\{=\}3 R=3R\{=\}3轮次×\\timesK=2K\{=\}2个块 ⇒\\;\\Rightarrow\\;展开深度N=KR=6N\{=\}KR\{=\}6, M=2N=12M\{=\}2N\{=\}12个子层访问 \(c\) 残差子层: xi\+1=Norm\(αxi\+fj\(xi;φj\)\)\\mathbf\{x\}\_\{i\+1\}=\\mathrm\{Norm\}\\bigl\(\\alpha\\,\\mathbf\{x\}\_\{i\}\+f\_\{j\}\(\\mathbf\{x\}\_\{i\};\\phi\_\{j\}\)\\bigr\) xi\\mathbf\{x\}\_\{i\} ++ Norm xi\+1\\mathbf\{x\}\_\{i\+1\} fj\(⋅;φj\)f\_\{j\}\(\\,\\cdot\\,;\\phi\_\{j\}\) 跳跃连接×α\\times\\,\\alpha 初始化增益β\\beta \(d\) 缩放规则 DeepLoop (p=12p=\\tfrac\{1\}\{2\}) α=\(2N\)1/2\\alpha=\(2N\)^\{1/2\} β=\(8N\)−1/2\\beta=\(8N\)^\{\-1/2\}
图1: DeepLoop框架概述。\(a\) 物理块: 一个循环Transformer存储K=2K\{=\}2个物理块（每个包含一个注意力和一个FFN子层），块1 \(φ1\\phi\_\{1\}\)和块2 \(φ2\\phi\_\{2\}\)。\(b\) 展开执行: 整个KK块堆栈每轮应用一次，共R=3R\{=\}3轮，得到N=KR=6N\{=\}KR\{=\}6个展开块和M=2N=12M\{=\}2N\{=\}12个子层访问；φj\\phi\_\{j\}标签指示哪个物理块被复用。\(c\) 残差子层: 每次访问应用xi\+1=Norm\(αxi\+fj\(xi;φj\)\)\\mathbf\{x\}\_\{i\+1\}\{=\}\\mathrm\{Norm\}\(\\alpha\\,\\mathbf\{x\}\_\{i\}\+f\_\{j\}\(\\mathbf\{x\}\_\{i\};\\phi\_\{j\}\)\)，其中α\\alpha缩放跳跃连接，β\\beta是每个矩阵的初始化增益。\(d\) 缩放规则: DeepLoop设置α=\(2N\)1/2\\alpha\{=\}\(2N\)^\{1/2\}和β=\(8N\)−1/2\\beta\{=\}\(8N\)^\{\-1/2\}。

## 1 引言

深度是提高Transformer表达能力最可靠的方法之一，但在标准架构中，深度和参数数量同时增长：增加一层也会增加一组新的注意力和前馈权重（vaswani2017attention; kaplan2020scaling; hoffmann2022training）。循环Transformer通过将KK个物理块堆叠应用RR轮来解耦这些轴，产生展开深度N=KRN=KR，同时只存储KK个块（图1）。这种机制将经典深度共享（Universal Transformers、ALBERT和Subformer（dehghani2018universal; lan2019albert; reid2021subformer））与最近的循环深度和测试时计算模型（这些模型在不等比增加参数的情况下增加额外的顺序计算（giannou2023looped; yang2023looped; gatmiry2024can; geiping2025scaling; saunshi2025reasoning））联系起来。然而，要使循环深度成为实用的缩放轴，其残差参数化必须在相同物理块被多次回访时保持稳定。

困难在于，标准的残差缩放分析是针对非绑定深度编写的。残差参数化对于深度Transformer优化至关重要（xiong2020layer; nguyen2019transformers; liu2020understanding; huang2020improving; bachlechner2021rezero; wang2024deepnet）。特别是，DeepNorm通过选择跳跃缩放α\\alpha和残差分支初始化增益β\\beta，使得参数更新的一阶效应在M=2NM=2N次残差子层访问中保持有界（wang2024deepnet）。该计算假设每个展开的残差子层拥有一个独立参数张量，因此贡献一个独立的更新项。

权重共享恰恰违反了这一假设。当一个物理子层被访问RR次时，其优化器更新会聚合来自所有轮次的逐次梯度。然后，更新后的张量会被所有那些访问在下一个线性化前向计算中复用。因此，绑定深度创建了两条耦合的聚合路径：更新由多次访问写入，然后由多次访问读取。这种效应的大小取决于访问对齐。如果逐次梯度和敏感性在不同轮次中近似正交，那么循环模型的行为与非绑定深度在常数上类似。如果它们相干，那么共享更新可能会获得一个额外的因子RR。

本文使绑定深度效应变得明确。我们引入了一个访问对齐系数κR\\kappa\_\{R\}，并展示了循环Transformer的一阶稳定性条件变为

MκR\(βα\)2=O\(1\),M\\kappa\_\{R\}\\left\(\\frac\{\\beta\}\{\\alpha\}\\right\)^\{2\}=O\(1\),

而不是非绑定DeepNorm条件M\(β/α\)2=O\(1\)M\(\\beta/\\alpha\)^\{2\}=O\(1\)。对于缩放族α=\(cN\)p\\alpha=\(cN\)^\{p\}和β=\(dN\)−p\\beta=\(dN\)^\{\-p\}，当访问解相关时，恢复通常的DeepNorm阈值p=1/4p=1/4。在保守的对齐模式下，其中κR=Θ\(R\)\\kappa\_\{R\}=\\Theta\(R\)且KK固定而RR增长，阈值变为p=1/2p=1/2。

由此产生的方法是DeepLoop：保留DeepNorm Post-LN架构，但使用循环感知的缩放规则

α=\(2N\)1/2,β=\(8N\)−1/2。\\alpha=\(2N\)^\{1/2\},\\qquad\\beta=\(8N\)^\{\-1/2\}。DeepLoop不引入门控、学习到的残差系数、辅助损失或架构特定的调优常数。它是一种确定性的残差参数化，适用于通过回访共享块来增加有效深度的场景。

我们在受控的GPT风格循环语言建模消融实验中评估DeepLoop，规模为GPT-2小规模和GPT-2中等规模。在R=1R=1时（没有物理块被回访），DeepLoop基本保持中性；在更大的循环次数下，它持续改进最终验证损失。这种优势转移到了八任务语言模型评估套件上：在R=1R=1时下游平均分接近，随着循环深度增加，结果倾向于DeepLoop。这些结果支持了分析的核心预测：残差缩放应取决于深度是如何实现的，而不仅仅是名义上的展开层数。

我们的贡献是：

1.  \[leftmargin=\*, itemsep=1pt, topsep=1pt\]
2.  1. 我们识别了非绑定残差缩放分析中缺失的绑定深度聚合机制：共享更新在多次访问中累积，然后通过相同的访问被回读；
3.  2. 我们推导了一个带有显式访问对齐系数κR\\kappa\_\{R\}的循环感知一阶扰动界，在解相关状态下恢复DeepNorm，并在对齐的固定物理深度状态下产生p=1/2p=1/2指数阈值；
4.  3. 我们提出DeepLoop，缩放规则α=\(2N\)1/2\\alpha=\(2N\)^\{1/2\}和β=\(8N\)−1/2\\beta=\(8N\)^\{\-1/2\}，作为Post-LN循环Transformer的一行保守参数化；
5.  4. 我们在GPT-2小规模和GPT-2中等规模上提供了受控的预训练和下游消融实验，表明当循环次数大于一时，修正恰好变得有用。

## 2 背景

### 2.1 循环Transformer与有效深度

一个标准深度为NN的Transformer不同块依次应用一次。而循环Transformer选择KK个物理块并应用RR轮，

xr,k\+1=Bk\(xr,k;φk\),k=1,…,K,r=1,…,R,xr\+1,1=xr,K\+1,\\mathbf\{x\}\_\{r,k\+1\}=B\_\{k\}\(\\mathbf\{x\}\_\{r,k\};\\phi\_\{k\}\),\\qquad k=1,\\ldots,K,\\quad r=1,\\ldots,R,\\qquad\\mathbf\{x\}\_\{r\+1,1\}=\\mathbf\{x\}\_\{r,K\+1\},(1)

因此有效深度为N=KRN=KR，而块参数φ1,…,φK\\phi\_\{1\},\\ldots,\\phi\_\{K\}只存储一次。Universal Transformer（dehghani2018universal）对应完全循环的极端情况，其中相同的转换在深度上复用；ALBERT（lan2019albert）同样跨层绑定Transformer参数。本文研究的场景是将KK保持固定或较小，并增加RR，从而在不增加物理块数量的情况下增加测试时计算和展开深度。

### 2.2 残差归一化

令glg\_\{\\ell\}表示一个残差子层，无论是注意力还是MLP。两种常见的归一化位置是：

Pre-LN:xl\+1=xl\+gl\(Norm\(xl\)\),Post-LN:xl\+1=Norm\(xl\+gl\(xl\)\)。\\displaystyle\\text\{Pre-LN:\}\\quad\\mathbf\{x\}\_\{\\ell\+1\}=\\mathbf\{x\}\_\{\\ell\}\+g\_\{\\ell\}\(\\mathrm\{Norm\}\(\\mathbf\{x\}\_\{\\ell\}\)\),\\qquad\\text\{Post-LN:\}\\quad\\mathbf\{x\}\_\{\\ell\+1\}=\\mathrm\{Norm\}\\left\(\\mathbf\{x\}\_\{\\ell\}\+g\_\{\\ell\}\(\\mathbf\{x\}\_\{\\ell\}\)\\right\)。

Pre-LN改善了深度Transformer的优化稳定性，而Post-LN可以保留更具表达力的残差流（xiong2020layer; nguyen2019transformers）。DeepNorm（wang2024deepnet）修改了Post-LN残差路径，在归一化之前缩放跳跃连接：

xl\+1=Norm\(αxl\+gl\(xl;θl\)\)。\\displaystyle\\mathbf\{x\}\_\{\\ell\+1\}=\\mathrm\{Norm\}\\left\(\\alpha\\,\\mathbf\{x\}\_\{\\ell\}\+g\_\{\\ell\}\(\\mathbf\{x\}\_\{\\ell\};\\theta\_\{\\ell\}\)\\right\)。(2)

对于具有NN个块和M=2NM=2N次残差子层应用的仅有编码器或仅有解码器的Transformer，DeepNorm设置

α=\(2N\)1/4,β=\(8N\)−1/4,\\alpha=\(2N\)^\{1/4\},\\qquad\\beta=\(8N\)^\{\-1/4\},(3)

这里β\\beta是一个初始化增益，而不是残差流上的额外运行时乘数。如果SlDN\\mathcal\{S\}^\{\\mathrm\{DN\}\}\_\{\\ell\}表示由DeepNet在子层l\\ell中缩放的残差分支矩阵，例如注意力中的值投影和输出投影以及前馈矩阵，那么DeepNorm初始化

Wl,q\(0\)=βDNW~l,q\(0\),q∈SlDN,βDN=\(8N\)−1/4,W\_\{\\ell,q\}^\{\(0\)\}=\\beta\_\{\\mathrm\{DN\}\}\\,\\widetilde\{W\}\_\{\\ell,q\}^\{\(0\)\},\\qquad q\\in\\mathcal\{S\}^\{\\mathrm\{DN\}\}\_\{\\ell\},\\qquad\\beta\_\{\\mathrm\{DN\}\}=\(8N\)^\{\-1/4\},(4)

其中通常的无缩放初始化器用于W~l,q\(0\)\\widetilde\{W\}\_\{\\ell,q\}^\{\(0\)\}。因此进入扰动论证的量是

βDNαDN=\(8N\)−1/4\(2N\)1/4=12N。\\frac\{\\beta\_\{\\mathrm\{DN\}\}\}\{\\alpha\_\{\\mathrm\{DN\}\}\}=\\frac\{\(8N\)^\{\-1/4\}\}\{\(2N\)^\{1/4\}\}=\\frac\{1\}\{2\\sqrt\{N\}\}。(5)

总结DeepNorm计算的有用方式不是α2β\\alpha^\{2\}\\beta的界，而是一阶更新条件

M\(βα\)2=Θ\(1\)。M\\left\(\\frac\{\\beta\}\{\\alpha\}\\right\)^\{2\}=\\Theta\(1\)。(6)

实际上，代入公式(5)得到2N\(β/α\)2=1/22N\(\\beta/\\alpha\)^\{2\}=1/2。

### 2.3 为什么权重共享改变了深度缩放

在非共享的深度为NN的Transformer中，每个子层参数每个输入接收一个梯度贡献。在循环Transformer中，相同的物理参数每轮被访问一次，因此其梯度是RR次访问的和。如果逐次贡献解相关，那么循环的行为在常数上类似于非绑定深度。如果访问是对齐的，绑定的更新可能大RR倍，并且相同的更新会被展开计算中的所有RR次访问读取。DeepLoop是对公式(6)的相应保守修正：它选择一个更小的更新与残差比值β/α\\beta/\\alpha，使得绑定参数更新的一阶效应在展开所有RR次访问后仍然有界。

## 3 DeepLoop Transformer

一个循环Transformer对相同的KK个物理块复用RR轮。其展开深度为N=KRN=KR，残差子层访问次数为M=2NM=2N。DeepLoop保留了DeepNorm架构，但在固定物理深度、增加循环次数的场景中使用指数p=1/2p=1/2而不是p=1/4p=1/4。推导基于两个事实：RMSNorm通过因子1/α1/\\alpha暴露残差分支，而参数绑定将一阶更新和从MM项变为MκRM\\kappa\_\{R\}项，其中κR\\kappa\_\{R\}衡量循环访问之间的对齐程度。

### 3.1 设置：循环后归一化块

每个物理块包含两个残差子层：注意力和MLP。令j=\(k,s\)j=\(k,s\)索引一个物理子层，其中k∈\{1,…,K\}k\\in\\\{1,\\ldots,K\\\}且s∈\{attn,ffn\}s\\in\\\{\\mathrm\{attn\},\\mathrm\{ffn\}\\\}，并令i=\(r,j\)i=\(r,j\)表示其第rr次展开访问。我们记J=2KJ=2K为物理残差子层的数量，M=JR=2KR=2NM=JR=2KR=2N为展开访问次数。DeepLoop使用

xi\+1=Norm\(αxi\+fj\(xi;φj\)\),i=1,…,M,\\displaystyle\\mathbf\{x\}\_\{i\+1\}=\\mathrm\{Norm\}\\left\(\\alpha\\,\\mathbf\{x\}\_\{i\}\+f\_\{j\}\(\\mathbf\{x\}\_\{i\};\\phi\_\{j\}\)\\right\),\\qquad i=1,\\ldots,M,

其中相同的物理参数φj\\phi\_\{j\}对所有r=1,…,Rr=1,\\ldots,R复用。在我们的实现中，Norm\\mathrm\{Norm\}是RMSNorm。增益β\\beta是应用于DeepNorm指定的残差分支矩阵的每个矩阵初始化增益。它不应被解读为整个子层输出的标准差，因为一个分支可能包含多个缩放矩阵。具体来说，对于每个物理子层jj，令Sj\\mathcal\{S\}\_\{j\}是接收DeepNorm增益的矩阵集合。DeepLoop初始化

Wj,q\(0\)=βDLW~j,q\(0\),q∈Sj,βDL=\(8N\)−1/2。W\_\{j,q\}^\{\(0\)\}=\\beta\_\{\\mathrm\{DL\}\}\\,\\widetilde\{W\}\_\{j,q\}^\{\(0\)\},\\qquad q\\in\\mathcal\{S\}\_\{j\},\\qquad\\beta\_\{\\mathrm\{DL\}\}=\(8N\)^\{\-1/2\}。(7)

然后前向递推在每次访问中使用这些初始化后的参数通过fj\(xi;φj\)f\_\{j\}\(\\mathbf\{x\}\_\{i\};\\phi\_\{j\}\)；β\\beta不会作为每次访问的独立乘法因子重新应用。使用αDL=\(2N\)1/2\\alpha\_\{\\mathrm\{DL\}\}=\(2N\)^\{1/2\}，DeepLoop的更新与残差比值为

βDLαDL=\(8N\)−1/2\(2N\)1/2=14N。\\frac\{\\beta\_\{\\mathrm\{DL\}\}\}\{\\alpha\_\{\\mathrm\{DL\}\}\}=\\frac\{\(8N\)^\{\-1/2\}\}\{\(2N\)^\{1/2\}\}=\\frac\{1\}\{4N\}。(8)

\{假设\}

\[DeepNorm尺度的局部敏感性\] 对于每个展开的访问i=\(r,j\)i=\(r,j\)，从该访问处的残差分支参数扰动出发的线性化输出映射的范数，以及对应