Long-Short Transformer: Efficient Transformers for Language and Vision

1 Introduction . Transformer-based models [ 1 ] have achieved great success in the domains of natural language processing ( NLP ) [ 2 , 3 ] and computer vision [ 4–6 ] . 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 [ 4 , 6 , 11 ] . As a result , it is crucial to design an efficient attention mechanism for long sequence modeling that generalizes well across different domains . ∗Work done during an internship at NVIDIA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 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–15 ] and random sparse patterns [ 16 ] . 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–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 [ 20 ] and vision [ 21 ] . 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 [ 22 ] , image synthesis [ 23 ] , and text-to-image synthesis [ 24 ] , 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 [ 17 , 18 ] on Long Range Arena benchmark [ 20 ] . • We identify a scale mismatch problem between the embeddings from the long-range and shortterm 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 [ 6 , 11 ] 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 . 2 Related Work . 2.1 Efficient Transformers . 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 [ 12 ] , Image Transformer [ 13 ] , Axial Transformer [ 25 ] for modeling images , and Longformer [ 14 ] , blockwise self-attention [ 26 ] , ETC [ 15 ] , Big Bird [ 16 ] for modeling language . ii ) Low-rank projection attention , including Linformer [ 17 ] , Nyströmformer [ 18 ] , Synthesizer [ 19 ] . 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 n ) and allows all query tokens ( Q ) to attend these compressed representations . iii ) Memory-based mechanisms like Compressive Transformer [ 10 ] and Set Transformer [ 27 ] , 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 [ 28 ] , Linear Transformer [ 29 ] , and Random Feature Attention [ 30 ] . vi ) Similarity and clustering based methods , including Reformer [ 31 ] , Routing Transformer [ 32 ] , and Sinkhorn Transformer [ 33 ] . 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 [ 17 ] . In the similar vein , a few other methods also try to combine the strengths of different methods . For example , Longformer [ 14 ] and ETC [ 15 ] augment local window attention with task motivated global tokens . Such global tokens may not be applicable for some tasks ( e.g. , autoregressive modelling ) . BigBird [ 16 ] 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 [ 34 ] combines convolution and self-attention , but it still has quadratic complexity for long sequences . 2.2 Vision Transformers . Vision Transformer ( ViT ) [ 4 ] 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 [ 35 ] ) with sufficient training data . DeiT [ 36 ] 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 [ 37 ] . Instead of applying transformer at a single low resolution of patches ( e.g. , 16× 16 patches ) , very recent works , including Pyramid Vision Transformer ( PVT ) [ 5 ] , Swin-Transformer [ 38 ] , T2T-ViT [ 39 ] , Vision Longformer ( ViL ) [ 11 ] and Convolutional Vision Transformer ( CvT ) [ 6 ] , 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× 56 = 3136 patches for images with 224× 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 [ 38 ] achieves linear complexity by limiting the computation of self-attention only within each local window . HaloNet [ 40 ] applies local attention on blocked images and only has quadratic complexity with respect to the size of the block . Perceiver [ 41 ] uses cross-attention between data and latent arrays to replace the self-attention on data to remove the quadratic complexity bottleneck . Vision Longformer ( ViL ) [ 11 ] , another concurrent work , achieves linear complexity by adapting Longformer [ 14 ] 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 [ 23 , 24 ] ) . 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 . 3 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 . 3.1 Preliminaries and Notations . Multi-head attention is a core component of the Transformer [ 1 ] , which computes contextual representations for each token by attending to the whole input sequence at different representation subspaces . It is defined as MultiHead ( Q , K , V ) = Concat ( H1 , H2 , ... , Hh ) WO , ( 1 ) where Q , K , V ∈ Rn×d are the query , key and value embeddings , WO ∈ Rd×d is the projection matrix for output , the i-th head Hi ∈ Rn×dk is the scaled dot-product attention , and dk = d/h is the embedding dimension of each head , Hi = Attention ( QW Q i , KW K i , V W V i ) = softmax [ QWQi ( KWKi ) ᵀ √ dk ] VWVi = AiVW V i , ( 2 ) where WQi , W K i , W V i ∈ Rd×dk are learned projection matrices , and Ai ∈ Rn×n denotes the full attention matrix for each attention head . The complexity of computing and storing Ai is O ( n2 ) , which can be prohibitive when n is large . For simplicity , our discussion below is based on the case of 1D input sequences . It is straightforward to extend to the 2D image data given a predetermined order . 3.2 Short-term Attention via Segment-wise Sliding Window . We use the simple yet effective sliding window attention to capture fine-grained local correlations , where each query attends to nearby tokens within a fixed-size neighborhood . Similar techniques have also been adopted in [ 14 , 16 , 11 ] . Specifically , we divide the input sequence into disjoint segments with length w for efficiency reason . All tokens within a segment attend to all tokens within its home segment , as well as w/2 consecutive tokens on the left and right side of its home segment ( zero-padding when necessary ) , resulting in an attention span over a total of 2w key-value pairs . See Figure 5 in Appendix for an illustration . For each query Qt at the position t within the i-th head , we denote the 2w key-value pairs within its window as K̃t , Ṽt ∈ R2w×d . For implementation with PyTorch , this segment-wise sliding window attention is faster than the per-token sliding window attention where each token attends to itself and w tokens to its left and right , and its memory consumption scales linearly with sequence length ; see [ 14 ] and our Figure 3 for more details . 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 [ 14 ] . 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 [ 16 ] , 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 .

Summary: This work presents Long-Short Transformer (Transformer-LS), which integrates a dynamic projection based long-range attention and a local window short-term attention to capture both global and local features. Specifically, 1) the long-range attention is implemented by projecting the Key and Value embeddings into shorter ones (in terms of number of tokens) based on the projection matrix generated from the original Key embeddings and 2) the short-term attention is implemented by dividing the input sequences into disjoint segments. A scale mismatch is identified from the long-range and short-term attention and solved by the proposed DualLN. The experiments are conducted in both NLP and CV tasks.