PaperlayerView exact revision on arXiv

Source-linked AI summary

SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design

Seokju Yun, Youngmin Ro

arXiv:2401.16456v2cs.CV

TL;DR

Efficient Vision Transformers contain computational redundancy in their macro architecture and multi-head attention, motivating a memory-efficient redesign. SHViT uses larger-stride patchification and single-head attention to reduce this redundancy, achieving strong speed–accuracy tradeoffs across devices and vision tasks. On ImageNet-1K, SHViT-S4 is 1.3% more accurate and 2.4× faster on iPhone12 than MobileViTv2×1.0.

  • Problem

    Efficient Vision Transformers incur computational redundancy from high-resolution token processing and multi-head attention, limiting their speed on resource-constrained devices.

  • Method

    SHViT combines a 16×16 patchify stem with a 3-stage hierarchy and Single-Head Self-Attention that processes partial channels to reduce spatial, channel, and memory-access redundancy.

  • Results

    SHViT achieves state-of-the-art speed–accuracy performance across classification, detection, and segmentation; SHViT-S4 is 1.3% more accurate and 2.4× faster on iPhone12 than MobileViTv2×1.0.

  • Takeaways & Limitations

    The proposed macro and micro designs provide a fast, accurate SHViT family across diverse devices and vision tasks.

Abstract

from arXiv · show

Recently, efficient Vision Transformers have shown great performance with low latency on resource-constrained devices. Conventionally, they use 4x4 patch embeddings and a 4-stage structure at the macro level, while utilizing sophisticated attention with multi-head configuration at the micro level. This paper aims to address computational redundancy at all design levels in a memory-efficient manner. We discover that using larger-stride patchify stem not only reduces memory access costs but also achieves competitive performance by leveraging token representations with reduced spatial redundancy from the early stages. Furthermore, our preliminary analyses suggest that attention layers in the early stages can be substituted with convolutions, and several attention heads in the latter stages are computationally redundant. To handle this, we introduce a single-head attention module that inherently prevents head redundancy and simultaneously boosts accuracy by parallelly combining global and local information. Building upon our solutions, we introduce SHViT, a Single-Head Vision Transformer that obtains the state-of-the-art speed-accuracy tradeoff. For example, on ImageNet-1k, our SHViT-S4 is 3.3x, 8.1x, and 2.4x faster than MobileViTv2 x1.0 on GPU, CPU, and iPhone12 mobile device, respectively, while being 1.3% more accurate. For object detection and instance segmentation on MS COCO using Mask-RCNN head, our model achieves performance comparable to FastViT-SA12 while exhibiting 3.8x and 2.0x lower backbone latency on GPU and mobile device, respectively.

1. Introduction

SHViT addresses computational redundancy in efficient Vision Transformers at both macro and micro design levels. It combines memory-efficient architecture and single-head attention to improve the speed–accuracy tradeoff across devices and vision tasks.

  • Motivation: Efficient Vision Transformers combine convolutional architecture with cost-efficient attention to meet real-time constraints, but global attention remains computationally expensive.ViTs model long-range dependencies effectively, while vanilla global attention has quadratic complexity with image size.
  • Macro design: 4×4 patch embeddings create early-stage bottlenecks because high-resolution stages process many tokens.At 224×224 resolution, stage 1 processes 3136 tokens and stage 2 processes 784 tokens.
  • Macro design: A 16×16 patchify stem with a 3-stage design is 3.0× faster on GPU and 2.8× faster on CPU than the compared 4-stage design, with a 1.5% accuracy decrease.At 256×256 training resolution, the alternative design is comparable in performance while remaining significantly faster.
  • Micro design: Multi-head attention contains noticeable channel redundancy, particularly in later stages, motivating Single-Head Self-Attention that processes only a subset of channels.The remaining channels stay unchanged, reducing computational redundancy and memory access cost while allowing wider and deeper models within the same budget.
  • Results: SHViT-S4 reaches 79.4% ImageNet top-1 accuracy, with throughput of 14283 images/s on an Nvidia A100 GPU and 509 images/s on an Intel Xeon CPU.It outperforms EfficientNet-B0 by 2.3% in accuracy, 69.4% in GPU inference speed, and 90.6% in CPU speed.
  • Results: SHViT achieves a favorable speed–accuracy tradeoff across classification, detection, and segmentation on GPU, CPU, mobile, and other inference platforms.The paper reports 1.3% higher accuracy and 2.4× faster iPhone12 inference than MobileViTv2×1.0, plus lower-latency COCO detection and segmentation backbones than EfficientViT-M4.

