Post-LayerNorm 回归：稳定、表达力强且支持深度扩展

通过 Highway 式残差连接取代 ResNet 式路径，使 Post-LN 在超深模型中稳定训练

前置知识

Layer Normalization (层归一化)

LayerNorm 是深度神经网络中用于稳定训练的关键技术。给定输入向量 $x$，其计算公式为 $\text{LN}(x) = \frac{x}{\|x\|_2} \odot \gamma$，其中 $\gamma \in \mathbb{R}^d$ 是可学习的仿射变换参数，$\odot$ 表示逐元素乘法。LayerNorm 的核心作用是对每一层的激活值进行标准化，使其均值和方差保持在合理范围内，从而防止梯度爆炸或消失。在 Transformer 中，LayerNorm 的放置位置（Pre-LN 或 Post-LN）对模型的训练稳定性和表达能力有深远影响。本文重点关注 Post-LN 的恢复与改进。

本文的核心创新正是围绕 LayerNorm 展开的，通过分析 Post-LN 中梯度消失的根源，提出了新的架构设计。理解 LayerNorm 的工作原理是理解本文理论分析和方法设计的基础。

Post-LayerNorm 与 Pre-LayerNorm

Post-LN 和 Pre-LN 是 Transformer 中两种主要的归一化放置策略。Post-LN 的公式为 $x_{l+1} = \text{LN}(x_l + F(x_l))$，即先做残差加法再归一化，这是原始 Transformer 的设计。Pre-LN 的公式为 $x_{l+1} = x_l + F(\text{LN}(x_l))$，即先归一化再做变换，残差连接绕过归一化。Pre-LN 因训练稳定性好而成为现代 LLM（如 GPT-3、LLaMA）的默认选择，但其梯度主要通过恒等连接传播，导致深层网络的有效深度不足，表达能力受限。Post-LN 虽然在深层保持更强的梯度信号和层间耦合，但在大规模训练中容易出现不稳定。

本文要复兴 Post-LN 架构，理解两种归一化策略的优劣是理解本文动机的关键。Pre-LN 的有效深度问题和 Post-LN 的不稳定性正是本文要解决的核心矛盾。

残差连接（Residual Connection）

残差连接是 ResNet 提出的深度学习基础组件，标准形式为 $x_{l+1} = x_l + F(x_l)$。这种设计通过恒等映射路径让信号可以直接跨层传播，缓解了深度网络的训练困难。在 Transformer 中，残差连接同时用于自注意力（MHA）和前馈网络（FFN）两个子层。残差连接与 LayerNorm 的交互方式（加法前还是加法后归一化）决定了梯度的传播路径和激活的累积方式，是影响深度扩展能力的关键设计选择。

本文的核心技术贡献是用 Highway 式连接替换 ResNet 式残差连接，这是解决 Post-LN 梯度消失的根本手段。理解标准残差连接的工作原理是理解本文改进策略的前提。

Highway Network（高速公路网络）

Highway Network 是一种早期用于训练深度前馈网络的机制，其核心思想是让隐藏层自适应地调节输入被变换还是直接传递的比例。给定输入 $x$ 和变换 $F(x)$，Highway 层计算 $x_{l+1} = T(x_l) \odot F(x_l) + C(x_l) \odot x_l$，其中 $T(x_l)$ 是可训练的门控，$C(x_l)$ 通常设为 $1 - T(x_l)$。这种门控机制确保梯度在必要时可以绕过变换，防止梯度衰减。与 ResNet 的固定恒等路径不同，Highway Network 通过可学习的门动态控制信息流。

本文将 Highway Network 的门控思想引入 Post-LN Transformer，用简化的 Highway 式连接替换 ResNet 式残差路径，这是本文方法的技术灵感来源。

梯度消失与梯度爆炸

