FLatten Transformer: Vision Transformer using Focused Linear Attention

Introduction

Recent years have witnessed the vast development of Transformer and self-attention in the field of computer vision. With the advent of Vision Transformer , self-attention techniques have shown great potential in a variety of vision tasks including image classification , semantic segmentation , object detection , and multi-modal tasks .

However, applying Transformer to vision models is a non-trivial task. Unlike lightweight convolution neural networks , the quadratic computation complexity $\mathcal{O}(n^{2})$ with respect to sequence length $n$ leads to high computation costs when employing self-attention with a global receptive field. Previous works have sought to mitigate this challenge by confining the global receptive field to a smaller region, such as designing sparse global attention patterns or applying smaller attention windows . Albeit effective, these methods are either prone to disregarding informative features in other regions due to their attention patterns or inevitably sacrifice the ability to model long-range dependencies.

Linear attention, on the other hand, has been considered a simple yet effective alternative to address the computation dilemma by reducing the general complexity. Early research leverages a locally-sensitive hashing scheme that compresses the computation complexity from $\mathcal{O}(n^{2})$ to $\mathcal{O}(n{\rm log}(n))$. Nevertheless, it introduces a large constant before the complexity term, which makes it still unaffordable under common cases. More recent studies have noticed that the utilization of Softmax function in the self-attention operation practically compels a pairwise computation between all queries and keys, resulting in the predominant $\mathcal{O}(n^{2})$ complexity. To tackle this, several approaches adopt simple activation functions or tailored mapping functions to approximate the original Softmax function. As illustrated in Fig. 1, by changing the computation order from (query·key)·value to query·(key·value), the overall computation complexity can be reduced to $\mathcal{O}(n)$. However, compared to Softmax attention, current linear attention approaches still suffer from severe performance drop and may involve additional computation overhead from the mapping function, thereby constraining their practical application.

In this paper, we target on the limitations of current linear attention approaches and propose a novel Focused Linear Attention module, which achieves both high efficiency and expressiveness. Specifically, we undertake a dual-pronged analysis of the factors contributing to the performance decline in linear attention and subsequently propose corresponding solutions. First, the distribution of attention weight in the former linear attention modules is relatively smooth, lacking the focus ability to address the most informative features. As a remedy, we propose a simple mapping function to adjust the feature direction of queries and keys, making the attention weights more distinguishable. Second, we notice that the diminished rank of the attention matrix curtails the diversity of features in linear attention. To address this, we propose a rank restoration module by applying an additional depthwise convolution (DWC) to the original attention matrix, which helps to restore the matrix rank and keeps the output feature of different positions diversified. Leveraging these improved techniques, our module demonstrates comparable or superior performance to its Softmax counterparts, while enjoying the benefits of low computation complexity.

We empirically validate the effectiveness of our module on image classification, semantic segmentation, and object detection tasks using five advanced vision Transformer models. The results demonstrate consistent improvements over all baselines and other linear attention approaches.

Related Works

Transformer and self-attention mechanism are first introduced in the field of natural language processing and have earned wide research interest in computer vision. Nevertheless, the high computation complexity of self-attention set constraints on the direct application to vision tasks. Previous works have attempted to address this concern from several perspectives. The pioneer Vision Transformer considers reducing the input resolution by merging neighbouring pixels into a single token. Similar insights have been adopted in the following researches and also extend to downstream tasks . Another line of research reduces the feature resolution gradually and adopts carefully designed attention patterns to constrain the number of attentive tokens. For instance, PVT uses a sparse attention pattern and selects attentive tokens from a global perspective. DAT follows the path and designs a deformable attention module to achieve data-dependent attention pattern. Swin Transformer selects attentive tokens locally by dividing input into isolated windows. NAT follows the query-centric pattern in convolution and designs independent attentive tokens for all queries. Some researches also notice that convolution operations are valuable to Transformer models and may help to improve the overall efficiency . CMT combines Transformer blocks with efficient convolution operators like depthwise convolution , and achieves better efficiency-performance trade-off. ACmix shares the computation overhead of convolution and self-attention, and integrates both modules with limited cost. Methods have also been proposed for the efficient training of Transformers . In application scenarios demanding high efficiency, MobileFormer maintains two paths for convolution and Transformer respectively and enjoys the benefit from both modules. Dyn-Perceiver achieves efficient visual recognition through dynamic early exiting . MobileViT takes advantage of the success of MobileNets and uses the combination of mobilenet blocks and Transformer blocks to achieve light-weight and low latency.

