[Paper Review] EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention

Xinyu Liu, Houwen Peng+4 more|arXiv (Cornell University)|May 11, 2023

Advanced Memory and Neural Computing36 citations

TL;DR

EfficientViT introduces a memory-efficient vision transformer with a sandwich layout and cascaded group attention to reduce memory access and computation, yielding higher speed and accuracy than prior efficient models.

ABSTRACT

Vision transformers have shown great success due to their high model capabilities. However, their remarkable performance is accompanied by heavy computation costs, which makes them unsuitable for real-time applications. In this paper, we propose a family of high-speed vision transformers named EfficientViT. We find that the speed of existing transformer models is commonly bounded by memory inefficient operations, especially the tensor reshaping and element-wise functions in MHSA. Therefore, we design a new building block with a sandwich layout, i.e., using a single memory-bound MHSA between efficient FFN layers, which improves memory efficiency while enhancing channel communication. Moreover, we discover that the attention maps share high similarities across heads, leading to computational redundancy. To address this, we present a cascaded group attention module feeding attention heads with different splits of the full feature, which not only saves computation cost but also improves attention diversity. Comprehensive experiments demonstrate EfficientViT outperforms existing efficient models, striking a good trade-off between speed and accuracy. For instance, our EfficientViT-M5 surpasses MobileNetV3-Large by 1.9% in accuracy, while getting 40.4% and 45.2% higher throughput on Nvidia V100 GPU and Intel Xeon CPU, respectively. Compared to the recent efficient model MobileViT-XXS, EfficientViT-M2 achieves 1.8% superior accuracy, while running 5.8x/3.7x faster on the GPU/CPU, and 7.4x faster when converted to ONNX format. Code and models are available at https://github.com/microsoft/Cream/tree/main/EfficientViT.

Motivation & Objective

• Analyze factors affecting Vision Transformer (ViT) inference speed, focusing on memory access, computation redundancy, and parameter usage.
• Design a memory-efficient ViT building block to improve throughput without sacrificing accuracy.
• Develop a new attention mechanism to reduce redundancy and enhance feature diversity in heads.
• Reallocate parameters to improve efficiency of critical modules while pruning less important components.
• Demonstrate transferability of EfficientViT to downstream vision tasks.

Proposed method

• Introduce sandwich layout blocks where a single memory-bound MHSA is sandwiched between FFN layers to reduce memory-bound operations.
• Propose Cascaded Group Attention (CGA) that feeds attention heads with different feature splits and cascades outputs across heads to increase diversity and reduce computation.
• Apply Taylor structured pruning to identify important channels and guide parameter reallocation, expanding critical module width while shrinking less important parts.
• Use overlapping patch embedding and a three-stage hierarchical architecture with BN throughout to improve speed and practicality on hardware.
• Evaluate throughput and accuracy on ImageNet-1K across GPU/CPU/ONNX and test transferability to downstream tasks.

Figure 1 : Speed and accuracy comparisons between EfficientViT (Ours) and other efficient CNN and ViT models tested on an Nvidia V100 GPU with ImageNet-1K dataset [ 17 ] .

Experimental results

Research questions

• RQ1 How can ViTs be redesigned to minimize memory-bound operations during inference?
• RQ2 Can feeding attention heads with split feature channels (CGA) reduce redundancy and increase attention diversity without sacrificing accuracy?
• RQ3 What parameter reallocation strategies yield better speed-accuracy trade-offs for lightweight ViTs?
• RQ4 Do the proposed EfficientViT blocks generalize to downstream vision tasks and deployments on real hardware?

Key findings

• EfficientViT-M5 achieves 77.1% Top-1 accuracy with 10,621 images/s throughput on Nvidia V100, outperforming MobileNetV3-Large by 1.9% in accuracy and 40.4% (GPU) / 45.2% (CPU) in throughput.
• EfficientViT-M2 reaches 70.8% Top-1 accuracy, beating MobileViT-XXS by 1.8% and delivering 5.8x GPU and 3.7x CPU speedups, with 7.4x faster ONNX performance.
• EfficientViT-M4 outperforms several efficient models on ImageNet with higher throughput and competitive accuracy (e.g., vs. EdgeViT-XXS, 4.4x GPU and 3.0x CPU faster).
• Replacing memory-bound MHSA with the sandwich layout reduces memory time consumption while increasing FFN-based channel communication.
• CGA reduces attention computation by distributing heads across feature splits and cascading outputs, improving efficiency and accuracy (ablation shows gains over MHSA).
• Parameter reallocation, increasing V-channel width and reducing Q/K dimensions and FFN expansion, yields measurable accuracy and throughput gains.

Figure 2 : Runtime profiling on two standard vision transformers Swin-T and DeiT-T. Red text denotes memory-bound operations, i.e. , the time taken by the operation is mainly determined by memory accesses, while time spent in computation is much smaller.