PaperlayerView exact revision on arXiv

Source-linked AI summary

Locality-Attending Vision Transformer

Sina Hajimiri, Farzad Beizaee, Fereshteh Shakeri, Christian Desrosiers, Ismail Ben Ayed, Jose Dolz

arXiv:2603.04892v1cs.CV

TL;DR

ViTs’ global attention supports classification but can miss the spatial granularity needed for segmentation. The paper proposes LocAt, which combines a soft Gaussian locality bias with patch-representation refinement. Across benchmarks, LocAt improves segmentation while preserving or improving classification, with scope limited to natural images and small foundation models.

  • Problem

    Global self-attention helps ViTs capture long-range context but can obscure the fine-grained spatial detail required for dense prediction.

  • Method

    LocAt adds Gaussian-Augmented attention for soft local biasing and Patch Representation Refinement for more meaningful patch representations before classification.

  • Results

    LocAt consistently improves segmentation across ViT-based baselines while preserving or improving classification performance.

  • Takeaways & Limitations

    The modular add-on provides a simple upgrade for classification-trained ViTs that strengthens dense-prediction representations without changing the training regime.

  • Takeaways & Limitations

    Evaluation covered natural-image benchmarks and a small foundation model, leaving other domains and large foundation models for future work.

Abstract

from arXiv · show

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/.

1 INTRODUCTION

ViTs capture global context effectively for classification, but that global focus can obscure the fine-grained spatial detail required for segmentation. LocAt addresses this tension with a modular locality-aware add-on while retaining image-level capabilities.

  • ViTs use global self-attention to capture long-range dependencies and support holistic high-level representations for image classification.
  • Global attention can dilute local cues needed for precise localization, while purely local attention may miss long-range dependencies.
  • Foundation-model features such as CLIP’s often lack the spatial granularity needed for precise dense prediction because they are not trained for quality local representations.
  • LocAt combines a learnable Gaussian locality bias with Patch Representation Refinement to improve segmentation while preserving or improving classification accuracy.

2 RELATED WORK

Prior work introduces locality and spatial detail through hierarchical backbones, convolutional hybrids, attention modifications, positional encodings, and token-representation improvements. LocAt instead refines the ViT backbone with a modular approach aimed at preserving both local detail and global context.

  • Hierarchical ViT backbones for dense prediction: Hierarchical backbones add multi-scale features, token merging, or parallel local-global branches, but typically require substantial architectural changes.
  • Convolution-based hybrids: Convolution-based hybrids inject local priors through additional projections, gates, coupled features, or feed-forward convolutions, potentially reducing plug-and-play compatibility.
  • Locality mechanisms inside attention: Attention-locality methods use windows, neighborhoods, content-based key sampling, sparse routing, or global-local mixers to balance receptive field and locality.
  • Positional encodings that strengthen locality: Relative and rotary positional encodings strengthen spatial awareness and are presented as complementary to attention-locality mechanisms.
  • Foundation models for dense prediction: LocAt refines the standard ViT backbone to pursue dense prediction without extra adaptation layers or specialized fine-tuning.

3 PRELIMINARIES

A ViT layer processes one [CLS] token together with spatial patch tokens through normalized self-attention and a feed-forward network. Self-attention derives queries, keys, and values, computes scaled attention logits, and projects the resulting representation.

  • Each layer receives a sequence containing one [CLS] token and hw spatial patch tokens, each represented by a C-dimensional vector.
  • Self-attention uses Wq, Wk, and Wv to produce d-dimensional queries, keys, and values from the input tokens.
  • A patch’s attention logits are given by the corresponding row of the scaled query-key product qk^⊤/√d.

4 METHOD

