🤖 AI

Stability of Transformers under Layer Normalization

This paper presents a principled theoretical and numerical analysis of how different layer normalization placements affect the forward and backward stability of Transformers, deriving explicit bounds on hidden state growth and gradient propagation to guide architectural design and residual scaling for improved training dynamics.

Original authors: Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

Published 2026-08-10

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write a story. You give it a massive brain made of thousands of tiny layers, like a skyscraper with hundreds of floors. Every time the robot reads a word, it passes that information up the elevator, floor by floor, changing its understanding a little bit at each stop. This is how modern AI, known as "Transformers," works. They are the engines behind the smartest chatbots and image generators today. But there's a catch: as these buildings get taller, they become unstable. Sometimes, the information gets so distorted on the way up that the robot starts screaming nonsense numbers, or the math used to teach it gets so wild that the whole system crashes. To stop this, engineers use a safety device called "Layer Normalization." Think of it like a speed bump or a governor on a car engine; it keeps the information from speeding out of control. For years, people just guessed where to put these speed bumps—sometimes before the engine, sometimes after—and hoped for the best. But nobody really knew why one spot worked better than another, or if the robot was secretly building a house of cards that would collapse under its own weight.

This paper is like a group of detectives who decided to stop guessing and start using the laws of physics to solve the mystery. The authors treat the AI's learning process like a journey through time, using a branch of math called "optimal control theory" (which is usually used to figure out the best way to steer a rocket or a drone). They ask: "If we want this robot to learn perfectly, what happens to the information as it travels up the tower?" They discovered that the old way of placing the safety device (called "Pre-LN") is actually a trap. Their math proves that even if the robot tries its absolute hardest to learn, the information will eventually grow so huge that it explodes, leading to a chaotic mess. It's like trying to fill a bucket with a firehose; no matter how hard you try to aim, the water will overflow.

However, they found a better way. By moving the safety device to both the entrance and exit of each floor (a new method they call "Peri-LN"), the information stays calm and controlled. Their equations show that instead of exploding, the data grows in a predictable, gentle line, like a well-behaved plant. They also found a simple trick to make this even better: slowing down the steps the robot takes at each floor. Imagine if the robot took tiny, careful steps instead of giant leaps; this keeps the whole building steady. When they tested this on real AI models, the results matched their math perfectly. The "Peri-LN" models didn't crash, and they learned just as well, if not better, than the old ones. The paper doesn't just say "this feels right"; it provides a mathematical guarantee that this new design keeps the AI stable, offering a clear blueprint for building the next generation of super-smart machines without them falling apart.

Technical Summary: Stability of Transformers under Layer Normalization

Problem Statement

Despite their dominance in modern deep learning, training deep Transformers remains prone to instability, particularly as model depth increases. These instabilities manifest in two critical areas: forward instability (uncontrolled growth of hidden states leading to numerical issues) and backward instability (exploding gradients during backpropagation).

A primary architectural component used to mitigate these issues is Layer Normalization (LN). However, the placement of LN within Transformer blocks has largely evolved through empirical heuristics rather than rigorous theoretical justification. The paper identifies three main configurations:

Post-LN: LN applied after residual connections. Often requires delicate optimization scheduling and exhibits suboptimal performance.
Pre-LN: LN applied before attention and feedforward modules. While standard for its improved training stability over Post-LN, it is empirically observed to produce excessively large hidden states, leading to potential numerical instability and gradient explosion in deep networks.
Peri-LN: A newer configuration applying LN to both the input and output of modules. Empirical studies suggest it offers superior stability and more regular hidden states compared to Pre-LN, but its theoretical properties remain poorly understood.

The core problem addressed is the lack of a principled theoretical framework to explain why these different placements lead to distinct training dynamics and stability profiles, and whether training drives models toward regular solutions or pathological behaviors.

Methodology

The authors propose a novel theoretical framework grounded in discrete-time optimal control theory to analyze Transformer training dynamics.

