Linformer: Self-Attention with Linear Complexity

1 Linformer: Self-Attention with Linear ComplexitySinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, Hao MaFacebook AI, Seattle, WA{sinongwang, belindali, hanfang, mkhabsa, transformer models have shown extraordinary success in achieving state-of-the-art results in many natural language processing applications. However, trainingand deploying these models can be prohibitively costly for long sequences, asthe standard Self-Attention mechanism of the Transformer usesO(n2)time andspace with respect to sequence length. In this paper, we demonstrate that theself-attention mechanism can be approximated by a low-rank matrix. We furtherexploit this finding to propose a new Self-Attention mechanism, which reducesthe overall Self-Attention Complexity fromO(n2)toO(n)in both time and resulting Linear transformer, theLinformer, performs on par with standardTransformer models, while being much more memory- and IntroductionTransformer models (Vaswani et al.)}

2 , 2017) have become ubiquitous for wide variety of problemsin natural language processing (NLP), including translation (Ott et al., 2018), text classification,question answering, among others (Raffel et al., 2019; Mohamed et al., 2019). Over the last couple ofyears, the number of parameters in state-of-the-art NLP transformers has grown drastically, from theoriginal 340 million introduced in BERT-Large to 175 billion in GPT-3 (Brown et al., 2020). Althoughthese large-scale models yield impressive results on wide variety of tasks, training and deployingsuch model are slow in practice. For example, the original BERT-Large model (Devlin et al., 2019)takes four days to train on 16 Cloud TPUs, and the recent GPT-3 (Brown et al., 2020) consumedorders of magnitude more petaflops / day to train compared to its predecessor, GPT-2 (Radford et al.,2019). Beyond training, deploying Transformer models to real world applications is also expensive,usually requiring extensive distillation (Hinton et al.

3 , 2015) or main efficiency bottleneck in Transformer models is its Self-Attention mechanism. Here, eachtoken s representation is updated by attending toallother tokens in the previous layer. This operationis key for retaining long-term information, giving Transformers the edge over recurrent models onlong sequences. However, attending to all tokens at each layer incurs a Complexity ofO(n2)withrespect to sequence length. Thus, in this paper, we seek to answer the question:can Transformermodels be optimized to avoid this quadratic operation, or is this operation required to maintainstrong performance?Prior work has proposed several techniques for improving the efficiency of Self-Attention . One populartechnique is introducing sparsity into attention layers (Child et al., 2019; Qiu et al., 2019; Beltagyet al., 2020) by having each token attend to only a subset of tokens in the whole sequence.

4 Thisreduces the overall Complexity of the attention mechanism toO(n n)(Child et al., 2019). However,as shown in Qiu et al. (2019), this approach suffers from a large performance drop with limitedefficiency gains, , a 2% drop with only 20% speed up. More recently, the Reformer (Kitaev et al.,2020) used locally-sensitive hashing (LSH) to reduce the Self-Attention Complexity toO(nlog(n)).However, in practice, the Reformer s efficiency gains only appear on sequences with length>2048(Figure 5 in Kitaev et al. (2020)). Furthermore, the Reformer s multi-round hashing approach actuallyincreasesthe number of sequential operations, which further undermines their final efficiency Under [ ] 14 Jun 2020In this work, we introduce a novel approach for tackling the Self-Attention bottleneck in approach is inspired by the key observation thatself-attention is low rank.

5 More precisely, weshow both theoretically and empirically that the stochastic matrix formed by Self-Attention can beapproximated by a low-rank matrix. Empowered by this observation, we introduce a novel mechanismthat reduces Self-Attention to anO(n)operation in both space- and time- Complexity : we decomposethe original scaled dot-product attention into multiple smaller attentions through Linear projections,such that the combination of these operations forms a low-rank factorization of the original summary of runtimes for various Transformer architectures, including ours, can be found in Table predominant application of Transformers, that has seen the most gains, is using them as pretrainedlanguage models, whereby models are first pretrained with a language modeling objective on a largecorpus, then finetuned on target tasks using supervised data (Devlin et al.)

