Technical Summary: Laplacian Heads Improve Transformers by Smoothing Token Representations

Problem Statement

The standard Transformer architecture updates token representations through multi-head attention and residual connections. In this formulation, each attention head $i$ computes a softmax attention matrix $P^{(i)}$ and updates the input $X$ via a weighted sum of value projections: $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

While effective, this mechanism lacks explicit control over the within-sequence variance of token representations. The prevailing view in the literature, particularly regarding Graph Neural Networks (GNNs) and deep Transformers, treats "oversmoothing"—where token representations become too similar, losing distinctiveness—as a detrimental failure mode. Consequently, most prior work focuses on preventing smoothing through normalization, edge dropping, or diversity losses. This paper challenges the assumption that token oversmoothing is inherently harmful, proposing instead that controlled smoothing can be beneficial for representation learning.

Methodology

The authors propose a parameter-free architectural modification: replacing a subset of standard attention heads with Laplacian heads.

The Laplacian Operator

In the standard update, a head computes $P^{(i)}X$, which is an attention-weighted average of tokens. The proposed Laplacian head replaces the attention matrix $P^{(i)}$ with the graph Laplacian operator $I - P^{(i)}$, where $I$ is the identity matrix. The update equation becomes:
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
where $A$ is the set of standard attention heads and $L$ is the set of Laplacian heads.

Motivations

The paper provides two theoretical motivations for this modification:

1. Direct Control of Mean and Variance: Standard attention heads update the mean of token representations. In contrast, the term $(I - P^{(i)})X$ computes the difference between each token and its attention-weighted mean. Adding this difference allows the model to directly control the within-sequence variance of token representations, a capability absent in standard Transformers.
2. Graph Diffusion: If tokens are viewed as nodes in a graph with edge weights defined by $P^{(i)}$, then $I - P^{(i)}$ is the random-walk normalized graph Laplacian. Applying this operator corresponds to a single step of heat diffusion on the token graph, progressively smoothing the signal (token representations) across the sequence.

Key Contributions

1. Architectural Proposal: A simple, parameter-free modification to the Transformer architecture that integrates Laplacian heads alongside standard attention heads.
2. Empirical Validation: Demonstration that this modification improves performance across three distinct learning paradigms: supervised classification, language modeling, and self-supervised learning.
3. Geometric Characterization: A detailed analysis of how Laplacian heads alter the geometry of learned token representations, showing that they induce beneficial forms of "collapse" and smoothing rather than harmful oversmoothing.

Experimental Results

1. Image Classification (Supervised Learning)

Using the DeiT-3 family (ViT-B) on CIFAR-10, CIFAR-100, and ImageNet-1k:

• Performance: Adding Laplacian heads consistently improved Top-1 accuracy. The improvement was nearly monotonic with the number of Laplacian heads.
  • On ImageNet, the model with 12 Laplacian heads (ViT-B-12L) achieved 82.32% accuracy compared to the baseline's 81.45%.
  • On CIFAR-100, the gain was nearly 2% (77.04% vs. 75.12%).
• Geometry:
  • ANOVA: Laplacian heads significantly reduced Within-Sequence Variance while increasing Between-Class Variance. This indicates that tokens within the same sequence collapse toward a common mean, while the means of different sequences separate more effectively.
  • Neural Collapse: The models exhibited stronger convergence to the Neural Collapse geometry. Token representations formed well-separated, simplex-like clusters (Simplex ETF) for each class, with sequence means aligning with the global class geometry.

2. Language Modeling (Autoregressive)

Using GPT-2 style models (561M parameters) pre-trained on FineWeb-Edu and fine-tuned on reasoning/code benchmarks:

• Performance: Models with Laplacian heads outperformed the baseline on average across ARC-Easy, ARC-Challenge, MMLU, GSM8K, and HumanEval.
  • The model with 5 Laplacian heads achieved the highest average performance (30.06%), with significant gains on ARC-Easy (+16%) and ARC-Challenge (+7%).
• Geometry:
  • Spectral Decay: Laplacian heads produced token representations with faster-decaying spectra. The number of principal components required to explain 99% of the variance dropped by 21% (from 547.7 to 430.4) when all heads were Laplacian.
  • Linguistic Collapse: The models showed increased separability among token representations that share the same next-token prediction, suggesting a stronger form of "linguistic collapse" where tokens with similar semantic roles cluster together.

3. Self-Supervised Learning (DINO)

Using ViT-S models on ImageNet-1k for DINO pre-training:

• Performance: Laplacian heads improved linear probing accuracy on ImageNet and segmentation performance on ADE20K.
  • ViT-S-5L achieved the best ADE20K mIoU (32.27%) compared to the baseline (29.81%).
• Geometry:
  • Segmentation: PCA visualizations showed that Laplacian heads produced sharper object boundaries and preserved finer details (e.g., helmet shapes, jersey numbers) compared to the baseline.
  • Spectral Decay: Similar to language modeling, these models exhibited faster spectral decay, concentrating within-image token representations into a lower-dimensional subspace.

Significance and Claims

The paper argues that the prevailing view—that token oversmoothing is universally harmful—is incomplete. Instead, the authors claim that smoothing is a mechanism for representation learning that can be beneficial when controlled.

• Task-Specific Benefits: The "smoothing" induced by Laplacian heads is not a failure mode but a feature that aligns with task objectives:
  • In supervised classification, it facilitates Neural Collapse by collapsing tokens within sequences and aligning sequence means with class geometry.
  • In language modeling, it enhances Linguistic Collapse, improving the separability of tokens with identical next-token targets.
  • In self-supervised learning, it concentrates representations in lower-dimensional subspaces, aiding downstream tasks like segmentation.
• Mechanism: By explicitly incorporating graph diffusion (via the Laplacian) into the architecture, the model gains the ability to directly control within-sequence variance, leading to representations that are more robust and geometrically structured.

The authors conclude that while their fixed-strategy implementation is likely suboptimal, the results suggest that smoothing should be investigated further as a fundamental component of representation learning, rather than solely as a problem to be mitigated.