cs.CVMar 5, 2026
SaveLocality-Attending Vision Transformer
Organizations: ÉTS Montreal, LIVIA, ILLS
Abstract
Vision transformers have demonstrated remarkable success in classification by leveraging global self-attention to capture long-range dependencies. However, this same mechanism can obscure fine-grained spatial details crucial for tasks such as segmentation. In this work, we seek to enhance segmentation performance of vision transformers after standard image-level classification training. More specifically, we present a simple yet effective add-on that improves performance on segmentation tasks while retaining vision transformers' image-level recognition capabilities. In our approach, we modulate the self-attention with a learnable Gaussian kernel that biases the attention toward neighboring patches. We further refine the patch representations to learn better embeddings at patch positions. These modifications encourage tokens to focus on local surroundings and ensure meaningful representations at spatial positions, while still preserving the model's ability to incorporate global information. Experiments demonstrate the effectiveness of our modifications, evidenced by substantial segmentation gains on three benchmarks (e.g., over 6% and 4% on ADE20K for ViT Tiny and Base), without changing the training regime or sacrificing classification performance. The code is available at https://github.com/sinahmr/LocAtViT/.
Figures & tables
Figure 1: Qualitative evaluation on the attention maps. The final attention maps (before the classification head) of ViT and LocAtViT for the [CLS] token and three patches are illustrated for an image with label school bus .
Figure 2: LocAt considerably enhances different baselines in segmentation, while preserving or even improving classification.
Figure 3: Illustration of the Gaussian-Augmented attention for a grid. (a) The Gaussian addition is obtained based on the query and is added to the attention logits. The row in the attention logits matrix presents the attention of patch to all patch tokens. The reshaped matrix illustrates that with GAug, both local and global attentions are integrated. (b) The supplement matrix encourages attending to the locality and is computed using the pairwise squared difference tensor from Eq. 7 . For simplicity, the [CLS] token is not shown in this visualization, and Gaussian variances and scaling coefficients are set to a constant value for all patches.
| Method | Segmentation mIoU (%) | Top-1 (%) | #Params | FLOPs | ||||||
| ADE | P-Context | C-Stuff | ImageNet | (M) | (G) | |||||
| Tiny | ViT | 17.30 | 33.71 | 20.29 | 72.39 | 6 | 1.26 | |||
| + LocAt | 23.47 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+6.17}}}} | 38.57 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+4.86}}}} | 26.15 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+5.86}}}} | 73.94 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+1.55}}}} | 6 | 1.27 | ||||
| Swin | 25.58 | 36.78 | 28.34 | 81.18 | 28 | 4.50 | ||||
| + LocAt | 26.52 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.94}}}} | 37.65 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.87}}}} | 29.09 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.75}}}} | 81.43 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.25}}}} | 28 | 4.51 | ||||
| RegViT | 15.98 | 33.45 | 19.58 | 72.90 | 6 | 1.29 | ||||
Table 1: Segmentation performance of models and their counterparts with our LocAt extension (in gray), along with their classification performance on ImageNet-1K, which the models are initially trained on. Results demonstrate that (i) LocAt substantially boosts segmentation performance ( our primary focus ), while preserving or even improving the classification performance, and (ii) this effect holds for a variety of methods, for different backbone sizes. Furthermore, (iii) the segmentation gains appear not only in weaker baselines, but also in strong, high-performing models, where classification improvements are harder to achieve.
| Tiny | Base | |||||||||
| Method | PASCAL | ADE20K | PASCAL | ADE20K | ||||||
| Vanilla | + LocAt | Vanilla | + LocAt | Vanilla | + LocAt | Vanilla | + LocAt | |||
| ViT | 39.2 | 50.3 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+11.1}}}} | 12.0 | 15.2 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+3.2}}}} | 55.8 | 58.7 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+2.9}}}} | 19.5 | 21.5 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+2.0}}}} | ||
| Swin | 45.2 | 45.3 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.1}}}} | 16.1 | 16.3 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.2}}}} | 57.6 | 62.8 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+5.2}}}} | 23.3 | 24.6 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+1.3}}}} | ||
| RegViT | 39.4 | 52.3 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+12.9}}}} | 12.5 | 15.9 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+3.4}}}} | 55.5 | 60.3 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+4.8}}}} | 19.4 | 22.8 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+3.4}}}} | ||
| RoPEViT | 50.7 | 54.7 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+4.0}}}} | 16.0 | 17.5 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+1.5}}}} | 61.0 | 61.4 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+0.4}}}} | 22.4 | 23.7 {}_{{\textbf{{\color[rgb]{0.4,0.55,0.4}+1.3}}}} | ||
Table 4: Hummingbird dense nearest-neighbor retrieval (mIoU %) of models and their counterparts with our LocAt extension for different backbone sizes on PASCAL VOC and ADE20K.
Table 5: Ablation study on model’s architecture. We report segmentation performance (mIoU %) over three benchmarks and classification accuracy (top-1 %) on ImageNet-1K. PE and GAP stand for positional embeddings and global average pooling.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Segmentation mIoU (%) | Top-1 (%) | ||||
| ADE | P-Context | C-Stuff | ImageNet | |||
| CvT-21 | 21.40 | 40.91 | 29.29 | 82.50 | ||
| Conformer | 22.11 | 40.03 | 26.37 | 83.83 | ||
| ConViT | 23.08 | 44.82 | 25.20 | 82.30 | ||
| Twins | 30.47 | 44.55 | 32.27 | 82.71 | ||
| DaViT | 30.68 | 44.87 | 32.38 | 84.64 | ||
Table 6: Segmentation and classification performance of Base backbones from prior work and the proposed LocAtViT.
Figure 4: Qualitative evaluation on the attention maps. The final attention maps (before the classification head) of ViT and LocAtViT for the [CLS] token and three different patches are illustrated for three different images from mini-ImageNet with labels: orange , Komondor , and corn .
| Tiny | Base | #Params | ||||
| LocAtViT ( Section 4 ) | 78.47 | 84.86 | - | |||
| Gaussian from | 79.10 | 85.18 | , | |||
| Isotropic Gaussian | 78.71 | 84.66 | ||||
| Fixed width | 75.20 | 82.81 | ||||
| 76.41 | 82.65 | |||||
| 75.53 | 82.42 | |||||
Table 7: Ablations on GAug attention components . #Params shows the difference in the number of the parameters of each model compared to LocAtViT (first row). Experiments are conducted on mini-ImageNet, and the classification accuracy (top-1 %) is reported.
| Kernel | Tiny | Base | ||||||||
| ADE | P-Context | C-Stuff | ImageNet | ADE | P-Context | C-Stuff | ImageNet | |||
| No (ViT) | 17.30 | 33.71 | 20.29 | 72.39 | 28.40 | 43.10 | 30.43 | 80.99 | ||
| Gaussian | 23.47 | 38.57 | 26.15 | 73.94 | 32.64 | 45.35 | 33.62 | 82.31 | ||
| Inv-dist | 22.18 | 38.16 | 25.25 | 74.00 | 28.42 | 43.48 | 30.82 | 81.94 | ||
| Laplace | 21.67 | 37.80 | 25.56 | 74.01 | 29.74 | 44.10 | 31.95 | 82.24 | ||
Table 8: Effect of different distance-based kernels. Segmentation performance (mIoU %) over three benchmarks and classification accuracy (top-1 %) on ImageNet-1K are reported.
Figure 5: Degradation of local features in vanilla ViT. Features in ViT collapse to the global information in the last layers while in LocAtViT, patch features encode local information.
Figure 6: Layer-wise statistics of the learned Gaussian standard deviation in LocAtViT. For each layer, we summarize the distribution of learned standard deviation values using percentile ribbons (10–90% and 30–70%) and overlay the median (solid) and mean (dashed).