Formulation as Optimal Control: The Transformer training problem is cast as a discrete-time mean-field control problem. The layer index $i$ represents discrete time, and the hidden state $X_i$ represents the system state. The residual connections act as the velocity of the dynamics.
Hamilton-Jacobi-Bellman (HJB) Analysis: The authors analyze the optimality conditions of the training problem using the HJB partial differential equation (PDE) coupled with mass transport equations.
- They examine the Hamiltonian, which acts as a feedback control mechanism defining the optimal velocity field for hidden states.
- For Pre-LN, they demonstrate that the Hamiltonian is unbounded, leading to ill-posed optimality conditions where hidden states can grow arbitrarily large.
- For Peri-LN, the output normalization constrains the velocity field to an ellipsoid (based on Lemma 1), ensuring the Hamiltonian is bounded and the problem is well-posed.
Gradient Analysis: The paper analyzes backward stability by examining the propagation of gradients through the product of local sensitivities (Jacobian matrices) across layers.
Residual Scaling: The framework is extended to include a time increment parameter $\Delta t < 1$ , effectively scaling the magnitude of residual steps to control signal propagation.

Key Contributions

1. Theoretical Framework for Trained Models

Unlike prior work that often focuses on model initialization or relies on empirical evidence, this framework analyzes trained models at convergence. It provides a systematic assessment of whether training dynamics lead to regular solutions or pathologies.

2. Rigorous Analysis of LN Placements

Pre-LN Instability: The paper proves that for Pre-LN, the optimal solution to the training problem is unbounded in magnitude. Even with weight decay, hidden states can grow exponentially ( $O(e^D)$ ), and gradients grow proportionally with activations, leading to training instability.
Peri-LN Stability: The authors derive explicit bounds showing that Peri-LN guarantees linear growth for the mean absolute value of hidden states ( $O(D)$ ) and quadratic growth for variance ( $O(D^2)$ ). Crucially, the local sensitivity of Peri-LN is shown to be invariant to activation magnitude, preventing gradient explosion even when activations are large.

3. Uncertainty Quantification

The analysis establishes that Peri-LN provides controlled propagation of uncertainty (measured via Wasserstein distance) from input to terminal hidden states, whereas Pre-LN can amplify input uncertainty to unbounded levels.

4. Residual Step Scaling

Guided by the stability analysis, the authors introduce a residual step scaling mechanism ( $\Delta t < 1$ ). They theoretically demonstrate that this scaling:

Sharpens bounds on hidden state growth and output uncertainty.
Mitigates gradient explosion by scaling local sensitivity.
Improves both forward and backward stability without additional computational cost.

Experimental Results

The authors validate their theoretical findings using GPT-2 variants (124M to 1.5B parameters) trained on the OpenWebText dataset.

Training Stability: In experiments with 5 trials, Pre-LN models diverged in 1 run when weight decay was enabled and in 3 runs when weight decay was disabled. In contrast, Peri-LN models remained stable in all 5 trials under both weight decay settings.
Performance: Peri-LN achieved performance comparable to Pre-LN even when hyperparameters were tuned for Pre-LN.
Impact of Scaling: Combining Peri-LN with residual scaling ( $\Delta t = 0.1$ ) consistently yielded the best performance across all metrics (Val Loss, Perplexity, ROUGE, BERT scores) and effectively controlled the growth of hidden state moments (mean absolute value and variance) across layers, aligning with theoretical predictions.

Significance and Claims

The paper claims to provide the first rigorous theoretical justification for the empirical observations regarding Layer Normalization placements in Transformers.

Explanation of Instability: It explains why Pre-LN leads to unbounded hidden states and gradient explosion via the ill-posedness of the underlying HJB PDE.
Justification for Peri-LN: It provides a theoretical basis for the stability of Peri-LN, showing it maintains controlled growth and invariant sensitivity.
Design Guidance: The framework offers a principled method to "sanity-check" new architectural modifications before expensive empirical training.
Practical Improvement: The introduction of residual step scaling is presented as a simple, cost-free modification that significantly enhances stability and performance, validated by experiments showing that Peri-LN with scaling outperforms standard configurations.

The authors emphasize that their work bridges the gap between empirical heuristics and theoretical understanding, offering a continuous-time perspective that yields both theoretical insights and practical architectural improvements.