Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Normalization in Attention Dynamics

Nikita Karagodin, Shu Ge, Yury Polyanskiy, Philippe Rigollet

TL;DR

The work reframes normalization in deep transformers as speed-regulation on token directions, modeling dynamics on the unit sphere via a unified interacting-particle ODE across six schemes: Post-LN, Pre-LN, Mix-LN, Peri-LN, nGPT, and sqrt-scaling. It demonstrates that, despite differing speed controls, all schemes share a common velocity field, enabling a rigorous analysis of asymptotic clustering and velocity evolution. The key findings show that Peri-LN and tunable nGPT can achieve favorable initial and terminal dynamics, mitigating representation collapse while sustaining meaningful updates across layers. These insights provide a principled basis for comparing normalization schemes and guide architectural choices to balance depth, stability, and expressive capacity in large transformer models.

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.

Normalization in Attention Dynamics

TL;DR

The work reframes normalization in deep transformers as speed-regulation on token directions, modeling dynamics on the unit sphere via a unified interacting-particle ODE across six schemes: Post-LN, Pre-LN, Mix-LN, Peri-LN, nGPT, and sqrt-scaling. It demonstrates that, despite differing speed controls, all schemes share a common velocity field, enabling a rigorous analysis of asymptotic clustering and velocity evolution. The key findings show that Peri-LN and tunable nGPT can achieve favorable initial and terminal dynamics, mitigating representation collapse while sustaining meaningful updates across layers. These insights provide a principled basis for comparing normalization schemes and guide architectural choices to balance depth, stability, and expressive capacity in large transformer models.

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.
Paper Structure (23 sections, 8 theorems, 99 equations, 5 figures, 3 tables)

This paper contains 23 sections, 8 theorems, 99 equations, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work.
  3. Our contributions.
  4. Normalized Attention Dynamics
  5. Attention
  6. Normalization.
  7. Speed regulation formulation
  8. Asymptotic clustering
  9. Initial and terminal token velocities
  10. Prelude: Symmetric initialization
  11. Initial velocity
  12. Terminal velocity
  13. Limitations
  14. Theoretical limitations
  15. Practical limitations
  16. ...and 8 more sections

Key Result

Theorem 3.1

Consider the normalized attention dynamics NA with $Q = K = V = I_d$. Then for uniformly sampled initializations $\Theta(0) \in ({\mathcal{S}}^{d-1})^{\otimes n}$Post-LN, nGPT, sqrt-scaling cluster asymptotically whereas for a standard Gaussian sample of $X(0):= r(0)\cdot \Theta(0)$ with $\Theta(0) \in ({\mathcal{S}}^{d-1})^{\otimes n}, r(0) \in \mathbb{R}^{\otimes n}$ for Pre-LN, Mix-LN, Peri-LN

Figures (5)

  • Figure 1: Normalization layer placements in various architectures.
  • Figure 2: (a) Evolution of cosine similarity $\gamma(t)$, (b) its speed $\dot{\gamma}(t)$ over time, (c) phase-plot of $\dot{\gamma}(t)$ vs. $\gamma(t)$, for introduced normalization strategies. Here nGPT has $\alpha_t \equiv 1$, to showcase the significance of that parameter. Pre-LN and Peri-LN are the last to converge, mitigating representation collapse. On the other hand, Post-LN, nGPT and Peri-LN move faster in early layers, effectively utilizing them. In the phase-plot (c) we see how at the same position the speed is defined by a known speed control parameter, ranking different methods.
  • Figure 3: Convergence in nGPT from orthogonal initialization for different choices of $\alpha_t$ -- constant, root, linear, combination of linear and constant with weights $\sin(4t)$ and $\cos(4t$).
  • Figure 4: Evolution of average cosine similarity $\gamma(t)$ with 90% confidence interval with randomly initialized weights (Kaiming init), $d = 512, n_{\textrm{heads}}=1, \beta = \sqrt{d}, d>n$ and random initial $X$. We set $\alpha_t \equiv 1$ for nGPT. We see that Peri-LN and nGPT initially move faster, and that Post-LN and nGPT eventually collapse tokens faster than Pre-LN and Peri-LN. See Appendix \ref{['ap:pics']} for more studies, including multi-head, untied weights and more.
  • Figure 5: Evolution of average cosine similarity for tokens under the pure attention update.

Theorems & Definitions (15)

  • Theorem 3.1
  • Corollary 3.2
  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3
  • proof
  • Remark
  • Lemma D.1
  • proof
  • Proposition D.2
  • ...and 5 more