Multi-Scale Vision Longformer: A New Vision Transformer for High-Resolution Image Encoding

Pengchuan Zhang, Xiyang Dai, Jianwei Yang, Bin Xiao, Lu Yuan, Lei Zhang, Jianfeng Gao

Introduction

Vision Transformer (ViT) has shown promising results on image classification tasks for its strong capability of long range context modeling. But its quadratic increase of both computational and memory complexity hinders its application on many vision tasks that require high-resolution feature maps computed on high-resolution imagesIn this paper, encoding a high-resolution image means generating high-resolution feature maps for high-resolution images. , like object detection , segmentation , and human pose estimation . Vision-language tasks, like VQA, image captioning, and image-text retrieval, also benefit from high-resolution feature maps , which are extracted with pre-trained CNN models. Developing a vision Transformer that can process high-resolution feature maps is a critical step toward the goal of unifying the model architecture of vision and language modalities and improving multi-modal representation learning.

In this paper, we propose a new vision Transformer architecture Multi-Scale Vision Longformer, which significantly enhances the baseline ViT for encoding high-resolution images using two techniques: (1) the multi-scale model structure, and (2) the attention mechanism of Vision Longformer.

Models with multi-scale (pyramid, hierarchical) structure provide a comprehensive encoding of an image at multiple scales, while keeping the computation and memory complexity manageable. Deep convolutional networks are born with such multi-scale structure, which however is not true for the conventional ViT architecture. To obtain a multi-scale vision Transformer, we stack multiple (e.g., four) vision Transformers (ViT stages) sequentially. The first ViT stage operates on a high-resolution feature map but has a small hidden dimension. As we go to later ViT stages, the feature map resolution reduces while the hidden dimension increases. The resolution reduction is achieved by performing patching embedding at each ViT stage. In our experiments, we find that with the same number of model parameters and the same model FLOPs, the multi-scale ViT achieves a significantly better accuracy than the vanilla ViT on image classification task. The results show that the multi-scale structure not only improves the computation and memory efficiency, but also boosts the classification performance. The proposed multi-scale ViT has the same network structure as conventional (multi-scale) CNN models such as ResNet , and can serve as a replace-and-plug-in choice for almost all ResNet applications. In this paper, we demonstrate this plausible property in image classification, object detection and instance segmentation.

The multi-scale structure alone is not sufficient to scale up ViT to process high-resolution images and feature maps, due to the quadratic increase of the computation and memory complexity with respect to the number of tokens in the self-attention layers. Compared to natural language tasks where data is 1-D, this problem is more severe in vision tasks where the increase in complexity is quartic (fourth order) with the increase of image resolution. For example, the computational complexity of a $4\times$ higher resolution multi-head self attention (MSA) layer (hidden dimension reduced by 4, i.e., $4H\times 4W\times\frac{D}{4}$ ) equals to that of 64 layers in the original size (i.e., $H\times W\times D$ ). To address this challenge, we develop a 2-D version of Longformer, called Vision Longformer, to achieve a linear complexity w.r.t. the number of tokens (quadratic w.r.t. resolution). Our experiments show that compared to the baseline ViT, Vision Longformer shows no performance drop while significantly reduces the computational and memory cost in encoding images. The result indicates that the “local attention + global memory” structure in Vision Longformer is a desirable inductive bias for vision Transformers. We also compare Vision Longformer with other efficient attention mechanisms. The result again validates its superior performance on both image classification and object detection tasks.

The main contributions of this paper are two-fold: (1) We propose a new vision Transformer that uses the multi-scale model structure and the attention mechanism of 2-D Longformer for efficient high-resolution image encoding. (2) We perform a comprehensive empirical study to show that the proposed ViT significantly outperforms strong baselines, including previous ViT models, their ResNet counterparts, and a model from a concurrent work, on image classification, object detection and segmentation tasks.

Related Work

The Vision Transformer (ViT) applies a standard Transformer, originally developed for natural language processing (NLP), for image encoding by treating an image as a word sequence, i.e., splitting an image into patches (words) and using the linear embeddings of these patches as an input sequence. ViT has shown to outperform convolution neural network (CNN) models such as the ResNet , achieving state-of-the-art performance on multiple image classification benchmarks, where training data is sufficient. DeiT is another computer vision model that leverages Transformer. It uses a teacher-student strategy specific to Transformers to improve data efficiency in training. Thus, compared to ViT, it requires much less training data and computing resources to produce state-of-the-art image classification results. In addition to image classification, Transformers have also been applied to other compute vision tasks, including object detection , segmentation , image enhancement , image generation , video processing , and vision-language tasks .