梯度消失和梯度爆炸是深度网络训练中的两个核心挑战。梯度消失指反向传播时梯度信号逐层衰减，导致浅层参数几乎无法得到有效更新，模型无法学习层次化的特征表示。梯度爆炸则相反，梯度信号逐层放大导致数值不稳定。在 Post-LN Transformer 中，梯度必须通过 LayerNorm 的 Jacobian 传播，当残差输出和变换特征被求和后再归一化时，梯度会表现出极端的变异性，尤其在深层网络中。本文通过理论推导证明了 Post-LN 中梯度累积量随深度 $L$ 呈指数衰减：$\prod_{l=1}^{L} J^*_{\text{LN}_{l,1}}(z_l) = \mathcal{O}\left(\frac{1}{L^{L/2}}\right)$。

梯度消失是 Post-LN 被 Pre-LN 取代的根本原因，也是本文要解决的核心技术问题。理解梯度传播的动力学是理解 Keel 架构设计原理的关键。

DeepNorm

DeepNorm 是一种为深度 Transformer 设计的归一化方法，通过引入深度相关的缩放因子来稳定 Post-LN 训练。其公式为 $x_{l+1} = \text{LN}(\alpha x_l + F(x_l))$，其中 $\alpha$ 根据理论公式设定（对于 decoder-only 架构 $\alpha = L^{0.25}$），同时将权重初始化缩放 $\beta = L^{-0.25}$。DeepNorm 的核心思路是通过约束前向输出幅度来防止爆炸，但它依赖特殊的初始化方案，在 LLM 训练中随着权重从初始化漂移，其稳定效果会减弱。本文的 Keel 架构从梯度流角度进行分析，采用了不同的稳定策略。

DeepNorm 是本文的直接对比基线和改进对象。理解 DeepNorm 的设计思路和局限性有助于理解 Keel 的创新之处。

研究动机

大语言模型（LLM）的性能提升主要依赖扩展（scaling），但传统的扩展轴正在遇到瓶颈。宽度扩展（增加隐藏维度）迅速饱和，上下文长度扩展成本越来越高，单纯增加参数也无法解锁质的新能力。深度扩展（增加网络层数）理论上可以表示指数级更丰富的函数并支持更多层次的推理，但当前的 Transformer 架构在极端深度下训练极不稳定。现代 LLM 普遍采用 Pre-LN 归一化策略，它通过在变换前归一化输入来稳定早期训练，避免了 Post-LN 常见的发散问题。然而，Pre-LN 引入了结构性限制：梯度主要通过恒等连接传播，导致深层网络的有效贡献逐层递减。当模型深度增加时，Pre-LN 的深度扩展回报远低于宽度扩展，限制了深度作为新扩展轴的潜力。具体实验表明，在 512 层配置中，Pre-LN 的最大可容忍学习率仅为 $4.67 \times 10^{-3}$，而标准 Post-LN 更低至 $2.8 \times 10^{-4}$。

本文的目标是本文的具体目标是恢复 Post-LayerNorm 架构在深度 LLM 训练中的可行性，使其能够发挥理论上优于 Pre-LN 的深度扩展能力和表达能力，同时克服训练不稳定性问题。作者希望通过一种简单而有效的架构修改，使 Post-LN 能够在超过 1000 层的深度下稳定训练，不需要特殊的初始化方案或复杂的优化技巧。具体来说，目标包括：1）证明 Post-LN 的不稳定性根源在于 ResNet 式残差路径而非归一化本身；2）设计一种能保证梯度在任意深度不消失的架构；3）在保持训练稳定性的同时，获得比 Pre-LN 更强的模型表达能力，特别是在数学和代码等推理密集型任务上。

与已有工作不同的是，本文的独特切入角度是从梯度动力学的角度分析 Post-LN 的失败模式，并提出针对性的架构解决方案。现有方法（DeepNorm、Admin、HybridNorm、Mix-LN 等）主要通过修改归一化放置位置或添加深度相关的缩放因子来缓解不稳定性，但这些方法没有从根本上解决 Post-LN 的梯度退化问题。本文首次正式推导了 Post-LN 中梯度信号的衰减上界，证明了 ResNet 式残差路径（而非归一化本身）是梯度消失的主要来源。基于这一分析，作者没有继续在归一化位置上做文章，而是将注意力转向残差连接的设计，用 Highway 式门控连接替换 ResNet 式恒等路径，直接从梯度流动角度解决问题。此外，本文还建立了深度传播与序列建模中 Test-Time Training（TTT）之间的结构同构关系，为未来研究提供了理论桥梁。

