关于循环变换器中残差缩放：稳定性与可迁移性

# 循环Transformer的残差缩放：稳定性与可迁移性

**来源：** https://arxiv.org/html/2606.18524  
1\]清华大学 2\]字节跳动Seed 3\]M-A-P (2026年6月16日)

###### 摘要

循环（权重共享）Transformer 使用一个共享的残差块重复 \(N\) 次 (即每一步都使用相同的 \(f\)，更新规则为 \(h \leftarrow h + \varepsilon f(h)\))，在不增加参数的情况下增加了有效深度。先前的深度缩放分析建议，对于深度为 \(L\) 的残差网络，应使用 \(\varepsilon=1/\sqrt{L}\)。我们表明，对于循环架构而言这并不足够：权重共享使得不同迭代之间的残差更新产生关联，因此需要更强的缩放 \(\varepsilon=1/N\)。对于多层块（\(L\) 个独特层各循环 \(N\) 次），我们推导出一个分解式参数化 \(\varepsilon=\lambda/(N\sqrt{L})\)，该参数化将两种增长来源分开：\(1/N\) 控制层内循环相关性，\(1/\sqrt{L}\) 控制跨层方差。一个关键结果是，最优学习率仅取决于独特层的数量 \(L\)，而与循环次数 \(N\) 无关，从而实现了从小 \(N\) 到大 \(N\) 的直接超参数迁移，无需重新调参。在循环 Transformer 上的实验证实，与 \(1/\sqrt{N}\) 缩放相比，\(1/N\) 缩放提高了可训练性，并在不同的循环次数下获得了更优的损失。

## 1 引言