Developing an efficient attention mechanism for high-resolution image encoding is the focus of this work. Our model is inspired by the efficient attention mechanisms developed for Transformers, most of which are for NLP tasks. These mechanisms can be grouped into four categories. The first is the sparse attention mechanism, including content-independent sparsity and content-dependent sparsity . Axial Transformer and Image Transformer are among few sparsity-based efficient attentions that are developed for image generation. The second is the memory-based mechanism, including Compressive Transformers and Set Transformer . These models use some extra global tokens as static memory and allow all the other tokens to attend only to those global tokens. The third is the low-rank based mechanism. For example the Linformer projects the input key-value pairs into a smaller chunk, and performs cross-attention between the queries and the projected key-value pairs. The fourth is the (generalized) kernel-based mechanism, including Performer and Linear Transformers. Many models utilize hybrid attention mechanisms. For example, Longformer, BigBird and ETC combine the sparsity and memory mechanisms; Synthesizers combines the sparsity and low-rank mechanisms. Readers may refer to and for a comprehensive survey and benchmarks, respectively.

In this paper, we developed a 2-D version of Longformer, called Vision Longformer, which utilizes both the sparsity and memory mechanisms. Its conv-like sparsity mechanism is conceptually similar to the sparsity mechanism used in the Image Transformer.

The multi-scale vision Transformer architecture is another technique we use in our proposed high-resolution Vision Longformer. The hierarchical Transformers for NLP contain two stages, with the first stage processing overlapping segments and the second stage using the embeddings of the CLS tokens from all segments as input. In our proposed Vision Longformer, size reduction is performed by the patch embedding at the beginning of each stage, by merging all tokens in a patch from previous stage into a single token at the current stage. We typically use 4 stages for our model since we have empirically verified that using 4 stages is better than using 2 or 3 stages, especially for object detection tasks. Informer takes a similar stacked multi-stage approach to encoding long sequences, where the size reduction between stages is achieved by max-pooling.

Pyramid Vision Transformer (PVT) , Swin Transformer and HanoNet are concurrent works of ours. All these works use a multi-scale architecture where multiple (slightly modified) ViTs are stacked. The authors of PVT propose the spatial-reduction attention (SRA) to alleviate the cost increase in self-attention layers. However, the computation and memory complexity of PVT still increases quartically w.r.t. resolution (with a much smaller constant). Swin Transformer and HanoNet utilizes similar local attention mechanism as our Vision Longformer, but from different perspectives and implementations.

Multi-Scale Stacked Vision Transformers

As shown in Figure 1 (Bottom), we improve the encoding efficiency of vision Transformer by making the following modifications to the vanilla ViT. The modified ViT is referred to as Efficient ViT (E-ViT).

We add a Layer Normalization (LayerNorm) after the patch embedding.

We define a number of global tokens, including the CLS token. Correspondingly, the tokens associated with image and feature patches are referred to as local tokens afterwards.

We replace the vanilla full self-attention with an efficient attention mechanism, denoted by ${\color[rgb]{1,0,0}a}$ , which will be described in detail in Sections 3.2 and 3.3.

We use either an Absolute 2-D Positional Embedding (APE for short, separately encoding $x$ and $y$ coordinates and concatenating them) or a Relative Positional Bias (RPB for short) to replace the original absolute 1-D positional embedding.

Except for attention ${\color[rgb]{1,0,0}a}$ , E-ViT has the following architecture parameters inherited from the vanilla ViT : input patch size ${\color[rgb]{0,1,0}p}$ , number of attention blocks $n$ , hidden dimension $d$ and number of heads $h$ , denoted as $\text{E-ViT}(a\times n/p~{};h,d,n_{g})$ . Using the full attention mechanism (i.e., $a$ = full) and one global token (i.e., the CLS token with $n_{g}=1$ ), the deficient $\text{E-ViT}(\text{full}\times 12/16~{};h,d,1)$ models still achieve better ImageNet classification performance than the baseline ViT for both tiny ( $h=3,d=192$ ) and small ( $h=6,d=384$ ) model sizes, as shown in Table 2. The performance gain is attributed to the added LayerNorm, as we show in the Supplementary.

Mathematically, an $\text{E-ViT}(a\times n/p~{};h,d,n_{g})$ encoding module can be written as:

Stack multiple E-ViT modules as multi-scale vision Transformers. As illustrated in Figure 1 (Top), a multi-scale Vision Transformer is built by stacking multiple E-ViT modules (or stages). In what follows, we describe several design choices we have made when building the multi-scale ViT.

What are the patch size and hidden dimension at each stage? As required in object detection and human pose estimation, for models with 4-scale feature maps, the first feature map needs to down-sample the image by 4 and thus stage 1 can be written as $\text{E-ViT}(a_{1}\times n_{1}/4~{};h_{1},d_{1},n_{g,1})$ . We typically use only one attention block, i.e., $n_{1}=1$ . The first stage generates the highest-resolution feature map, which consumes lots of memory, as shown in Table 2. We also construct several 3-stage models, whose first stage patch size is 8. For later stages, the patch sizes are set to 2, which downsizes the feature map resolution by 2. Following the practice in ResNet, we increase the hidden dimension twice when downsizing the feature map resolution by 2. We list a few representative model configurations in Table 1. Different attention types ( ${\color[rgb]{1,0,0}a}$ ) have different choices of number of global tokens $n_{g}$ . But they share the same model configurations. Thus we do not specify ${\color[rgb]{1,0,0}a}$ and $n_{g}$ in Table 1. Please refer to the Supplementary for the complete list of model configurations used in this paper,

