The Devil in Linear Transformer (2022)

Abstract: Linear transformers aim to reduce the quadratic space-time complexity of vanilla transformers. However, they usually suffer from degraded performances on various tasks and corpus. In this paper, we examine existing kernel-based linear transformers and identify two key issues that lead to such performance gaps: 1) unbounded gradients in the attention computation adversely impact the convergence of linear transformer models; 2) attention dilution which trivially distributes attention scores over long sequences while neglecting neighbouring structures. To address these issues, we first identify that the scaling of attention matrices is the devil in unbounded gradients, which turns out unnecessary in linear attention as we show theoretically and empirically. To this end, we propose a new linear attention that replaces the scaling operation with a normalization to stabilize gradients. For the issue of attention dilution, we leverage a diagonal attention to confine attention to only neighbouring tokens in early layers. Benefiting from the stable gradients and improved attention, our new linear transformer model, transNormer, demonstrates superior performance on text classification and language modeling tasks, as well as on the challenging Long-Range Arena benchmark, surpassing vanilla transformer and existing linear variants by a clear margin while being significantly more space-time efficient. The code is available at https://github.com/OpenNLPLab/Transnormer .

## introduction transformer models show great performance on a wide range of natural language processing and computer vision tasks @xcite @xcite . one issue of the vanilla transformer model lies in its quadratic space-time complexity with respect ⋆ equal contribution. the corresponding author (email: zhongyiran@gmail.com). this work was done when weixuan sun and yiran zhong were in the sensetime research. to the input length. various prior works attempt to alleviate this inefficiency @xcite @xcite @xcite @xcite . in this work, we focus on a particular subset of these methods, known as kernel-based linear transformers @xcite @xcite @xcite considering their desirable linear space-time complexity. despite their space-time efficiency, linear transformers are not always in favor for practical adoption, largely due to the degraded performance than the vanilla model. to address this issue, we take a close look at existing kernel-based linear transformers and identify two deficiencies that lead to such a performance gap. unbounded gradients. most existing linear transformers inherit attention formulation from the vanilla transformer, which scales attention scores to ensure they are bounded within [0, 1]. however, we theoretically show that such a scaling strategy renders unbounded gradients for linear transformer models. as a result, the unbounded gradients em-pirically lead to unstable convergence as our preliminary experiments suggest. attention dilution. previous works @xcite @xcite @xcite suggest that in vanilla transformer, softmax attention maps tend to be local. in contrast, as shown in fig 2 , we observe that linear transformers often trivially distribute attention scores over the entire sequence even in early layers. due to this issue, which we refer as attention dilution, important local information is less well preserved in linear models, resulting in inferior performance. this negative impact of attention dilution is also evidenced by the performance drop in our controlled experiments if partly replacing vanilla attention in transformer layers with linear attention ones. to mitigate these issues, we propose a linear transformer model, called transnormer, which shows better performance than vanilla transformer on a wide range of task while being significantly faster during runtime, as shown in fig. 1 . to avoid the unbounded gradients, we introduce normattention, which gets rid of scaling over attention matrices while appending an additional normalization only after the attention layer. the choice of the normalization operator is unrestricted, for example, layernorm @xcite or rm-snorm @xcite both serve the purpose. we show empirical results demonstrating that with normattention, the gradients are more stable during training, which in turn leads to more consistent convergence. to alleviate the attention dilution issue, we modify the vanilla attention and allow each token to only attend to its neighbouring tokens, resulting in a diagonal attention. to mimic the behaviors on local semantics of the vanilla transformer, we employ the diagonal attention on early layers while using normattention for later ones. in this way, we encourage the model to capture both local and global language context. note that our diagonal attention can be efficiently computed such that the overall linear space-time complexity of transnormer is preserved. we perform extensive experiments on standard tasks, where transnormer demonstrates lower language modeling perplexities on wikitext-103 and overall higher text classification accuracy on glue than vanilla model and other competing methods. in addition, on the challenging long-range arena benchmark, transnormer also shows favorable results while being faster and more scalable with longer inputs during both training and inference time. ## background and related work we first briefly review vanilla transformer @xcite and its efficient variants. the key component of transformers is the self-attention, which operates on query q, key k and value v matrices; each of them is the image of a linear projection taking @xmath0: with n the input length, d the hidden dimension. the output @xmath1: in the kernel-based methods @xcite @xcite @xcite , a kernel function ϕ(•) maps queries and keys to their hidden representations. then the output of the linear attention can be rewritten as: choromanski et al. (2020) and katharopoulos et al. (2020) adopt activation function 1 + elu to process query and key. wang et al. (2020) assumes attention matrices are lowrank. peng et al. (2020) and zheng et al. (2022) approximate softmax under constrained theoretical bounds. qin et al. (2022) propose a linear alternative to the attention based on empirical properties of the softmax function. these methods focus on either approximating or altering the softmax operator while preserving its properties. compared with the vanilla transformer, these methods often trade performance for efficiency, usually resulting in worse task performance. in this paper, we argue that there are two essential reasons leading to such a performance gap, discussed in detail as follows. ## the devil in linear attention in this section, we motivate the design principles of transnormer by providing theoretical evidence for the unbounded gradients, and empirical results showing the adverse influence of attention dilution. ## method based on the aforementioned observations, we propose a new linear transformer network called transnormer that addresses the above two limitations of current linear transformers. the overall architecture is shown in fig. 3 . ## experiments in this section, we compare our method to other linear transformers and the vanilla transformer on autoregressive language modeling, bidirectional language modeling as well as the long range arena benchmark @xcite . we also provide an extensive ablation study to vindicate our choice in designing the transnormer. we validate our method on two variants of the transnormer. the transnormer t1 uses the rela attention @xcite in the dia-gattention and the elu as the activation function in the normattention. the transnormer t2 uses the softmax attention @xcite in the diagattention and the 1+elu as the activation function in the normattention. for experiments, we first study the autoregressive language modeling on wikitext-103 @xcite in section 5.2. then in section 5.2 we test our method on bidirectional language modeling, which is pre-trained on wikitext-103 @xcite and then fine-tuned on several downstream tasks from the glue benchmark @xcite . finally, we test transnormer on the long-range arena benchmark @xcite to evaluate its ability in modeling longrange dependencies and efficiencies in section 5.2. ## conclusion in this paper, we identified two key issues that cause the inferior performance of existing linear transformer models: 1) unbounded gradients; 2)