Long-Short Transformer: Efficient Transformers for Language and Vision

Introduction

Transformer-based models have achieved great success in the domains of natural language processing (NLP) and computer vision . These models benefit from the self-attention module, which can capture both adjacent and long-range correlations between tokens while efficiently scaling on modern hardware. However, the time and memory consumed by self-attention scale quadratically with the input length, making it very expensive to process long sequences. Many language and vision tasks benefit from modeling long sequences. In NLP, document-level tasks require processing long articles [e.g., 7, 8], and the performance of language models often increases with sequence length [e.g., 9, 10]. In computer vision, many tasks involve high-resolution images, which are converted to long sequences of image patches before being processed with Transformer models . As a result, it is crucial to design an efficient attention mechanism for long sequence modeling that generalizes well across different domains.

Various methods have been proposed to reduce the quadratic cost of full attention. However, an efficient attention mechanism that generalizes well in both language and vision domains is less explored. One family of methods is to sparsify the attention matrix with predefined patterns such as sliding windows [e.g., 12, 13, 14, 15] and random sparse patterns . These methods leverage strong inductive biases to improve both computational and model performance, but they limit the capacity of a self-attention layer because each specific token can only attend to a subset of tokens. Another family of methods leverages low-rank projections to form a low resolution representation of the input sequence, but the successful application of these methods has been limited to certain NLP tasks [e.g., 17, 18, 19]. Unlike sparse attention, this family of methods allows each token to attend to the entire input sequence. However, due to the loss of high-fidelity token-wise information, their performance sometimes is not as good as full attention or sparse attention on tasks that require fine-grained local information, including standard benchmarks in language and vision .

Despite the rapid progress in efficient Transformers, some proposed architectures can only be applied to bidirectional models [e.g., 15, 16, 18]. Transformer-based autoregressive models have achieved great successes in language modeling , image synthesis , and text-to-image synthesis , which also involve long texts or high-resolution images. It is desirable to design an efficient transformer that can be applied to both autoregressive and bidirectional models.

In this work, we unify a local window attention and a novel long-range attention into a single efficient attention mechanism. We show that these two kinds of attention have complementary effects that together yield the state-of-the-art results on a range of tasks in language and vision, for both autoregressive and bidirectional models. Specifically, we make the following contributions:

We propose Long-Short Transformer (Transformer-LS), an efficient Transformer that integrates a dynamic projection based attention to model long-range correlations, and a local window attention to capture fine-grained correlations. Transformer-LS can be applied to both autoregressive and bidirectional models with linear time and memory complexity.

We compute a dynamic low-rank projection, which depends on the content of the input sequence. In contrast to previous low-rank projection methods, our dynamic projection method is more flexible and robust to semantic-preserving positional variations (e.g., insertion, paraphrasing). We demonstrate that it outperforms previous low-rank methods on Long Range Arena benchmark .

We identify a scale mismatch problem between the embeddings from the long-range and short-term attentions, and design a simple but effective dual normalization strategy, termed DualLN, to account for the mismatch and enhance the effectiveness of the aggregation.

We demonstrate that Long-Short Transformer, despite its low memory and runtime complexity, outperforms the state-of-the-art models on a set of tasks from Long Range Arena, and autoregressive language modeling on enwik8 and text8. In addition, the proposed efficient attention mechanism can be easily applied to the most recent vision transformer architectures and provides state-of-the-art results, while being more scalable to high-resolution images. We also investigate the robustness properties of the Transformer-LS on diverse ImageNet datasets.

Related Work

In recent years, many methods have been introduced for dealing with the quadratic cost of full attention. In general, they can be categorized as follows: i) Sparse attention mechanism with predefined patterns (e.g., sliding window), including Sparse Transformer , Image Transformer , Axial Transformer for modeling images, and Longformer , blockwise self-attention , ETC , Big Bird for modeling language. ii) Low-rank projection attention, including Linformer , Nyströmformer , Synthesizer . For example, Linformer uses linear layers to project the original high resolution keys ($K$) and values ($V$) with length $n$ to low resolution with size $r$ ($r\ll n$) and allows all query tokens ($Q$) to attend these compressed representations. iii) Memory-based mechanisms like Compressive Transformer and Set Transformer , which use extra memories for caching global long-range information for use in computing attention between distant tokens. iv) Kernel-based approximation of the attention matrix, including Performer , Linear Transformer , and Random Feature Attention . vi) Similarity and clustering based methods, including Reformer , Routing Transformer , and Sinkhorn Transformer .

