Post-LayerNorm Is Back: Stable, ExpressivE, and Deep

📄 arXiv: 2601.19895v1 📥 PDF

作者: Chen Chen, Lai Wei

分类: cs.LG, cs.CL

发布日期: 2026-01-27

💡 一句话要点

Keel：基于Highway连接的Post-LN Transformer，实现深度LLM的稳定训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 深度学习 Transformer Post-LayerNorm Highway连接 梯度消失 语言模型 深度扩展 模型训练

📋 核心要点

  1. 现有Transformer架构在极深层训练时面临梯度消失问题，限制了模型深度的扩展和表达能力。
  2. Keel通过引入Highway连接替代ResNet风格的残差连接，从而保持梯度在残差分支中的流动，避免信号消失。
  3. 实验表明，Keel能够在超过1000层的深度下稳定训练，并在困惑度和深度扩展方面优于Pre-LN Transformer。

📝 摘要（中文）

大型语言模型（LLM）的扩展正面临瓶颈。模型宽度增加的回报递减，而扩展上下文长度并不能提高根本的表达能力。相比之下，深度扩展在理论上提供了更优越的表达能力，但当前的Transformer架构难以在极端深度下进行可靠的训练。本文重新审视了Post-LayerNorm（Post-LN）公式，其在规模上的不稳定性导致其被现代LLM中的Pre-LN所取代。我们表明，Post-LN的主要失效模式源于ResNet风格的残差路径，这会在深度网络中引入梯度消失。我们提出了Keel，一种Post-LN Transformer，它用Highway风格的连接取代了这种残差路径。这种修改保留了通过残差分支的梯度流，防止了信号从顶层到低层的消失。与先前的方法不同，Keel能够在极端深度下进行稳定训练，而无需专门的初始化或复杂的优化技巧。Keel在超过1000层的深度下进行稳健的训练，并且在困惑度和深度扩展特性方面始终优于Pre-LN。这些发现表明，Post-LN与Highway风格的连接相结合，为构建深度可扩展的LLM提供了一个简单而有效的基础，为未来无限深度架构开辟了可能性。

🔬 方法详解

问题定义：现有Transformer架构，特别是采用Post-LayerNorm的架构，在深度扩展时面临训练不稳定的问题。ResNet风格的残差连接导致梯度在深层网络中消失，限制了模型深度的进一步增加，从而影响了模型的表达能力。现有方法通常需要复杂的初始化或优化技巧来缓解这个问题，但效果有限，且增加了训练的复杂性。

核心思路：Keel的核心思路是用Highway连接替换Post-LN Transformer中的ResNet风格残差连接。Highway连接允许信息以可控的方式通过，从而更好地保持梯度在网络中的流动，避免梯度消失。这种设计旨在解决Post-LN架构在深度扩展时的训练不稳定问题，而无需引入复杂的优化技巧。

技术框架：Keel采用标准的Transformer架构，但关键区别在于LayerNorm的位置和残差连接的类型。具体来说，Keel使用Post-LayerNorm，并将传统的ResNet风格残差连接替换为Highway连接。Highway连接包含一个门控机制，控制有多少信息从输入直接传递到输出，有多少信息通过残差分支。整个架构保持了Transformer的基本结构，包括自注意力机制和前馈网络。

关键创新：Keel的关键创新在于使用Highway连接替代ResNet风格的残差连接，从而解决了Post-LN Transformer在深度扩展时的梯度消失问题。与现有方法相比，Keel无需复杂的初始化或优化技巧，即可实现深度模型的稳定训练。这种简单的修改显著提高了模型的深度可扩展性。

关键设计：Highway连接的关键设计在于门控机制。设输入为x，残差分支的输出为F(x)，则Highway连接的输出为T * F(x) + (1 - T) * x，其中T是一个门控向量，控制信息通过残差分支的比例。门控向量T通常由sigmoid函数生成，其输入是x的线性变换。这种设计允许网络动态地调整信息通过残差分支和直接连接的比例，从而更好地保持梯度流动。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Keel在深度超过1000层的Transformer模型上实现了稳定训练，并在困惑度指标上优于传统的Pre-LN Transformer。实验结果表明，Keel能够有效解决Post-LN架构的梯度消失问题，显著提升深度模型的性能和可扩展性。具体性能数据（如困惑度降低的具体数值）在论文中给出。

🎯 应用场景

Keel的潜在应用领域包括自然语言处理、机器翻译、文本生成等。通过实现深度LLM的稳定训练，Keel可以提升模型在复杂任务中的表现，例如长文本理解、多轮对话等。此外，Keel的设计思想也可以应用于其他深度学习模型，例如计算机视觉领域的深度卷积神经网络。

📄 摘要（原文）

Large language model (LLM) scaling is hitting a wall. Widening models yields diminishing returns, and extending context length does not improve fundamental expressivity. In contrast, depth scaling offers theoretically superior expressivity, yet current Transformer architectures struggle to train reliably at extreme depths. We revisit the Post-LayerNorm (Post-LN) formulation, whose instability at scale caused its replacement by Pre-LN in modern LLMs. We show that the central failure mode of Post-LN arises from the ResNet-style residual pathway, which introduces gradient vanishing in deep networks. We present Keel, a Post-LN Transformer that replaces this residual path with a Highway-style connection. This modification preserves the gradient flow through the residual branch, preventing signal vanishing from the top layers to the bottom. Unlike prior methods, Keel enables stable training at extreme depths without requiring specialized initialization or complex optimization tricks. Keel trains robustly at depths exceeding 1000 layers and consistently improves perplexity and depth-scaling characteristics over Pre-LN. These findings indicate that Post-LN, when paired with a Highway-style connection, provides a simple and effective foundation for building deeply scalable LLMs, opening the possibility for future infinite-depth architectures.