[Paper Review] Fast Vision Transformers with HiLo Attention

Zizheng Pan, Jianfei Cai+1 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications112 citations

TL;DR

The paper introduces LITv2, a fast Vision Transformer backbone using HiLo attention that disentangles high- and low-frequency patterns for local high-frequency and global low-frequency interactions, achieving faster throughput and competitive accuracy across tasks.

ABSTRACT

Vision Transformers (ViTs) have triggered the most recent and significant breakthroughs in computer vision. Their efficient designs are mostly guided by the indirect metric of computational complexity, i.e., FLOPs, which however has a clear gap with the direct metric such as throughput. Thus, we propose to use the direct speed evaluation on the target platform as the design principle for efficient ViTs. Particularly, we introduce LITv2, a simple and effective ViT which performs favourably against the existing state-of-the-art methods across a spectrum of different model sizes with faster speed. At the core of LITv2 is a novel self-attention mechanism, which we dub HiLo. HiLo is inspired by the insight that high frequencies in an image capture local fine details and low frequencies focus on global structures, whereas a multi-head self-attention layer neglects the characteristic of different frequencies. Therefore, we propose to disentangle the high/low frequency patterns in an attention layer by separating the heads into two groups, where one group encodes high frequencies via self-attention within each local window, and another group encodes low frequencies by performing global attention between the average-pooled low-frequency keys and values from each window and each query position in the input feature map. Benefiting from the efficient design for both groups, we show that HiLo is superior to the existing attention mechanisms by comprehensively benchmarking FLOPs, speed and memory consumption on GPUs and CPUs. For example, HiLo is 1.4x faster than spatial reduction attention and 1.6x faster than local window attention on CPUs. Powered by HiLo, LITv2 serves as a strong backbone for mainstream vision tasks including image classification, dense detection and segmentation. Code is available at https://github.com/ziplab/LITv2.

Motivation & Objective

• Motivate evaluation of ViTs by direct speed measurements on target hardware, not only FLOPs.
• Propose HiLo attention to separately model high-frequency local details and low-frequency global structure.
• Show that HiLo-based LITv2 achieves faster throughput with competitive accuracy on ImageNet, COCO, and ADE20K.
• Demonstrate reduced memory usage and efficiency advantages over state-of-the-art ViTs and some CNNs.

Proposed method

• Introduce HiLo attention that splits attention heads into high-frequency (Hi-Fi) and low-frequency (Lo-Fi) groups.
• Hi-Fi uses local window self-attention on high-resolution maps (e.g., 2x2 windows).
• Lo-Fi applies average pooling to each window to obtain low-frequency signals and performs attention with down-sampled keys/values from these pooled signals.
• Split heads by a ratio α, with (1−α)N_h allocated to Hi-Fi and αN_h to Lo-Fi, then concatenate the outputs.
• Replace fixed relative positional encoding with a depthwise 3x3 convolution in FFN to inject position information and enlarge receptive fields.
• Adopt HiLo in LITv2 across three model sizes (S/M/B) with training on ImageNet-1K and evaluation on downstream tasks.

Experimental results

Research questions

• RQ1 Can HiLo attention provide a speedup on CPU/GPU while maintaining competitive accuracy compared to standard MSAs and other efficient attentions?
• RQ2 How does HiLo’s high/low-frequency disentangling affect performance on high-resolution and dense prediction tasks?
• RQ3 Do architectural changes (e.g., convolutional FFN, removing RPE) further improve speed and accuracy across vision tasks?
• RQ4 What is the optimal head split α and window size s for different workloads and resolutions?
• RQ5 Do HiLo-based backbones improve throughput and memory usage in downstream detection/segmentation tasks compared to SoTA ViTs?

Key findings

• LITv2 with HiLo outperforms LITv1 in Top-1 accuracy while using fewer FLOPs and achieving faster throughput.
• HiLo yields inference speedups of 13% (S), 27% (M), and 36% (B) on ImageNet-1K over LITv1 when measured on GPUs.
• For comparable FLOPs, LITv2-S achieves faster throughput than PVT-S and Twins-PCPVT-S on ImageNet-1K.
• On COCO object detection/instance segmentation, LITv2 backbones deliver higher AP than ResNet backbones and competitive or superior results to SoTA ViTs, with favorable FPS and FLOPs.
• On ADE20K semantic segmentation, LITv2 variants achieve higher mIoU and FPS than several baselines (e.g., LITv2-S/LITv2-M/LITv2-B surpass Swin-Ti/Swin-S/Swin-B in mIoU with increased FPS).
• HiLo consistently reduces FLOPs and improves speed versus other efficient attention mechanisms (SRA, W-MSA, T-MSA) on ImageNet-1K and downstream tasks.