循环（权重共享）Transformer 重复使用单个块 \(f\) 共 \(N\) 次（每一步都使用相同的 \(f\)，更新规则为 \(h \leftarrow h + f(h)\)），在不增加参数的情况下增加了有效深度。这种设计出现在 Universal Transformers [5](https://arxiv.org/html/2606.18524#bib.bib5)、ALBERT [11](https://arxiv.org/html/2606.18524#bib.bib11) 以及近期关于算法推理和隐式计算的工作中 [27, 7, 20, 8, 17, 30]。

(a) 深度网络：\(L\) 个不同的权重矩阵 \(W_0, W_1, W_2, \cdots, W_{L-1}\)；\(h_{\ell+1} = h_\ell + \varepsilon r_\ell\)；\(\varepsilon\)：残差缩放因子；\(r_\ell = W_\ell \phi(h_\ell)\)；\(\|\sum_\ell r_\ell\| = \Theta(\sqrt{L})\)；\(\|\sum_\ell r_\ell\|^2 = \Theta(L)\)；随机游走范数增长；标准缩放足够：\(\varepsilon = \lambda/\sqrt{L}\)  
(b) 循环网络：单个共享权重 \(W\)，重复使用 \(N\) 次 (\(W \times N\))；\(h_{n+1} = h_n + \varepsilon r_n\)；\(\varepsilon\)：残差缩放因子；\(r_n = W \phi(h_n)\)；\(\|\sum_n r_n\| = \Theta(N)\)；\(\|\sum_n r_n\|^2 = \Theta(N^2)\)；线性范数增长；线性缩放需要：\(\varepsilon = \lambda/N\)

**图 1：** 权重共享改变了残差累积的方式。(a) 在具有独立权重的深度网络中，残差更新指向不同方向，像随机游走一样累积，范数为 \(\Theta(\sqrt{L})\)。标准缩放 \(\varepsilon = 1/\sqrt{L}\) 使输出保持有界。(b) 在循环网络中，单个块被重复使用 \(N\) 次。共享权重使得后续更新方向对齐，因此它们的和以 \(\Theta(N)\) 增长，这需要更强的缩放 \(\varepsilon = 1/N\)。

在实践中，增加 \(N\) 常常导致训练不稳定，例如隐藏状态爆炸和对学习率高度敏感 [30]。一种标准补救措施是将每个残差分支乘以一个随深度缩小的因子 \(\varepsilon\)，得到 \(h \leftarrow h + \varepsilon f(h)\)。先前的深度缩放分析建议，对于深度残差网络，\(\varepsilon = 1/\sqrt{N}\) 是一个合适的选择 [2, 6]。但在循环架构中，每一步都重复使用相同的 \(f\)，这一规则是否仍然适用尚不明确。我们发现，对于循环模型，\(1/\sqrt{N}\) 缩放确实是不够的。考虑经过 \(N\) 次迭代后的残差流范数 \(\|h_N\|\)。对于标准（非共享）深度网络，\(\varepsilon = 1/\sqrt{L}\) 成功地使 \(\|h_L\|\) 在深度 \(L\) 增长时保持有界（图 2，上排）。然而，对于循环网络，\(\varepsilon = 1/\sqrt{N}\) 无法控制 \(\|h_N\|\)，后者会随着 \(N\) 迅速增长；相比之下，\(\varepsilon = 1/N\) 则能使它保持有界（图 2，下排）。我们的理论分析（第 3 节）解释了这一差异：\(1/\sqrt{N}\) 规则依赖于每层具有独立权重的假设，但权重共享使得后续更新产生相关性，将残差流范数的增长从 \(\Theta(\sqrt{N})\) 放大到 \(\Theta(N)\)。

**图 2：** 线性缩放稳定了循环网络；平方根缩放则无效。归一化残差流范数 \(R = d^{-1/2} \|h\|_2\)（对数尺度）与深度 \(L\)（上排）或循环次数 \(N\)（下排）的关系，显示在 Llama 风格预归一化 Transformer 诊断程序的前 10 个训练步中；线条按步数着色。\(1/\sqrt{L}\) 缩放稳定了深度网络（面板 b），但 \(1/\sqrt{N}\) 缩放对循环网络无效（面板 e）。线性缩放 \(\varepsilon = 1/N\) 使残差流范数在不同循环次数下保持有界（面板 f）。

除了稳定前向传播，\(1/N\) 缩放还能固定学习率。导致平方范数呈 \(\Theta(N^2)\) 增长的相同构造性累积也会放大权重更新：一次优化器步骤引起的输出变化量级为 \(\eta \varepsilon N\)，因此设置 \(\varepsilon = 1/N\) 可以使稳定学习率在 \(N\) 变化时保持不变（第 3 节）。这使得超参数迁移成为可能：在 \(N=1\) 时调优的学习率在更大的 \(N\) 下仍然保持接近最优，无需重新调参。

我们进一步将分析扩展到实用的多层块，其中 \(L\) 个不同层各重复使用 \(N\) 次（第 4 节）。这引入了第二个方差来源：在独立层之间，更新像标准深度网络一样，以 \(L\) 的随机游走方式累积。分解式参数化 \(\varepsilon = \lambda/(N\sqrt{L})\) 可以独立处理这两个来源：\(1/N\) 消除了层内二次增长，\(1/\sqrt{L}\) 控制了跨层随机游走。由此得到的学习率准则 \(\eta \lesssim 1/(\lambda \sqrt{L})\) 仅取决于独特深度 \(L\)，与 \(N\) 无关，因此对于多层块，超参数迁移仍然成立。

在 FineWeb-Edu [18] 上训练的仅解码器 Transformers 上的实验证实了这些预测（第 5 节）。支撑二次累积的成对相关结构（循环步更新之间密集的正余弦相似度）在初始化之后及整个训练过程中持续存在，表明 \(\Theta(N^2)\) 增长机制并非初始化伪影。线性残差缩放提高了可训练性，并在 \(N \in \{1,2,4,8\}\) 范围内实现了学习率的一致迁移：最优学习率几乎不变，而相比之下，在 \(1/\sqrt{N}\) 缩放下最优学习率会随着 \(N\) 变化而偏移。分解式参数化 \(\varepsilon = \lambda/(N\sqrt{L})\) 进一步将这种迁移扩展到多个深度 \(L \in \{12,24,48\}\)，单个学习率在所有测试的 \((N,L)\) 组合下保持接近最优。

总体而言，我们的工作将深度缩放框架 [26, 6] 扩展到了权重共享架构，表明在参数重复使用的情况下，标准的独立性假设不再成立，并推导出了修正的缩放规则。除了理论贡献外，由此得到的参数化直接解决了两个实际问题：它提高了大循环次数下的可训练性，并消除了在改变 \(N\) 时重新调超参数的需要。通过使循环次数稳定且可调，这些结果确立了重复使用作为权重共享 Transformer 的一个实用缩放轴。

## 2 相关工作

##### 循环和参数共享的 Transformer。

重复使用一个块作为循环深度已在 Universal Transformers [5] 中作为跨层参数共享在 ALBERT [11] 中被探索过。更近期的循环模型在迭代算法学习和多步上下文过程 [27, 8]、长度泛化 [7] 以及类似潜推理风格的计算时缩放 [20, 30] 方面表现出色。相关的参数共享形式也出现在循环神经网络 [17] 中。我们的工作解决了一个互补的问题：如何参数化残差缩放，使得训练保持稳定且超参数可以迁移。

##### 深度残差堆叠和深度 Transformer 的稳定性。

大规模深度训练拥有一系列稳定技术，包括残差重参数化和初始化规则（如 Fixup 和 ReZero [29, 1]）以及针对 Transformer 的稳定器（如 DeepNet/DeepNorm [23]）。在理论方面，先前对深度非共享残差网络的分析表征了残差缩放何时能在深度极大限制下控制信号 [15]。另一条互补的研究路线研究连续深度模型及其 ResNet 类比 [4]：Marion [14] 推导了一个基于 Lipschitz 的界，其复杂度项依赖于连续权重矩阵之间的差异。这些工作大多将深度视为不同层的堆叠，而在我们的设置中，循环步骤重复使用相同的参数。

##### 超参数迁移与参数化。

张量程序和 \(\mu\)P 风格的分析建立了跨模型规模的迁移原则，并推动了系统化的参数化选择 [26]。近期关于深度迁移的分析在非共享假设下研究了学习率和初始化选择如何随深度变化 [2, 10]。CompleteP 和后续工作将这一方向扩展到深度 Transformer 和更广泛的迁移轴 [6, 16]。我们的工作与之互补，针对**循环轴**，并表明共享权重会创建跨步相关性，从而改变稳定性阈值和迁移机制。

## 3 共享层的循环缩放

我们分析仅仅一个共享 MLP 在初始化时的缩放，以隔离权重共享的影响；第 4 节将其扩展到多层块。

### 3.1 设置

考虑以下简化的残差模型，它抽象了 Transformer 块的一个残差分支：
\[
h_{n+1} = h_n + \varepsilon W \phi(h_n), \qquad \varepsilon = N^{-\alpha}, \quad n=0,\dots,N-1.
\]
这里 \(N\) 是循环次数（共享层被应用的次数），\(\alpha > 0\) 是缩放指数，控制残差分支随 \(N\) 增长的缩减力度。隐藏状态 \(h_n \in \mathbb{R}^d\) 从给定的输入 \(h_0\) 初始化，满足 \(\|h_0\|_2^2/d = \Theta(1)\)；模型输出为 \(h_N\)。共享权重矩阵 \(W \in \mathbb{R}^{d \times d}\) 独立同分布地取自 \(W_{ij} \sim \mathcal{N}(0,1/d)\)，\(\phi\) 是 ReLU 激活函数。我们的目标是确定使 \(h_N\) 在 \(N\) 增长时保持有界的最小 \(\alpha\)，以及由此导出的学习率缩放法则。我们记 \(u_n \triangleq \phi(h_n)\) 为激活后向量，\(r_n \triangleq W u_n\) 为每一步的残差，\(R_n \triangleq d^{-1/2} \|h_n\|_2\) 为归一化残差流范数。

**图 3：** 权重共享导致持续的跨步相关性。块级增量 \(\delta_i = h_i - h_{i-1}\) 之间的成对余弦相似度，比较非共享深度堆叠（面板 a；64 个独立的 12 层块副本，有效深度 \(12 \times 64\)）与循环网络（面板 b；\(L=12, N=64, d=768\)）。两种配置具有相同的有效深度；唯一区别是块权重是否共享。两者均在无残差缩放的情况下，经过 10 个训练步后测量。在非共享情况下，非对角相关性可忽略（范围 \([-0.037, 0.034]\)）。在循环情况下，共享权重产生密集的正向对齐（范围 \([0.027, 0.995]\)），与 \(\Theta(N^2)\) 累积一致（定理 1）。

### 3.2 二次方差累积

展开 (1) 得 \(h_N = h_0 + \varepsilon \sum_{n=0}^{N-1} r_n\)。为了了解 \(R_N\) 如何随 \(N\) 缩放，我们展开 \(R_N^2\)：
\[
R_N^2 = R_0^2 + \underbrace{\frac{2\varepsilon}{d} \sum_{n=0}^{N-1} \langle h_0, r_n \rangle}_{B_N} + \underbrace{\frac{\varepsilon^2}{d} \sum_{n=0}^{N-1} \sum_{m=0}^{N-1} \langle r_n, r_m \rangle}_{C_N}.
\]
交叉项 \(B_N\) 求和了 \(N\) 个固定输入 \(h_0\) 与残差 \(r_n\) 的内积，因此 \(B_N = O(\varepsilon N)\)。然而，二次项 \(C_N\) 求和了 \(N^2\) 对相互作用，等于 \(\frac{\varepsilon^2}{d} \left\| \sum_n r_n \right\|_2^2\)。