Research map / Architectures and efficiency / Attention and transformers
SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SHViT, an efficient vision transformer that uses a larger-stride (16x16 overlapping) patchify stem, a 3-stage macro design, and a single-head self-attention applied to a partial channel subset, achieving strong speed-accuracy tradeoffs on GPU, CPU, and iPhone.
Problem
Efficient vision transformers carry computational and memory-access redundancy at both macro design (4x4 patch embedding, 4 stages) and micro design (multi-head attention with redundant heads), limiting speed on resource-constrained devices.
Approach
Combine a larger-stride overlapping patchify stem with 3 stages to reduce early-stage token counts, use convolutions in the first stage, and replace multi-head attention with a single-head self-attention applied to only a partial channel subset (r=1/4.67) in parallel with depthwise convolution, which removes head redundancy and cuts memory-bound operations.
Key results
- SHViT-S4 reaches 79.4% top-1 on ImageNet-1K at 14283 images/s on A100 GPU, outperforming EfficientNet-B0 by 2.3% accuracy and MobileViTv2x1.0 by 1.3% accuracy while 3.3x faster on GPU.
- On COCO with Mask R-CNN, SHViT-S4 achieves 39.0 APb and 35.9 APm with 0.28 ms GPU backbone latency, versus FastViT-SA12 at 38.9 APb and 1.06 ms.
- In the single-head ablation, replacing SHSA with MHSA lowers GPU throughput from 20522 to 18036 images/s and CPU from 62 to 50 while accuracy drops from 77.4% to 77.7% (MHSA is higher accuracy but slower).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SHViT-S4 instead of EfficientViT-M4 | ImageNet classification · ImageNet-1K | Top-1 accuracy: 74.3 -> 79.4 | Helped |
| Use SHViT-S4 instead of EfficientViT-M4 | RetinaNet · COCO 2017 | AP: 32.7 -> 38.8 | Helped |
| Use SHViT-S4 instead of FastViT-SA12 | Mask R-CNN · COCO 2017 | APb: 38.9 -> 39.0 | Helped |
| Use SHViT-S4 instead of FastViT-SA12 | Mask R-CNN · COCO 2017 | APm: 35.9 -> 35.9 | No clear effect |
| Replace single-head attention with MHSA (ablation) | SHViT-S3 · ImageNet-1K | Top-1 accuracy: 77.7 -> 77.4 | Helped |
| Replace single-head attention with MHSA (ablation) | SHViT-S3 · ImageNet-1K | GPU throughput: 20522 -> 18036 | Hurt |
| Add self-attention (ablation) | SHViT-S3 without self-attention · ImageNet-1K | Top-1 accuracy: 76.3 -> 77.4 | Helped |
| Set SHSA partial ratio to 1/8 instead of 1 / 4.67 (ablation) | SHViT-S3 · ImageNet-1K | Top-1 accuracy: 77.4 -> 77.1 | Hurt |
| Set SHSA partial ratio to 1/2 instead of 1 / 4.67 (ablation) | SHViT-S3 · ImageNet-1K | Top-1 accuracy: 77.4 -> 77.5 | Helped |
Surprising findings
- Removing the attention layer entirely (variant 2) costs only about 1.1 points of accuracy relative to the SHSA model while being faster, showing that macro design dominates speed-accuracy tradeoffs.
- Attention heads in DeiT-T and Swin-T show high similarity, and removing a single head in Swin-T can slightly improve accuracy.
Limitations
- Speed advantages depend on the inference platform; at low resolutions SHViT-S4 is slightly slower than FastViT on mobile.
- Ablation of the single-head design is limited to SHViT-S3 and the partial ratio search is coarse.
Experimental setup
Datasets: ImageNet-1K MS COCO 2017
Benchmarks: ImageNet-1K MS COCO 2017
Metrics: top-1 accuracy top-5 accuracy throughput (images/s) latency (ms) AP APmask FLOPs
Baselines: MobileViTv2 EfficientViT-M4 FastViT EfficientFormer MobileOne EfficientNet FasterNet PoolFormer EdgeViT MobileNetV3
Builds on: MetaFormer EfficientViT DeiT Swin Transformer MobileViTv2
Closest papers by meaning
- FasterViT: Fast Vision Transformers with Hierarchical Attention (ICLR 2024 · similarity 0.88)
- You Only Need Less Attention at Each Stage in Vision Transformers (CVPR 2024 · similarity 0.88)
- RepViT: Revisiting Mobile CNN From ViT Perspective (CVPR 2024 · similarity 0.88)
- HydraViT: Stacking Heads for a Scalable ViT (NeurIPS 2024 · similarity 0.88)
- Mobile Attention: Mobile-Friendly Linear-Attention for Vision Transformers (ICML 2024 · similarity 0.88)
- Real-time Core-Periphery Guided ViT with Smart Data Layout Selection on Mobile Devices (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.