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 paper proposed Long-Short Transformer (Transformer-LS), which aggregates sliding window attention (i.e. short-term attention) with linear attention via low-rank projection (i.e. long-term attention), by simply concatenating the two sets of keys. Due to the scale mis-match of the keys from the long and short term attentions, the authors further proposed the dual layernorm (DualLN) method.  Experiments were conducted on Long-range Arena (LRA), auto-regressive language modeling and ImageNet classification. Transformer-LS achieves improvements over previous efficient transformer models on both accuracy and efficiency.