The Laplacian Mechanism Improves Transformers by Reshaping Token Geometry

Transformers leverage attention, the residual connection, and layer normalization to control the variance of token representations. We propose to modify attention into a Laplacian mechanism that gives the model more direct control over token variance. We conjecture that this helps transformers achieve the ideal token geometry. To investigate our conjecture, we first show that incorporating the Laplacian mechanism into transformers induces consistent improvements across benchmarks in computer vision and language. Next, we study how the Laplacian mechanism impacts the geometry of token representations using various tools: 1) principal component analysis, 2) cosine similarity metric, 3) analysis of variance, and 4) Neural Collapse metrics. Our investigation shows that the Laplacian mechanism reshapes token embeddings toward a geometry of maximal separability: tokens collapse according to their classes, and the class means exhibit Neural Collapse.

💡 Research Summary

The paper introduces a simple yet powerful modification to the standard self‑attention mechanism in transformers, called the Laplacian mechanism, which directly controls the variance of token embeddings. In the conventional attention pipeline, token representations are first projected into queries, keys, and values; the softmax of the scaled dot‑product yields attention weights that compute a weighted average of the values (PV). This average is added back to the original tokens via a residual connection, and layer normalization subsequently projects the result onto a sphere, indirectly regulating token variance. The authors observe that this indirect control may limit the model’s ability to reach an optimal geometric configuration of token embeddings, especially the “Neural Collapse” (NC) geometry identified in supervised classification networks, where class features collapse to their means, class means form a simplex equiangular tight frame (ETF), classifier weights align with class means, and predictions are made by nearest‑class‑mean decision boundaries.

To address this, the authors define the Laplacian mechanism as L(X) = V – PV, i.e., the difference between the value vectors and their attention‑weighted mean. By replacing a subset of attention heads with Laplacian heads, the model can explicitly increase or decrease token variance without relying on the subsequent normalization step. This modification introduces no additional trainable parameters and can be applied uniformly across all layers; the number of Laplacian heads per layer (k) is a hyper‑parameter that the authors vary in experiments.

Experimental Evaluation – Vision
The authors integrate the Laplacian heads into the DeiT‑3 family of Vision Transformers, focusing on the ViT‑B architecture (12 blocks, 12 heads per block). They evaluate models with k ∈ {0, 3, 6, 9, 11, 12} Laplacian heads on CIFAR‑10, CIFAR‑100, and ImageNet‑1k, using identical training recipes and three drop‑path rates (0.1, 0.3, 0.4). Across all datasets, adding Laplacian heads yields consistent improvements in top‑1 accuracy. On CIFAR‑10 and CIFAR‑100 the gain grows monotonically with k, reaching up to +0.53 % and +1.92 % respectively. On ImageNet‑1k the improvement is more modest (≈ +0.5 % to +1 %) and does not increase monotonically, suggesting that very large tasks benefit from a balanced mix of standard and Laplacian heads. The authors also show that the gains persist across different drop‑path rates, confirming robustness.

Experimental Evaluation – Language
For autoregressive language modeling, a GPT‑2‑style decoder (561 M parameters, 10 heads per block) is trained with k ∈ {0, 1, 3, 5, 7, 10} Laplacian heads. Training follows a three‑stage schedule: large‑scale pre‑training on 11.2 B tokens, mid‑training, and supervised fine‑tuning (SFT) on a mixture of chat and task data. The SFT model is evaluated on six benchmarks: ARC‑Easy, ARC‑Challenge, MMLU, GSM8K, and HumanEval, reporting zero‑shot accuracy for the first three and pass@10 for the latter two. Adding Laplacian heads improves the average score from 26.97 % (baseline) to up to 30.94 % (k = 5). Notably, ARC‑Easy improves by +9.18 % and ARC‑Challenge by +6.32 % with five Laplacian heads, while GSM8K sees a +2.58 % gain. HumanEval performance remains comparable, with the best results at k = 7 or 9. These findings suggest that direct variance control benefits not only classification but also generative reasoning and code synthesis tasks.

Geometric Analysis of Token Representations
To understand why the Laplacian mechanism improves performance, the authors conduct a multi‑faceted geometric analysis on the vision models:

1. Principal Component Analysis (PCA) – Projecting the last‑layer token embeddings onto the top two principal components reveals that models with more Laplacian heads form well‑separated, class‑specific clusters, whereas the baseline shows overlapping clouds.

2. Variance Decomposition (ANOVA‑style) – The authors define three variance components: Within‑Sequence (tokens around their sequence mean), Within‑Class (sequence means around the class mean), and Between‑Class (class means around the global mean). Total variance decomposes additively into these three. As k increases, the fraction of total variance attributed to Within‑Sequence drops sharply, while Between‑Class variance rises, indicating that tokens collapse within each sequence but become more discriminative across classes.

3. Layer‑wise Cosine Similarity – Measuring average cosine similarity among tokens within the same sequence shows a steeper increase with depth for models containing Laplacian heads, confirming that the mechanism accelerates within‑sequence alignment.

4. Neural Collapse Extension – Neural Token Collapse (NTC) – Building on the four NC properties, the authors define NTC: (NTC0) tokens in the same sequence collapse; (NTC1) tokens of the same class collapse to a class mean; (NTC2‑NTC4) class means and classifier weights satisfy the ETF geometry and nearest‑class‑mean decision rule. Empirical metrics (e.g., class‑mean alignment, simplex ETF loss) demonstrate that models with many Laplacian heads closely approach NTC, whereas the baseline remains far from this ideal.

Interpretation and Implications
The Laplacian heads give the model a direct “variance knob”: by subtracting the attention‑weighted mean from the values, the residual update can shrink or expand the spread of token embeddings before normalization. This reduces the “redundant radial movement” that standard residual + LN must later cancel, making the learning dynamics more efficient. Consequently, the model more readily drives token embeddings toward a geometry that maximizes inter‑class separability (high Between‑Class variance) while minimizing intra‑class dispersion (low Within‑Class and Within‑Sequence variance). The resulting NTC geometry mirrors the Neural Collapse observed in deep classification networks, suggesting a universal principle: optimal supervised learning pushes representations toward a simplex ETF configuration, regardless of whether the backbone is convolutional or transformer‑based.

Limitations and Future Work
The paper focuses on classification and next‑token prediction; other tasks such as dense prediction, multimodal fusion, or reinforcement learning remain unexplored. The optimal proportion of Laplacian to standard heads appears dataset‑dependent; overly aggressive replacement may harm performance on very large or highly heterogeneous datasets. Moreover, the analysis is primarily empirical; a theoretical framework linking the Laplacian update to convergence properties or generalization bounds would strengthen the contribution.

Conclusion
By replacing a subset of attention heads with a Laplacian variant that directly manipulates token variance, the authors achieve consistent accuracy gains across vision and language benchmarks and uncover a clear geometric shift toward Neural Token Collapse. This work highlights token variance control as a valuable design lever for future transformer architectures, opening avenues for more principled, geometry‑aware modifications that could further close the gap between empirical performance and theoretical optimality.