Our method seamlessly integrates both low-rank projection and local window attentions, to leverage their strengths for modeling long-range and short-term correlations. In particular, our long-range attention uses a dynamic low-rank projection to encode the input sequence, and outperforms the previous low-rank projection method used by the Linformer . In the similar vein, a few other methods also try to combine the strengths of different methods. For example, Longformer and ETC augment local window attention with task motivated global tokens. Such global tokens may not be applicable for some tasks (e.g., autoregressive modelling). BigBird further combines local window and global token attention with random sparse attention. It is not applicable in autoregressive tasks because the global token and random sparse pattern are introduced. To compress the model footprint on edge devices, Lite Transformer combines convolution and self-attention, but it still has quadratic complexity for long sequences.

2 Vision Transformers

Vision Transformer (ViT) splits images as small patches and treats the patches as the input word tokens. It uses a standard transformer for image classification and has shown to outperform convolutional neural networks (e.g., ResNet ) with sufficient training data. DeiT has applied the teacher-student strategy to alleviate the data efficiency problem of ViT and has shown strong comparable performance using only the standard ImageNet dataset . Instead of applying transformer at a single low resolution of patches (e.g., $16\times 16$ patches), very recent works, including Pyramid Vision Transformer (PVT) , Swin-Transformer , T2T-ViT , Vision Longformer (ViL) and Convolutional Vision Transformer (CvT) , stack a pyramid of ViTs to form a multi-scale architecture and model long sequences of image patches at much higher resolution (e.g., $56\times 56=3136$ patches for images with $224\times 224$ pixels). Most of these methods have quadratic complexity of self-attention with respect to the input image size.

To reduce the complexity, Swin-Transformer achieves linear complexity by limiting the computation of self-attention only within each local window. HaloNet applies local attention on blocked images and only has quadratic complexity with respect to the size of the block. Perceiver uses cross-attention between data and latent arrays to replace the self-attention on data to remove the quadratic complexity bottleneck. Vision Longformer (ViL) , another concurrent work, achieves linear complexity by adapting Longformer to Vision. ViL augments local window attention with task-specific global tokens, but the global tokens are not applicable for decoding task (e.g., image synthesis ). In contrast, our method reduces the quadratic cost to linear cost by combining local window attention with global dynamic projection attention, which can be applied to both encoding and decoding tasks.

Long-Short Transformer

Transformer-LS approximates the full attention by aggregating long-range and short-term attentions, while maintaining its ability to capture correlations between all input tokens. In this section, we first introduce the preliminaries of multi-head attention in Transformer. Then, we present the short-term attention via sliding window, and long-range attention via dynamic projection, respectively. After that, we propose the aggregating method and dual normalization (DualLN) strategy. See Figure 1 for an illustration of our long-short term attention.

Multi-head attention is a core component of the Transformer , which computes contextual representations for each token by attending to the whole input sequence at different representation subspaces. It is defined as

2 Short-term Attention via Segment-wise Sliding Window

The sliding window attention can be augmented to capture long-range correlations in part, by introducing different dilations to different heads of sliding window attention . However, the dilation configurations for different heads need further tuning and an efficient implementation of multi-head attention with different dilations is non-trivial. A more efficient alternative is to augment sliding window attention with random sparse attention , but this does not guarantee that the long-range correlations are captured in each layer as in full attention. In the following section, we propose our long-range attention to address this issue.

3 Long-range Attention via Dynamic Projections

Previous works have shown that the self-attention matrix can be well approximated by the product of low-rank matrices . By replacing the full attention with the product of low-rank matrices , each query is able to attend to all tokens. Linformer is one of the most representative models in this category. It learns a fixed projection matrix to reduce the length of the keys and values, but the fixed projection is inflexible to semantic-preserving positional variations.