核心方法

Keel 的整体思路是通过一个看似简单但影响深远的架构修改来解决 Post-LN 的梯度消失问题。直觉上，Post-LN 的不稳定性不是归一化本身的缺陷，而是残差和变换特征在归一化前被相加的方式导致的。当这两个信号对齐不佳时，LayerNorm 的 Jacobian 会导致梯度极端变异。Keel 的解决方案是：将 ResNet 式残差连接替换为 Highway 式连接，通过一个缩放因子 $\alpha$ 来加权跳过连接，同时在残差分支的输入处注入额外的 LayerNorm。技术路线包括：1）理论推导 Post-LN 梯度衰减的数学上界；2）证明 Highway 式连接可以保证梯度范数在任意深度保持一致（约等于 1）；3）在残差分支内加入输入归一化以稳定残差分支的梯度流；4）通过消融实验验证从朴素缩放到最终 Keel 公式的设计演化过程。

Keel 的核心创新点是将 Post-LN Transformer 的 ResNet 式残差路径替换为 Highway 式门控连接，并引入残差分支归一化。具体来说，Keel 的前向传播公式为 $x_{l+1} = \text{LN}_{l,1}(\alpha x_l + F_l(\text{LN}_{l,2}(x_l)))$，其中 $\alpha = L$（$L$ 为总子层数）。与已有方法的本质区别在于：DeepNorm 依赖特殊的初始化方案（$\beta = L^{-0.25}$）并在后期训练中效果减弱，而 Keel 将稳定机制内化到架构本身，通过结构设计而非初始化来确保全程稳定。与 HybridNorm/Mix-LN 等混合归一化方案不同，Keel 不改变 Post-LN 的本质拓扑（快捷路径仍携带归一化后的信号），而是通过 Highway 式缩放控制梯度流动。关键理论结果是，当 $\alpha = L$ 时，梯度累积量的极限为 $\lim_{L \to \infty} \prod_{l=1}^{L} J^*_{\text{LN}_{l,1}}(z_l) = \lim_{L \to \infty} \frac{\alpha}{\sqrt{\alpha^2 + 1}} = 1$，证明了梯度不会消失。

方法步骤详情

Keel 的计算过程包含以下关键步骤。首先，对于第 $l$ 层的输入 $x_l$，先通过一个额外的 LayerNorm $\text{LN}_{l,2}$ 进行归一化，作为残差分支 $F_l$ 的输入。这一步的目的是稳定残差分支的方差，防止未归一化的输入导致梯度衰减。然后，将归一化后的输入送入标准的 Transformer 子层（注意力或 FFN），得到变换输出 $F_l(\text{LN}_{l,2}(x_l))$。接下来，将快捷路径 $\alpha x_l$ 与变换输出相加，其中 $\alpha = L$ 是 Highway 式缩放因子。最后，对求和结果应用 Post-LayerNorm $\text{LN}_{l,1}$ 得到该层输出 $x_{l+1}$。实现细节方面：1）第一个注意力层移除 Post-LN 以确保从嵌入层开始的稳定信号初始化；2）所有 LayerNorm 使用可学习的仿射权重 $\gamma$ 但省略偏置项 $\beta = 0$；3）$\alpha = L$ 的设置对超深模型的训练稳定性至关重要，对于较小的架构可将 $\alpha$ 作为可调超参数（$\alpha > 1$）。这种设计使得 Keel 可以使用比 Pre-LN 基线更大的学习率，加速收敛。

技术新颖性