However, these approaches still relied on the Softmax operator, whose inherit high computation complexity inevitably results in the inconvenience in model architecture design and practical application.

2 Linear Attention

Apart from the above methods, another line of research addresses high computation complexity with linear attention . Specifically, linear attention replaces the Softmax function in self-attention with separate kernel functions. In this case, linear attention does not have to compute the pairwise similarity $QK^{T}$ first. As illustrated in Fig. 1, based on the associative property of matrix multiplication, linear attention can change the computation order by computing $K^{T}V$ first, thus reducing the computation complexity from $\mathcal{O}(N^{2}d)$ to $\mathcal{O}(Nd^{2})$. Though efficient, how to design linear attention module as effective as softmax attention is a nontrivial problem. Performer approximates the Softmax operation with orthogonal random features. Efficient attention applies Softmax function to $Q$ and $K$ respectively, which naturally ensures each row of $QK^{T}$ sums up to 1. Nyströmformer and SOFT approximate the full self-attention matrix via matrix decomposition. Hydra attention replaces Softmax with cosine similarity and proposes hydra trick which reduces the computation complexity to $\mathcal{O}(Nd)$. EfficientVit uses depth-wise convolution to improve linear attention’s local feature extraction capacity. Castling-ViT proposes linear angular kernel to measure spectral similarity between each $Q_{i}$ and $K_{j}$.

Nevertheless, current linear attention designs either do not have enough expressive capability to catch up with Softmax attention or involve additional computation overhead from the complex kernel function. In this work, we analyze the reasons for the performance drop of linear attention from the focus ability and feature diversity perspectives. Based on these analyses, we propose a novel linear attention module called focused linear attention which achieves better performance than Softmax attention with lower computation complexity (Fig. 2).

Preliminaries

Due to the quadratic computation complexity, simply using self-attention with global receptive field becomes intractable, which usually leads to excessive computation costs. Previous works either addressed this concern by designing sparse global attention pattern or applying smaller attention windows . Though effective, these approaches become susceptible to the carefully-designed attention patterns, or inevitably sacrifice the ability to model long-range dependencies.

2 Linear Attention

Comparably, linear attention is considered as an effective alternative which restricts the computation complexity from $\mathcal{O}(N^{2})$ to $\mathcal{O}(N)$. Specifically, carefully designed kernels are introduced as the approximation of the original similarity function, i.e.,

where the self-attention module can be rewritten as:

In this way, we can change the computation order from $(QK^{T})V$ to $Q(K^{T}V)$ based on the associative property of matrix multiplication (as illustrated in Fig. 1):

where the computation complexity with respect to token number is reduced to $\mathcal{O}(N)$.

However, current linear attention approaches also face the dilemma between model complexity and expressiveness. On one hand, simple approximations, e.g., using ReLU activation , are too loose and lead to significant performance drop. On the other hand, carefully designed kernel functions or matrix decomposition approaches may incur additional computation overhead. In general, there is still a gap between the practical performance of linear attention and Softmax attention.

Focused Linear Attention

Although enjoying linear computational complexity, various previous works have also proved that simply replacing Softmax attention with linear attention usually results in severe performance drop . In this section, we first perform a detailed analysis of the inferior performances of linear attention from two perspectives: focus ability and feature diversity. Then, we introduce our Focused Linear Attention which adequately addresses these concerns and achieves high efficiency and expressive capability.

Softmax attention practically provides a nonlinear re-weighting mechanism, which makes it easy to concentrate on important features . As shown in Fig. 3, the distribution of attention map from Softmax attention is especially sharp on certain regions, e.g., foreground objects. Comparably, the distribution in linear attention is relatively smooth, making its output closer to the average of all features and failing to focus on more informative regions.

As a remedy, we propose a simple yet effective solution by adjusting the direction of each query and key features, driving similar query-key pairs closer while pushing dissimilar query-key pairs away. Specifically, we present a simple mapping function $f_{p}$ called Focused Function:

and $x^{\ast\ast p}$ represents element-wise power $p$ of $x$. We follow previous linear attention modules to use the ReLU function first to ensure the non-negativity of input and validity of denominator in Eq.(4). A direct observation is that the norm of the feature is preserved after the mapping, i.e., $\left\|x\right\|\!=\!\left\|f_{p}(x)\right\|$, indicating that only feature direction is adjusted.

On this basis, we show that under mild assumptions, the proposed mapping function $f_{p}$ practically affects the distribution of attention.

For a pair of feature $\{x,y\}$ with $m\!\neq\!n$:

Please refer to Appendix for complete proof. ∎

