Source-linked AI summary
Shunted Self-Attention via Multi-Scale Token Aggregation
Sucheng Ren, Daquan Zhou, Shengfeng He, Jiashi Feng, Xinchao Wang
TL;DR
Vision Transformers’ uniform receptive fields and token-merging strategies limit simultaneous modeling of objects at different scales while self-attention remains computationally costly. The paper introduces shunted self-attention, which combines fine- and coarse-grained attention through multi-scale token aggregation, and reports state-of-the-art results across classification, detection, and segmentation.
Problem
Uniform receptive fields and aggressive token merging limit Vision Transformers’ ability to capture different object scales while self-attention has quadratic memory cost in the number of tokens.
Method
Shunted self-attention splits heads across attention granularities, selectively aggregating tokens for coarse features while preserving tokens and local details for fine-grained features.
Results
The proposed model outperforms prior Transformers and achieves state-of-the-art results on classification, detection, and segmentation, including better performance than Focal Transformer with similar model size.
Takeaways & Limitations
SSA jointly models multi-scale objects within one self-attention layer while reducing computation and retaining fine-grained details.
Abstract
from arXiv · showhide
Recent Vision Transformer~(ViT) models have demonstrated encouraging results across various computer vision tasks, thanks to their competence in modeling long-range dependencies of image patches or tokens via self-attention. These models, however, usually designate the similar receptive fields of each token feature within each layer. Such a constraint inevitably limits the ability of each self-attention layer in capturing multi-scale features, thereby leading to performance degradation in handling images with multiple objects of different scales. To address this issue, we propose a novel and generic strategy, termed shunted self-attention~(SSA), that allows ViTs to model the attentions at hybrid scales per attention layer. The key idea of SSA is to inject heterogeneous receptive field sizes into tokens: before computing the self-attention matrix, it selectively merges tokens to represent larger object features while keeping certain tokens to preserve fine-grained features. This novel merging scheme enables the self-attention to learn relationships between objects with different sizes and simultaneously reduces the token numbers and the computational cost. Extensive experiments across various tasks demonstrate the superiority of SSA. Specifically, the SSA-based transformer achieves 84.0\% Top-1 accuracy and outperforms the state-of-the-art Focal Transformer on ImageNet with only half of the model size and computation cost, and surpasses Focal Transformer by 1.3 mAP on COCO and 2.9 mIOU on ADE20K under similar parameter and computation cost. Code has been released at https://github.com/OliverRensu/Shunted-Transformer.
1. Introduction
Vision Transformers model long-range dependencies effectively, but uniform receptive fields and token-merging strategies limit simultaneous handling of objects at different scales. Shunted Self-Attention addresses this by combining coarse- and fine-grained attention within one layer, improving efficiency and performance across vision tasks.
- Motivation: Vision Transformers use global self-attention to model long-range dependencies, but its quadratic token cost drives down-sampling and related efficiency compromises.Down-sampling can lose feature information, while token merging can mix small-object features with background noise.
- Motivation: Existing attention mechanisms rely on static receptive fields and uniform information granularity, limiting simultaneous capture of objects at different scales.This limitation makes models less effective in scenes containing objects with distinct sizes.
- Method: SSA maintains light computation while attending to hybrid-scale objects, including large objects and small or remote objects that PVT may miss.The qualitative comparison highlights attention to small objects such as lights and fans alongside large objects such as sofas.
- Method: Shunted Self-Attention lets different heads within one layer handle coarse-grained and fine-grained features through multi-scale token aggregation.Fine-grained groups aggregate fewer tokens, while coarse-grained groups aggregate more tokens to reduce computation while retaining large-object modeling.
- Results: The Shunted Transformer achieves state-of-the-art results across classification, detection, and segmentation tasks.The model is built by stacking multiple SSA-based blocks and is evaluated as a backbone across downstream tasks.
2. Related Work
Related ViT variants reduce self-attention cost through low-resolution features, local attention, or token merging, but these strategies can weaken global dependency modeling or small-object representation. SSA is introduced to preserve coarse- and fine-grained details while maintaining global modeling.
- Vision Transformers: ViT models apply self-attention directly to image tokens, but quadratic complexity makes large-token inputs difficult to process.Consequently, many models down-sample feature maps and operate on low-resolution features.
- Vision Transformers: Low-resolution self-attention reduces computation but compromises fine-grained features and impedes dense prediction tasks such as detection and segmentation.The limitation is especially relevant when spatial detail must be retained.
- Efficient ViT Variants: Efficient ViT variants use local attention or token merging to make self-attention applicable to large feature maps.Local attention partitions feature maps into regions, whereas token merging reduces the number of tokens.
- Efficient ViT Variants: Local-window methods require shifted windows or many layers to obtain global receptive fields, while PVT-style merging can mix small-object information with background.These trade-offs motivate an approach that preserves both global dependencies and multi-scale details.
- Efficient ViT Variants: SSA is proposed to preserve coarse- and fine-grained details while maintaining global dependency modeling over image tokens.Its token aggregation is designed to avoid the excessive merging that harms small-object representation.
3. Method
Shunted Transformer blocks combine multi-granularity self-attention with detail-specific feed-forward processing, while the architecture uses staged feature maps and overlapping patch embedding. SSA varies token aggregation across heads to balance fine detail, large-object modeling, and computation.
- Shunted Self-Attention: SSA blocks capture multi-granularity information by using different attention granularities across heads within the same layer.Fine-grained groups aggregate fewer tokens, while coarse-grained groups aggregate more tokens to reduce computation and capture large objects.
- Shunted Self-Attention: SSA reduces key and value lengths differently across attention heads, unlike PVT's identical spatial reduction across heads.Multi-scale token aggregation uses different down-sampling sizes to obtain key and value representations at varied granularities.
- Shunted Self-Attention: SSA adds local enhancement to values through depth-wise convolution, preserving more fine-grained and low-level details than spatial reduction.The local enhancing component is applied to the value representation after multi-scale token aggregation.
- Shunted Self-Attention: The down-sampling rate r controls the trade-off between computation and detail preservation in SSA.Larger r merges more tokens and lowers computation, whereas smaller r preserves more detail at higher computational cost.
- Detail-specific Feed-Forward: The detail-specific feed-forward layer inserts a cross-token operation between two fully connected layers to complement local information.The detail-specific layer is implemented with a depth-wise convolution in practice.
- Architecture Details and Variants: The model uses four stages connected by stride-2 convolutional embeddings, with feature-map resolution halved and channel dimension doubled between stages.Its patch embedding begins with overlapping convolutions and produces an input sequence of length H/4 × W/4.
4. Experiments
Experiments evaluate the Shunted Transformer across ImageNet classification, COCO detection and instance segmentation, ADE20K segmentation, and component ablations. Across these settings, it generally improves performance while maintaining comparable or smaller parameter and computation budgets.
- The experiments cover ImageNet-1K classification, COCO detection and instance segmentation, ADE20K semantic segmentation, and ablation studies.
- Image Classification: 22M→11M parameters and 4.6G→2.1G FLOPs let the tiny model match DeiT-S, while a similarly sized model exceeds DeiT-S by 3%.
- Object Detection and Instance Segmentation: 2.3 points higher bbox mAP than Focal Transformer is obtained with 85% of its model size under the 1× COCO schedule.
- Object Detection and Instance Segmentation: 1.5 points higher mask mAP than Focal Transformer is obtained under the 1× COCO schedule, alongside gains over ResNet-50.
- Semantic Segmentation: 2.9% higher performance than Focal Transformer is reported with UpperNet on ADE20K, while Semantic FPN exceeds Swin Transformer by 6.7 mIOU with 20% less model size.
- Ablation Studies: The token aggregation function preserves global and local information with computation similar to convolutional spatial reduction, while detail-specific feed-forward layers improve ImageNet accuracy.
5. Conclusion
The paper presents SSA to model multi-scale features within one self-attention layer. Experiments report state-of-the-art performance across classification, detection, and segmentation tasks.
- SSA maintains various-scale feature maps within one self-attention layer to attend to multi-scale objects.
- The proposed model outperforms prior Transformers and achieves state-of-the-art results on classification, detection, and segmentation tasks.