Application to Autoregressive Models: In autoregressive models, each token can only attend to the previous tokens, so the long-range attention should have a different range for different tokens. A straightforward way to implement our global attention is to update $\bar{K}_{i},\bar{V}_{i}$ for each query recurrently, but this requires re-computing the projection in Eq. (3) for every token due to the nonlinearity of softmax, which results in $O(rn^{2})$ computational complexity. To preserve the linear complexity, for autoregressive models, we first divide the input sequence into equal-length segments with length $l$, and apply our dynamic projection to extract $\bar{K}_{i},\bar{V}_{i}$ from each segment. Each token can only attend to $\bar{K}_{i},\bar{V}_{i}$ of segments that do not contain its future tokens. Formally, let $Q_{t}$ be the query at position $t$, $K_{(l-1)s:ls},V_{(l-1)s:ls}$ be the key-value pairs from the $s$-th segment, and $s_{t}=\lfloor{t/l}\rfloor$. For autoregressive models, we compute the long-range attention of $Q_{t}$ by attending to $K_{i,t},V_{i,t}$, defined as

In this way, the dynamic low-rank projection is applied to each segment only once in parallel, preserving the linear complexity and the high training speed. By comparison, Random Feature Attention is slow at training due to the requirement for recurrence.

4 Aggregating Long-range and Short-term Attentions

To avoid such issues, we add two sets of Layer Normalizations after the key and value projections for the local window and global low-rank attentions, so that their scales are aligned at initialization, but the network can still learn to re-weight the norms after training. Specifically, the aggregated attention is now computed as

where $\text{LN}{}_{L}(\cdot),\text{LN}{}_{G}(\cdot)$ denote the Layer Normalizations for the local and global attentions respectively. In practice, to maintain the consistency between the local attention and dynamic projection, we use $\text{LN}_{L}(K),\text{LN}_{L}(V)$ instead of $K,V$ to compute $\bar{K}_{i},\bar{V}_{i}$ in Eq. 3. As illustrated in Figure 2 Right, the Transformer-LS models trained with DualLN has consistently lower validation loss than the models without DualLN.

Experiments

In this section, we demonstrate the effectiveness and efficiency of our method in both language and vision domains. We use PyTorch for implementation and count the FLOPs using fvcore .

To evaluate Long-Short Transformer as a bidirectional encoder for long text, we train our models on the three NLP tasks, ListOps, Text, and Retrieval, from the recently proposed Long Range Arena (LRA) benchmark , following the setting of Peng et al. and Tay et al. . For fair comparisons, we use the PyTorch implementation and the same data preprocessing/split, training hyperparameters and model size from , except for Retrieval where we accidentally used more warmup steps and improved the results for all models. See Appendix B for more details. The results on these three tasks are given in Table 1. Results of the other two image-based tasks of LRA, as well as models implemented in JAX, are given in Appendix C and C.2.

In addition, we follow the pretraining procedure of Longformer to pretrain our models based on RoBERTa-base and RoBERTa-large , and fine-tune it on the IMDb sentiment classification dataset. The results are given in Table 3.

Results. From Table 3, our base model outperforms Longformer-base, and our large model achieves improvements over RoBERTa-large, demonstrating the benefits of learning to model long sequences. Comparisons with models on LRA are given in Table 1. Transformer-LS (best) with the best configurations of $w,r$ for each task are given in Table 7 in Appendix B. We also report the results of using fixed hyperparameter $w=8,r=32$ on all tasks. Overall, our Transformer-LS (best) is significantly better than other efficient Transformers, and the model with $w,r=8,32$ performs favorably while using only about 50% to 70% computation compared to other efficient Transformers on all three tasks. The advantage of aggregating local and long-range attentions is the most significant on ListOps, which requires the model to understand the tree structures involving both long-term and short-term relations. On Retrieval, where document-level encoding capability is tested, we find our global attention more effective than window attention. The test accuracy of using only dynamic projection is about 10% higher than Linformer on Text (i.e., 66.28 vs. 56.12), which has the highest variance in sequence length (i.e. standard deviation 893). This demonstrates the improved flexibility of dynamic projection at learning representations for data with high variance in sequence length, compared to the learned but fixed projection of Linformer. Similarly, Linformer, Nyströmformer and our model outperform full attention on ListOps, indicating they may have better inductive bias, and efficient Transformers can have better efficacy beyond efficiency.

Robustness of Dynamic Projection. In Table 2, we compare the robustness of Linformer and the proposed Dynamic Projection (DP) against insertion and deletion on Text and Retrieval tasks of LRA. We train the models on the original, clean training sets and only perturb their test sets. For insertion, we insert 10 random punctuations at 10 random locations of each test sample. For deletion, we delete all punctuations from the test samples. Both transforms are label-preserving in most cases. By design, dynamic projection is more robust against location changes.

