Normalization in Attention Dynamics

AI summary

Overview

• Research area: Theory of deep learning — specifically the mathematical analysis of transformer attention as an interacting particle system, and how layer-normalization placement shapes the geometry of token representations.
• Technical level: Advanced. The paper is built on Riemannian geometry of the sphere, continuous-time ODE limits of residual networks, gradient-flow arguments (Łojasiewicz-type convergence), and random-matrix / high-dimensional probability bounds.
• Scope: The paper gives a unified dynamical-systems account of six normalization schemes by recasting each as a different speed regulation of the same attention-driven particle flow on the sphere.

What This Paper Is About

Transformers apply normalization layers in different positions (before attention, after attention, before and after, and so on), and these choices quietly change how token representations evolve across depth — sometimes making deep layers nearly useless, sometimes causing representations to collapse together. The authors ask how these schemes influence deep representations, and answer by showing that once you look only at the direction of each token, every scheme obeys the same attention-driven motion on the sphere, differing only in a per-token speed factor. Comparing those speed factors across schemes yields a principled ranking, with Peri-LN emerging as the strongest overall choice.

Key Contributions

1. A unified "speed regulation" formulation. The paper rewrites Post-LN, Pre-LN, Mix-LN, Peri-LN, nGPT and sqrt-scaling as a single normalized attention dynamics (NA), where each scheme contributes only a scalar multiplier (a speed factor) to a common velocity field; RMS normalization of a token is written as Norm(x) = x/||x||.
2. An extension of asymptotic clustering beyond Post-LN. The authors generalize the Post-LN clustering theorem of Geshkovski et al. (2025) to every scheme in their framework, proving convergence of trajectories to a critical point and giving conditions under which tokens synchronize to one cluster.
3. A quantitative comparison of initial and terminal velocities. Using a symmetric orthogonal initialization, they derive closed-form initial velocities and terminal decay rates for all six schemes, then sharpen the early-layer picture with a high-dimensional bound on the attention vector.
4. Identification of Peri-LN as a particularly effective scheme. Combining the initial-velocity and terminal-velocity analyses, Peri-LN is the only scheme that behaves well at both ends — order-one movement in early layers and slow polynomial convergence in deep layers — while nGPT offers explicit control through its trainable parameter α_t.

Main Findings