LocAtViT adds Gaussian-Augmented attention and Patch Representation Refinement to classification-trained ViTs. The first softly biases patch attention toward nearby positions, while the second improves spatial-token supervision before classification.

  • Method overview: LocAtViT combines GAug attention and PRR while retaining ViT’s standard classification objective.
  • Gaussian-Augmented attention: GAug adds a supplement matrix to attention logits so each patch attends more strongly to nearby patches without eliminating global interactions.
  • Gaussian-Augmented attention: The Gaussian kernel is centered on each source patch and decays smoothly with distance, with σ controlling the effective receptive field.
  • Gaussian-Augmented attention: Patch-specific variances are predicted from spatial queries, allowing attention spans to differ across the two coordinate axes.
  • Gaussian-Augmented attention: Per-query scaling controls the strength of the Gaussian supplement, producing a soft data-dependent locality mechanism that can approach global attention.
  • Patch Representation Refinement: Classification supervision directly targets only the [CLS] output, leaving final patch representations without direct supervision despite their importance for dense prediction.
  • Patch Representation Refinement: PRR applies parameter-free non-uniform attention before the classification head, preserving distinct patch contributions and diversifying gradient flow.
  • Method overview: GAug and PRR act at different stages and can be attached independently to a ViT backbone.

5 EXPERIMENTS

Experiments evaluate LocAt across classification, segmentation, self-supervised learning, qualitative attention, and architectural ablations. LocAt improves dense prediction while preserving or improving recognition performance under the paper’s evaluation settings.

  • Segmentation performance: LocAtViT Tiny improves ViT segmentation by +6.17% on ADE20K, +4.86% on PASCAL Context, and +5.86% on COCO Stuff.The gains are reported against ViT across three segmentation benchmarks.
  • Segmentation performance: LocAt improves segmentation across five ViT-based model families while preserving or improving ImageNet-1K classification performance.The evaluated backbones include ViT, Swin Transformer, RegViT, RoPEViT, and Jumbo.
  • Classification performance: LocAt improves ViT classification by 3-6% on mini-ImageNet and 4-7% on CIFAR-100 with only 2,340 new parameters.The reported parameter increase is 0.003% for the Base backbone.
  • Foundation models: Replacing ViT with LocAtViT in DINO improves both linear-classifier and k-NN classification performance in the self-supervised regime.The study trains DINO ViT-S/16 and LocAtViT-S/16 on ImageNet-1K for 50 epochs.
  • Foundation models: LocAt consistently improves frozen-feature dense nearest-neighbor retrieval on PASCAL VOC and ADE20K across architectures.The Hummingbird protocol uses visual-token retrieval without fine-tuning or decoder training.
  • Qualitative analysis: Qualitative attention maps show more concentrated bus-feature activation and more interpretable patch-token attention for LocAtViT than vanilla ViT.Both models were trained solely for classification, while LocAtViT also highlights coherent scene structures such as background regions.
  • Ablation study: GAug and PRR each improve classification and segmentation, with their combination producing further gains.PRR also outperforms GAP for segmentation and improves classification accuracy more than GAP.
  • Ablation study: LocAt without positional embeddings surpasses ViT without or with positional embeddings, indicating that it captures spatial information represented by the default embeddings.The authors state that LocAt complements rather than replaces newer positional encoding approaches such as RoPE.

6 CONCLUSION

LocAtViT is a modular add-on for dense prediction that preserves image-level capabilities and improves segmentation through locality-biased attention and stronger patch representations. The evaluation is limited to natural-image benchmarks and small foundation models.

  • Contribution: LocAtViT combines Gaussian-augmented attention with patch representation refinement to improve vision transformers for dense prediction.GAug biases attention toward local regions, while PRR strengthens patch-token representations and gradient flow.
  • Results: The add-on delivers consistent segmentation gains across multiple ViT baselines without compromising classification accuracy.
  • Implication: The method is designed as a simple, largely orthogonal upgrade that integrates with existing classification-trained ViTs.
  • Limitations: Evaluation covered natural-image classification and segmentation benchmarks, leaving other domains such as medical imaging and remote sensing untested.
  • Limitations: Large foundation models at CLIP scale were not evaluated because of computational-budget constraints.

B LOCATVIT COMPARISON WITH RELATED WORK

The comparison evaluates LocAtViT against prior vision-transformer and related-backbone methods using a common segmentation pipeline and the reported classification results.

  • Compared methods: LocAtViT is compared with five baseline methods and multiple related works, including CvT-21, Conformer, ConViT, Twins, DaViT, and GCViT.
  • Evaluation setup: The evaluation uses publicly available implementations, code, and checkpoints, with models assessed on the same segmentation pipeline.
  • Reported metrics: Table 6 reports segmentation and classification performance for Base backbones from prior work and LocAtViT.