2 Autoregressive Language Modeling

We compare our method with other efficient transformers on the character-level language modeling where each input token is a character.

Setup. We train and evaluate our model on enwik8 and text8, each with 100M characters and are divided into 90M, 5M, 5M for train, dev, test, following . Our smaller 12-layer and larger 30-layer models are Pre-LN Transformers with the same width and depth as Longformer , except that we add relative position encoding to the projected segments in each layer. We adopt the cache mechanism of Transformer-XL , setting the cache size to be the same as the input sequence length. We follow similar training schedule as Longformer, and train our model in 3 phases with increasing sequence lengths. The input sequence lengths are 2048, 4096 and 8192 respectively for the 3 phases. By comparison, Longformer trains their model in 5 phases on GPUs with 48GB memory (The maximal of ours is 32GB) where the sequence length is 23,040 in the last phase. The window size of Longformer increases with depth and its average window size is 4352 in phase 5, while our effective number of attended tokens is 1280 on average in the last phase. Each experiment takes around 8 days to finish on 8 V100 GPUs. Detailed hyperparameters are shown in Appendix D. For testing, same as Longformer, we split the dataset into overlapping sequences of length 32K at a step size of 512, and evaluate the BPCs for predicting the next 512 tokens given the previous 32K characters.

Results Table 4 shows comparisons on text8 and enwik8. Our method has achieved state-of-the-art results. On text8, we achieve a test BPC of 1.09 with the smaller model. On enwik8, our smaller model achieves a test BPC of 0.99, and outperforms the state-of-the-art models with comparable number of parameters. Our larger model obtains a test BPC of 0.97, on par with the Compressive Transformer with 2$\times$ parameters. Our results are consistently better than Longformer which is trained on longer sequences with 5 stages and 48 GPU memory. In Figure 3, we show our model is much more memory and computational efficient than full attention.

3 ImageNet Classification

We train and evaluate the models on ImageNet-1K with 1.3M images and 1K classes. We use CvT and ViL , state-of-the art vision transformer architectures, as the backbones and replace their attention mechanisms with our long-short term attention, denoted as CvT∗-LS and ViL-size-LS in Table 5. CvT uses overlapping convolutions to extract dense patch embeddings from the input images and feature maps, resulting in a long sequence length in the early stages (e.g., $56\times 56=3136$ patches for images with $224^{2}$ pixels). For ViL, our sliding window uses the same group size $w$, but each token attends to at most $2w\times 2w$ (rounding when necessary) tokens inside the window, instead of $3w\times 3w$ as ViL, which allows adding our dynamic projection without increasing the FLOPs. We set $r=8$ for the dynamic projections for both ViL-LS-Medium and ViL-LS-Base. Note that, our efficient attention mechanism does not depend on the particular architecture, and it can be applied to other vision transformers [e.g., 4, 36, 5]. Please refer to Appendix E for more details.

Classification Results. The results are shown in the Table 5, where we also list test accuracies on ImageNet Real and ImageNet V2. Except for CvT, we compare with the original ViT and the enhanced DeiT , PVT that also uses multi-scale stragey, ViL that uses window attention and global tokens to improve the efficiency. Training at high-resolution usually improves the test accuracy of vision transformer. With our long-short term attention, we can easily scale the training to higher resolution, and the performance of CvT∗-LS and ViL-LS also improves. Our best model with CvT (CvT∗-LS-21 at $448^{2}$) achieves 0.3% higher accuracy than the best reported result of CvT while using the same amount of parameters and 76% of its FLOPs. In CvT architecture, the spatial dimension of feature maps in earlier stages are large, representing more fine-grained details of the image. Similar to training with high-resolution images, the model should also benefit from denser feature maps. With our efficient long-short term attention, we can better utilize these fine-grained feature maps with less concerns about the computational budget. In this way, our CvT∗-LS-17 achieves better result than CvT-21 at resolution 224 using fewer parameters and FLOPs, and our CvT∗-LS-21S model further improves our CvT∗-LS-21 model.