• Normalization acts as speed regulation, not as a change of direction. Tracking only directions θ_1, …, θ_n on the sphere S^{d−1}, all schemes follow θ̇_j(t) = (1/s_j(t)) P_{θ_j(t)} A_j^t(Θ(t)), i.e. the same projected attention field scaled by a scheme-specific factor s_j(t). The speed factor does change the trajectories, since it depends on time and position.
• Radial behavior separates the schemes. Post-LN, nGPT and sqrt-scaling have ṙ_j(t) = 0 (tokens pinned to the sphere). Pre-LN has ṙ_j(t) = ⟨θ_j(t), A_j^t(Θ(t))⟩, Mix-LN switches that on only for t > τ, Peri-LN has ṙ_j(t) = ⟨θ_j(t), A_j^t(Θ(t))⟩ / ||A_j^t(Θ(t))||, and nGPT again gives ṙ_j(t) = 0. Pre-LN and Peri-LN are regulated not by an explicit mechanism but by token magnitude.
• Asymptotic clustering holds more broadly than Post-LN (Theorem 3.1). With Q = K = V = I_d and uniformly sampled initial directions on (S^{d−1})^{⊗n}, Post-LN, nGPT and sqrt-scaling synchronize to one cluster with probability 1. For Pre-LN, Mix-LN and Peri-LN, started from a standard Gaussian X(0) = r(0) · Θ(0), the probability is 1 that either the tokens synchronize to one cluster or min_{j∈[n]} liminf_{t→∞} ṙ_j(t) = 0. The result also holds for Q^t = Q, K^t = K, V^t = V = Q^⊤K = K^⊤Q.
• Unconditional synchronization under a simple condition (Corollary 3.2). For Pre-LN and Peri-LN with n ≤ e^β, synchronization is unconditional, following from the lower bound ṙ_j ≥ (1/(n e^β))(e^β − (n − 1)) ≥ 1/(n e^β).
• Terminal velocities differ qualitatively (Theorem 4.1). From the symmetric orthogonal initialization, initial velocities are: 2/(e^β + n − 1) for Post-LN, 2/(r_0(e^β + n − 1)) for Pre-LN, 2/(e^β + n − 1) for Mix-LN, 2/(r_0 sqrt(e^{2β} + n − 1)) for Peri-LN, 2α_0/sqrt(e^{2β} + n − 1) for nGPT, and 2/(e^β + n − 1) for sqrt-scaling. Terminal behavior: C e^{−2t} for Post-LN, C/t³ for Pre-LN, Mix-LN and Peri-LN, C α_t e^{−2∫^t α_s ds} for nGPT, and C e^{−4√t}/√t for sqrt-scaling. Pre-LN, Mix-LN and nGPT (with constant α_t) therefore show a polynomial slowdown, whereas the others decay faster — meaning they use intermediate layers more effectively and resist representation collapse.
• The temperature β exponentially damps the initial velocity, which the authors read as a suggestion that initializing Q and K with smaller magnitudes in early layers may be beneficial.
• Early layers move much more under Peri-LN and nGPT (Theorem 4.2). With max{||Q^⊤K||_op, ||V||_op} ≤ 1, β = 1, i.i.d. uniform initial directions, and e^{√d} ≥ n log n ≥ d, the attention vector obeys ||A_j(0)|| ≤ C(√(log n / n) + log n / d) with probability 1 − n^{−C}. Because Peri-LN and nGPT damp the angular velocity by a factor proportional to ||A_j(0)||, their first-layer angular displacement exceeds that of Post-LN, Pre-LN, Mix-LN and sqrt-scaling by a factor Ω(min(d/log n, √(n/log n))). Post-LN and Pre-LN advance more slowly, with step sizes on the order of O(log n / d).
• Clustering rates in a local cone (Theorem 4.3). For tokens initialized in a cone with ⟨θ_j(0), θ_k(0)⟩ ≥ 1 − δ and δ < 1/(100 n² β²), with V = I_d and ||Q^⊤K|| ≤ 1: radial growth satisfies r_k(t) ≥ (1 − δ)t for both Pre-LN and Peri-LN, and the intra-cluster variance Var(t) evolves as d/dt Var(t) = −Θ(Var(t)) for Post-LN, −Θ(Var(t)/t) for Pre-LN, −Θ(Var(t)/t) for Peri-LN, −Θ(Var(t)/α_t) for nGPT, −Θ(Var(t)/t) for Mix-LN, and −Θ(Var(t)/√t) for sqrt-scaling.
• Linear versus diffusive radial growth matters. In the aligned regime studied here the token norm grows linearly, r_j(t) ∼ t, whereas empirical work on randomly initialized transformers reports r_j(t) ∼ √t; because the angular update is scaled by 1/r_j(t), linear growth turns exponential clustering into polynomial decay.
• Theory and experiment line up. Figure 2 plots the ODE-based evolution of γ(t) with β = 5 and n = 256; Figure 4 shows average cosine similarity with a 90% confidence interval under randomly initialized weights (Kaiming init), d = 512, n_heads = 1, β = √d, d > n, random initial X, and α_t ≡ 1 for nGPT. Peri-LN and nGPT move faster initially, while Post-LN and nGPT eventually collapse tokens faster than Pre-LN and Peri-LN.
• The nGPT parameter α_t is powerful. It influences both initial and terminal velocity, and Figure 3 compares constant, root, linear, and a combination of linear and constant weighted by sin(4t) and cos(4t).

Methodology in Plain English

The authors strip a transformer down to a pure attention stack, dropping MLP layers and setting the query, key and value matrices to the identity (Q = K = V = I_d), following the simplified setup of Geshkovski et al. (2025). Each token embedding is split into a direction θ_k = x_k/||x_k|| on the unit sphere and a magnitude r_k = ||x_k||. Because the final decoding layer is normally preceded by a normalization step, only the directions matter for prediction, so the analysis focuses there.

They then treat residual connections as a forward-Euler discretization of a continuous ODE, replacing X^{t+1} − X^t with Ẋ(t). Writing the update rules of all six schemes in this continuous form reveals that, after projecting onto the tangent space of the sphere with P_x y = y − ⟨y, x⟩x, every scheme yields the same attention-driven velocity field, divided by a scheme-specific speed factor s_j(t).