C ADDITIONAL QUALITATIVE EXPERIMENTS

Additional qualitative experiments visualize final attention maps for ViT and LocAtViT across classification tokens and spatial patches on mini-ImageNet images.

  • Visualization setup: The additional qualitative experiment provides three mini-ImageNet images together with attention maps for the [CLS] token and several patches in both models.
  • Attention-map comparison: Figure 4 compares final attention maps from ViT and LocAtViT before the classification head.
  • Visualization setup: The visualization includes the [CLS] token and three different patches across three mini-ImageNet images.The images are labeled orange, Komondor, and corn.

D ABLATION STUDY ON SELF-ATTENTION

The ablations examine how GAug’s variance construction, kernel width, scaling, and automatic scaling affect classification performance and parameter count.

  • Study scope: The ablation study evaluates design choices inside the GAug self-attention module.
  • Variance construction: Computing Gaussian variance from the spatial query matrix is compared with computing it from the self-attention input x.Using x improves performance but significantly increases the number of parameters.
  • Kernel variance: An isotropic Gaussian kernel performs comparably, whereas a fixed kernel width substantially diminishes performance.The comparison indicates that dynamic, input-dependent kernel width is important.
  • Supplement scaling: Removing learnable scaling consistently reduces accuracy, showing that unscaled addition of the locality term is suboptimal.The scaling balances the Gaussian prior with the original global attention logits.
  • Automatic scaling: Auto α performs close to learnable α with slightly fewer parameters, but the main model retains learnable scaling for formulation simplicity and flexibility.

E ABLATION STUDY ON ALTERNATIVE DISTANCE-BASED KERNELS

The study compares Gaussian, Laplace, and inverse-distance locality kernels within the same GAug formulation. All improve over vanilla ViT, while Gaussian provides the strongest segmentation gains across three benchmarks and remains competitive in ImageNet-1K accuracy.

  • Laplace and inverse-distance kernels predict scale parameters from query tokens using spatial-grid distances between patches.
  • Alternative kernels replace G in Eq. (10), while the rest of GAug, including scaling with α, remains unchanged.
  • All three locality-augmented kernels improve performance over the baseline ViT.
  • The Gaussian kernel delivers the strongest segmentation gains on all three benchmarks.
  • Gaussian remains competitive in ImageNet-1K accuracy compared to Laplace and inverse-distance kernels.

F LOCAL FEATURE ANALYSIS ACROSS LAYERS

Layer-wise analysis shows that vanilla ViT patch tokens lose distinct local information and become increasingly [CLS]-like with depth. LocAtViT preserves local patch structure without collapsing representations onto the global [CLS] content, while its learned locality scales remain nondegenerate.

  • Figure 5 contrasts vanilla ViT feature collapse with LocAtViT patch features that encode local information.
  • The locality score measures cosine similarity between each patch and its eight immediate neighbors in a surrounding 3 × 3 window.
  • Patch-[CLS] similarity distinguishes genuine local structure from globally collapsed patch representations.
  • Vanilla ViT patch tokens gradually lose distinct local information and become dominated by global [CLS]-like content as depth grows.
  • LocAtViT preserves strong locality in patch features without collapsing them onto the [CLS] token.
  • Learned Gaussian standard deviations remain inside the allowed interval and do not cluster near its bounds across layers.

H LIMITATIONS OF THE GAUSSIAN BIAS

GAug and PRR are most effective when attention is globally connected and locality is not already enforced by the backbone. Gains are smaller on windowed-attention models, and no improvement was observed for GCViT.

  • The largest GAug and PRR improvements occur on backbones with unrestricted patch-patch attention.
  • Gains are noticeably smaller on Swin, whose windowed attention already hard-codes locality.
  • No performance improvement was observed when applying the approach to GCViT with attention confined to small grids.
  • The authors attribute GCViT’s negative result to limited room for the Gaussian bias to reshape narrow-window locality patterns.
Loading 2603.04892v1…