Our ViL-LS-Medium and ViL-LS-Base with long-short term attention improve the accuracies of ViL-Medium and ViL-Base from 83.5 and 83.7 to 83.8 and 84.1 respectively, without an increase in FLOPs. When increasing the resolution for training ViL-LS-Medium from $224^{2}$ to $384^{2}$, the FLOPs increased (approximately) linearly and the accuracy improved by 0.6%, showing our method still benefits greatly from increased resolution while maintaining the linear complexity in practice.

Short-term Attention Suppresses Oversmoothing. By restricting tokens from different segments to attend to different windows, our short-term sparse local attention encourages diversity of the feature representations and helps to alleviate the over-smoothing problem (where all queries extract similar information in deeper layers and the attention mechanism is less important), thus can fully utilize the depth of the network. As in , we provide the cosine similarity of patch embeddings of our CvT∗-LS-13 and re-implemented CvT-13 (81.1 accuracy) in Figure 6 within Appendix. This is one of the reasons why our efficient attention mechanism can get even better results than the full attention CvT model in the same setting.

Robustness evaluation on Diverse ImageNet Datasets.

As vision models have been widely used in safety-critical applications (e.g. autonomous driving), their robustness is vital. In addition to out-of-distribution robustness (ImageNet-Real and Imageet-v2), we further investigate the robustness of our vision transformer against common corruption (ImageNet-C), semantic shifts (ImageNet-R), Background dependence (ImageNet-9) and natural adversarial examples (ImageNet-A). We compare our methods with standard classification methods, including CNN-based model (ResNet ) and Transformer-based models (DeiT ) with similar numbers of parameters. As shown in Table 6, we observe that our method significantly outperforms the CNN-based method (ResNet-50). Compared to DeiT, our models also achieve favorable improvements. These results indicate that the design of different attention mechanisms plays an important role for model robustness, which sheds new light on the design of robust vision transformers. More details and results can be found in Appendix E.

Conclusion

In this paper, we introduced Long-Short Transformer, an efficient transformer for long sequence modeling for both language and vision domain, including both bidirectional and autoregressive models. We design a novel global attention mechanism with linear computational and memory complexity in sequence length based on a dynamic projection. We identify the scale mismatch issue and propose the DualLN technique to eliminate the mismatch at initialization and more effectively aggregate the local and global attentions. We demonstrate that our method obtains the state-of-the-art results on the Long Range Arena, char-level language modeling and ImageNet classification. We look forward to extending our methods to more domains, including document QA, object detection and semantic segmentation on high-resolution images.

References

Appendix A Details of Norm Comparisons

Appendix B Details for Experiments on Long Range Arena

We compare our method with the following three tasks:

ListOps. ListOps is designed to measure the parsing ability of models through hierarchically structured data. We follow the setting in in which each instance contains 500-2000 tokens.

Text. This is a binary sentiment classification task of predicting whether a movie review from IMDb is positive or negative . Making correct predictions requires a model to reason with compositional unsegmented char-level long sequences with a maximum length of 4k.

Retrieval. This task is based on the ACL Anthology Network dataset . The model needs to classify whether there is a common citation between a pair of papers, which evaluates the model’s ability to encode long sequences for similarity-based matching. The max sequence length for each byte-level document is 4k and the model processes two documents in parallel each time.

Architecture.

On all tasks, the models have 2 layers, with embedding dimension $d=64$, head number $h=2$, FFN hidden dimension 128, smaller than those from . Same as , we add a CLS token as a global token and use its embedding in the last layer for classification. We re-implement the methods evaluated by Xiong et al. , and report the best results of our re-implementation and those reported by Xiong et al. . For our method, the results we run a grid search on the window size $w$ and the projected dimension $r$, and keep $2w+r\leq 256$ to make the complexity similar to the other methods. The maximum sequence length for ListOps and Text are 2048 and 4096. For Retrieval, we set the max sequence for each of the two documents to 4096.

Hyperparameters for Training.

Our hyperparameters are the same as Nyströmformer unless otherwise specified. Specifically, we follow and use Adam with a fixed learning rate of $10^{-4}$ without weight decay, batch size 32 for all tasks. The number of warmup training steps $T_{w}$ and total training steps $T$ are different due to the difference in numbers of training samples. For Retrieval, we accidentally found using $T_{w}=8000$ rather than the default $T_{w}=800$ of improves the results for all models we have evaluated. See Table 8 for the configurations of each task.

