arXiv:2608.09417cs.LG2026-08

揭示后归一化变压器崩溃的根源:注意力放大相似性,梯度修复失效。

Why Post-Norm Transformers Collapse: Attention Amplification and Gradient Repair Failure

论文配图:Why Post-Norm Transformers Collapse: Attention Amplification and Gradient Repair Failure
图 1 · 摘自论文原文
  • 用词元相似度分析两阶段崩溃机制:初始化时注意力放大相似性,训练中梯度几何级衰减。
  • 48层模型实验验证:初始相似度上升、后期梯度收缩,崩溃状态损失趋近预测值。
  • 适合研究模型训练稳定性、注意力机制与归一化设计的读者参考。

深度解码器仅使用后归一化(Post-Norm)的Transformer在常规初始化下对预热和学习率极为敏感。尽管已有研究指出秩坍缩与梯度消失是相关症状,但因果注意力如何生成高相似性表示,以及训练动态为何无法修复这些现象仍不清楚。本文通过将词元相似度作为标量状态变量,提出两阶段分析:首先,在初始化时,因果注意力近似为前缀平均算子,随深度增加词元相似度;而SwiGLU分支仅起较小阻尼作用。其次,当训练进入高相似度状态后,预归一化残差范数增长使RMSNorm反向传播因子收缩,在温和条件下,早期层梯度呈几何级衰减。补充结果表明,坍缩网络的最佳预测指标是频率分布,其损失存在相对较高的基线,且坍缩层梯度在频率分布上完全消失。在C4数据集上训练的48层解码器仅模型实验验证了预测的初始化相似度增长和坍缩时梯度收缩,并显示崩溃运行始终接近预测的频率损失。这些结果区分了前向相似性放大与后向修复失败,同时刻画了坍缩网络的行为。

原文摘要 · Abstract (English)

Deep decoder-only Transformers often replace the original Post-Norm architecture with Pre-Norm variants because Post-Norm training is highly sensitive to warmup and learning rate under conventional initialization schemes. Although prior work has identified rank collapse and gradient vanishing as related symptoms, it remains poorly understood how causal attention creates high-similarity representations and why training dynamics fail to repair them. We give a two-stage analysis of Post-Norm rank collapse using token similarity as a scalar state variable. First, at initialization, causal attention acts approximately as a prefix-averaging operator that increases token similarity across depth, while the SwiGLU branch contributes only a smaller damping effect. Second, once training enters a high-similarity regime, growth of pre-normalization residual norms makes the RMSNorm backward factor contractive; under mild conditions, gradients to earlier layers decay geometrically. As a complementary result, we characterize the properties of a collapsed network: its best predictor is frequency distribution with relatively high loss floor, and gradients in collapsed layers vanish at frequency distribution. Experiments on 48-layer decoder-only Transformers trained on C4 dataset match the predicted initialization-time similarity growth and collapse-time gradient contraction, and show that collapsed runs stay near the predicted frequency loss. Together, these results distinguish the forward similarity amplification and backward repair incapacity in Post-Norm collapse, while also characterizing the behavior of collapsed networks.

Transformer训练崩溃注意力机制梯度消失

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。