How to connect global tokens between consecutive stages? The choice varies at different stages and among different tasks. For the tasks in this paper, e.g., classification, object detection, instance segmentation, we simply discard the global tokens and only reshape the local tokens as the input for next stage. In this choice, global tokens only plays a role of an efficient way to globally communicate between distant local tokens, or can be viewed as a form of global memory. These global tokens are useful in vision-language tasks, in which the text tokens serve as the global tokens and will be shared across stages.

Should we use the average-pooled layer-normed features or the LayerNormed CLS token’s feature for image classification? The choice makes no difference for flat models. But the average-pooled feature performs better than the CLS feature for multi-scale models, especially for the multi-scale models with only one attention block in the last stage (including all models in Table 1). Please refer to the Supplementary for an ablation study.

As reported in Table 2, the multi-scale models outperform the flat models even in low-resolution classification tasks, demonstrating the importance of multi-scale structure. However, the full self-attention mechanism suffers from the quartic computation/memory complexity w.r.t. the resolution of feature maps, as shown in Table 2. Thus, it is impossible to train 4-stage multi-scale ViTs with full attention using the same setting (batch size and hardware) used for DeiT training.

2 Vision Longformer: A “Local Attention + Global Memory” Mechanism

We propose to use the ”local attention + global memory” efficient mechanism, as illustrated in Figure 2 (Left), to reduce the computational and memory cost in the E-ViT module. The 2-D Vision Longformer is an extension of the 1-D Longformer originally developed for NLP tasks. We add $n_{g}$ global tokens (including the CLS token) that are allowed to attend to all tokens, serving as global memory. Local tokens are allowed to attend to only global tokens and their local 2-D neighbors within a window size. After all, there are four components in this ”local attention + global memory” mechanism, namely global-to-global, local-to-global, global-to-local, and local-to-local, as illustrated in Figure 2 (Left). In Equation (2), a Multi-head Self-Attention (MSA) block with the Vision Longformer attention mechanism is denoted as $MSA_{{\color[rgb]{1,0,0}\text{ViL}}}$ , i.e., $a=\text{ViL}$ in Equation (2).

Relative positional bias for Vision Longformer. Following , we add a relative positional bias $B$ to each head when computing the attention score:

where $Q,K,V$ are the query, key and value matrices and $d$ is the query/key dimension. This relative positional bias makes Vision Longformer translational invariant, which is a desired property for vision models. We observe significant improvements over the absolute 2-D positional embedding, as shown in Table 2 for ImageNet classification and Section 4.4 for COCO object detection.

Theoretical complexity. Given the numbers of global and local tokens, denoted by $n_{g}$ and $n_{l}$ respectively, and local attention window size $w$ , the memory complexity of the $MSA_{{\color[rgb]{1,0,0}\text{ViL}}}$ block is $\mathcal{O}(n_{g}(n_{g}+n_{l})+n_{l}w^{2})$ . Although points out that separating the attention parameters for global and local tokens is useful, we do not observe obvious gain in our experiments and thus simply let them share the same set of attention parameters. We empirically set the window size $w$ to 15 for all E-ViT stages, which makes our model comparable with the global attention window size 14 of ViT/16 acted on $224\times 224$ images. With such a window size, only attentions in the first two stages (in 4-stage multi-scale ViTs) are local. The attentions in the later two stages are equivalentEquivalent in our sliding chunks implementation, which is our default choice. to full attention. In our experiments, we find that it is sufficient to use only one global token ( $n_{g}=1$ ) for ImageNet classification problems. So, the effective memory complexity of the $MSA_{{\color[rgb]{1,0,0}\text{ViL}}}$ block is $\mathcal{O}((15^{2}+1)n_{l})$ , which is linear w.r.t. the number of tokens.

Results in Table 2 show that in comparison with the full attention models, the proposed multi-scale Vision Longformer achieves a similar or even better performance, while saving significant memory and computation cost. The memory saving is significant for feature maps with resolution $56\times 56$ (i.e., the feature maps in the first stage of a 4-stage multi-scale model). The savings are even more significant for higher resolution feature maps. This makes Vision Longformer scalable to high-resolution vision tasks, such as object detection and segmentation. When equipped with relative positional bias, Vision Longformer outperforms the full attention models with absolute positional embedding. This indicates that the “local attention + global memory” mechanism is a good inductive bias for vision Transformers.