Therefore, with a proper $p$, our focused function $f_{p}(\cdot)$ practically achieves a more distinguished difference between similar query-key pairs (Eq. 7) and dissimilar query-key pairs (Eq. 8), restoring the sharp attention distribution as the original Softmax function.

For better understanding, we give an example to show the effects of $f_{p}$ in Fig. 4. It can be seen that $f_{p}$ actually “pulls” each vector to its nearest axis, and $p$ determines the degree of this “pulling”. By doing so, $f_{p}$ helps divide the features into several groups according to their nearest axes, improving the similarity within each group while reducing the similarity between the groups. The visualizations are in accordance with our analysis above.

2 Feature diversity

Apart from focus ability, feature diversity is also one of the factors that set restriction on the expressive power of linear attention. One of the possible reasons may give credit to the rank of the attention matrix , where a significant difference can be seen. Take one of the Transformer layers from DeiT-Tiny with $N\!=\!14\!\times\!14$ for example, we can see from Fig. 5 (a) that the attention matrix has the full rank (196 out of 196), showing the diversity when aggregating features from values.

Nevertheless, this can be hardly achieved in the case of linear attention. As a matter of fact, the rank of the attention matrix in linear attention is bounded by the number of tokens $N$ and the channel dimension $d$ for each head:

where $d$ is usually smaller than $N$ in common vision Transformer designs, e.g., $d\!=\!64,N\!=\!196$ in DeiT and $d\!=\!32,N\!=\!49$ in Swin Transformer . In this case, the upper bound of attention matrix rank is restricted at a lower ratio, which indicates that many rows of the attention map are seriously homogenized. As the output of self-attention is the weighted sum of the same set of $V$, the homogenization of attention weights inevitably leads to the resemblance among the aggregated features.

To better illustrate, we substitute the original Softmax attention in DeiT-Tiny with linear attention, and show the rank of the attention map in Fig. 5 (b). It can be observed that the rank is greatly decreased (54 out of 196) and many rows of the attention matrix are similar.

As a remedy, we present a simple yet effective solution to address this limitation of linear attention. Specifically, a depthwise convolution (DWC) module is added to the attention matrix and the output can be formulated as:

To better understand the effect of this DWC module, we can consider it as a kind of attention, in which each query will only focus on several adjacent features in space instead of all features $V$. This locality ensures that even if the linear attention values corresponding to two queries are the same, we can still get different outputs from different local features, thus maintaining feature diversity. The effect of DWC can also be explained from the perspective of matrix rank. Based on Eq.(10), we have:

where we denote $M_{\rm DWC}$ as the sparse matrix corresponding to the depthwise convolution function, and denote $M_{eq}$ as the equivalent full attention map. As $M_{\rm DWC}$ has the potential to be a full rank matrix, we practically increase the upper bound of the rank of the equivalent attention matrix, which incurs little computation overhead while greatly improving the linear attention’s performance.

To better illustrate, we conduct similar modifications on DeiT-Tiny. With the additional DWC module, the rank of the attention map in the linear attention can be restored to full rank (196 out 196 as shown in Fig. 5 (c)), which keeps the feature diversity as the original Softmax attention.

3 Focused linear attention module

Based on the aforementioned analysis, we propose a novel linear attention module, dubbed focused linear attention, which reduces the computation complexity while maintaining the expressive power. Specifically, we first design a novel mapping function to imitate the sharp distribution of the original Softmax attention. On this basis, we focus on the low-rank dilemma in previous linear attention modules, and adopt a simple depthwise convolution to restore feature diversity. In this way, our new module can enjoy benefits from both linear complexity and high expressiveness. Specifically, our module can be formulated as:

In general, our module has the following advantages:

(1) Low computation complexity as linear attention. By changing the computation order of self-attention, the complexity is transformed from $\mathcal{O}(N^{2}d)$ to $\mathcal{O}(Nd^{2})$, where $N$ and $d$ denote the token number and channel dimension of each head respectively. $d$ is usually smaller than $N$ in common vision Transformer designs, e.g., $d\!=\!64,N\!=\!196$ in DeiT and $d\!=\!32,N\!=\!49$ in Swin Transformer , the overall computation is practically decreased. Also, compared to previous linear attention modules that design complex kernel function, our proposed focused function $f_{p}$ only adopts simple operators which achieves approximation with minimum computation overhead.

(2) High expressive capability as Softmax attention. As we have analyzed above, previous kernel-based linear attention designs are generally inferior to the Softmax counterpart from the focus ability and feature diversity perspective. With the proposed focused function $f_{p}$ and depthwise convolution, our focused linear attention can achieve even better performance than Softmax attention.

