Fast Vision Transformers with HiLo Attention
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 performs the
attention to model the global relationship between the average-pooled
low-frequency keys from each window and each query position in the input
feature map. Benefit from the efficient design for both groups, we show that
HiLo is superior to the existing attention mechanisms by comprehensively
benchmarking on FLOPs, speed and memory consumption on GPUs. 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/zip-group/LITv2.