Three implementations of Vision Longformer and its random-shifting training strategy. Vision Longformer is conceptually similar to conv-like local attention. We have implemented Vision Longformer in three ways: (1) using Pytorch’s unfold function (nn.unfold or tensor.unfold), (2) using a customized CUDA kernel and (3) using a sliding chunk approach. The unfold implementation is simple but very slow, i.e., 24 times slower than full attention on $40\times 40\times 768$ feature map. The implementation using the customized CUDA kernel is about 20% faster than the full attention in the same setting, while achieving the theoretical memory complexity. The sliding-chunk approach is the fastest, which is 60% faster than the full attention with a cost of consuming slightly more memory than the theoretical complexity. With the sliding chunk implementation, we also propose a random-shifting training strategy for Vision Longformer, which further improves the training speed and memory consumption during training. Please refer to the Supplementary for details of these implementations and the random-shifting training strategy.

3 Other Efficient Attention Mechanisms

We compare Vision Longformer with the following alternative choices of efficient attention methods. We put detailed descriptions of these methods and their experimental setup in the Supplementary.

Pure global memory ( ${\color[rgb]{1,0,0}a}=\text{global}$ ). In Vision Longformer, see Figure 2 (Left), if we remove the local-to-local attention, then we obtain the pure global memory attention mechanism (called Global Attention hereafter). Its memory complexity is $\mathcal{O}(n_{g}(n_{g}+n_{l}))$ , which is also linear w.r.t. $n_{l}$ . However, for this pure global memory attention, $n_{g}$ has to be much larger than 1. We gradually increase $n_{g}$ (by 2 each time) and its performance gets nearly saturated at 128. Therefore, $n_{g}=128$ is the default for this Global attention.

Linformer ( ${\color[rgb]{1,0,0}a}=\text{LIN}$ ) projects the $n_{l}\times d$ dimensional keys and values to $K\times d$ dimensions using additional projection layers, where $K\ll n_{l}$ . Then the $n_{l}$ queries only attend to these projected $K$ key-value pairs. The memory complexity of Linformer is $\mathcal{O}(Kn_{l})$ . We gradually increase $K$ (by 2 each time) and its performance gets nearly saturated at 256. Therefore, $K=256$ is the default for this Linformer attention, which turns out to be the same with the recommended value. Notice that Linformer’s projection layer (of dimension $K\times n_{l}$ ) is specific to the current $n_{l}$ , and cannot be transferred to higher-resolution tasks that have a different $n_{l}$ .

Spatial Reduction Attention (SRA) ( ${\color[rgb]{1,0,0}a}=\text{SRA}$ ) is similar to Linformer, but uses a convolution layer with kernel size $R$ and stride $R$ to project the key-value pairs, hence resulting in $n_{l}/R^{2}$ compressed key-value pairs. Therefore, The memory complexity of SRA is $\mathcal{O}(n_{l}^{2}/R^{2})$ , which is still quadratic w.r.t. $n_{l}$ but with a much smaller constant $1/R^{2}$ . When transferring the ImageNet-pretrained SRA-models to high-resolution tasks, SRA still suffers from the quartic computation/memory blow-up w.r.t. the feature map resolution. Pyramid Vision Transformer uses this SRA to build multi-scale vision transformer backbones, with different spatial reduction ratios ( $R_{1}=8,R_{2}=4,R_{3}=2,R_{4}=1$ ) for each stage. With this PVT’s setting, the key and value feature maps at all stages are essentially with resolution $H/32\times W/32$ .

Performer ( ${\color[rgb]{1,0,0}a}=\text{performer}$ ) uses random kernels to approximate the Softmax computation in MSA, and achieves a linear computation/memory complexity with respect to $n_{l}$ and the number of random features. We use the default 256 orthogonal random features (OR) for Performer, and provide other details in the Supplementary.

Compare Vision Longformer with other attention mechanisms. On the ImageNet classification task in Table 3, all efficient attention mechanisms above show a large performance gap from Vision Longformer. Linformer performs very competitively. Global attention and Performer have a similar performance with the DeiT model (72.2 for tiny and 79.8 for small). We use spatial reduction ratios $16,8,4,2$ from stage1 to stage4 for the multi-scale SRA model, which is different from the reduction ratios $8,4,2,1$ in PVT . This more aggressive spatial reduction makes the classification performance worse in Table 3, but makes the memory cost manageable when transfer to detection tasks for input image size $8000\times 1333$ . For a more complete comparison of these models, including model parameters, FLOPs and memory usage, please refer to the Supplementary.

Why is Longformer better? One possible reason is that the conv-like sparsity is a good inductive bias for vision transformers, compared with other attention mechanisms. This is supported by the visualization of the attention maps from pretrained DeiT models . Another explanation is that Vision Longformer keeps the key and value feature maps high resolution. However, low resolution-based attention mechanims like Linformer and SRA and pure global attention lose the high-resolution information in the key and value feature maps.