2. Analysis and Method

The paper analyzes redundancy in macro and micro Vision Transformer design, then introduces SHViT with larger-stride tokenization, convolutional early stages, and single-head attention in later stages. Its design combines partial-channel attention with convolution to improve the speed–accuracy tradeoff while reducing memory overhead.

  • Macro Design: 3.0× / 2.8× faster on GPU / CPU, respectively, model (b) is 1.5% worse than model (a) at the original resolution, while its 256×256 variant is comparable and significantly faster.The comparison uses models with similar channel counts for equivalent feature-map sizes.
  • Micro Design: Early-stage attention heads often behave convolutionally, while depthwise convolution provides a more favorable speed–accuracy tradeoff than attention in the first stage.This motivates using convolutions with spatial inductive bias as the initial token mixer.
  • Micro Design: 78.3% average head similarity in latter-stage DeiT-T, or 64.8% with 6 heads, indicates substantial multi-head redundancy; using one of 12 or 24 Swin-T heads reduces performance by only 0.95% points on average.Most heads can be removed with limited accuracy change, and some removals slightly improve Swin-T scores.
  • Single-Head Self-Attention: SHSA applies single-head attention to only partial input channels and leaves the remainder untouched, with r = 1/4.67 by default.The attention output and residual channels are concatenated before a projection applied to all channels.
  • Single-Head Self-Attention: The partial-channel approach combines complementary local and global features in parallel, while reducing memory-bound operations such as reshaping and normalization.This design is intended to avoid head redundancy during both training and inference and better exploit GPU/CPU compute.

3. Experiments

Experiments show that SHViT delivers a strong accuracy–throughput tradeoff across classification, mobile inference, and downstream vision tasks. Its large-stride macro design and single-head attention support efficient performance across hardware platforms.

  • ImageNet-1K Classification: SHViT achieves a better accuracy–throughput/latency tradeoff across GPU, CPU, ONNX, and mobile platforms.The comparison covers ImageNet-1K classification and multiple efficient model families.
  • ImageNet-1K Classification: 79.1–79.4% top-1 accuracy accompanies 2.9×/3.3× higher speed than EfficientNet-B1 on the A100 GPU/Intel CPU with ONNX format.
  • ImageNet-1K Classification: 1.9% higher accuracy and 2.6×, 1.7×, and 4.2× faster inference than EdgeViT-XS are achieved by SHViT-S4 on A100 GPU, Intel CPU, and ONNX, respectively.The authors attribute the ONNX advantage partly to fewer reshape operations in the single-head design.
  • Mobile Latency Evaluation: At 1024 × 1024 resolution, SHViT-S4 has 34.4% and 69.7% lower latency than FastViT and EfficientFormer on iPhone 12.At low resolutions, SHViT-S4 is slightly slower than those models.
  • Downstream Tasks: SHViT-S4 is 2.3× faster than MobileNetV3 on mobile detection and exceeds it by +8.9 AP.Against MobileFormer, it is 3.2× faster on the A100 GPU and 8.2× faster on mobile while achieving better performance.
  • Downstream Tasks: SHViT-S4 gains 1.7 APb and 1.3 APm over PoolFormer while running 4.3×, 8.1×, and 2.1× faster on GPU, CPU, and mobile.
  • Ablation and Design Analysis: The 16×16 patchify stem and 3-stage hierarchy reduce computational cost while producing meaningful high-resolution token representations.The authors connect the remaining performance gap against EfficientViT to their micro-design choices.

4. Related Work

Efficient vision models combine CNN-like local processing with selective attention to reduce inference cost. Related work also identifies redundancy among attention heads, motivating alternatives to conventional multi-head designs.

  • Efficient CNNs: CNN efficiency strategies include convolution decomposition, channel shuffling, cheap linear transformations, compound scaling, and structural re-parameterization.
  • Efficient ViT Architectures: Efficient ViT architectures commonly use convolution for high-resolution features and attention for low-resolution, high-level features.
  • Attention Redundancy: Several studies find that attention heads often behave similarly and can be pruned without notably affecting performance.EfficientViT instead feeds heads different channel splits, while other work regularizes or measures head similarity.