Keel 的技术新颖性体现在多个层面。从理论角度，本文首次正式推导了 Post-LN 中梯度信号衰减的精确数学上界，证明了 ResNet 式残差路径是梯度消失的根源，而非归一化本身。这个理论洞察为后续架构设计提供了明确的方向。从方法角度，将 Highway Network 的门控思想与 Post-LN 结合是一个全新的视角。之前的 Highway Network 主要用于早期的前馈网络，本文将其简化为一个标量缩放因子 $\alpha$ 并引入 Transformer，既保留了 Highway 式梯度流动的优势，又避免了完整门控的计算开销。从设计演化角度，本文展示了从朴素残差缩放（Attempt 1）到可学习输入缩放（Attempt 2）再到解耦缩放与方差（Attempt 3）最终到 Keel 最终公式的设计路径，每一步都有清晰的理论动机和实验验证。从深度-序列同构角度，本文建立了深度维度上的 Test-Time Training 与序列维度上 TTT 之间的结构对应关系，为跨领域的技术迁移提供了理论基础。

实验结果

实验结果全面验证了 Keel 在训练稳定性、表达能力和深度扩展三个维度的优势。在训练稳定性方面，Keel 在 64 层配置下最大可容忍学习率达到 $1.01 \times 10^{-2}$，比 Pre-LN（$7.65 \times 10^{-3}$）高出 32%，比标准 Post-LN（$3.0 \times 10^{-4}$）高出近两个数量级。在 512 层配置下，Keel 的 Max LR 为 $6.31 \times 10^{-3}$，显著优于 Pre-LN 的 $4.67 \times 10^{-3}$。在表达能力方面，512 层 3B 参数模型在 250B token 预训练后，Keel 在所有能力域全面超越 Pre-LN，整体平均分 55.5 vs 52.3（+3.2），特别是在数学和代码任务上提升显著：GSM-8K +5.7（43.8 vs 38.1）、MBPP +3.2（26.0 vs 22.8）。在深度扩展方面，Keel 的优势随深度增加而扩大：64 层时平均分提升 +1.7，128 层 +1.2，512 层 +3.8，1024 层 +3.0。1T token 预训练实验中，512 层 Keel 在 GSM-8K 上达到 60.9（Pre-LN 为 51.0，+9.9），MMLU-Pro 上 35.6（Pre-LN 为 26.6，+9.0）。SFT 阶段同样验证了 Keel 的优势持续存在，BBH 上 51.7 vs 46.4（+5.3），CMMLU 上 71.8 vs 66.5（+5.3）。在「更深 vs 更宽」实验中，3B 参数预算下 512 层 Keel（平均 55.5）同时超越了 512 层 Pre-LN（52.3，+3.2）和 128 层 Pre-LN（52.2，+3.3），证明深度扩展在 Keel 架构下真正可行。

**Figure 1**: KEEL 实现稳定、高表达力和深度可扩展的 LLM 训练

**Figure 5**: Pre-LN 和 Keel 在 FineWeb-EDU 数据集上的训练损失曲线

**Figure 7**: 开源模型（Qwen2.5-72B-Instruct 和 LLaMA-3.3-70B-Instruct）的层移除分析

查看结构化数据

任务	指标	本文	基线	提升
GSM-8K (数学推理)	5-shot Accuracy	60.9%	51.0%	+9.9%
MMLU-Pro (综合知识)	5-shot Accuracy	35.6%	26.6%	+9.0%
CMMLU (中文理解)	5-shot Accuracy	72.0%	66.6%	+5.4%
AGI-Eval (综合评测)	0-shot Accuracy	46.5%	37.9%	+8.6%
CommonsenseQA (常识推理)	0-shot Accuracy	69.8%	64.5%	+5.3%
MBPP (代码生成)	0-shot Accuracy	40.6%	35.0%	+5.6%
HellaSwag (常识推理)	0-shot Accuracy	69.8%	68.2%	+1.6%
HumanEval (代码生成)	0-shot Accuracy	33.5%	29.9%	+3.6%
BBH (困难推理)	3-shot Accuracy	51.7%	46.4%	+5.3%

局限与改进

