DeepLoop：循环 Transformer 的深度缩放方法 DeepLoop: Depth Scaling for Looped Transformers
为循环 Transformer 设计 p=1/2 的残差缩放规则，稳定权重共享下的深层训练
前置知识
循环 Transformer（Looped Transformer）
循环 Transformer 把一个由 $K$ 个物理块组成的紧凑堆栈反复执行 $R$ 轮，从而在不增加存储参数量的前提下获得展开深度 $N=KR$。也就是说同一个物理块（包含 attention 和 FFN 两个残差子层）会被多次复用，等价于 Universal Transformer、ALBERT 这类深度方向权重共享的架构。它的吸引力在于把深度变成了一种可控的计算资源：测试时多算几轮就能提升表达能力，而参数预算保持不变。
整篇论文都是围绕'同一组参数被多次访问'这一特殊性展开的，不理解循环 Transformer 的核心机制就无法理解为什么标准的残差缩放分析在这里失效。
DeepNorm 与 Post-LN 残差缩放
DeepNorm 是为超深 Post-LN Transformer 设计的初始化方案：它在残差路径上引入跳跃缩放 $\alpha$ 和残差分支初始化增益 $\beta$，使得 $N$ 个块、$M=2N$ 个残差子层访次的模型在第一步优化后输出变化仍保持有界。DeepNorm 推导得到 $\alpha=(2N)^{1/4}$、$\beta=(8N)^{-1/4}$，即指数 $p=1/4$。其核心稳定性条件是 $M(\beta/\alpha)^2=O(1)$。这里的 $\beta$ 不是运行时乘子，而是对 value/output 投影和 FFN 矩阵的初始化增益。
DeepLoop 完全沿用 DeepNorm 的 Post-LN 架构，只改指数；不理解 DeepNorm 就无法体会 DeepLoop'从 $p=1/4$ 升到 $p=1/2$'这一改动背后的稳定性论证。
一阶扰动分析（First-order Perturbation）
一阶扰动分析研究在一次优化器更新 $\theta\leftarrow\theta+\delta\theta$ 之后，前向输出 $F(x;\theta)$ 的变化量 $\Delta F=F(x;\theta+\delta\theta)-F(x;\theta)$。对深度网络做泰勒展开并取一阶项，可以把每个残差子层对最终输出的灵敏度求和，得到 $\|\Delta F\|\le C\cdot\sum_i \|U_i G_i\|$ 的形式。DeepNorm 的 $\alpha$、$\beta$ 推导正是基于这种逐访次求和的一阶界。这种方法的优势是把'深层网络能否稳定训练'转化为对单访次灵敏度的代数约束。
DeepLoop 的核心贡献就是把 DeepNorm 的一阶扰动求和从 $M$ 项改写为 $M\cdot\kappa_R$ 项，理解一阶扰动分析是看懂 tied-depth bound 推导的前提。
权重共享 / 参数绑定（Parameter Tying）
权重共享指同一个参数张量被多个计算位置复用。在循环 Transformer 中，物理块 $j$ 的参数 $\phi_j$ 会在 $R$ 轮中被访问 $R$ 次。这导致两件事：第一，该参数的梯度是 $R$ 次访次的梯度之和；第二，更新后的参数又被这 $R$ 次访次读取。这种'写入和读取都跨多次访次'的耦合，在标准 DeepNorm 假设（每个子层参数独立）下并不存在，是 DeepLoop 必须额外建模的关键机制。
论文的全部新颖性都建立在这个 tied-depth 现象上：理解参数绑定如何改变梯度聚合与读取路径，才能理解为什么要引入访问对齐系数 $\kappa_R$。
研究动机
循环 Transformer 想把'深度'变成一种不增加参数的扩展轴：存 $K$ 个物理块、跑 $R$ 轮、得到 $N=KR$ 的展开深度。然而要让这条扩展轴在工程上稳定可用，残差参数化必须能在'同一组参数被反复访问'的情况下保持训练稳定。问题在于，现有的标准残差缩放分析——尤其是为超深 Post-LN Transformer 设计的 DeepNorm——是针对 untied depth 写的：它假设 $M=2N$ 个残差子层访次中的每一个都拥有自己独立的参数张量，因此只贡献一个独立的一阶更新项。在权重共享的循环 Transformer 中这个假设被打破：同一物理子层的参数在 $R$ 轮里被反复读写，其优化器更新会聚合 $R$ 次访次的梯度，然后在下一次前向中被这 $R$ 次访次重新读取。如果各轮访次的梯度和灵敏度方向一致（aligned），共享更新会被放大 $\Theta(R)$ 倍，DeepNorm 的 $p=1/4$ 设置在固定物理深度 $K$ 下会留下 $\Theta(R)$ 的残差增长，导致深层循环训练不稳定。
本文的目标是本文的目标是为循环 Transformer 推导一个'循环感知'的残差缩放规则，使得在权重共享、物理块被多次访问的设定下，一次优化步之后输出的变化仍保持一阶有界。具体而言，作者希望给出缩放族 $\alpha=(cN)^p$、$\beta=(dN)^{-p}$ 中指数 $p$ 的正确阈值，并保留 DeepNorm 简洁的 Post-LN 架构与初始化范式：不引入门控、不引入可学习残差系数、不引入辅助损失，只通过确定性的 $\alpha$ 和 $\beta$ 来稳定训练。作者还要验证这个规则在 GPT-2 small/medium 的循环语言建模和层次化推理模型上都能带来实际收益。
与已有工作不同的是，以往工作要么关注循环 Transformer 的表达能力（Universal Transformer、ALBERT、recurrent-depth LMs），要么关注残差参数化本身（DeepNorm、$\mu$P、ReZero），但没有把'参数跨轮共享'这一事实显式纳入残差缩放的稳定性推导。本文的独特切入角度是引入一个显式的访问对齐系数 $\kappa_R$ 来量化各轮访次的梯度和灵敏度之间的相关性，从而把 untied DeepNorm 的界 $M(\beta/\alpha)^2=O(1)$ 推广为循环版 $M\cdot\kappa_R\cdot(\beta/\alpha)^2=O(1)$。这一公式把 DeepNorm 作为一个特例（$\kappa_R=O(1)$ 时 $p=1/4$）恢复出来，同时在最保守的对齐场景下给出新的指数阈值 $p=1/2$。这填补了'参数共享如何改变深度缩放'这一长期被忽略的理论缺口。
核心方法
DeepLoop 的整体思路是：保留 DeepNorm 的 Post-LN 残差块结构与初始化方式，只修改缩放指数。直觉上，RMSNorm 会在每个子层把信号归一到单位 RMS，所以前向信号尺度不是约束瓶颈，真正的约束是一次优化步后输出对参数扰动的一阶灵敏度。在 untied 情形下，每个子层访次独立贡献一个 $O(\beta/\alpha)$ 的灵敏度项，求和得到 $M(\beta/\alpha)^2=O(1)$。但在循环情形下，共享参数的更新被 $R$ 次访次写入、又被 $R$ 次访次读取，形成双求和结构。引入访问对齐系数 $\kappa_R$ 把这个双求和吸收掉后，界变成 $M\cdot\kappa_R\cdot(\beta/\alpha)^2=O(1)$，其中 $0\le\kappa_R\le R$：访次独立时 $\kappa_R=O(1)$，访次完全对齐时 $\kappa_R=\Theta(R)$。把这一界套到缩放族 $\alpha=(cN)^p$、$\beta=(dN)^{-p}$ 上，作者证明阈值是 $p\ge(1+\gamma)/4$，其中 $\kappa_R=\Theta(R^\gamma)$；$\gamma=0$ 退化为 DeepNorm 的 $p=1/4$，$\gamma=1$ 给出 $p=1/2$。DeepLoop 取后者并沿用 DeepNorm 常数 $(c,d)=(2,8)$，最终规则是 $\alpha=(2N)^{1/2}$、$\beta=(8N)^{-1/2}$。
核心创新点是显式地把'tied-depth 聚合效应'写成访问对齐系数 $\kappa_R=\max_j \big[(\sum_r U_{r,j})(\sum_r G_{r,j})\big]/(R\cdot C_U\cdot C_G\cdot(\beta/\alpha)^2)$。这一系数同时刻画了两个耦合路径：共享参数的更新由多次访次写入（梯度的 $R$ 项求和），更新后的参数又被这 $R$ 次访次读取。与已有方法的本质区别在于：DeepNorm 把 $M=2N$ 个访次当成统计独立的更新项求和，DeepLoop 则承认这些项之间存在结构性相关，并用 $\kappa_R$ 把相关性的上界显式代入界里。其结果不是推翻 DeepNorm，而是把它推广成一个特例——当 $\kappa_R=O(1)$（访次近正交）时 $p=1/4$ 仍然正确；只有当访次对齐、$\kappa_R=\Theta(R)$ 时才需要把指数从 $1/4$ 提到 $1/2$。这种'同一公式覆盖两种情形'的设计使方法既是理论上的最小修正，也是工程上的稳妥默认。
方法步骤详情
方法推导分八步。第 1 步（Setup）：把循环块写成 $x_{i+1}=\mathrm{Norm}(\alpha\cdot x_i+f_j(x_i;\phi_j))$，其中 $j=(k,s)$ 索引物理子层（$K$ 个块、每块 attn 和 ffn 两个子层，共 $J=2K$ 个物理子层），$i=(r,j)$ 表示第 $r$ 轮的展开访次，总访次数 $M=JR=2KR=2N$。第 2 步（局部灵敏度假设 Assumption 3.1）：每个访次对参数扰动的一阶输出灵敏度和有效更新范数都是 $O(\beta/\alpha)$，宽度、Norm 增益、注意力头数、学习率等常数被吸收进 $O(\cdot)$。第 3 步（untied 界）：在参数不共享时求和 $M$ 项得到 $\|\Delta F\|\le C'\cdot M\cdot(\beta/\alpha)^2$，即 Proposition A.2，DeepNorm 的 $p=1/4$ 满足 $M(\beta/\alpha)^2=1/2$。第 4 步（定义 $\kappa_R$）：把循环下的双求和写成 $\sum_j\sum_r U_{r,j}\cdot\sum_t G_{t,j}$，提取出 $\kappa_R$ 满足 $0\le\kappa_R\le R$。第 5 步（tied 界）：代入得到 $\|\Delta F\|\le C''\cdot M\cdot\kappa_R\cdot(\beta/\alpha)^2$（Proposition A.3），充分条件为 $M\cdot\kappa_R\cdot(\beta/\alpha)^2=O(1)$。第 6 步（指数阈值 Proposition 3.2）：对缩放族 $\alpha=(cN)^p$、$\beta=(dN)^{-p}$，在固定 $K$、$N=KR$、$\kappa_R=\Theta(R^\gamma)$ 下，条件当且仅当 $p\ge(1+\gamma)/4$。第 7 步（选定规则）：保守取 $\gamma=1$、$p=1/2$，并沿用 DeepNorm 常数得到 DeepLoop 规则 $\alpha=(2N)^{1/2}$、$\beta=(8N)^{-1/2}$，此时 $MR(\beta/\alpha)^2=1/(8K)$ 为 $O(1)$。第 8 步（推广到层次化推理器 Section 4）：把访次数替换为梯度可见访次 $M_g=2(K_H+C_L K_L)$，并允许每个模块使用独立的 $(\alpha_H,\beta_H)$ 与 $(\alpha_L,\beta_L)$；在内循环对齐时仍得出 $p=1/2$ 的约束。
技术新颖性
技术新颖性体现在三方面。其一，论文首次显式识别了'tied-depth 聚合机制'：共享更新由多次访次写入、再由同样这几次访次读取，这一双耦合路径是 untied 残差缩放分析里完全不存在的现象。其二，作者用一个标量 $\kappa_R$ 把这个机制装进闭式界里，使得 DeepNorm 的 $p=1/4$ 与循环情形的 $p=1/2$ 能从同一个公式里派生出来，理论上是 DeepNorm 的最小一致性推广。其三，论文把这个规则扩展到层次化推理器（如 HRM）：通过引入'梯度可见访次数' $M_g$ 来刻画 one-step 梯度截断，证明在固定物理深度下沿内循环轴增长有效深度时仍然落到 $p=1/2$ 这一阈值，并且允许高低层模块使用非对称缩放。整体而言，DeepLoop 不引入任何门控、可学习残差系数、辅助损失或架构常数，是 Post-LN 循环 Transformer 的'一行修正'，这是它在工程上最大的简洁性优势。
实验结果
核心发现按三组实验逐一分析。第一组（Table 1，FineWeb-Edu 50B token、GPT-2 small 124M 与 medium 350M、上下文 1024）：在 $R=1$（没有物理块被复用）时 DeepLoop 与基线基本持平——small 上 +0.0004、medium 上 +0.0011 nats，处于噪声内；一旦 $R\ge3$ 启用递归深度，DeepLoop 在两种规模的所有 $R$ 上都严格更好。具体差距在 small 上：$R=3$ 时 -0.0160、$R=5$ 时 -0.0231、$R=7$ 时 -0.0186；在 medium 上差距随 $R$ 单调扩大：$R=3$ 时 -0.0153、$R=5$ 时 -0.0196、$R=7$ 时达到 -0.0278 nats，medium 上递归深度信号更强。这印证了理论预测：收益恰好出现在'物理块被复用'时。第二组（Table 2，lm-evaluation-harness 八任务、medium 规模）：0-shot 平均上 DeepLoop 在 $R=3$、5、7 分别领先 +0.43、+1.06、+0.93 个百分点；1-shot 平均上 $R=7$ 时 DeepLoop 达到 55.20%，相对基线 54.62% 领先 +0.58；$R=7$ 下 DeepLoop 在八任务中的 7 个上同时领先（仅 PIQA 略低），其中 WinoGrande 0-shot 跳升 +1.74 pp（59.04 vs 57.30）。第三组（Table 3，ARC-AGI-1 + 层次化推理模型 HRM）：在保持 backbone、AdamATan2 优化器、100K epoch 与同一份 hash 校验的 arc-aug-1000 数据完全一致的前提下，仅把残差参数化换成 DeepLoop（$M_g=24$、$N_g=12$），论文协议的 $K=2$ 投票准确率从 36.50% 提升到 39.75%（+3.25 pp），并且 $K=1/2/10/100/1000$ 整条投票阶梯全部改善；四种子对照显示 $K=2$ 标准差约 0.5 pp，因此 +3.25 pp 大约是 $6\sigma$ 效应。此外 Appendix C 的 p-sweep 在 $R=3$ 上经验性地定位训练稳定性边界在 $p\approx1/2$ 附近，与 Proposition 3.2 的理论阈值一致。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| FineWeb-Edu 验证损失（GPT-2 medium，50B token，step 100K，R=7） | 交叉熵（nats，越低越好） | 2.5280（DeepLoop） | 2.5558（baseline：tying+input-embed RMSNorm） | -0.0278 nats |
| FineWeb-Edu 验证损失（GPT-2 small，R=7） | 交叉熵（nats，越低越好） | 2.7514（DeepLoop） | 2.7700（baseline） | -0.0186 nats |
| lm-evaluation-harness 八任务平均（medium，1-shot，R=7） | acc 平均（%，越高越好） | 55.20%（DeepLoop R=7） | 54.62%（baseline R=7） | +0.58 pp |
| ARC-AGI-1（HRM，K=2 投票，N=400 题） | 投票准确率（%） | 39.75%（DeepLoop p=1/2） | 36.50%（vanilla HRM） | +3.25 pp（约 6σ） |
局限与改进
作者明确承认三点局限。其一，全部主要结果（Table 1、Table 2）都是单种子运行，作者写道'Multi-seed runs would be needed to quantify run-to-run variance'，所以 medium 上 -0.0278 nats、ARC-AGI +3.25 pp 等数字的统计置信度尚未严格量化（尽管 ARC-AGI 上做了四种子 $K=2$ 对照）。其二，论文没有直接测量 $\kappa_R$ 或跨轮梯度对齐度，$p=1/2$ 的'保守'取法是基于最坏情形假设；真实模型很可能并未达到完全对齐，因此存在指数过度保守的可能。其三，实验只覆盖到 GPT-2 small 与 medium（124M 与 350M）规模，未在更大尺度上验证阈值是否依旧适用；层次化推理器方面也只在 ARC-AGI-1 上做了单一任务评测。我自己的观察是：DeepLoop 仅针对 Post-LN+RMSNorm 推导，对 Pre-LN 或其他归一化放置的分析并不直接适用；且实验中 $R$ 仅到 7，未触及非常大 $R$（如 $R\ge16$）下可能出现的别的不稳定性。
独立分析的弱点
独立分析有四个弱点。第一，访问对齐系数 $\kappa_R$ 在训练中是动态变化的、且未被显式监测，$p=1/2$ 是基于最坏情形的对齐假设；改进方向是在训练中实时估计跨轮梯度余弦相似度，并据此自适应地选择更小的 $p$（例如 $p=1/4$ 到 $1/2$ 之间）以减少过度保守带来的容量浪费。第二，主要语言建模实验是单种子，medium 规模 -0.0278 nats 的差距在某些任务（如 $R=5$ 的 1-shot Avg）上甚至被基线反超 0.23 pp，说明增益并非在所有 (R, shot, task) 单元上一致；改进方向是多种子重复并报告均值±方差。第三，理论与实验之间的'$p$ 阈值'验证只在 $R=3$ 的一个网格上做了 sweep（Appendix C），未在多个 $R$ 上系统刻画稳定性边界随 $R$ 的移动；改进方向是把 p-sweep 扩展到 $R\in\{3,5,7,10\}$ 并拟合经验阈值曲线与理论 $(1+\gamma)/4$ 对照。第四，方法只覆盖 Post-LN，而现代大模型常采用 Pre-LN 或 RMSNorm+SWIGLU 等组合，DeepLoop 的常数 $(c,d)=(2,8)$ 是针对两子层解码块的，对 FFN 子层数不为 1 的变体需要重新推导；改进方向是把 perturbation 推广到任意子层组合并给出对应的 $c$、$d$。
未来方向
作者明确点名的未来方向有三条：直接测量 $\kappa_R$ 或跨轮梯度对齐度以验证对齐假设是否成立；测试同一稳定性边界在更大尺度或别的参数化（如 $\mu$P、Depth-$\mu$P）下是否依旧适用；探索通过训练目标或正则化鼓励访次去相关，从而安全地使用比 $p=1/2$ 更小的指数。基于本文成果可延伸的方向包括：把 tied-depth 分析推广到 Pre-LN 与其他归一化放置；研究 $R$ 取很大（如 $R\ge16$）时是否存在 $\kappa_R$ 之外的新不稳定性来源；把层次化推理器的非对称 $(\alpha_H,\beta_H)/(\alpha_L,\beta_L)$ 参数化进一步发展为'每模块独立 $p$'的自动调参方案；以及在 MoE、长上下文 Attention、状态空间模型等同样存在参数复用的架构上检验 $\kappa_R$ 框架的适用性。
复现评估
复现评估总体较好。论文提供了项目页 https://github.com/lszeshu/DeepLoop，承诺开源代码；实验设置描述相当详细：GPT-2 small 用 4×H200 141GB、medium 用 8×H200 141GB，FineWeb-Edu 数据、100K optimizer steps、context 1024、global batch 480；ARC-AGI 部分强调了与 baseline 完全一致的 backbone、AdamATan2 优化器、100K epoch 以及 hash 校验的 arc-aug-1000 数据，并以 published HRM ARC-2 检查点（5.00% on $K=2$）作为 harness 校验。算法本身极简——只改 $\alpha$ 与 $\beta$ 两个标量——所以代码改动量小。复现难点主要在三方面：一是大规模训练成本（50B token、最多 8 卡 H200）对独立研究者不友好；二是 baseline 已经做了 input-embedding RMSNorm 与 tying，要严格复现 $\Delta$ 需要精确对齐这些细节；三是 ARC-AGI-1 评测需要专用 pipeline 与 count-first 投票协议，作者公开了 hash 校验细节以减少数据歧义。
论文图表
