A Mechanistic Diagnostic of Rank Collapse in Post-Norm Decoder Transformers
2026-08-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study why a certain type of Transformer model (Post-Norm deep decoder-only) struggles during training, often leading to 'rank collapse,' where token representations become too similar. They find that initially, the model’s causal attention makes tokens more alike, and later, gradients shrink so much that the network can’t fix this problem. Their analysis shows that a collapsed network ends up predicting based mainly on token frequency, which limits performance. Experiments confirm their theoretical findings about how similarity and gradients behave in these models during training.
Post-NormPre-NormTransformercausal attentionrank collapsetoken similaritygradient vanishingRMSNormSwiGLUfrequency distribution
Authors
Xingjian Wang, Qingyu Han, Xiaodong Luo, Yin Zhang
Abstract
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.