Mixed attention mechanisms (Partial X-former) for classification tasks. For classification tasks with $224\times 224$ image size as input, the feature map size at Stage3 in multi-scale ViTs is $14\times 14$ . This is the same as the feature map size in ViT and DeiT, and best suits for full attention. A natural choice is to use efficient attention in the first two stages (with high-resolution feature map but with small number of blocks) and to use full attention in the last two stages. Multi-scale ViTs with this mixed attention mechanisms are called “Parital X-former”. We also report these Partial X-formers’ performance in Table 3. All these Partial X-formers perform well on ImageNet classification, with very little (even no) gap between Full Attention and Vision Longformer. These Partial X-forms achieve very good accuracy-efficiency performance for low-resolution classification tasks. We do not have “Partial ViL” for classification because ViL’s window size is 15, and thus its attention mechanism in the last two stages is equivalent to the full attention.

4 Transfer to High-resolution Vision Tasks

Similar to the transfer-ability of ImageNet-pretrained CNN weights to downstream high-resolution tasks, such as object detection and segmentation, multi-scale Vision Longformer pretrained on ImageNet can be transferred to such high-resolution tasks, as we will show in Section 4.3.

However, Linformer is not transferable because the weights of the linear projection layer is specific to a resolution. The Partial X-formers and Multi-scale ViT with full attention are not transferable due to its prohibitively large memory usage after transferred to high-resolution tasks. In Table 8, we also show the superior performance of Vision Longformer over other attention mechanisms, on the object detection and segmentation tasks.

Experiments

In this section, we show the final performance of Multi-scale Vision Longformer (short for ViL) on ImageNet classification in Section 4.1 & 4.2 and downstream high-resolution detection tasks in Section 4.3. We follow the DeiT training configuration for ImageNet classification training, and use the standard “ $\times 1$ ” and “ $\times 3$ +MS” training schedules with the “AdamW” optimizer for detection tasks. We refer to the Supplementary for detailed experimental settings.

Following DeiT and PVT, we build multi-scale ViLs with four different sizes, i.e., tiny, small, medium and base. The detailed model configuration is specified in Table 1. We train multi-scale ViLs purely on ImageNet1K, following the setting in DeiT .

In Table 4, we report our results and compare with ResNets, ViT , DeiT and PVT . Our models outperform other models in the same scale by a large margin. We again confirm that the relative positional bias (RPB) outperforms the absolute 2-D positional embedding (APE) on Vision Longformer. When compared with Swin Transformers , our models still performs better with fewer parameters.

2 ImageNet-21K pretraining and ImageNet-1K finetuning

When trained purely on ImageNet-1K, the performance gain from ViL-Medium to ViL-Base is very marginal. This is consistent with the observation in ViT : large pure transformer based models can be trained well only when training data is sufficient.

Therefore, we conducted experiments in which ViL-Medium/Base models are first pre-trained on ImageNet-21k with image size $224^{2}$ and finetuned on ImageNet-1K with image size $384^{2}$ . For ViT models on image size $384^{2}$ , there are in total $24\times 24$ tokens with full attention. For ViL models on image size $384^{2}$ , we set the window sizes to be $(13,17,25,25)$ from Stage1 to Stage4. Therefore, in the last two stages, the ViL models’ attention is still equivalent to full attention.

As shown in In Table 5, the performance gets significantly boosted after ImageNet-21K pretraining for both ViL medium and base models. We want to point out that the performance of ViL-Medium model has surpassed that of ViT-Base/16, ViT-Large/16 and BiT-152x4-M, in the ImageNet-21K pretraining setting. The performance of ViL-Base models are even better. This shows the superior performance and parameter efficiency of ViL models.

3 Detection Tasks

We apply our ViL to two representative object detection pipelines including RetinaNet and Mask-RCNN . We follow the conventional setting to use our Vision Longformer as the backbone to generate feature maps for both detection pipelines. Similar to , we extract the features from all four scales and then feed them to the detection and/or instance segmentation head. To adapt the learned relative positional bias to the higher image resolution in detection, we perform bilinear interpolation on it prior to the training. In our experiments, all models are evaluated on COCO dataset , with 118k images for training and 5k images for evaluation. We report the results for both 1 $\times$ and 3 $\times$ +MS training schedules, and compare them with two backbone architectures: ResNet and PVT .

As shown in Table 6, our ViL achieves significantly better performance than the ResNet and PVT architecture. The improvements are uniform over all model sizes (tiny, small, medium, base) and over all object scales ( $AP_{S}$ , $AP_{M}$ , $AP_{L}$ ). The improvement is so large that ViL-Tiny with “3x+MS” schedule already outperforms the ResNeXt101-64x4d and the PVT-Large models. A similar trend is observed with the Mask R-CNN pipeline. As shown in Table 7, our ViL backbone significantly surpasses ResNet and PVT baselines on both object detection and instance segmentation. When compared with the concurrent Swin Transformer , our model also outperforms it with fewer parameter and FLOPs. More specifically, our ViL-Small achieves 47.1 $AP^{b}$ with 45M parameters, while Swin-Tiny achieves 46.0 $AP^{b}$ with 48M parameters. These consistent and significant improvements with both RetinaNet and Mask R-CNN demonstrate the promise of our proposed ViL when using it as the image encoder for high-resolution dense object detection tasks.

