Perceptrons and localization of attention's mean-field landscape

TL;DR

The paper casts Transformer dynamics as a mean-field interacting particle system on the unit sphere and analyzes how a perceptron block reshapes the energy landscape under a Wasserstein gradient-flow framework. It shows that, in contrast to pure attention, the coupled system drives stationary measures to be atomic and spatially localized, with precise results in low dimensions and generic sparsity in higher dimensions. The authors derive sharp bounds on cluster masses and the number of atoms (scaling like \\sqrt{\\beta}) and provide concrete verifications via simulations in multiple dimensions, including both ReLU and GeLU activations and both unnormalized and normalized attention. Overall, the work reveals a fundamental mechanism by which perceptrons counteract attention-driven collapse and enforce discrete, interpretable structure in the mean-field landscape, with implications for understanding representation geometry in long-context Transformer regimes.

Abstract

The forward pass of a Transformer can be seen as an interacting particle system on the unit sphere: time plays the role of layers, particles that of token embeddings, and the unit sphere idealizes layer normalization. In some weight settings the system can even be seen as a gradient flow for an explicit energy, and one can make sense of the infinite context length (mean-field) limit thanks to Wasserstein gradient flows. In this paper we study the effect of the perceptron block in this setting, and show that critical points are generically atomic and localized on subsets of the sphere.

Figures (11)

• Figure 2: Cluster masses (in blue, the largest being the thickest) at final time across $\sqrt{\upbeta}$ for gradient descent with GeLU perceptron, initialized with $N=1000$ points of mass $10^{-3}$ (see \ref{['sec:numerics']} for setup). The horizontal and red dashed lines represent the numerical term and the full upper bound in \ref{['eq: conclusion']}, respectively.
• Figure 3: Gradient ascent on $\mathbb{S}^1$ with ReLU perceptron. Left: pure self-attention. Middle: self-attention with a ReLU perceptron. Right: measure at final time. Background shading represents the potential landscape (green: positive; orange: negative values).
• Figure 4: Gradient ascent on $\mathbb{S}^2$ with $\upbeta=1$. Top row: pure self-attention. Bottom row: self-attention with ReLU perceptron. An animation is available at https://github.com/antonioalvarezl/2026-MLP-Attention-Energy/blob/main/examples/USAS2.gif.
• Figure 5: Gradient descent on $\mathbb{S}^1$ with ReLU perceptron.
• Figure 6: Histograms at final time for gradient descent with GeLU perceptron.

Theorems & Definitions (24)

• Theorem 3.1
• Theorem 3.2
• Theorem 3.3
• Remark 3.4
• Theorem 3.5
• Corollary 3.6
• Proposition 3.7
• Lemma 5.1
• proof
• Remark 5.2