Fast Vision Transformers with HiLo Attention
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 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.
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.4 × faster than spatial reduction attention and 1.6 × 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 .
Introduction
Real-world applications usually require a model to have an optimal speed and accuracy trade-off under limited computational budget, such as UAV and autonomous driving.
This motivates substantial works toward efficient vision Transformer (ViT) design, such as PVT [51], Swin [32] and Focal
Transformer [60], among others.
To measure the computational complexity, a widely adopted metric in recent ViT design is the number of float-point operations, i .
e ., FLOPs.
However, FLOPs is an indirect metric, which can not directly reflect the real speed on the target platform.
For example,
Focal-Tiny is much slower than Swin-Ti on GPUs although their FLOPs are comparable.
In general, the discrepancy between the indirect metric (FLOPs) and the direct metric (speed) in recent ViTs can be attributed to two main reasons.
First, although self-attention is efficient on low-resolution feature maps, the quadratic complexity in both memory and time makes it much slower on high-resolution images due to intensive memory access cost [34], where fetching data from off-chip DRAM can be speed-consuming.
Second, some efficient attention mechanisms in ViTs have low theoretical complexity guarantee but are actually slow on GPUs due to particular operations
that
[Figure]
Attention
C on c a t
P r o j e c t i on
QKV
Scaled Dot-Product
Attention
C on c a t
P r o j e c t i on
High Frequency Attention (Hi-Fi) are not hardware-friendly or cannot be parallelized, such as the multi-scale window partition [60], recursion [44] and dilated window [20].
With these observations, in this paper we propose to evaluate ViT by the direct metric, i .
e ., throughput, not only FLOPs.
Based on this principle, we introduce LITv2, a novel efficient and accurate vision
Transformer that outperforms most state-of-the-art (SoTA) ViTs on standard benchmarks while being practically faster on GPUs.
LITv2 is bulit upon LITv1 [36], a simple ViT baseline which removes all multi-head self-attention layers (MSAs) in the early stages while applying standard MSAs in the later stages.
Benefit from this design, LITv1 is faster than many existing works on ImageNet classification due to no computational cost from the early MSAs while the later MSAs only need to process downsampled low-resolution feature maps.
However, the standard MSA still suffers from huge computational cost on high-resolution images, especially for dense prediction tasks.
To address this problem, we propose a novel efficient attention mechanism, termed HiLo .
HiLo is motivated by the fact that natural images contain rich frequencies where high/low frequencies play different roles in encoding image patterns, i .
e ., local fine details and global structures, respectively.
A typical MSA layer enforces the same global attention across all image patches without considering the characteristics of different underlying frequencies.
This motivates us to propose to separate an
MSA layer into two paths where one path encodes high-frequency interactions via local self-attention with relatively high-resolution feature maps while the other path encodes low-frequency interactions via global attention with down-sampled feature maps, which leads to a great efficiency improvement.
Specifically, HiLo employs two efficient attentions to disentangle Hi gh/ Lo w frequencies in feature maps.
As shown in Figure 1, in the upper path, we allocate a few heads to the high frequency attention (Hi-Fi) to capture fine-grained high frequencies by local window self-attention ( e .
g .,
2 × 2 windows), which is much more efficient than standard MSAs.
The lower path, implementing the low-frequency attention (Lo-Fi), first applies average pooling to each window to obtain low-frequency signals.
Then, we allocate the remaining heads for Lo-Fi to model the relationship between each query position in the input feature map and the average-pooled low-frequency keys from each window.
Benefit from the reduced length of keys and values, Lo-Fi also achieves significant complexity reduction.
Finally, we concatenate the refined high/low-frequency features and forward the resulting output into subsequent layers.
Since both Hi-Fi and Lo-Fi are not equipped with time-consuming operations such as dilated windows and recursion, the overall framework of HiLo is fast on both CPUs and GPUs.
We show by comprehensive benchmarks that HiLo achieves advantage over the existing attention mechanisms in terms of performance, FLOPs, throughput and memory consumption.
Besides, we find the fixed relative positional encoding in LITv1 dramatically slows down its speed on dense prediction tasks due to the interpolation for different image resolutions.
For better efficiency, we propose to adopt one 3 × 3 depthwise convolutional layer with zero-padding in each FFN to incorporate the implicitly learned position information from zero-padding [27].
Moreover, the 3 × 3 convolutional filters simultaneously help to enlarge the receptive field of the early multi-layer perceptron (MLP) blocks in LITv1.
Finally, we conduct extensive experiments on ImageNet, COCO and ADE20K to evaluate the performance of LITv2.
Comprehensive comparisons with SoTA models show that our architecture achieves competitive performance with faster throughput, making ViTs more feasible to run low-latency applications for real-world scenarios.
Related Work
Vision Transformers.
Vision Transformers are neural networks that adopt self-attention mechanisms into computer vision tasks.
In [18], Dosovitskiy et al .
propose a ViT for image classification, which inherits the similar architecture from a standard Transformer [48] in natural language processing (NLP) tasks.
Since then, subsequent works have been proposed to improve ViT by incorporating more convolutional layers [54, 61], introducing pyramid feature maps [51, 32], enhancing the locality [62], as well as automatically searching a well-performed architecture [5, 3] with neural architecture search (NAS).
Some others also seek for token pruning to accelerate the inference speed of ViTs [37] or applying ViT into low-level vision tasks [47].
Compared to existing works, this paper focuses on a general ViT-based backbone for computer vision (CV) tasks and aims to achieve better efficiency on
GPUs while maintaining competitive performance.
Efficient attention mechanisms.
Efficient attention mechanisms aim to reduce the quadratic complexity of standard MSAs.
Existing efforts in NLP can be roughly categories into low-rank decomposition [50], kernelization [28, 39], memory [40] and sparsity mechanism [10].
However, simply adopting these method usually performs suboptimally in CV tasks [32, 63].
In CV, representative efficient self-attention mechanisms includes spatial reduction attention (SRA)
[51], local window attention [32, 26] and Twins attention [12].
However, they only focus on either local or global attention at the same layer.
To address this problem, TNT [21] introduced additional global tokens and MixFormer [6] mixed local window attention with depthwise convolutional layers.
Some other attention mechanisms consider both simultaneously, such as Focal [60] and QuadTree [44].
However, due to the inefficient operations which are not hardware-friendly and cannot be reflected in FLOPs ( e . g ., multi-scale window partition, recursion), they are slow on GPUs even compared to standard
MSA.
To this end, the proposed HiLo attention simultaneously captures rich local-global information at the same MSA layer and is faster and more memory-efficient compared to the existing works.
Frequency domain analysis in vision.
The frequency domain analysis in CV has been well studied in the literature.
According to [13, 16], the low frequencies in an image usually capture global structures and color information while the high frequencies contain fine details of objects ( e .
g ., sharp edges).
Based on this insight, a plethora of solutions have been proposed for image super-resolution [66, 19], generalization [25], image re-scaling [56] and neural network compression [59, 7].
Furthermore,
Octave convolution [9] targeted convolutional layers and proposed to locally applies convolution on high/low-resolution feature maps, separately.
Different from it, the proposed HiLo is a novel attention mechanism that captures both local and global relationships with self-attention.
Background
Multi-head self-attention.
Transformers are built upon multi-head self-attention, which enables to capture long-range relationships for tokens at different positions.
Specifically, let X ∈ R N × D be the input sequence into a standard MSA layer, where N is the length of the input sequence and D refers to the number of hidden dimensions.
Each self-attention head calculates the query Q , key K and value V matrices with a linear transformation from X ,
[Equation]
where W q , W k , W v ∈ R D
× D h are learnable parameters and D h is the number of hidden dimensions for a head.
Next, the output of a self-attention head is a weighted sum over N value vectors,
[Equation]
For an MSA layer with N h heads, the final output is computed by a linear projection of the concatenated outputs from each self-attention head, which can be formulated by
[Equation]
where W o ∈ R ( N h × D h ) × D is a learnable parameter.
In practice, D is usually equal to N h × D h .
Overall, a standard MSA layer have the computational cost of 4 ND 2 +
Transformer blocks.
A standard vision Transformer as described in [18] consists of a patch embedding layer, several blocks and a prediction head.
Let l be the index of a block.
Then each block contains an MSA layer and a position-wise feed-forward network (FFN), which can expressed as
[Equation]
where LN denotes the LayerNorm [2] and an FFN consists of two FC layers with GELU [24] nonlinearity in between.
Recent works on ViT have proposed to divide the blocks into several stages (typically 4 stages) to generate pyramid feature maps for dense prediction tasks.
Furthermore, to reduce the computational cost on high-resolution feature maps in the early stages, the MSA in Eq.
(4) has been replaced with efficient alternatives, such as SRA [51] and W-MSA [32].
Bottlenecks of LITv1.
Recent studies have shown that the MSA layers in the early stages in a model still focus on local patterns [14].
With the same observation, LITv1 [36] removes all early
MSAs ( i .
e ., exclude Eq.
(4) in each block) while applying standard MSAs at the later stages.
This design principle has achieved better efficiency with competitive performance on ImageNet compared to PVT [51] and Swin [32].
However, LITv1 still has two main bottlenecks in speed: 1) Given a high-resolution image, the standard MSAs in the later stages still result in huge computational cost.
The fixed relative positional encoding [32] dramatically slows down the speed when dealing with different image resolutions.
This is due to interpolating the fixed-size positional encoding for each different image resolution.
In the next section, we describe a novel attention mechanism with zero padding positional encoding to comprehensively accelerate LITv1.
Method 4.1 HiLo Attention
We propose to separately process high/low frequencies in a feature map at an attention layer.
We name the new attention mechanism as HiLo, which is depicted in Figure 1.
Essentially, the low-frequency attention branch (Lo-Fi) is to capture the global dependencies of the input (image/features), which does not need a high-resolution feature map but requires global attention.
On the other hand, the high-frequency attention branch (Hi-Fi) is to capture the fine detailed local dependency, which requires a high-resolution feature map but can be done via local attention.
In the next, we describe the two attentions in detail.
High-frequency attention.
Intuitively, as high frequencies encode local details of objects, it can be redundant and computationally expensive to apply global attention on a feature map.
Therefore, we propose to design Hi-Fi to capture fine-grained high frequencies with local window self-attention ( e .
g .,
2 × 2 windows), which saves significant computational complexity.
Furthermore, we employ the simple non-overlapping window partition in Hi-Fi, which is more hardware-friendly compared to the time-consuming operations such as window shifting [32] or multi-scale window partition [60].
Low-frequency attention.
Recent studies have shown that the global attention in MSA helps to capture low frequencies [38].
However, directly applying MSA to high-resolution feature maps requires huge computational cost.
As averaging is a low-pass filter [49], Lo-Fi firstly applies average pooling to each window to get low-frequency signals in the input X .
Next, the average-pooled feature maps are projected into keys K ∈ R N/s 2 × D h and values V ∈ R N/s 2 × D h , where s is the window size.
The queries Q in Lo-Fi still comes from the original feature map X .
We then apply the standard attention to capture the rich low-frequency information in feature maps.
Note that due to the spatial reduction of K and V , Lo-Fi simultaneously reduces the complexity for both Eq.
(1) and Eq.
Head splitting.
A naive solution for head assignment is to allocate both Hi-Fi and Lo-Fi the same number of heads as the standard MSA layer.
However, doubling heads results in more computational cost.
In order to achieve better efficiency, HiLo separates the same number of heads in an MSA into two groups with a split ratio α , where (1 − α ) N h heads will be employed for Hi-Fi and the other αN h heads are used for Lo-Fi.
By doing so, as each attention has a lower complexity than a standard MSA, the entire framework of HiLo guarantees a low complexity and ensures high throughput on GPUs.
Moreover, another benefit of head splitting is that the learnable parameter W o can be decomposed into two smaller matrices, which helps to reduce model parameters.
Finally, the output of HiLo is a
[Figure]
concatenation of the outputs from each attention
[Equation]
where [ · ] denotes the concatenation operation.
Complexity Analysis.
Without loss of generality, we assume Hi-Fi and Lo-Fi have an equal number of heads ( i .
e .,
α = 0 .
and the feature map has equal width and height.
Then, Hi-Fi and Lo-Fi have a computational cost of 74 ND 2 + s 2 ND and ( 34 + 1 s 2 ) ND 2 + 1 s 2 N 2 D , respectively.
Derivation for this result can be found in the supplementary material.
As shown in Figure 2-(a) and (b), under a small input image resolution and a small value of s ( e .
g .
, s = 2 ), both Hi-Fi and Lo-Fi are comparably efficient.
However, with a much higher resolution, Lo-Fi will result in a huge computational cost as it still has a quadratic complexity in terms of N in Eq.
(2), i .
e ., 1 s 2 N 2 D .
In this case, slightly increasing s ( e . g ., s = 4 ) helps Lo-Fi achieve better efficiency while preserving the accuracy.
Combining the two attentions together, a larger window size also helps the overall framework of HiLo to reduce more FLOPs on high-resolution images, as shown in Figure 2-(c).
Thus, we suggest a practical guideline for adopting HiLo into existing frameworks: increasing the window size in order to get better efficiency on high-resolution images.
We further show in Section 5.2 that this principle helps
LITv2 achieve a better speed and accuracy trade-off on downstream tasks, e .
g ., dense object detection.
4.2 Positional Encoding
Positional encoding is essential to self-attention due to its permutation-invariant property.
In LITv1, the later MSAs adopt the same relative positional encoding (RPE) scheme as Swin [32].
This approach has significantly improves Swin by 0.7% in Top-1 accuracy on ImageNet compared to using absolute positional encoding [32].
However, on dense prediction tasks, the fixed RPE has to be interpolated for different image resolutions, which dramatically slows down the training/inference speed of LITv1.
As a recent study [27] has shown that position information can be implicitly learned from zero-padding in CNNs, we propose to adopt one layer of 3 × 3 depthwise convolutional layer with zero-padding in each FFN to replace the time-consuming RPE.
Notably, due to the elimination of early MSAs, the early blocks in LITv1 only have FFNs left, which results in a tiny receptive field of 1 × 1 .
To this end, we show in Section 5.4 that the 3 × 3 convolutional filters adopted in each FFN also improve LITv2 by simultaneously enlarging the receptive field in the early stages.
4.3 Model Architecture
LITv2 has three variants: LITv2-S, LITv2-M and LITv2-B, corresponding to the small, medium and base settings in LITv1, respectively.
For a fair comparison, we keep the network width and depth as the same as LITv1.
The overall modifications are simply in two steps: 1) Adding one layer of depthwise convolution with zero-padding in each FFN and removing all relative positional encodings in all MSAs.
2) Replacing all attention layers with the proposed HiLo attention.
Detailed architecture configurations can be found in the supplementary material.
Experiment
In this section we conduct experiments to validate the effectiveness of the proposed LITv2.
Following common practice [51, 32, 12, 60], we experiment LITv2 on three tasks, including image classification on ImageNet-1K [43], object detection and instance segmentation on COCO [31] and semantic segmentation on ADE20K [65].
[Table] [Table]
LITv2-S 3.
ResNet-101 [53] 7.
ConvNext-S [33] 8.
PVT-M [51] 6.
Twins-SVT-B [12] 8.
Swin-S [32] 8.
LITv1-M [36] 8.
LITv2-M 7.
ResNet-152 [53] 11.
ConvNext-B [33] 15.
Twins-SVT-L [12] 14.
Swin-B [32] 15.
LITv1-B [36] 15.
LITv2-B 13.
DeiT-B ↑ 384 [45] 55.
Swin-B ↑ 384 [32] 47.
OOM 6.
LITv2-B ↑ 39.
We conduct image classification experiments on ImageNet-1K [43], a large-scale image dataset which contains ∼ 1.2M training images and 50K validation images from 1K categories.
We measure the model performance by Top-1 accuracy.
Furthermore, we report the FLOPs, throughput, as well as training/test memory consumption on GPUs.
We compare with two CNN-based models [53, 33] and several representative SoTA ViTs [51, 32, 54, 60, 12].
Note that this paper does not consider mobile-level architectures [8, 35].
Instead, we focus on models with the similar model size.
Besides, we are also not directly comparable with NAS-based methods [3, 5] as LITv2 is manually designed.
Implementation details.
All models are trained for 300 epochs from scratch on 8 V100 GPUs.
At training time, we set the total batch size as 1,024.
The input images are resized and randomly cropped into 224 × 224 .
The initial learning rate is set to 1 × 10 − 3 and the weight decay is set to 5 × 10 − 2 .
We use AdamW optimizer with a cosine decay learning rate scheduler.
All training strategies including the data augmentation are same as in LITv1.
For HiLo, the window size s is set to 2.
The split ratio α is set to 0.9, which is chosen from a simple grid search on ImageNet-1K. The depthwise convolutional layers in FFNs are set with a kernel size of 3 × 3 , stride of 1 and zero padding size of 1 .
[Table]
Results.
In Table 1, we report the experiment results on ImageNet-1K. First, compared to LITv1 baselines, LITv2 achieves consistent improvement on Top-1 accuracy while using less FLOPs.
Moreover, benefit from HiLo, LITv2 achieves faster throughput and significant training time memory reduction ( e .
g ., 13% , 27% , 36% inference speedup for the small, medium and base settings, respectively)
compared to LITv1.
Second, compared to CNNs, LITv2 models outperform all counterparts of
ResNet and ConvNext in terms of FLOPs, throughput and memory consumption while achieving comparable performance.
Last, compared to SoTA ViTs, LITv2 surpasses many models in terms of throughput and memory consumption with competitive performance.
For example, under the similar amount of FLOPs, LITv2-S achieves faster inference speed than PVT-S and Twins-PCPVT-S with better performance.
Although Focal-Tiny achieves better Top-1 accuracy than LITv2-S, it runs much slower ( i .
e ., 384 vs. 1,471 images/s) and requires a large amount of memory to train.
Besides, when finetuning on a higher resolution, LITv2-B outperforms both DeiT-B and Swin-B with a faster throughput and lower complexity.
5.2 Object Detection and Instance Segmentation on COCO
In this section, we conduct experiments on COCO 2017, a common benchmark for object detection and instance segmentation which contains ∼ 118K images for the training set and ∼ 5K images for the validation set.
Following common practice [12, 51], we experiment with two detection frameworks:
RetinaNet [30] and Mask R-CNN [22].
We measure model performance by Average Precision (AP).
Implementation details.
All backbones are initialized with pretrained weights on ImageNet-1K.
We train each model on 8 GPUs with 1 × schedule (12 epochs) and a total batch size of 16.
For a fair comparison, we adopt the same training strategy and hyperparameter settings as in LITv1 [36].
Note that we pretrain LITv2 with a local window size of 2 and α = 0 .
9 on ImageNet-1K.
Under the same α , a larger window size helps to achieve lower complexity and thus improves the speed at high resolution, as explained in Section 4.1.
In this case, we also train models with a slightly larger window size of s = 4 for better efficiency, which we denote with “*”.
By default, FLOPs is evaluated based on the input resolution of 1280 × 800 .
FPS is measured on one RTX 3090 GPU based on the mmdetection [4] framework.
Results.
In Table 2, we report the experimental results on COCO.
In general, LITv2 outperforms
LITv1 by a large margin in almost all metrics.
Besides, our LITv2 significantly surpasses ResNet in terms of AP, though it runs slightly slower in some cases.
More importantly, our LITv2 beats all the compared SoTA ViTs, achieving the best AP with compelling fast inference speed.
Furthermore, by adopting a larger window size ( i . e ., s = 4 ), LITv2 achieves better efficiency with a slightly performance drop.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
5.3 Semantic Segmentation on ADE20K
[Table]
In this section
[Table]
we evaluate LITv2 on the semantic segmentation task.
We conduct experiments on ADE20K [65], a widely adopted dataset for semantic segmentation which has ∼ 20K training images, ∼ 2K validation images and ∼ 3K test images.
Following prior works, we adopt the framework of Semantic FPN [29] and measure the model performance by mIoU.
We train each model on 8 GPUs with a total batch size of 16 with 80K iterations.
All backbones are initialized with pretrained weights on ImageNet1K.
The stochastic depth for the small, medium and base models of LITv2 are 0.2, 0.2 and 0.3, respectively.
All other training strategies are the same as in LITv1 [36].
Results.
In Table 3, we compare LITv2 with
ResNet and representative ViTs on ADE20K.
In general, LITv2 achieves fast speed while outperforming many SoTA models.
For example, our
LITv2-S, LITv2-M and LITv2-B surpass SwinTi, Swin-S and Swin-B by 2.8%, 0.5% and 1.2% in mIoU with higher FPS, respectively.
5.4 Ablation Study
In this section, we provide ablation studies for LITv2, including the comparison with other efficient attention variants, the effect of α in HiLo, as well as the effect of architecture modifications.
By default, the throughput and memory consumption are measured on one RTX 3090 GPU with a batch size of 64 under the resolution of 224 × 224 .
Comparing HiLo with other attention mechanisms.
Based on LITv2-S, we compare the performance of HiLo with other efficient attention mechanisms on ImageNet-1K, including spatial reduction attention (SRA) in PVT [51], shifted-window based attention (W-MSA) in Swin [32] and alternated local and global attention (T-MSA) in Twins [12].
In our implementation, we directly
[Figure]
[Figure]
[Figure]
T opA cc u r a cy ( % ) 3.
F L O P s ( G )
“RPE”
[Table]
[Figure]
The magnitude is averaged over 100 samples.
The lighter the color, the larger the magnitude.
A pixel that is closer to the centre means a lower frequency.
Visualization code can be found in Section A.6.
replace HiLo with each compared method.
The results are reported in Table 4.
In general, HiLo reduces more FLOPs while achieving better performance and faster speed than the compared methods.
Furthermore, in Figure 3, we provide comprehensive benchmarks for more attention mechanisms based on different image resolutions, including Focal [60], QuadTree [44] and Performer [11].
Suffering from weak parallelizability, they are even slower than that of using standard MSAs on GPUs.
Compared to them, HiLo achieves competitive results in terms of the FLOPs, throughput and memory consumption.
Moreover, we conduct experiments based on ADE20K and Semantic FPN and show that HiLo achieves more performance gain than other attention mechanisms on the downstream dense prediction task.
Effect of α .
As shown in Figure 4, since the complexity of Lo-Fi is lower than Hi-Fi under the resolution of 224 × 224 and the window size of 2, a larger α helps to reduce more FLOPs as we allocate more heads to Lo-Fi.
Moreover, we found HiLo performs badly with α = 0 , in which case only the Hi-Fi is left and HiLo only focuses on high frequencies.
We speculate that low frequencies play an important role in self-attention.
For other values of α , we find the performance difference is around 0.2%, where α = 0 .
9 achieves the best performance.
However, it is worth noting that although the pure Lo-Fi branch ( α = 1 . 0 ) can achieve competitive results on ImageNet-1K, high-frequency signals play an important role in capturing fine object details, which is particularly important for dense prediction tasks such as semantic segmentation.
For example, with α = 0 . 9 , LITv2-S based
Semantic FPN achieves more performance gain (+0.6%) than that of using α = 1 .
Effect of architecture modifications.
Based on LITv2-S, we explore the effect of architecture modifications.
As shown in Table 5, benefit from the enlarged receptive field in the early stages, the adoption of depthwise convolutions improves the performance on both ImageNet and COCO.
Next, by removing the relative positional encoding, we significantly improve FPS on dense prediction tasks with a slightly performance drop on both datasets.
Also note that since depthwise convolutions have encoded positional information by zero paddings [27], the elimination of RPE does not result in a significant performance drop compared to prior works [32].
Finally, benefit from HiLo, we achieve more gains in model efficiency on both ImageNet and COCO.
Spectrum analysis of HiLo.
In Figure 5, we visualize the magnitude of frequency component [42] by applying Fast Fourier Transform (FFT) to the output feature maps from Hi-Fi and Lo-Fi attentions, respectively.
The visualisation indicates that Hi-Fi captures more high frequencies and Lo-Fi mainly focuses on low frequencies.
This strongly aligns with our aim of disentangling high and low frequencies in feature maps at a single attention layer.
[Table]
[Figure]
H i L o S R A
L o c a l W i n d o w X C A
L i n f o r m e r M S
A
S h
i f t e d
W i n d o w
C r o ss W i n d o w D A T
P e
r
f o
[Figure]
e
[Figure]
F
[Figure]
c a
[Figure]
H
[Figure]
r N
[Figure]
Q u
[Figure]
T
[Figure]
g
[Figure]
t ( i
[Figure]
Speed and performance comparisons with more ViTs on different GPUs.
We compare the inference speed with more models and on more types of GPUs.
Table 6 reports the results.
It shows that LITv2-S still achieves consistent faster throughput (images/s) than many ViTs on NVIDIA A100,
Tesla V100, RTX
[Figure]
and RTX 3090.
It is also worth noting that under similar performance (82.0%),
LITv2-S is 2.1 × faster than PVTv2-B2 [52], 1.7 × faster than XCiT-S12 [1] and ConvNext-Ti [33], and 3.5 × faster than Focal-Tiny [60] on V100, which is another common GPU version for speed test in previous works [32, 45, 33, 64].
Throughput comparisons with more attention mechanisms on CPUs and GPUs.
In Figure 6, we show that HiLo is consistently faster than many attention mechanisms [51, 32, 1, 50, 11, 17, 55, 60, 41, 44, 20, 18] on both CPUs and GPUs.
In particular, under CPU testing, HiLo is 1.4 × faster than
SRA [51], 1.6 × faster than local window attention [32] and 17.4 × faster than VAN [20].
Detailed benchmark configurations can be found in Section A.5.
Conclusion and Future Work
In this paper, we have introduced LITv2, a novel efficient vision Transformer backbone with fast speed on GPUs and outperforms most SoTA models on ImageNet and downstream tasks.
We have also presented HiLo attention, the core of LITv2 which helps to achieve better efficiency especially on high-resolution images.
With competitive performance, HiLo achieves great advantage over the existing attention mechanisms across FLOPs, throughput and memory consumption.
Future work may include incorporating convolutional stem [57] and overlapping patch embedding [52] for better performance, or extending HiLo on more tasks such as speech recognition and video processing.
Limitations and societal impact.
HiLo adopts a head splitting ratio to assign different numbers of heads into Hi-Fi and Lo-Fi.
In our experiments, this ratio is determined by a grid search on
ImageNet ( i .
e .,
α = 0 . 9 )
However, different tasks may have different importance on high and low frequencies.
Thus, the optimal value of α is task-specific and needs to be set manually.
Besides, our work potentially brings some negative societal impacts, such as the huge energy consumption and carbon emissions from large-scale training on GPU clusters.
A Appendix
We organize our supplementary material as follows.
• In Section A.1, we describe the architecture specifications of LITv2.
• In Section A.2, we provide the derivation for the computational cost of HiLo attention.
• In Section A.3, we study the effect of window size based on CIFAR-100.
• In Section A.4, we provide additional study on the Lo-Fi branch where we directly compute the queries from pooled feature maps.
• In Section A.5, we describe more details of the throughput benchmark for more attention mechanisms on CPUs and GPUs.
• In Section A.6, we provide more visualisation examples for spectrum analysis of HiLo attention.
[Table]
DTM
P 3 =
C 3 =
P 3 =
C 3 =
P 3 =
C 3 =
[Table]
  α 3 = 0 .s 3 =
H 3 =
E 3 =
 ×  α 3 = 0 .s 3 =
H 3 =
E 3 =
 ×  α 3 = 0 .s 3 =
H 3 =
E 3 =
 ×
[Table]
H
[Table]
DTM
P 4 =
C 4 =
P 4 =
C 4 =
P 4 =
C 4 =
[Table]
  α 4 = 1 .s 4 =
H 4 =
E 4 =
 ×  α 4 = 1 .s 4 =
H 4 =
E 4 =
 ×  α 4 = 1 .s 4 =
H 4 =
E 4 =
 ×
A.1 Architecture Specifications of LITv
The overall framework of LITv2 is depicted in Figure 7.
We also provide detailed architecture specifications of LITv2 in Table 7.
In general, we set the same network depth and width as LITv1.
It is worth noting that recent works [51, 32, 12, 60, 54] usually adopt standard MSAs at the last stage, including LITv1.
Following common practice, we set α = 1 .
and s = 1 at the last stage to make
HiLo behave as a standard MSA.
LITv2 also excludes MSAs in the first two stages due to the tiny receptive field of attention heads, as visualized in Figure 3 of LITv1 [36].
A.2 Computational Cost of HiLo Attention
Let N and D be the number of tokens and the number of hidden dimensions in an HiLo attention layer.
We denote s as the window size.
For simplicity, we assume Hi-Fi and Lo-Fi have an equal number of heads and the feature map has equal width and height.
Then, the computational cost of each attention comes from three parts: 1)
The projections of Q , K , V matrices.
The attention
L i n e a r
E m b e dd i n g
ConvFFNBlock
P a t c h
P a r t i t i o n
D T M
ConvFFNBlock
D T M
TransformerBlock
D T M
TransformerBlock 𝐻 4 ×𝑊4×𝐶 !
𝐻 8 ×𝑊8×𝐶 " 𝐻 16 ×𝑊16×𝐶 # 𝐻 32 ×𝑊32×𝐶 $
Stage
Stage Stage
Stage𝐻×𝑊×
C o n v FF
N
N o r m
H i
L o N o r m
C o n v FF
N
N o r m × 𝐿 !
× 𝐿 " × 𝐿 # × 𝐿 $
[Table]
[Table]
computation and weighted-sum of values.
The final linear projection of the weighted-sum values.
For Hi-Fi, the computational cost for each part is
[Equation] [Equation]
respectively.
Overall, this gives rise to a total computational cost of 74 ND 2 + s 2 ND for Hi-Fi.
Next, the computational cost for each part in Lo-Fi is
[Equation] [Equation]
respectively.
Thus, the total computational cost of Lo-Fi is ( 34 + 1 s 2 ) ND 2 + 1 s 2 N 2 D .
A.3 Effect of Window Size
Based on LITv2-S, we study the effect of window size in HiLo by experimenting on CIFAR-100.
As shown in Table 8, the window size does not affect the model parameters since the parameters of HiLo do not depend on it.
Moreover, as both Hi-Fi and Lo-Fi are comparably efficient under the small resolution of image classification ( i .
e ., 224 × 224), all settings have the comparable FLOPs, speed and
[Table] [Table]
memory footprint, where the difference is mainly due to the extra cost from padding on feature maps for window partition [32].
Overall, we find the window size of 2 performs the best, which therefore serves as our default setting in LITv2 for image classification.
Also note that as discussed in the main manuscript, a slightly larger window size ( e .
g ., 4) can help LITv2 achieve better efficiency on larger resolutions with a slightly performance drop.
A.4 Additional Study on the Lo-Fi Branch
In the proposed HiLo attention, the Lo-Fi branch computes queries from the original input feature maps.
An alternative approach is to directly compute the queries from the average-pooled feature maps.
However, in self-attention, the number of queries determines the spatial size of the output feature maps.
When computing the queries from pooled feature maps, the spatial size of the output feature maps is inconsistent with that of the original input.
One solution is to use interpolation (e.g. bilinear) and concatenate the interpolated feature maps with the outputs from Hi-Fi.
However, as shown in Table 9, this approach (denoted as "pooled queries") brings inferior performance and much slower throughput than our proposed design.
Note that although computing queries from pooled feature maps can slightly achieve a lower theoretical model complexity, frequently applying interpolation on GPUs results in a high memory access cost (MAC).
Therefore, it instead slows down the inference speed on GPUs.
A.5 Details of Throughput Benchmark for Different Attention Mechanisms
To evaluate the inference speed of HiLo on CPUs and GPUs, we benchmark the throughput based on a single attention layer and the standard settings of training ViT-B [18] on ImageNet.
Specifically, under the input resolution of 224 × 224, attention layers in ViT-B need to handle 14 × 14 (1/16 scale) feature maps, where each attention layer has 12 heads and each head has 64 dimensions.
For a fair comparison, we adopt the aforementioned configurations for all compared methods by default.
Besides, since different methods have distinct hyperparameters, we adopt their default settings for dealing with 1/16 scale feature maps.
For example, HiLo adopts a window size of 2 and alpha of 0.
when processing 1/16 scale feature maps.
In Table 10, we report more benchmark results.
Overall, we show that under a similar amount of parameters, a single layer of HiLo uses less FLOPs than compared methods, meanwhile it is faster on both CPUs and GPUs.
A.6 More Visualisations on Spectrum Analysis
In Figure 8 and Figure 9, we provide frequency magnitude visualisations for Hi-Fi and Lo-Fi attention outputs, respectively.
Clearly, the results indicate that Hi-Fi captures more high frequencies in
LITv2 while Lo-Fi mainly focuses on low frequencies.
We also provide the PyTorch-style code in
Algorithm 1 to explain our visualisation.
Algorithm 1 PyTorch-style Code for Visualising Frequency Magnitude.
[Table] [Table]
[Figure] [Figure]