A Mechanistic Diagnostic of Rank Collapse in Post-Norm Decoder Transformers

2026-08-10

AI Summary

This work addresses the instability and vanishing gradients in Post-Norm decoder-only Transformers caused by rank collapse, a phenomenon whose underlying mechanisms remain poorly understood. The authors model token similarity as a scalar state variable and, for the first time, disentangle the roles of causal attention—amplifying representation similarity during forward propagation—and the lack of corrective capacity during backward propagation in driving rank collapse, analyzing both initialization and training phases. By integrating theoretical analysis, modeling the backward dynamics of RMSNorm, and incorporating the branch effect of SwiGLU alongside a prefix-averaging operator approximation, they characterize the behavioral signatures of collapsing networks. Experiments on 48-layer decoders confirm that similarity grows at initialization, gradients contract upon collapse, and loss converges toward the prediction based on token frequency distribution.

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.

Problem

• rank collapse
• Post-Norm Transformers
• causal attention
• token similarity
• gradient vanishing

Innovation

• rank collapse
• Post-Norm Transformer
• token similarity
• gradient vanishing
• causal attention