4 Ablation Study for Detection Tasks

Compare with other efficient attention mechanisms. Similar to Sec 4.4, we study SRA , Global Transformer and Performer and their corresponding partial version with Mask R-CNN pipeline (trained with the 1 $\times$ schedule). As we can see in Table 8, when efficient attention mechanisms are used in all stages, ViL achieves much better performance than the other three mechanisms. Specifically, our ViL achieves 42.9 $AP^{b}$ while the other three are all around 36.0 $AP^{b}$ . When efficient attention mechanisms are only used in the first two stages (Par-Xformer), the gaps between different mechanisms shrink to around 1.0 point while our ViL still outperform all others. Moreover, the ViL model outperforms the partial models of all other attention mechanisms and has a very small gap (0.4 $AP^{b}$ ) from the Partial-ViL model. These results show that the “local attention + global memory” mechanism in Vision Longformer can retain the good performance of the full attention mechanism in ViT, and that it is a clear better choice than other efficient attention mechanisms for high-resolution vision tasks.

The effects of window size and number of global tokens are not obvious in ImageNet classification, as long as the last two stages use full attention. For different window sizes in $and different number of global tokens in$ , the final top-1 accuracy differs by at most 0.2 for ViL-Small models. Meanwhile, their effects are significant in high-resolution tasks, where ViL models use local attention in all stages. In Figure 3, we report their effects in COCO object detection with Mask R-CNN. We notice that the window size plays a crucial role and the default window size 15 gives the best performance. Smaller window sizes lead to serious performance drop. As shown in Figure 3 (Right), as long as there is one global token, adding more global tokens does not improve the performance any more.

Conclusions

In this paper, we have presented a new Vision Transformer (ViT) architecture Multi-Scale Vision Longformer to address the computational and memory efficiency that prevents the vanilla ViT model from applying to vision tasks requiring high-resolution feature maps. We mainly developed two techniques: 1) a multi-scale model structure designed for Transformers to provide image encoding at multiple scales with manageable computational cost, and 2) an efficient 2-D attention mechanism of Vision Longformer for achieving a linear complexity w.r.t. the number of input tokens. The architecture design and the efficient attention mechanism are validated with comprehensive ablation studies. Our experimental results show that the new ViT architecture effectively addresses the computational and memory efficiency problem and outperforms several strong baselines on image classification and object detection.

References

Appendix A Settings

We listed the model configuration of all models used in this paper in Table 9. We do not specify the attention mechanism here, because the model configuration is the same for all attention mechanisms and the attention-specific parameters are specified in Table 15.

A.2 Experimental settings

Table 10 summarizes our training setups for our different models.

For the ImageNet classification task, our setting mainly follow that in DeiT . For example, we do not use dropout but use random path. We use all data augmentations in DeiT , except that we apply Repeated Augmentation only on Medium and Base models. When fine-tuning from a ImageNet-21K pretrained checkpoint, we mainly follow the practice of ViT , train on image size $384\times 384$ , use SGD with momentum $0.9$ , use no weight decay, and use only random cropping for data augmentation.

For COCO object detection/segmentation tasks, we follow the standard “ $1\times$ ” and “ $3\times+\text{MS}$ ” schedules. We only change the optimizer from SGD to AdamW and search for good initial learning rate and weight decay. For the “ $1\times$ ” schedule, the input image scale is fixed to be $(800,1333)$ for the min and max sizes, respectively. For the “ $3\times+\text{MS}$ ” schedule, the input image is randomly resized to have min size in $\{640,672,704,736,768,800\}$ . We found that there is obvious over-fitting in Training ViL-Medium and ViL-Base models on COCO, mainly because that these two models are relatively large but they are only pretrained on ImageNet. Therefore, we are taking the best checkpoint (one epoch per checkpoint) along the training trajectory to report the performance.

Appendix B More experimental results

In this section, we present two ablation studies on the model architecture of multi-scale Vision Longformer.

Ablation of the effects of LayerNorm and 2-D positional embedding in the patch embedding. In Table 2, we show that our flat model $\text{E-ViT}(\text{full}\times 12/16)$ , which only differs from the standard ViT/DeiT model by an newly-added LayerNorm after the patch embedding and the 2-D positional embedding, has better performance than the standard ViT/DeiT model. In Table 11, we show that this better performance comes from the newly-added LayerNorm.

Feature from the CLS token or from average pooling? As shown in Table 12, for ViL models that has only one attention block in the last stage (ViL 1-2-8-1), the average pooled feature from all tokens works better than the feature of the CLS token. However, when there are more than 2 attention blocks in the last stage (ViL 1-1-8-2), the difference between these two features disappears. The ViL 1-1-8-2 model has better performance than the ViL 1-2-8-1 model because it has more trainable parameters.

