Vision Transformers Need Better Token Interaction

Technical Summary: Vision Transformers Need Better Token Interaction

Problem Statement
Vision Transformers (ViTs) are renowned for learning strong image-level representations, yet their patch-level representations often degrade during prolonged training, becoming less effective for dense prediction tasks such as semantic segmentation and depth estimation. While previous research has attributed this "dense degradation" to high-norm artifacts or outlier tokens, this paper argues that these are symptoms rather than the root cause. The authors identify the core issue as semantic diffusion: an optimization shortcut where global semantic information spreads indiscriminately through patch tokens beyond what is locally justified. This phenomenon allows the model to improve global recognition by exploiting unrestricted token interactions, but at the cost of degrading the spatial structure required for dense tasks.

Methodology
The authors propose that the solution is not to eliminate global context or suppress background tokens entirely, but to make token interactions more selective. They introduce sparse attention as a minimal intervention to modify the attention geometry.

• Mechanism: The standard softmax attention mechanism is replaced with entmax-1.5. Unlike softmax, which produces dense probability distributions, entmax-1.5 is a sparse probability mapping that can assign exact zero weights to low-utility interactions.
• Implementation: The method preserves global token connectivity (unlike local window attention) and introduces no additional parameters, register tokens, or separate processing branches (unlike specialization or register-based approaches).
• Hypothesis: By encouraging sparse attention distributions, the model is forced to be more selective in its token mixing. This reduces the "optimization shortcut" of semantic diffusion, preventing semantically irrelevant background patches from diluting the local structure of foreground patches while maintaining the ability to capture global semantics.

Key Contributions

1. Reframing Dense Degradation: The paper characterizes dense degradation not merely as an artifact issue but as a consequence of semantic diffusion, an optimization shortcut where global semantics diffuse into patch tokens in a locally inconsistent manner.
2. Analysis of Token Roles: Through Point-in-Box (PiB) analysis and layer-wise probing, the authors demonstrate that while shallow layers maintain better spatial alignment (higher PiB), they lack the semantic richness of deeper layers. Crucially, they show that the [CLS] token retains complementary global information that, when concatenated with patch features, significantly aids dense prediction, suggesting that global context should be preserved but better managed.
3. Minimal Intervention: The paper introduces a simple architectural change—swapping softmax for entmax-1.5—that acts as a representational bias toward selective token interaction without requiring additional supervision, teacher models, or complex architectural modifications.

Experimental Results
The method was evaluated on DINOv1 ViT-S/16 models trained for 200 epochs on ImageNet-1K. The results indicate that sparse attention improves dense representations without sacrificing global recognition:

• Image-Level Recognition: The Entmax-1.5 variant achieved a Top-1 accuracy of 70.06% on ImageNet linear probing, slightly outperforming the CLS baseline (69.50%) and other variants like Register-4 and Specialization.
• Semantic Segmentation: The most significant gains were observed here. Compared to the CLS baseline, Entmax-1.5 improved:
  • VOC mIoU: from 42.80 to 48.78 (+5.98).
  • ADE20K mIoU: from 19.85 to 21.97 (+2.12).
  • Cityscapes mIoU: from 36.79 to 37.87 (+1.08).
  • It outperformed both Register-4 and CLS/patch specialization on all three segmentation benchmarks.
• Depth Estimation: Improvements were more moderate but consistent across NYUv2, SUN RGB-D, and KITTI benchmarks (reducing RMSE by 0.004–0.005). The authors note that the baseline depth probing protocol already concatenates [CLS] and patch features, which may partially compensate for semantic diffusion, making gains harder to achieve.
• Qualitative Analysis: PCA visualizations of patch features showed that Entmax-1.5 produces cleaner, more spatially coherent feature maps with better separation between foreground and background regions compared to baselines.

Significance and Claims
The paper claims that selective token mixing is a simple and effective bias for improving dense ViT representations. The significance lies in demonstrating that the goal of dense prediction is not to remove global context, but to refine how tokens interact with it. By replacing softmax with entmax-1.5, the authors show that a lightweight modification to attention normalization can mitigate semantic diffusion, leading to substantial improvements in semantic segmentation while maintaining or slightly improving image-level recognition.

The authors remain modest regarding their claims, noting that this study is limited to DINOv1 ViT-S/16 and has not yet evaluated larger backbones or longer pretraining schedules. They position sparse attention not as a definitive final architecture, but as evidence supporting the broader principle that making token interactions more selective is a viable path to resolving dense degradation.