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

• URL: http://arxiv.org/abs/2608.09417v2
• Date: Tue, 11 Aug 2026 06:00:23 GMT
• Title: Why Post-Norm Transformers Collapse: Attention Amplification and Gradient Repair Failure
• Abstract summary: Deep decoder-only Transformers often replace the original Post-Norm architecture with Pre-Norm variants.<n>We show how causal attention creates high-similarity representations and why training dynamics fail to repair them.
• Score: 8.63727823701713
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• NICE: Scale-Stable Perturbations for Graph Neural Network Explanations via Noise Corruption [67.57304101271096]
  Element-wise Masking (EM) suppresses edge-induced messages toward zero.<n>Noise Corruption (NC) perturbs each message through matched-norm random-direction corruption.<n>NICE, a Noise Corruption-based explanation framework, learns a Restoration Boundary (SRB) under NC-induced uncertainty.
  arXiv  Detail & Related papers  (2026-08-17T03:08:14Z)
• The Structural Origin of Attention Sink: Variance Discrepancy, Super Neurons, and Dimension Disparity [22.040532283910522]
  This work provides a textitmechanistic explanation for the phenomenon.<n>We trace its root to the value aggregation process inherent in self-attention.<n>As a proof of concept, we propose textithead-wise RMSNorm, an architectural modification that stabilizes value aggregation outputs during pre-training.
  arXiv  Detail & Related papers  (2026-05-07T17:28:55Z)
• Monitoring Neural Training with Topology: A Footprint-Predictable Collapse Index [51.56484100374058]
  We propose an online, topology-aware monitor for evolving neural representations that couples Modular Morse Homology Maintenance (MMHM) with a composite Collapse Index (CI)<n>Instead of rebuilding complexes each epoch, we apply sparse edits at a fixed scale and maintain a discrete Morse matching, yielding fast, incremental updates.
  arXiv  Detail & Related papers  (2026-04-28T19:21:21Z)
• Representation Collapse in Machine Translation Through the Lens of Angular Dispersion [12.526791012139883]
  We show that a standard next-token prediction training strategy may lead to overlooked artifacts such as representation collapse.<n>We incorporate an existing regularization method based on angular dispersion and demonstrate empirically that it not only mitigates collapse but also improves translation quality.
  arXiv  Detail & Related papers  (2026-02-19T11:46:38Z)
• Generalizing GNNs with Tokenized Mixture of Experts [75.8310720413187]
  We show that improving stability requires reducing reliance on shift-sensitive features, leaving an irreducible worst-case generalization floor.<n>We propose STEM-GNN, a pretrain-then-finetune framework with a mixture-of-experts encoder for diverse computation paths.<n>Across nine node, link, and graph benchmarks, STEM-GNN achieves a stronger three-way balance, improving robustness to degree/homophily shifts and to feature/edge corruptions while remaining competitive on clean graphs.
  arXiv  Detail & Related papers  (2026-02-09T22:48:30Z)
• A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training [86.64715217940274]
  Outliers function jointly with normalization.<n>Outliers serve more as rescale factors rather than contributors.<n>Outliers can be absorbed into learnable parameters or mitigated via explicit gated rescaling.
  arXiv  Detail & Related papers  (2026-01-30T13:29:45Z)
• Shift Before You Learn: Enabling Low-Rank Representations in Reinforcement Learning [56.87989363424]
  We show that a low-rank structure naturally emerges in the shifted successor measure.<n>We quantify the amount of shift needed for effective low-rank approximation and estimation.
  arXiv  Detail & Related papers  (2025-09-05T15:48:20Z)
• Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation [8.973965016201822]
  Finding the right initialisation for neural networks is crucial to ensure smooth training and good performance.<n>In transformers, the wrong initialisation can lead to one of two failure modes of self-attention layers: rank collapse, where all tokens collapse into similar representations, and entropy collapse, where highly concentrated attention scores lead to instability.<n>Here, we provide an analytical theory of signal propagation through deep transformers with self-attention, layer normalisation, skip connections and gradients.
  arXiv  Detail & Related papers  (2025-05-30T08:18:23Z)
• Initialization Matters: On the Benign Overfitting of Two-Layer ReLU CNN with Fully Trainable Layers [20.25049261035324]
  We extend the analysis to two-layer ReLU convolutional neural networks (CNNs) with fully trainable layers. Our results show that the scaling of the output layer is crucial to the training dynamics. In both settings, we provide nearly matching upper and lower bounds on the test errors.
  arXiv  Detail & Related papers  (2024-10-24T20:15:45Z)
• Generalization vs. Specialization under Concept Shift [12.196508752999797]
  We analyze ridge regression under concept shift.<n>We derive an exact expression for prediction risk in the thermodynamic limit.<n>Our experiments on MNIST and FashionMNIST suggest that this intriguing behavior is present also in classification problems.
  arXiv  Detail & Related papers  (2024-09-23T22:30:28Z)
• Temporal Difference Learning with Compressed Updates: Error-Feedback meets Reinforcement Learning [47.904127007515925]
  We study a variant of the classical temporal difference (TD) learning algorithm with a perturbed update direction. We prove that compressed TD algorithms, coupled with an error-feedback mechanism used widely in optimization, exhibit the same non-asymptotic approximation guarantees as their counterparts. Notably, these are the first finite-time results in RL that account for general compression operators and error-feedback in tandem with linear function approximation and Markovian sampling.
  arXiv  Detail & Related papers  (2023-01-03T04:09:38Z)
• Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse [11.486545294602697]
  We shed new light on the causes and effects of rank collapse in Transformers. We show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish.
  arXiv  Detail & Related papers  (2022-06-07T09:07:24Z)
• Extreme Memorization via Scale of Initialization [72.78162454173803]
  We construct an experimental setup in which changing the scale of initialization strongly impacts the implicit regularization induced by SGD. We find that the extent and manner in which generalization ability is affected depends on the activation and loss function used. In the case of the homogeneous ReLU activation, we show that this behavior can be attributed to the loss function.
  arXiv  Detail & Related papers  (2020-08-31T04:53:11Z)
• Revisiting Initialization of Neural Networks [72.24615341588846]
  We propose a rigorous estimation of the global curvature of weights across layers by approximating and controlling the norm of their Hessian matrix. Our experiments on Word2Vec and the MNIST/CIFAR image classification tasks confirm that tracking the Hessian norm is a useful diagnostic tool.
  arXiv  Detail & Related papers  (2020-04-20T18:12:56Z)
• The Break-Even Point on Optimization Trajectories of Deep Neural Networks [64.7563588124004]
  We argue for the existence of the "break-even" point on this trajectory. We show that using a large learning rate in the initial phase of training reduces the variance of the gradient. We also show that using a low learning rate results in bad conditioning of the loss surface even for a neural network with batch normalization layers.
  arXiv  Detail & Related papers  (2020-02-21T22:55:51Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.