5. Conclusion

SHViT addresses spatial and channel redundancy with 16×16 patch embeddings, 3-scale hierarchical representations, and Single-Head Self-Attention. The resulting networks target fast inference and strong performance across devices and vision tasks.

  • Conclusion: SHViT combines 16×16 patch embeddings, 3-scale hierarchical representations, and Single-Head Self-Attention to address spatial and channel redundancy.
  • Conclusion: The model family achieves ultra-fast inference and high performance across diverse devices and vision tasks.
  • Future Work: Future work targets cost-effective fine-grained features for high-resolution detail and small-object recognition.The authors also identify integrating single-head design into sophisticated attention methods as an open direction.

A. Comparison with Tiny Variants of Large-scale Models

At higher input resolutions, SHViT remains competitive with tiny variants of larger-scale models while offering substantial throughput advantages. Its speed advantage extends across GPU and CPU inference.

  • High-Resolution Comparison: SHViT-S4r384 is 0.3% less accurate than Swin-T but 2.3× faster on the A100 GPU and 9.5× faster on the Intel CPU.
  • High-Resolution Comparison: Increasing resolution improves SHViT speed relative to recent models on mobile and other inference platforms.The authors present this as evidence that SHViT can be competitive in real-world applications.

B. Memory Efficiency Analysis

SHViT addresses deployment memory costs despite using more parameters than lightweight models. Its macro and micro designs reduce feature-map memory usage, yielding lower test memory than EfficientNet-B0.

  • SHViT has more parameters than lightweight models, with SHViT-S3 using 2.7× more parameters than EfficientNet-B0.The paper emphasizes that parameter count is not the only deployment consideration.
  • Memory access cost is especially important for I/O-bound devices, larger batch sizes, and high-resolution inputs.The first term in the memory-access equation becomes more critical in these settings.
  • SHViT’s macro and micro designs reduce memory usage by eliminating redundancy in feature-map spatial dimensions and channels.
  • Despite having more parameters than EfficientNet-B0, SHViT consumes less test memory.

C. Further Results on COCO Detection

The paper extends single-head attention to detection and reports speed advantages across inference settings. In Deformable DETR, the module improves test speed without significant performance degradation, while related comparisons cover classification and latency benchmarks.

  • Single-head attention is applied to encoder self-attention and decoder cross-attention, which carry significant computational costs.The design targets these layers to enhance detection speed.
  • SHViT-S4 latency comparisons with FastViT and EfficientFormer are measured across image resolutions on an A100 GPU and Intel CPU.
  • Single-head modules are used in the 2nd, 3rd, and 4th encoder and decoder layers, while multi-head designs remain in the initial and latter layers considered vital.The head dimension is increased from 32 to 64 to minimize performance degradation.
  • 24% test-speed improvement is achieved with the single-head attention module in Deformable DETR without significant performance degradation.

D. More Details on Redundancy Experiments

The redundancy experiments quantify similarity and test-time importance among attention heads. They identify substantial latter-stage redundancy while showing that early-stage head removal can significantly reduce performance.

  • HeadSim_i averages pairwise cosine similarity among heads within layer i and then averages the value across batches.The measure uses N_h heads and excludes self-comparisons.
  • Head ablation sets binary mask variables δ_i to zero for removed heads while retaining the remaining MHSA computation.With all masks equal to one, the layer is equivalent to MHSA.
  • The ablation study reports each layer’s best accuracy after reducing that layer to its single most important head.
  • DeiT-S-Distill shows significant computational redundancy among many heads in the latter stages.
  • Removing heads in the early stages causes a relatively substantial performance decline, where many heads operate similarly to convolution.

E. Further Discussions on Related Works

The paper distinguishes its macro-design motivation from prior patch-embedding work by analyzing spatial and channel redundancy directly. It also positions its treatment of partial channels as differing in motivation from prior FLOPs- or throughput-focused approaches.

  • SHViT analyzes patch-embedding redundancy separately in spatial and channel dimensions rather than deriving patch size through experiment grafting.
  • The paper states that existing works had not analyzed the speed and memory effects of resolving spatial redundancy across diverse devices and tasks.
  • Prior partial-channel designs primarily focused on FLOPs or throughput and used depthwise or vanilla convolutions on partial channels.
Loading 2401.16456v2…