Stabilizing Transformer Training by Preventing Attention Entropy Collapse

Shuangfei Zhai; Tatiana Likhomanenko; Etai Littwin; Dan Busbridge; Jason Ramapuram; Yizhe Zhang; Jiatao Gu; Josh Susskind

Source-linked AI summary

Stabilizing Transformer Training by Preventing Attention Entropy Collapse

Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, Josh Susskind

arXiv:2303.06296v2Mar 11, 2023cs.LGcs.AIcs.CLcs.CVstat.ML

TL;DR

Transformer training can become unstable when attention entropy collapses, but the relationship between collapse and instability remains partly unresolved. The paper studies this dynamic, introduces σReparam to control spectral norms, and evaluates it across several tasks. The method commonly prevents entropy collapse while enabling stable, competitive training with simplified recipes and deep architectures.

Problem
The paper addresses recurring Transformer instability associated with attention entropy collapse and asks how to prevent it and improve training stability.
Method
σReparam reparameterizes weight matrices with spectral normalization and a learned scalar, decoupling spectral-norm updates from weight dimensionality while retaining model capacity.
Results
Across vision, self-supervised learning, translation, speech recognition, and language modeling, σReparam prevents or reduces entropy collapse and supports stable, competitive training, including 16% shorter Vision Transformer training.
Takeaways & Limitations
σReparam provides a broadly applicable way to improve Transformer training stability and robustness while simplifying training recipes and supporting deep architectures.
Takeaways & Limitations
The causal relationship between entropy collapse and training instability remains unclear, and σReparam may still benefit from combination with other techniques.

Abstract

from arXiv · show

