Clustering in Deep Stochastic Transformers

TL;DR

This work analyzes token dynamics in deep Transformers under intrinsic stochasticity from standard random initialization of value matrices. By deriving a diffusion limit with RMS normalization, it shows that initialization noise qualitatively alters clustering behavior, permitting antipodal configurations and enabling explicit phase transitions governed by token dimension and attention temperature. The key contributions include a sphere-valued SDE limit for deep stochastic Transformers, an explicit two-token phase boundary, and a noise-induced clustering regime, all corroborated by numerical experiments. The findings indicate initialization noise is a structural factor in signal propagation and clustering, with practical implications for training stability and representation diversity in deep attention stacks.

Abstract

Transformers have revolutionized deep learning across various domains but understanding the precise token dynamics remains a theoretical challenge. Existing theories of deep Transformers with layer normalization typically predict that tokens cluster to a single point; however, these results rely on deterministic weight assumptions, which fail to capture the standard initialization scheme in Transformers. In this work, we show that accounting for the intrinsic stochasticity of random initialization alters this picture. More precisely, we analyze deep Transformers where noise arises from the random initialization of value matrices. Under diffusion scaling and token-wise RMS normalization, we prove that, as the number of Transformer layers goes to infinity, the discrete token dynamics converge to an interacting-particle system on the sphere where tokens are driven by a \emph{common} matrix-valued Brownian noise. In this limit, we show that initialization noise prevents the collapse to a single cluster predicted by deterministic models. For two tokens, we prove a phase transition governed by the interaction strength and the token dimension: unlike deterministic attention flows, antipodal configurations become attracting with positive probability. Numerical experiments confirm the predicted transition, reveal that antipodal formations persist for more than two tokens, and demonstrate that suppressing the intrinsic noise degrades accuracy.

Figures (9)

• Figure 1: Effect of RMS normalization on attention field. Left: unnormalized $A^{\mathrm{U}}_\beta(X, \boldsymbol{X})$. Right: tangent projection $P^{\perp}_X[A^{\mathrm{U}}_\beta(X, \boldsymbol{X})]$ constrains flow to sphere. Here $\bm{X}=(X, X^1, X^2)$ and $\beta=1$, $Q=K=\operatorname{Id}(2)$.
• Figure 2: Effect of value-matrix initialization on CIFAR-10 validation accuracy. Constant Weights (magenta) correspond to identity initialization of the matrices $V$ shared across all layers, while Random Weights (violet) correspond to independent Gaussian initialization of $V$ in each layer. Curves show the mean validation accuracy over five independent training runs; shaded regions indicate the min. and max. accuracy across runs at each epoch.
• Figure 3: Noise-induced phase transition for two tokens. Empirical probability of antipodal convergence as a function of $\beta$ for dimensions $d\in\{4,\ldots,10\}$. Each point is estimated from $40{,}000$ independent trajectories of the discrete dynamics \ref{['eq:discrete_time']} with $L=100$, the time horizon $T=500$, and $\sigma=1$. The dashed curve shows the theoretical phase boundary $\beta_c(d)=\tfrac{1}{2}\cosh^{-1}(d-2)$.
• Figure 4: Noise-induced clustering threshold. Dominant clustering outcome of the discrete hybrid dynamics in the $(\beta,\varepsilon)$ plane. Colors encode the maximum of the empirical fractions of single-cluster (magenta) and antipodal (violet) configurations. Dashed curve $\varepsilon=\sqrt{2e^{-\beta}}$: theoretical threshold from Prop. \ref{['lem:hybrid_clustering']}.
• Figure 5: One layer of the Deep Transformer with added normalization and diffusion scaling ($\sim 1/\sqrt{L}$). The diagram highlights that self-attention acts globally across all tokens. The residual connection is shown with a dashed outline to indicate that it represents a mode of connectivity rather than an operation.

Theorems & Definitions (20)

• Theorem 1: Continuous-time limit of normalized attention
• Lemma 1
• Lemma 2
• Lemma 3
• Definition 1
• Theorem 2
• Proposition 1: Noise‑induced clustering transition
• Definition 2
• proof : Proof of Theorem \ref{['thm:ctl_transformer']}
• proof : Proof of Lemma \ref{['lemma:sphere_supported']}