Error bars.

We have already provided the average of 4 runs with different random seeds in Table 1. Here we also provide the standard deviations for these experiments in Table 9.

Appendix C Additional Results on LRA

We give the results of our model on the image-based tasks, implemented in PyTorch, in Table 10.

C.2 Compare models implemented in JAX

To compare the results with the implementations from the original LRA paper , we re-implement our method in JAX and give the comparisons with other methods in Table 11. The accuracies of other methods come from the LRA paper. We evaluate the per-batch latency of all models on A100 GPUs using their official JAX implementation from the LRA paper. Our method still achieves improvements while being efficient enough. We were unable to run Reformer with the latest JAX since JAX has deleted jax.custom_transforms, which is required by the Reformer implementation, from its API.https://github.com/google/jax/pull/2026 Note the relative speedups from the LRA paper are evaluated on TPUs.

Appendix D Details for Autoregressive Language Modeling

We give an illustration for the segment-wise dynamic projection for autoregressive models as discussed in Section 3.3. With the segment-wise formulation, we can first compute the low-rank projection for each segment in parallel, and each query will only attend to the tokens from segments that do not contain the future token or the query token itself. The whole process is efficient and maintain the $O(n)$ complexity, unlike RFA which causes a slow-down in training due to the requirement for cumulative sum. However, in this way, some of the most recent tokens are ignored, as shown in Figure 4 (left). The window attention (with segment size $w\geq l/2$) becomes an indispensable component in this way, since it fills the gap for the missing recent tokens, as shown in Figure 4.

Experimental Setup.

Throughout training, we set the window size $w=512$, the segment length $l=16$, and the dimension of the dynamic low-rank projection $r=1$, which in our initial experiments achieved better efficiency-BPC trade-off than using $l=32,r=1$ or $l=64,r=4$. Our small and large models have the same architecture as Longformer , except for the attention mechanisms. We use similar training schedules as Longformer . Specifically, for all models and both datasets, we train the models for 430k/50k/50k steps with 10k/5k/5k linear learning rate warmup steps, and use input sequence lengths 2048/4096/8192 for the 3 phases. We use constant learning rate after warmup. We compared learning rates from {1.25e-4, 2.5e-4,5e-4,1e-3} for 100k iterations and found 2.5e-4 to work the best for both models on enwik8, and 5e-4 to work the best on text8. The batch sizes for the 3 phases are 32, 32, 16 respectively. Unlike Longformer and Transformer-XL, we remove gradient clipping and found the model to have slightly faster convergence in the beginning while converging reliably. For smaller models, we use dropout rate 0.2 and weight decay 0.01. For the larger model, we use dropout 0.4 and weight decay 0.1.

Appendix E Details for ImageNet Classification

We implement the CvT model based on a public repository, https://github.com/rishikksh20/convolution-vision-transformers because this is a concurrent work with no official implementation when we conduct this work. In Table 5, since our CvT re-implementation gets worse test results than reported ones in their arxiv paper, we still list the best test accuracy from Wu et al. for fair comparisons. We report the FLOPs of CvT with our implementation for reasonable comparisons, because our CvT∗-LS implementation is based on that. Same as CvT, all the models have three stages where the first stage downsamples the image by a factor of 4 and each of the following stages downsamples the feature map by a factor of 2. CvT∗-LS-13 and CvT∗-LS-21 have the same configuration as CvT-13 and CvT-21. CvT∗-LS-17 and CvT∗-LS-21 are our customized models with more layers and higher embedding dimensions in the first two stages ($$,$$ layers respectively and $$dimensions). We train the model for 300 epochs using a peak learning rate of$5e-4$with the cosine schedule with 5 epochs of warmup. We use the same set of data augmentations and regularizations as other works including PVT and ViL . In general, CvT∗-LS-13 and CvT∗-LS-21 closely follow the architectural designs of CvT for fair comparisons. Specifically, in CvT∗-LS, we feed the token embeddings extracted by the depth-wise separable convolution of CvT to our long-short term attention. For dynamic projection, we replace$W_{i}^{P}$in Eq. (3) with a depth-wise separable convolution to maintain consistency with the patch embeddings, but we change its BN layer into a weight standardization on the spatial convolution’s weights for simplicity. We do not use position encoding. All of our models have 3 stages, and the feature map size is the same as CvT in each stage when the image resolutions are the same. CvT∗-LS-13 and CvT∗-LS-21 follow the same layer configurations as CvT-13 and CvT-21, i.e., the number of heads, the dimension of each head and the number of Transformer blocks are the same as CvT in each stage. For all models on resolution$224\times 224$, we set$r=$and$w=$. For higher resolutions, we scale up$r$and/or$w$to maintain similar effective receptive fields for the attentions. At resolution$384\times 384$, we use$r=$and$w=$for the 3 stages. At resolution$448\times 448$, we use$r=$and$w=$.

