Stochastic Scaling Limits and Synchronization by Noise in Deep Transformer Models

1. Problem Statement

The paper addresses a gap in the theoretical understanding of Deep Transformer architectures, specifically regarding the joint dynamics of Self-Attention and Multi-Layer Perceptron (MLP) blocks.

• Context: While recent literature has modeled the "residual stream" (the evolution of token representations through layers) as deterministic mean-field interacting particle systems (focusing primarily on attention), these models often neglect the stochastic nature introduced by the MLP blocks or treat them as purely deterministic in the limit.
• Goal: The authors aim to derive a rigorous stochastic scaling limit for finite-depth, finite-width Transformers. They seek to characterize the continuous-time evolution of tokens when both the depth ($L$) and the MLP width ($m$) tend to infinity, while the number of tokens ($N$) is either fixed or large.
• Key Question: How does the randomness inherent in the random initialization of MLP weights affect the global dynamics, and does this noise induce synchronization (clustering) of tokens?

2. Methodology

The authors employ a multi-step mathematical framework combining probability theory, stochastic analysis, and mean-field theory.

A. Model Setup

• Architecture: A Transformer with $L$ layers, $N$ tokens embedded on a $d$-dimensional sphere $S^{d-1}$.
• Dynamics: The update rule involves two steps per layer:
  1. Self-Attention: A deterministic interaction term based on softmax attention.
  2. MLP: A feed-forward block with random weights (Gaussian initialization) and a Lipschitz activation function $\phi$.
• Scaling: The authors introduce specific scaling factors $\alpha_1(L) = T/L$ for attention and $\alpha_2(L) = \sqrt{T/L}$ for the MLP. This scaling ensures that as $L \to \infty$, the attention drift remains $O(1)$ while the cumulative effect of the MLP noise converges to a Brownian motion.

B. Derivation of the Limiting Dynamics

1. Finite $N$ Limit (Deep & Wide):

   • The authors prove that the discrete-time token trajectories converge pathwise to a continuous-time Stochastic Differential Equation (SDE) system.
   • The Limiting SDE: The tokens evolve according to a system of interacting SDEs driven by a common noise term.
     • The drift is deterministic and depends on the empirical distribution of tokens (mean-field self-attention).
     • The diffusion term arises from the MLP. Due to the Central Limit Theorem applied to the wide MLP, the random weights converge to a Gaussian random field. Crucially, this noise is common (shared across all tokens) but non-interacting (the noise covariance does not depend on the current state of other tokens, unlike in some previous works).
   • Technical Challenge: The authors handle the non-linearity of Layer Normalization (radial projection onto the sphere) by constructing a careful coupling between the finite-width MLP and the limiting Gaussian field, using a Taylor expansion to control the projection error.

2. Large $N$ Limit (Mean-Field):

   • They derive the Stochastic Partial Differential Equation (SPDE) governing the evolution of the empirical distribution of tokens ($\mu_t$).
   • This SPDE is of McKean-Vlasov type with common noise.
   • They prove Propagation of Chaos: As $N \to \infty$, the tokens become independent conditional on the common noise, and the empirical measure converges to the solution of the SPDE.

C. Synchronization Analysis

• The authors analyze the long-term behavior of the system using a Lyapunov function (interaction energy) associated with the attention mechanism.
• They investigate whether the common noise can overcome metastable states (saddle points) of the deterministic dynamics to force the system into a synchronized state (all tokens collapsing to a single point).

3. Key Contributions

1. First Stochastic Scaling Limit with MLP: This is one of the first works to rigorously derive a stochastic limit for Transformers that explicitly includes the MLP component. Previous works largely focused on deterministic limits or ignored the MLP's contribution to the noise structure.
2. Identification of Common Non-Interacting Noise: The paper identifies that the randomness from the MLP initialization manifests as a common noise term in the limiting SDE. Unlike previous stochastic models where noise was idiosyncratic (independent per token) or state-dependent, here the noise is shared but acts as a transport term independent of the particle interactions.
3. Quantitative Convergence Bounds: The authors provide explicit, quantitative error bounds for the convergence of the discrete Transformer dynamics to the continuous SDE.
   • The error scales as $O(\frac{\log L}{L} + \frac{1}{m})$.
   • They prove that the limits $L \to \infty$ (depth) and $m \to \infty$ (width) commute.
4. Synchronization by Noise: They establish conditions under which the common noise induces exponential synchronization (clustering) of tokens.
   • They prove that if the noise is sufficiently "coercive" relative to the deterministic drift, the interaction energy decays exponentially.
   • This provides a theoretical mechanism for the empirically observed "collapse" of tokens in deep transformers.

4. Key Results

• Theorem 1.1 (Pathwise Convergence): For fixed $N$, the discrete token trajectories $X^\ell$ converge to the solution $\bar{X}_t$ of the limiting SDE in $L^2$ norm with a rate depending on $L$ and $m$.
• Proposition 1.2 (Propagation of Chaos): For large $N$, the empirical measure of the particle system converges to the unique solution of the SPDE (1.10) with a rate depending on the dimension $d$ (e.g., $N^{-1/2}$ for $d < 4$).
• Theorem 1.4 (Exponential Dissipation): Under specific assumptions on the activation function and the noise strength (Assumption 1 and 2), the expected interaction energy $E_\beta(\mu_t)$ decays exponentially:
  $$\mathbb{E}[E_\beta(\mu_t)] \leq \exp((\varepsilon + \bar{\lambda}')t) \mathbb{E}[E_\beta(\mu_0)]$$
  where $\bar{\lambda}' < 0$ is a rate determined by the noise structure.
• Proposition 1.5 (Explicit Rates): For ReLU activations, the authors derive explicit formulas for the Lyapunov exponent $\lambda_1$ and the dissipation rate $\bar{\lambda}$, showing they are negative (implying synchronization) for typical parameter choices, provided the dimension $d > 3$.

5. Significance and Implications

• Theoretical Unification: The paper bridges the gap between deterministic mean-field theories and stochastic dynamics in deep learning, offering a more complete picture of how information propagates through Transformers.
• Mechanism for Clustering: It provides a rigorous mathematical explanation for the "synchronization" or "collapse" of token representations observed in practice. It suggests that the random initialization of MLPs is not just a source of variance but a structural driver that can force the system into a synchronized state, potentially aiding in the formation of coherent internal representations.
• Noise as a Feature: The work reframes the "noise" from random initialization not as a nuisance to be averaged out, but as a fundamental component of the dynamics that interacts with the attention mechanism to shape the model's behavior.
• Foundation for Training Analysis: By establishing pathwise convergence and quantitative bounds, this work lays the groundwork for future analyses of the training dynamics of Transformers, where the pathwise perspective is essential.

In summary, this paper establishes that deep Transformers with random MLP weights converge to a stochastic interacting particle system driven by common noise. This noise, under specific conditions, acts as a synchronization force that drives the system toward clustered states, offering a new theoretical lens on the internal dynamics of Large Language Models.