B.2 A comprehensive comparison of different attention mechanisms on ImageNet classification

We compare different attention mechanisms with different model sizes and architectures in Table 13 and Table 14. In Table 13, we show their performance on ImageNet-1K classification problem, measured by Top-1 accuracy. In Table 14, we show their number of parameters and FLOPs. We would like to comment that FLOPs is just a theoretical estimation of computation complexity, and it may not fit well the space/time cost in practice.

Appendix C Implementations and Efficiency of Vision Longformer In Practice

There is a trivial implementation of the conv-like sliding window attention, in which we compute the full quadratic attention and then mask out non-neighbor tokens. This approach suffers from the quadratic complexity w.r.t. number of tokens (quartic w.r.t. feature map size), and is impractical for real use, as shown by the blue curve in Figure 5. We only use it to verify the correctness of our other implementations.

We have implemented Vision Longformer in three ways:

Using Pytorch’s unfold function. We have two sub-versions: one using nn.functional.unfold (denoted as “unfold/nn.F”) and the other using tensor.unfold (denoted as “unfold/tensor”). As shown in Figure 5, the “unfold/tensor” version (red solid line) is more efficient both in time and memory than the “unfold/nn.F” version (red dotted line). However, both of them are even slower and use more memory than the full attention!

Using a customized CUDA kernel, denoted as “cuda_kernel”. We make use of the TVM, like what has done in Longformer , to write a customized CUDA kernel for Vision Longformer. As shown in Figure 5, the “cuda_kernel” (green line) achieves the theoretical optimal memory usage. Its time complexity is also reduced to linear w.r.t. number of tokens (quadratic w.r.t. feature map size). However, since it’s not making use of the highly optimized matrix multiplication libraries in CUDA, it’s speed is still slow in practice.

Using a sliding chunk approach, illustrated in Figure 4. For this sliding chunk approach, we have two subversions: one using Pytorch’s autograd to compute backward step (denoted as “SCw/Autograd”) and the other writing a customized torch.autograd.Function with hand-written backward function (denoted as “SCw/Handgrad”). Both sub versions of this sliding chunk approach are fully implemented with Pytorch functions and thus make use of highly optimized matrix multiplication libraries in CUDA. As shown in Figure 5, both of them are faster than the “cuda_kernel” implementation.

In the sliding chunk approach, to achieve a conv-like local attention mechanism with window size $2w+1$ , we split the feature map into chunks with size $w\times w$ . Each chunk only attends to itself and its 8 neighbor chunks. The Pytorch Autograd will save 9 copies of the feature map (9 nodes in the computing graph) for automatic back-propagation, which is not time/memory efficient. The “SCw/Handgrad” version defines a customized torch.autograd.Function with hand-written backward function, which greatly saves the memory usage and also speeds up the algorithm, as shown in Figure 5. We would like to point out that the memory usage of the “SCw/Handgrad” version is nearly optimal (very close to that of the “cuda_kernel”). Similar speed-memory trade-off with different implementations of local attention mechanism has been observed in the 1-D Longformer , too; see Figure 1 in . We would like to point out that Image Transformer has an implementation of of 2-D conv-like local attention mechanism, which is similar to our “SCw/Autograd” version. The Image Transformer applies it to the image generation task.

This sliding-chunk implementation (Figure 4 Right) lets one token attends to more tokens than the exact conv-like local attention (Figure 4 Left). Our sliding-chunk implementation has the choice to be

exactly the same with the conv-like local attention (Left), by masking out tokens that should not be attended to,

sliding chunk without padding, in which the chunks on the boundary have less chunks to attend to,

sliding chunk with cyclic padding, in which the chunks on the boundary still attend to 9 chunks with cyclic padded chunks.

Since these three masking methods only differ by the attention masks to mask out invalid tokens, their speed and memory usage are nearly the same, as shown in Figure 6. For ImageNet classification, we observe no obvious difference in top1 accuracy between “exact sliding window” and “sliding chunk without padding”, while “sliding chunk with cyclic padding” performs slightly worse most of the time. For object detection, we observe that “sliding chunk without padding” performs consistently better than “exact sliding window”, as shown in Figure 8. Therefore, we make “sliding chunk without padding” as the default making method for Vision Longformer, although it sacrifices some translational invariance compared with “exact sliding window”.

In Figure 7, we show the running time (including forward and backward) and memory usage of our “SCw/Handgrad” implementation of conv-like local attention (sliding chunk attention without padding mode) with different window sizes. We can see that the speed is not sensitive to the window size for small window sizes ( $\leq 17$ ) and the memory usage monotonically increases.

Finally, both the “unfold/nn.F” and the “cuda_kernel” implementations support dilated conv-like attention. The customized CUDA kernel is even more flexible to support different dilations for different heads. The sliding-chunk implementation does not support this dilated conv-like attention. In this paper, we always use the sliding-chunk implementation due to its superior speed and nearly optimal memory complexity.