作者明确指出了几个局限性。首先，本文主要聚焦于深度扩展下的训练稳定性，但宽度扩展（如隐藏维度、专家数量、FFN 维度增加）同样会导致训练不稳定，Post-LN 在更宽的设置下可能需要更大的 $\alpha$ 或更强的稳定机制，这方面有待进一步研究。其次，Keel 主要解决 Pre-LN 的有效深度问题，当宽度-深度比已经较高时，这个问题不太明显，Keel 的收益可能不太显著。第三，Keel 通常需要大量训练数据才能有效，在低数据场景下不推荐使用。从我自己的观察来看，本文的实验规模相对受限（最大 3B 参数），与当前主流 LLM（70B-400B）的规模差距较大，Keel 在更大规模模型上的表现需要进一步验证。此外，本文所有实验都在 decoder-only 架构上进行，对 encoder-decoder 或 encoder-only 架构的适用性未被讨论。训练成本方面，虽然 Keel 允许使用更高的学习率加速收敛，但每层多了一个 LayerNorm 操作，对计算效率的影响未被详细分析。

独立分析的弱点

独立分析本文存在以下几个弱点。首先，实验规模偏小，最大模型仅为 3B 参数，而当前 LLM 研究的主流规模是 7B-70B 甚至更大。在这个规模上得出的结论能否直接推广到更大规模模型，存在不确定性。改进方向是在更大规模（如 13B、34B）上进行验证。其次，Keel 的设计引入了额外的超参数 $\alpha = L$，虽然作者提供理论推导支持，但在实际应用中不同架构可能需要不同的 $\alpha$ 值，缺乏自动调整机制。可以考虑设计自适应的 $\alpha$ 学习策略。第三，每层额外的 LayerNorm 操作增加了计算开销，虽然可能被更高的学习率带来的更快收敛所抵消，但缺乏详细的 FLOPS 分析和实际训练时间对比。第四，本文没有与最近的其他深度扩展方法（如 Universal Transformer、层共享等）进行对比，可能遗漏了一些有价值的参考基线。最后，论文中部分图表的可读性较差，特别是 Figure 1 的三个子图，数据点和标签重叠严重，影响了结果的清晰传达。

未来方向

作者提出了几个明确的未来研究方向。首先是宽度扩展下的训练稳定性研究，当模型宽度增加时 Post-LN 可能需要更强的稳定机制。其次是深度传播与序列建模之间的同构关系研究，作者建立了 TTT 框架下的对应关系，认为序列建模中改进长上下文递归的技术可以被迁移到深度传播中，反之亦然。基于本文成果可以延伸的研究方向包括：1）将 Keel 与 Mixture-of-Experts（MoE）结合，探索深度扩展在稀疏模型中的效果；2）将 Highway 式连接的思想应用到 Vision Transformer 等其他模态；3）研究 Keel 架构下深度与宽度的最佳分配比例，建立新的 Scaling Law；4）探索 Keel 在长上下文训练中的表现，特别是其与 Attention Sink 现象的关系；5）将 Keel 的设计思想应用到 State Space Model 等新兴架构中。

复现评估

从复现角度来看，本文提供了相对详细的实验配置。模型配置在 Appendix C 中给出了完整的超参数表，包括隐藏维度、中间维度、注意力头数、KV 头数、初始化方案、Dropout 设置等。训练配置方面，AdamW 优化器参数（$\beta_1=0.9, \beta_2=0.95$）、权重衰减（0.01）、梯度裁剪（最大范数 1.0）、学习率调度（线性预热 2500 步 + 余弦衰减）都有明确说明。然而，本文使用的训练数据为内部数据集，不公开可用，这对独立复现构成了主要障碍。FineWeb-EDU 数据集的实验可以部分复现，但主实验（1T token）无法复现。代码方面，论文未提供开源代码或模型权重，增加了复现难度。算力方面，512 层 3B 参数模型的训练需要大量 GPU 资源，对一般研究团队来说门槛较高。建议作者开源 Keel 的实现代码和预训练模型。

Post-LayerNorm 回归：稳定、表达力强且支持深度扩展 Post-LayerNorm Is Back: Stable, ExpressivE, and Deep