6 , 2019; Liu et al., 2019;Lewis et al., 2019). Following Devlin et al. (2019), we pretrain our model on BookCorpus (Zhuet al., 2015) plus English Wikipedia using masked-language-modeling objective. We observe similarpretraining performance to the standard Transformer model. We then finetune our pretrained modelson three tasks from GLUE (Wang et al., 2018) and one sentiment analysis task, IMDB reviews (Maaset al., 2011). On these tasks, we find that our model performs comparably, or even slightly better, thanthe standard pretrained Transformer, while observing significant training and inference ArchitectureComplexity per LayerSequential OperationRecurrentO(n)O(n)Transformer, (Vaswani et al., 2017)O(n2)O(1)Sparse Tansformer, (Child et al., 2019)O(n n)O(1)Reformer, (Kitaev et al., 2020)O(nlog(n))O(log(n))LinformerO(n)O(1 )Table 1: Per-layer time Complexity and minimum number of sequential operations as a function ofsequence length (n) for various Backgrounds and Related Transformer and Self-AttentionThe Transformer is built upon the idea of Multi-Head Self-Attention (MHA), which allows the modelto jointly attend to information at different positions from different representation subspaces.

7 MHAis defined asMultiHead(Q,K,V) =Concat(head1,head2,..,headh)WO,(1)where Q,K,V Rn dmare input embedding matrices,nis sequence length,dmis the embeddingdimension, andhis the number of heads. Each head is defined as:headi=Attention(QWQi,KWKi,V WVi) =softmax[QWQi(KWKi)T dk] PV WVi,(2)whereWQi,WKi Rdm dk,WVi Rdm dv,WO Rhdv dmare learned matrices anddk,dvarethe hidden dimensions of the projection subspaces. For the rest of this paper, we will not differentiatebetweendkanddvand just Self-Attention defined in (2) refers to a context mapping matrixP Rn n. The TransformerusesPto capture the input context for a given token, based on a combination of all tokens in thesequence. However, computingPis expensive. It requires multiplying twon dmatrices, which isO(n2)in time and space Complexity . This quadratic dependency on the sequence length has becomea bottleneck for Related worksThere has been much prior literature on improving the efficiency of Transformers, especially theself-attention bottleneck.

8 The most common techniques for model efficiency that can be applied toTransformers (some specific to Transformers, others more general-purpose) include:2 Mixed Precision(Micikevicius et al., 2017): Using half-precision or mixed-precision representationsof floating points is popular in deep learning, and is also widely used in training Transformers (Ottet al., 2019). This technique can be further improved through Quantization Aware Training (Jacobet al., 2018; Fan et al., 2020), where the weights are quantized during training and the gradients areapproximated with the Straight-Through Estimator. This line of work is orthogonal to our approach,and we use mixed-precision training by Distillation(Hinton et al., 2015): Knowledge distillation aims to transfer the knowl-edge" from a large teacher model to a lightweight student model. The student model is then usedduring inference.

9 However this approach has drawbacks: It does not address speeding up theteachermodel during training, and moreover, student models usually suffer performance degradation com-pared to the teacher model. For example, when distilling a 12-layer BERT to a 6-layer BERT, thestudent model experiences an average performance drop on several benchmark tasks (Sanhet al., 2019).Sparse Attention(Child et al., 2019): This technique improves the efficiency of Self-Attention byadding sparsity in the context mapping matrixP. For example, the Sparse Transformer (Child et al.,2019) only computesPijaround the diagonal of matrixP(instead of the allPij). Meanwhile,blockwise Self-Attention (Qiu et al., 2019) dividesPinto multiple blocks and only computesPijwithin the selected blocks. However, these techniques also suffer a large performance degradation,while having only limited additional speed-up, , 2% drop with 20% speed Attention(Kitaev et al.)

10 , 2020): Locally-sensitive hashing (LSH) attention utilizes a multi-roundhashing scheme when computing dot-product attention, which in theory reduces the self-attentioncomplexity toO(nlog(n)). However, in practice, their Complexity term has a large constant1282and it is only more efficient than the vanilla transformer when sequence length is extremely Optimizer Efficiency: Microbatching (Huang et al., 2019) splits a batch into smallmicrobatches (which can be fit into memory), and then separately runs forward and backward passeson them with gradient accumulation. Gradient checkpointing (Chen et al., 2016) saves memoryby only caching activations of a subset of layers. The uncached activations are recomputed duringbackpropagation from the latest checkpoint. Both techniques trade off time for memory, and do notspeed up we ve noted, most common techniques have limitations in reducing both the training and inferencetime/memory consumption, we investigate how to optimize the Self-Attention layers and introduceour approach Self-Attention is Low RankIn this section, we demonstrate that the Self-Attention mechanism, , the context mapping matrixP,is first provide a spectrum analysis of the context mapping matrixP.