In Figure 5, 6 and 7, the evaluation is performed on a single multi-head self-attention module (MSA) with the conv-like local attention mechanism, instead of on the full multi-scale Vision Longformer. With this evaluation, we can clearly see the difference among different implementations of the conv-like local attention mechanism.

Appendix D Random-shifting strategy to improve training efficiency

We propose the random-shifting training strategy for Vision Longformer, to further accelerate the training speed of Vision Longformer. More specifically, instead of attending to all 8 neighbor patches, one patch can only attend to itself and one random neighbor patch during training. To achieve this, we define 10 modes of the sliding-chunk local attention:

0 (default): attend to itself and all 8 neighbor chunks,

i ( $1<=i<=8$ ) : attend to itself chunk and the i’th neighbor chunk.

The ordering of the 8 neighbor patches is visualized in Figure 9. During training, we can randomly sample one mode from 1 to 8 and perform the corresponding random-shifting attention. We switch from the random-shifting mode to the default 8-neighbor mode after $x$ % training iterations, and this switch time $x$ % is a hyper-parameter with default value 75%. This switch, can be seen as fine-tuning, is necessary to mitigate the difference of model’s behavior during training and inference. As shown in Figure 10, this random-shifting training strategy accelerates the Vision Longformer training significantly, while not harming the final model performance.

Appendix E Other Efficient Attention Mechanisms utilized in this work

In this paper, we compare Vision Longformer with the following alternative choices of efficient attention methods.

Pure global memory ( ${\color[rgb]{1,0,0}a}=\text{global}$ ). In Vision Longformer, see Figure 2 (Left), if we remove the local-to-local attention, then we obtain the pure global memory attention mechanism (called Global Attention hereafter). Its memory complexity is $\mathcal{O}(n_{g}(n_{g}+n_{l}))$ , which is also linear w.r.t. $n_{l}$ . However, for this pure global memory attention, $n_{g}$ has to be much larger than 1. In practice, we set different numbers of global tokens for different stages, as shown in Table 15, with more global tokens in the first 2 stages and less in the last 2 stages. This setting makes the memory/computation complexity comparable with other attention mechanisms under the same model size.

Linformer ( ${\color[rgb]{1,0,0}a}=\text{LIN}$ ) projects the $n_{l}\times d$ dimensional keys and values to $K\times d$ dimensions using additional projection layers, where $K\ll n_{l}$ . Then the $n_{l}$ queries only attend to these projected $K$ key-value pairs. The memory complexity of Linformer is $\mathcal{O}(Kn_{l})$ . We gradually increase $K$ (by 2 each time) and its performance gets nearly saturated at 256. Therefore, $K=256$ is our choice for this Linformer attention, which turns out to be the same with the recommended value. Notice that Linformer’s projection layer (of dimension $K\times n_{l}$ ) is specific to the current $n_{l}$ , and cannot be transferred to higher-resolution tasks that have a different $n_{l}$ . It is possible to transfer Linformer’s weight by resizing feature maps of a different size to the original feature map size that Linformer is trained with and then applying the Linformer’s projection. We do not explore this choice in this work.

In this paper, we benchmarked the performance of SRA/32 with SR ratios $R_{1}=8,R_{2}=4,R_{3}=2,R_{4}=1$ (same as PVT ) and SRA/64 with SR ratios $R_{1}=16,R_{2}=8,R_{3}=4,R_{4}=2$ (two times more downsizing from that in PVT ), as shown in Table 15. The SRA/64 setting makes the memory usage comparable with other efficient attention mechanisms under the same model size, but introduces more parameters due to doubling the kernel size of the convolutional projection layer.

Performer ( ${\color[rgb]{1,0,0}a}=\text{performer}$ ) uses random kernel approximations to approximate the Softmax computation in MSA, and achieves a linear computation/memory complexity with respect to $n_{l}$ and the number of random features $K$ . We use the default $K=256$ orthogonal random features (OR) for Performer. The memory/space complexity of performer is $\mathcal{O}(Kd+n_{l}d+Kn_{l})$ while its computation/time complexity is $\mathcal{O}(Kn_{l}d)$ . For the time complexity, we ignore the complexity of generating the orthogonal random features, which in practice cannot be ignored during training. We refer to Section B.3 in for a detailed discussion of theoretical computation/memory complexity of Performer.

One important technique in training Performer is to redraw the random features during training. In our ImageNet classification training, we adopt a heuristic adaptive redrawing schedule: redraw every $1+5T$ iterations in Epoch $T$ ( $T=0,1,...,299$ ). In our COCO object detection/segmentation training, the Performer is initialized from ImageNet pretrained checkpoint and thus there is no need to redraw very frequently in the initial training stage.Therefore, we redraw the random features every 1000 iterations in COCO object detection/segmentation training.