Training stability is of great importance to Transformers. In this work, we investigate the training dynamics of Transformers by examining the evolution of the attention layers. In particular, we track the attention entropy for each attention head during the course of training, which is a proxy for model sharpness. We identify a common pattern across different architectures and tasks, where low attention entropy is accompanied by high training instability, which can take the form of oscillating loss or divergence. We denote the pathologically low attention entropy, corresponding to highly concentrated attention scores, as $\textit{entropy collapse}$. As a remedy, we propose $σ$Reparam, a simple and efficient solution where we reparametrize all linear layers with spectral normalization and an additional learned scalar. We demonstrate that $σ$Reparam successfully prevents entropy collapse in the attention layers, promoting more stable training. Additionally, we prove a tight lower bound of the attention entropy, which decreases exponentially fast with the spectral norm of the attention logits, providing additional motivation for our approach. We conduct experiments with $σ$Reparam on image classification, image self-supervised learning, machine translation, speech recognition, and language modeling tasks. We show that $σ$Reparam provides stability and robustness with respect to the choice of hyperparameters, going so far as enabling training (a) a Vision Transformer {to competitive performance} without warmup, weight decay, layer normalization or adaptive optimizers; (b) deep architectures in machine translation and (c) speech recognition to competitive performance without warmup and adaptive optimizers. Code is available at \url{https://github.com/apple/ml-sigma-reparam}.

1 Introduction

The paper identifies attention entropy collapse as a recurring marker of unstable Transformer training and proposes σReparam to control spectral norms and stabilize training across tasks. Experiments report simplified training, improved robustness, and competitive results in vision, translation, speech, and language modeling.

Training instability: Sharp entropy drops can push Hessian sharpness beyond a stability threshold, causing training instability, whereas later interventions may recover below the threshold.The intervention experiments reduce attention temperature by 10× and compare changes during and after warmup.
Training instability: Low attention entropy is associated with slow convergence, loss fluctuations, and, in severe cases, training divergence.The authors monitor attention entropy across query positions, heads, and examples as a stability indicator.
Proposed method: σReparam applies spectral normalization followed by a learned multiplicative scalar to weight matrices, controlling spectral-norm growth without reducing model capacity.The method targets query and key projections through a lower bound linking attention entropy to the spectral norm of attention logits.
Empirical findings: Entropy collapse is commonly observed in baseline models across multiple benchmarks.The evaluation covers image classification, self-supervised learning, machine translation, speech recognition, and language modeling.
Empirical findings: 16% shorter training accompanies equivalent or slightly better Vision Transformer performance after removing pre-LN, warmup, weight decay, and adaptive optimizers.This result is reported against baseline training strategies.
Empirical findings: σReparam stabilizes very deep post-LN translation models up to 100L-100L encoder-decoder layers and simplifies speech-recognition training by removing warmup and adaptive optimization.For language modeling, it remains compatible with causal Transformers and achieves state-of-the-art-competitive results without post-LN.

2 Related Works

The related-work discussion positions σReparam alongside normalization, initialization, weight-reparameterization, and rank-collapse research. It distinguishes entropy collapse from rank collapse and emphasizes that σReparam targets training dynamics without constraining representational capacity.

Normalization and initialization: Entropy collapse can occur even with extensive normalization, while σReparam can operate without specific normalization layers and smooth attention-entropy curves.This contrasts the method with reliance on post-LN or pre-LN configurations.
Normalization and initialization: σReparam complements initialization methods by targeting optimizer- and hyperparameter-driven training dynamics rather than initial conditions alone.The authors describe it as an orthogonal approach to initialization schemes.
Weight reparameterization: Unlike SpectralNorm, σReparam normalizes by spectral norm while retaining a learned scalar, so it changes optimization dynamics without explicitly constraining model capacity.The discussion frames σReparam as a weight reparameterization method related to WeightNorm.
Collapse phenomena: Rank collapse produces rank-1 attention and vanishing query-key gradients, whereas entropy collapse preserves high rank and tends to produce high gradient norms.The paper treats these as distinct Transformer failure patterns.

3 Method

The method links attention entropy to Transformer training stability and controls the spectral norms of attention projections through σReparam. It combines an entropy lower bound with a reparameterization that changes optimization dynamics while preserving representational capacity.

Attention entropy: The attention layer forms row-wise softmax weights from query-key logits and applies them to value projections.
Attention entropy: Attention entropy is monitored as a training-dynamics signal because low entropy accompanies slow convergence, loss fluctuations, and divergence.
Entropy bound: Theorem 3.1 connects attention entropy to the spectral norm of the query-key projection product and provides a tight lower bound.The bound is achievable for some inputs and weights.
Entropy bound: For large σ and T, the minimum attainable entropy behaves like Ω(Tσe^-σ), decreasing exponentially with σ.
σReparam: σReparam reparameterizes linear-layer weights using spectral normalization and a learned scalar, decoupling spectral-norm updates from weight dimensionality.It preserves representational capacity while imposing a different optimization dynamic from spectral normalization alone.
Optimization dynamics: The training analysis relates large ideal-update spectral norms to layer width, motivating explicit control of spectral norms rather than relying only on learning-rate adjustment.The passage notes that large updates can be counterproductive because small learning rates may reduce performance and layer sizes vary.

4 Experiments

Across supervised vision, self-supervised learning, machine translation, speech recognition, and language modeling, σReparam improves training stability while simplifying or strengthening training recipes. Its experiments associate bounded attention entropy and spectral norms with more reliable optimization and competitive performance.

Supervised Image Classification: All configurations in the hyperparameter grid converge with 81.4% (±0.52%) average top-1 accuracy after applying σReparam and removing pre-LN.The result demonstrates robustness to hyperparameter choices in the tested grid.
Supervised Image Classification: σReparam reduces supervised ViT-B/16 training time by 50 epochs while maintaining lower attention spectral norms and smoothly decreasing attention entropy.These dynamics accompany accelerated ImageNet1k test performance.
Self-Supervised Training of Visual Representations: SimCLR exhibits two instability types: large early-layer gradient norms and entropy collapse from growing spectral norms; σReparam protects against both.Frozen Patcher addresses the first type but remains susceptible to the second.
Machine Translation: σReparam bounds attention entropy and resolves divergence in deep 18L-18L and 50L-50L post-LN models, as well as 18L-18L, 50L-50L, and 100L-100L DeepNorm models.For the 50L-50L post-LN model, it also resolves vanishing gradients.
Speech Recognition and Language Modeling: In speech recognition, σReparam removes learning-rate warmup and adaptive optimization while improving post-LN training stability and hyperparameter robustness; in language modeling, it removes all LayerNorms with comparable performance.The speech-recognition result reports stable training and comparable performance without an adaptive optimizer.

5 Conclusion

The work identifies attention entropy collapse as a recurring Transformer failure pattern and presents σReparam as a simple reparameterization that often improves training stability and robustness. The authors note that entropy collapse may not be causally established as the source of instability, and σReparam is not a complete solution.

Attention entropy collapse is commonly observed across Transformer settings and is often associated with training instability.
σReparam effectively addresses entropy collapse and often improves training stability and robustness.
The causal relationship between entropy collapse and Transformer training instability remains unclear.
σReparam is not a panacea, so practical training may still benefit from initialization, feature normalization, and advanced optimizers.

A Proof of Theorem 3.1 and Proposition 3.2

This section develops the entropy minimization argument underlying the attention-entropy bound and describes the temperature and Hessian-based stability analysis. It also relates Hessian eigenvalues to an optimizer-dependent divergence threshold.

The lower bound is tight because inputs and query/key weights exist that attain it.
The entropy minimization proof analyzes a softmax probability vector generated by an attention-logit row under a norm constraint.
A minimizer can be taken to contain negative components, and the proof characterizes it using two distinct values, one positive and one negative.
Inducing entropy collapse: Temperature interventions divide attention logits by a global scalar, allowing entropy collapse to be induced by reducing temperature toward a value much smaller than one.
Eigenvalues of the Hessian: Sharpness is defined as the largest-magnitude Hessian eigenvalue, while explicit Hessian construction is avoided through Hessian-vector products and numerical spectral methods.
The Stability Threshold: Under a local quadratic assumption, exceeding an optimizer-dependent stability threshold causes optimization iterations to diverge.

B.2 Results

These experiments examine Vision Transformer stability under sharp temperature changes, tracking performance, attention dynamics, optimization settings, and leading Hessian singular values. They vary both the intervention epoch and the target temperature.

The experiments vary when a 10× temperature reduction is applied during Vision Transformer training.
The plots track training performance, the first attention projection’s spectral norm, first-block attention entropy, learning rate, temperature, and the five largest Hessian singular values.
Reducing the temperature below 0.15 causes a sharp drop in attention entropy.
Temperature modifications are applied at epoch 10 while the same performance, attention, optimization, and Hessian quantities are plotted.

C Implementation of σReparam

This section describes σReparam’s spectral-normalization implementation and the experimental configurations used to evaluate stability across self-supervised vision and speech recognition. It also reports stable SimCLR training with the original 10-epoch warmup.

Implementation of σReparam: The implementation initializes singular-vector estimates as random unit vectors and the learned spectral norm parameter to one.
Implementation of σReparam: σReparam estimates a weight matrix’s spectral norm with power iteration and rescales the matrix so its effective spectral norm equals a learned scalar.
Implementation of σReparam: One power-iteration step is used by default per gradient update, with no empirical performance difference reported for multiple steps.
SimCLR experiments: SimCLR stability experiments use predefined shared hyperparameters and augmentations including half-strength ColorJitter, random crops, flips, grayscale, blur, tensor conversion, and ImageNet normalization.
SimCLR experiments: σReparam enables stable SimCLR+ViT training with the original 10-epoch warmup, rather than the extended 40-epoch period previously reported as stabilizing.
Speech recognition: The speech-recognition setup uses LibriSpeech and a 36-layer vanilla Transformer encoder trained with CTC loss.

E.2 Training Stability, Robustness and Generalization

In ASR, baseline Transformers show instability, vanishing gradients, and attention entropy collapse under some hyperparameters. σReparam broadens stable training and, with post-LN, combines stability with competitive generalization.

Baseline stability: Post-LN baseline ASR models fail under varied learning rates, warmup, and gradient clipping, with vanishing gradients observed.
σReparam stability: σReparam without LayerNorm yields stable training, accepts a wider hyperparameter range, and avoids attention entropy collapse.
σReparam with post-LN: σReparam with post-LN achieves validation and test performance similar to post-LN, lower training loss, and no vanishing gradients.
Deep ASR: 72-layer post-LN models cannot train, whereas pre-LN, σReparam, and σReparam with post-LN train out of the box with bounded attention entropy.
Training with SGD: With LARS, σReparam remains stable as learning rates change and supports learning rates up to 1, while pre-LN and post-LN diverge.
Ablation: Separate reparameterization of keys, queries, and values is less stable than joint reparameterization for ASR with LARS and no warmup.

F.2 Training Stability of Deep Models

Deep machine-translation Transformers exhibit depth- and normalization-dependent instability, often involving attention entropy collapse. σReparam bounds entropy and resolves several divergence or gradient problems, though the deepest post-LN case remains unstable.

Pre-LN: Pre-LN models converge stably across depths with bounded attention entropy, although performance decreases as depth increases.
Measurement: The MT experiments track encoder self-attention, encoder-decoder cross-attention, and decoder self-attention separately across training and validation.
Post-LN: Post-LN models are stable at 6L-6L, but 18L-18L can diverge across seeds as attention entropy collapses in encoder and cross-attention.
Post-LN: At 50L-50L and 100L-100L, post-LN models cannot train because of vanishing gradients and attention entropy collapse.
σReparam with post-LN: σReparam bounds attention entropy in 18L-18L and 50L-50L post-LN models, resolving divergence and vanishing gradients in the 50L-50L model.
σReparam with DeepNorm: σReparam bounds entropy in 18L-18L, 50L-50L, and 100L-100L DeepNorm models, resolving 100L-100L divergence but yielding inferior performance at 50L-50L and 100L-100L.

G.2 Results

In language modeling, the baseline shows no observed training instability or entropy collapse, while σReparam still supports strong convergence and performance after removing post-LayerNorm.

Baseline behavior: Masked autoregressive attention likely makes entropy collapse less likely, consistent with its absence in the baseline language-model experiments.
σReparam: σReparam applied to all linear layers enables strong convergence and validation/test performance after removing all post-LNs.
Performance: With mild weight decay, σReparam outperforms the baseline on validation and test perplexity.

I Ablations

Ablations show that σReparam’s learned scalar and initialization choices matter, while joint key-query-value reparameterization is preferable for stable ASR training with LARS.

Initialization: Initializing γ as 1 is preferred because computing it from the initialized kernel can produce spectral norms above 1 and cause divergence or no training.
Initialization: ASR training is robust to initialization standard deviation, with larger values performing better and a reported sweet spot of 0.2-0.3.
Reparameterization variants: Separate key, query, and value normalization can lower training loss, but joint reparameterization is more stable for ASR with LARS.
Discussion: Across representative domains, architectures, and losses, entropy collapse accompanies instability and σReparam is reported as compatible with all tested settings.