In addition, our module also has the potential of adapting to larger receptive field and different model architectures. Modern Transformer models based on Softmax attention mainly use a limited number of key/value pairs because of the quadratic complexity towards token numbers. Nevertheless, the linear complexity of our module endows us to expand the receptive field to a larger region while maintaining the same amount of computation, and enjoying the advantage of modeling long-range dependencies. Also, our module can serve as a plug-in module and be easily adopted on a variety of modern vision Transformer architectures. We empirically implement our module on five advanced models including DeiT , PVT , PVT-v2 , Swin Transformer and CSwin Transformer . Considering the advantage of enlarged receptive field, we adopt the focused linear attention block at early stages of the vision Transformers, and keep the rest of blocks unchanged. Detailed model architectures are shown in Appendix.

Experiments

To verify the effectiveness of our method, we conduct experiments on ImageNet-1K classification , ADE20K semantic segmentation , and COCO object detection . We also provide a detailed comparison with other linear attention modules based on two representative model structures. In addition, we perform comprehensive ablation studies to analyze each important design element.

ImageNet-1K contains 1.28M images for training and 50K images for validation. We practically implement our module on five advanced Vision Transformer models, and report the Top-1 accuracy on the validation split to compare with various state-of-the-art models.

For fair comparison, we use the exact same settings as the corresponding baseline model to train our FLatten model. Specifically, we use AdamW optimizer to train all our models for 300 epochs with a cosine learning rate decay and 20 epochs of linear warm-up. The basic learning rate for a batch size of 1024 is set to $1\times{10}^{-3}$, and then linearly scaled w.r.t. the batch size. We follow DeiT and apply RandAugment , Mixup , CutMix and random erasing to avoid overfitting. In addition, a weight decay of 0.05 is used. To be consistent with , we also adopt EMA in the training of our FLatten-CSwin models. In terms of larger resolution finetuning, we follow the setting in that finetunes the models for 30 epochs.

The classification results are provided in Fig. 6. It is shown that our method achieves consistent improvements against baseline models under comparable FLOPs or parameters. For example, our FLatten-PVT-T/S surpass PVT-T/S by 2.7% and 1.9% respectively with similar FLOPs. Based on Swin, our model achieves comparable performance with 60% FLOPs. Our model based on PVT-v2 and CSwin also achieves a better trade-off between computation cost and model performance. These results demonstrate that our module has high expressive capability and is applicable to various model structures.

2 Semantic Segmentation

ADE20K is a widely adopted benchmark for semantic segmentation with 20K/2K training/validation images. We employ our model on two representative segmentation models, SemanticFPN and UperNet . As shown in Tab. 1, our model achieves consistently better results under all settings. Specifically, we can see a $0.5\!\sim\!1\%$ mIoU improvement with comparable computation cost and parameters. The improvements in mAcc are more significant.

3 Object Detection

COCO object detection and instance segmentation dataset has 118K training and 5K validation images. We use ImageNet pretrained model as the backbone in Mask R-CNN and Cascade Mask R-CNN frameworks to evaluate the effectiveness. We conduct experiments on 1x and 3x schedules with different detection heads and show results in Tab. 2. Taking advantage of larger receptive field, our model shows better results under all settings.

4 Comparison with Other Linear Attention

To show a fair comparison with other linear attention modules, we conduct experiments based on two representative Vision Transformer structures, DeiT and Swin Transformer respectively. Based on these two models, we compare our focused linear attention module with four previous linear attention designs, including hydra attention , efficient attention , linear angular attention and enhanced linear attention .

As shown in Tab. 3, previous linear attention modules are generally inferior to the Softmax counterpart, while our model significantly outperforms all other designs and the Softmax baseline. This indicates that our module has high expressive capability and can achieve better performance than Softmax attention with lower computation complexity.

5 Inference Time

We further evaluate the practical efficiency of our model and compare it with two competitive baselines. The results are presented in Fig. 7. We test the inference latency on multiple hardware platforms, including a desktop CPU (Intel i5-8265U) and two server GPUs (RTX2080Ti and RTX3090). It can be observed that our model achieves a significantly better trade-off between runtime and accuracy on both CPU and GPU, enjoying up to 2.1x faster inference speed with on par or even better performances.

6 Ablation Study

In this section, we ablate the key components in our focused linear attention to verify the effectiveness of these designs. We report the results on ImageNet-1K classification based on FLatten-DeiT-T and FLatten-Swin-T.

