Scratching Visual Transformer's Back with Uniform Attention

TL;DR

The paper investigates the density of spatial interactions in Vision Transformers and finds a strong propensity toward dense attention, which is difficult to learn with softmax. It proposes Context Broadcasting (CB), a lightweight module that injects uniform attention to supply dense global context, reducing attention density in subsequent MSAs while preserving long-range interactions. Across image classification, semantic segmentation, object detection, and vision-language tasks, CB and its dimension-scaled variant CB_S yield consistent accuracy and robustness gains with negligible computational cost. The findings suggest that guiding ViTs with explicit dense-context signals can reallocate MSA capacity toward sparser, more informative interactions, offering a practical and broadly applicable improvement strategy for transformer-based vision models.

Abstract

The favorable performance of Vision Transformers (ViTs) is often attributed to the multi-head self-attention (MSA). The MSA enables global interactions at each layer of a ViT model, which is a contrasting feature against Convolutional Neural Networks (CNNs) that gradually increase the range of interaction across multiple layers. We study the role of the density of the attention. Our preliminary analyses suggest that the spatial interactions of attention maps are close to dense interactions rather than sparse ones. This is a curious phenomenon, as dense attention maps are harder for the model to learn due to steeper softmax gradients around them. We interpret this as a strong preference for ViT models to include dense interaction. We thus manually insert the uniform attention to each layer of ViT models to supply the much needed dense interactions. We call this method Context Broadcasting, CB. We observe that the inclusion of CB reduces the degree of density in the original attention maps and increases both the capacity and generalizability of the ViT models. CB incurs negligible costs: 1 line in your model code, no additional parameters, and minimal extra operations.

Figures (12)

• Figure 1: Motivation of our work.Top: dense attention is hard to learn with softmax, but self-attention tends to learn it more than sparse one. Bottom: we infuse dense attention explicitly, named $\mathtt{CB}$, to split the responsibility of interactions; the burden of interactions of self-attention is reduced. Self-attention is now more likely to learn sparse interaction that is in favor of softmax.
• Figure 2: Type of spatial interactions. We categorize the spatial interactions of self-attention into three types. The anchor token interacts with reference tokens.
• Figure 3: Impact on the capacity of the ViT model with a single extra block. Training loss and top-1 accuracy ($y$-axis) versus epochs ($x$-axis) of 8-depth ViT with additional $\mathtt{MSA}$ and $\mathtt{MLP}$ blocks. The decrease in training loss and the increase in validation accuracy implies an increase in the model capacity.
• Figure 4: Entropy analysis. We use pre-trained ViTs to measure layer-wise entropy. We plot the average and $15^{th}$ percentile of entropy values. The red dot line stands for the maximum entropy upper bound.
• Figure 5: Gradient around uniform attention. Softmax operation has high gradients around uniform attention ($x_1=x_2$).