Normalization in Attention Dynamics

Nikita Karagodin, Shu Ge, Yury Polyanskiy, Philippe Rigollet

Keywords: Transformers, Self-Attention, normalization, continuous-time interacting particle systems, clustering, representation collapse

TL;DR: Theoretical study of impact of normalization layers in evolution of tokens representations as they propagate through layers of a transformer.

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**, and **LN-scaling**---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.

Primary Area: Theory (e.g., control theory, learning theory, algorithmic game theory)

Submission Number: 18290