Convergence is then approached as a modulated gradient flow: when KQ^⊤ = QK^⊤ = V, the flow is a gradient flow for the energy E(Θ) = −Σ_{j,k} e^{β⟨Qθ_k, Kθ_j⟩}, and the authors show the classical Łojasiewicz argument still applies. To compare schemes quantitatively they use two tractable configurations: a symmetric orthogonal initialization (all pairwise inner products zero, r_j(0) = r_0), which collapses the dynamics to two scalars γ(t) and r(t), and a "pre-clustered" local cone, which isolates a single well-formed cluster. These are complemented by high-dimensional probability bounds on the attention vector and by simulations.

Why This Matters

Impact on research. Prior theoretical work on attention dynamics focused almost entirely on Post-LN, which has become the standard paradigm for analysis. This paper shows the same machinery extends to Pre-LN, Mix-LN, Peri-LN, nGPT and sqrt-scaling, and shifts attention from gradient stability (the usual empirical/theoretical focus) to the forward evolution and expressive structure of representations. It also supplies concrete analytic criteria — large initial velocity plus slow terminal velocity — for judging a normalization scheme.

Real-world applications:

• Designing normalization placement in large language models. Pre-LN is the default for GPT (Radford et al., 2019) and LLaMA (Touvron, H. et al, 2023), while Peri-LN has been reported to be used in Gemma-3 (Gemma Team et al., 2025); the analysis ranks these choices on measurable dynamical criteria.
• Training very deep transformer stacks. The paper's results bear directly on the curse of depth identified by Sun et al. (2025) and Gromov et al. (2025), where deep layers degenerate into near-identity transformations, and on representation collapse as a barrier to increasing depth.
• Pruning or skipping deep layers for inference cost. The observation that deep layers can be pruned with minimal impact on performance (Muralidharan et al., 2024; Siddiqui et al., 2024) is exactly the pathology the terminal-velocity analysis quantifies.
• Choosing scaling and initialization tricks. The results speak to residual rescaling by √depth (sqrt-scaling, following Noci et al., 2022), to explicitly normalizing tokens onto the unit sphere (nGPT, Loshchilov et al., 2025), and to the temperature β used in the attention logits.

Industry relevance. Normalization placement is a low-cost architectural decision with outsized effects on training stability and depth scaling, so a principled basis for comparing options — and the identification of Peri-LN as effective — is directly actionable for teams building or adapting deep transformer models.

Future Directions

• Explicit convergence and metastability rates. Theorem 3.1 proves convergence but relies on transforming the system to a compact autonomous frame, giving neither an explicit rate nor metastability guarantees; empirical work shows tokens can remain trapped in metastable states for extended periods before clustering.
• Characterizing the intermediate regime. Only initial and terminal speeds are bounded. Comparing two flows in general — even when one has a larger speed-control factor — remains an open problem, since a flow can locally move faster despite a larger factor.
• Relaxing the weight assumptions. The analysis assumes restrictive conditions on Q, K and V; for instance it predicts linear norm growth in Pre-LN when the matrices are tied, whereas empirical work reports a √t trend at initialization with random weights. Reconciling these gaps calls for a stochastic analysis.
• Adding MLP layers and gradient flow. MLP layers are omitted, and optimization pathologies such as exploding updates in Pre-LN are out of scope because gradient propagation is not modeled. The authors state that a companion paper dedicated to gradient analysis is in preparation.
• Empirical verification. The paper does not give a specific model architecture to train and validate, which currently limits the direct architectural recommendations it can offer.

Target Audience

This paper is aimed at machine learning theorists and mathematically inclined architecture researchers — readers comfortable with ODEs, Riemannian gradients and high-dimensional probability. It is also valuable for practitioners working on deep transformer training who want a principled account of why normalization placement matters, and for anyone following the interacting-particle line of work on transformers initiated by Geshkovski et al. (2023) and Sander et al. (2022). Readers seeking empirical benchmarks or a ready-to-train architecture will find that the paper explicitly defers both.

Authors’ abstract

We study the effect of normalization schemes on token representations in deep transformers. Modeling their evolution as interacting particles on the sphere, we show that normalization acts as a form of speed regulation. This perspective enables a unified analysis of several schemes -- including Post-LN, Pre-LN, Mix-LN, Peri-LN, nGPT -- revealing how they influence clustering dynamics and representation collapse. Our framework clarifies how different schemes shape token representations across layers and provides a principled basis for comparing them, identifying Peri-LN as a particularly effective choice.