Focused function $f_{p}$ and DWC. We first evaluate the effectiveness of our proposed focused function $f_{p}$ and depth-wise convolution. We start from the vanilla linear attention and introduce $f_{p}$ and DWC in turn. As shown in Tab. 4, adopting the proposed focused function $f_{p}$ provides +1.3 improvement. Using DWC to maintain feature diversity further leads to an accuracy gain of +2.3, achieving an overall accuracy of 74.1. These results prove that our proposed $f_{p}$ and DWC can greatly improve the expressive capability of linear attention, thus helping our focused linear attention module achieve better performance than Softmax attention.

Ablation on different $p$. In Tab. 5, we study the effect of focused factor $p$ on the model performance. We find that when $p$ changes between 2 and 32, the Top-1 classification accuracy does not change much, implying the robustness of our module to this hyper-parameter. Practically, we choose $p=3$ for all models in the paper without additional tuning.

Receptive field. We also study the impact of receptive field based on FLatten-Swin-tiny. As illustrated in Tab. 6, with the increase of window size, the performance of our model improves progressively. This further proves that though window attention is effective, it inevitably sacrifices the long-range dependency of self-attention from the global perspective and is still inferior to global attention. With linear complexity, it is possible for our module to realize a large even global receptive field while maintaining the same amount of computation.

Focused linear attention at different stages. We replace the shift-window attention of Swin-T with our module at different stages. As shown in Tab. 7, we can see that replacing the first two stages leads to a performance gain of 0.8, while replacing the last two stages slightly decreases the overall accuracy. We attribute this result to the fact that the first two stages of Swin have larger resolutions and are more suitable for our module with large receptive field.

Conclusion

In this paper, we propose a novel focused linear attention module. By addressing the limitations of previous linear attention methods from focus ability and feature diversity perspectives, our module achieves an impressive combination of high efficiency and expressive capability. Extensive experiments on image classification, object detection and semantic segmentation demonstrated that our module can be widely applied to a variety of vision Transformers and achieve a better trade-off between computation efficiency and model performance.

Acknowledgement

This work is supported in part by National Key R&D Program of China (2021ZD0140407), the National Natural Science Foundation of China (62022048, 62276150) and THU-Bosch JCML. We appreciate the generous donation of computing resources by High-Flyer AI.

References

Appendix

A. Proof of Proposition 1

As mentioned in the main paper, with the aim to restore the sharp distribution in linear attention, we present our Focused Function $f_{p}$:

and $x^{\ast\ast p}$ represents the power $p$ of $x$ bit by bit. We follow previous linear attention modules to use the ReLU function first to ensure the non-negativity of input. Therefore, when proving the effects of $f_{p}$, we only consider $x,y\geq 0$.

For a pair of feature $\{x,y\}$ with $m\!=\!n$:

For a pair of feature $\{x,y\}$ with $m\!\neq\!n$:

Then we consider the following two cases: (1) $m\!=\!n$:

Therefore, with a proper $p$, our focused function $f_{p}(\cdot)$ practically achieves a more distinguished difference between similar query-key pairs (Eq. 15) and dissimilar query-key pairs (Eq. 16). Actually, $f_{p}$ divides the features into several groups according to their nearest axes, improving the similarity within each group while reducing the similarity between the groups, thus restoring the sharp attention distribution as the original Softmax function.

B. More Visualizations

We visualize more examples of attention weights in Fig. 8. To better show the contribution of our focused function and DWC, we start from the vanilla linear attention and introduce $f_{p}$ and DWC separately. As demonstrated in the last three rows, DWC improves local focus ability but cannot focus on any position, while $f_{p}$ practically enhances model’s focus ability, helping model focus on more informative regions. Combining $f_{p}$ and DWC, our focused linear attention module restores the sharp distribution as the original Softmax attention.

C. Full Classification Results

Due to the page limit, we only present representative ImageNet classification results in Figure 6 of main paper. Here, we give all the classification results when applying our focused linear attention module on various sizes of the five baseline models in Tab.8.

D. Model Architectures

We summarize the architectures of five Transformer models adopted in the main paper, including DeiT , PVT , PVTv2 , Swin Transformer , CSwin Transformer in Tab.9-15. In practice, we substitute the original self-attention blocks at all stages of the DeiT, PVT and PVTv2 with the focused linear attention block, but only adopt our module at early stages of Swin and CSwin. The model structure (width and depth) are kept unchanged, except for CSwin-T and CSwin-B, where we increase the depth of the first and second stages and correspondingly reduce the depth of the third stage to better reflect our module’s advantage of enlarged receptive field.