Besides maintaining the CvT architectures, we also try other architectures to further explore the advantage of our method. With the efficient long-short term attention, it becomes affordable to stack more layers on higher-resolution feature maps to fully utilize the expressive power of attention mechanisms. Therefore, we have created two new architectures, CvT∗-LS-17 and CvT∗-LS-21S, that have more and wider layers in the first two stages, as shown in Table 12. Compared with CvT-21, CvT∗-LS-17 has 25% fewer parameters, less FLOPs, but obtained the same level of accuracy. CvT∗-LS-21S has fewer parameters than CvT∗-LS-21, more FLOPs, and 0.4% higher accuracy, demonstrating the advantage of focusing the computation on higher-resolution feature maps.

The effect of DualLN.

We trained the CvT∗-LS-13 model without DualLN, which has a test accuracy of 81.3, lower than the 81.9 with DualLN.

Appendix F Evaluate the robustness of models trained on ImageNet-1k.

For a fair comparison, we choose models with similar number of parameters. We select two representative models, including the CNN-based model (ResNet) and the transformer-based model (DeiT). We give detailed results on all types of image transforms on ImageNet-C in Table 13. We evaluate our method on various ImageNet robustness benchmarks as follows:

ImageNet-C. ImageNet-C refers to the common corruption dataset. It consists of 15 types of algorithmically common corruptions from noise, blur, weather, and digital categories. Each type contains five levels of severity. In Table 4, we report the normalized mean corruption error (mCE) defined in Hendrycks and Dietterich . In Table 13, we report the corruption error among different types. In both tables, the lower value means higher robustness.

ImageNet-A. ImageNet-A is the natural adversarial example dataset. It contains naturally collected images from online that mislead the ImageNet classifiers. It contains 7,500 adversarially filtered images. We use accuracy as our evaluation metric. The higher accuracy refers to better robustness.

ImageNet-R. ImageNet-R (Rendition) aims to evaluate the model generalization performance on out-of-distribution data. It contains renditions of 200 ImageNet classes (e.g. cartoons, graffiti, embroidery). We use accuracy as the evaluation metric.

ImageNet-9. ImageNet-9 aims to evaluate the model background robustness. It designs to measure the extent of the model relying on the image background. Following the standard setting , we evaluate the two categories, including Mixed-Same and Mixed-Rand. Mixed-Same refers to replace the background of the selected image with a random background of the same class by GrabCut ; Mixed-Rand refers to replace the image background with a random background of the random class.

From table 6, we find that our method achieves significant improvement compared to CNN-based network (ResNet). For instance, our method improves the accuracy by 23.6%, 22.1%, 9.7% compared to ResNet on ImageNet-C, ImageNet-A, and ImageNet-R, respectively. For ImageNet-9, our method also achieves favorable improvement by 4.3% on average (Mixed-same and Mixed-rand). It indicates that our method is insensitive to background changes. We guess the potential reasons for these improvements are (1) the attention mechanism and (2) the strong data augmentation strategies during the training for vision transformer . The first design helps the model focus more on the global context of the image as each patch could attend to the whole image areas. It reduces the local texture bias of CNN. The latter design increases the diversity of the training data to improve model’s generalization ability. Compared to DeiT, we also surprisingly find that our method achieves slightly better performance. One plausible explanation is that our long-term attention has a favorable smoothing effect on the noisy representations. Such improvements also indicate that different designs of attention and network architecture can be essential to improve the robustness. As the goal of this paper is not to design a robust vision transformer, the robustness is an additional bonus of our method. We believe that our observation opens new directions for designing robust vision Transformers. We leave the in-depth study as an important future work.

The detailed results of ImageNet-C and ImageNet-9 are shown in Table 13 and Table 14 respectively.