Peri-LN: Revisiting Normalization Layer in the Transformer Architecture

📄 arXiv: 2502.02732v3 📥 PDF

作者: Jeonghoon Kim, Byeongchan Lee, Cheonbok Park, Yeontaek Oh, Beomjun Kim, Taehwan Yoo, Seongjin Shin, Dongyoon Han, Jinwoo Shin, Kang Min Yoo

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-02-04 (更新: 2025-06-06)

备注: ICML2025 Camera-ready version


💡 一句话要点

提出Peri-LN,一种新型Transformer归一化策略,提升大规模模型训练稳定性和收敛速度。

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

关键词: Transformer 层归一化 Peri-LN 大规模模型 训练稳定性 梯度传播 激活方差 深度学习

📋 核心要点

  1. 现有Pre-LN和Post-LN归一化策略在大规模Transformer训练中存在局限性,影响训练稳定性和收敛速度。
  2. 论文提出Peri-LN策略,将归一化层置于Transformer子层周围,旨在平衡激活方差增长和稳定梯度流动。
  3. 实验结果表明,Peri-LN在高达32亿参数的Transformer模型上,实现了更稳定的训练和更快的收敛。

📝 摘要(中文)

本文针对Transformer架构中层归一化(LN)策略选择的难题,提出了一个全面的分析框架,用于理解不同LN策略如何影响大规模Transformer的训练动态。尽管Pre-LN和Post-LN长期占据主导地位,但一些开源模型已开始采用一种新的策略,即Peri-LN,它将归一化层置于子层周围。本文深入分析了不同LN策略的行为,揭示了它们如何塑造激活方差和梯度传播。通过在高达32亿参数的Transformer上进行的大量实验,验证了理论分析,结果表明Peri-LN能够实现更平衡的方差增长、更稳定的梯度流动和收敛稳定性。研究结果表明,Peri-LN值得在大规模Transformer架构中得到更广泛的考虑。

🔬 方法详解

问题定义:现有Transformer架构中的层归一化(LN)策略,如Pre-LN和Post-LN,在大规模模型训练时存在训练不稳定、收敛速度慢等问题。选择合适的LN策略仍然是一个挑战,尤其是在大型语言模型(LLM)中。

核心思路:论文的核心思路是提出一种新的归一化策略,即Peri-LN,它将归一化层放置在Transformer子层的周围,而不是像Pre-LN那样放在子层之前,或者像Post-LN那样放在子层之后。这种设计旨在更有效地控制激活方差的增长,并改善梯度在网络中的传播,从而提高训练的稳定性和收敛速度。

技术框架:Peri-LN的核心在于重新设计了LN层的位置。在标准的Transformer块中,通常包含一个自注意力层和一个前馈网络层。Peri-LN将LN层分别放置在自注意力层和前馈网络层之后,相当于对每个子层的输出进行归一化。整体架构仍然是标准的Transformer结构,只是LN层的位置发生了改变。

关键创新:Peri-LN的关键创新在于其归一化层的位置。与Pre-LN和Post-LN相比,Peri-LN能够更直接地控制每个子层的输出方差,从而避免了激活值过大或过小的问题。这种设计有助于保持梯度在网络中的稳定流动,从而提高训练的稳定性和收敛速度。

关键设计:Peri-LN的关键设计在于将LN层放置在每个子层(自注意力层和前馈网络层)的输出之后。具体的参数设置和损失函数与标准的Transformer模型相同。网络结构也与标准的Transformer模型相同,只是LN层的位置发生了改变。论文通过理论分析和实验验证了Peri-LN的有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在高达32亿参数的Transformer模型上进行了实验,结果表明Peri-LN能够实现更平衡的方差增长、更稳定的梯度流动和收敛稳定性。相较于传统的Pre-LN和Post-LN,Peri-LN在训练过程中表现出更强的鲁棒性和更快的收敛速度,验证了其在大规模模型训练中的有效性。

🎯 应用场景

Peri-LN具有广泛的应用前景,可用于改进各种基于Transformer的深度学习模型,尤其是在大规模语言模型、机器翻译、语音识别等领域。该方法能够提升模型的训练稳定性和收敛速度,降低训练成本,并有可能提高模型的最终性能。未来,Peri-LN有望成为大规模Transformer模型训练的标准配置。

📄 摘要(原文)

Selecting a layer normalization (LN) strategy that stabilizes training and speeds convergence in Transformers remains difficult, even for today's large language models (LLM). We present a comprehensive analytical foundation for understanding how different LN strategies influence training dynamics in large-scale Transformers. Until recently, Pre-LN and Post-LN have long dominated practices despite their limitations in large-scale training. However, several open-source models have recently begun silently adopting a third strategy without much explanation. This strategy places normalization layer peripherally around sublayers, a design we term Peri-LN. While Peri-LN has demonstrated promising performance, its precise mechanisms and benefits remain almost unexplored. Our in-depth analysis delineates the distinct behaviors of LN strategies, showing how each placement shapes activation variance and gradient propagation. To validate our theoretical insight, we conduct extensive experiments on Transformers up to $3.2$B parameters, showing that Peri-LN consistently achieves more balanced variance growth, steadier gradient flow, and convergence stability. Our results suggest that Peri-LN warrants broader consideration for large-scale Transformer architectures, providing renewed insights into the optimal placement of LN.