The Devil in Linear Transformer

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025 - 7041 December 7-11, 2022 ©2022 Association for Computational Linguistics The Devil in Linear Transformer ⋆Zhen Qin1, ⋆Xiaodong Han1, Weixuan Sun2,3, Dongxu Li2, Lingpeng Kong4,5, Nick Barnes2, �Yiran Zhong4 1SenseTime Research, 2Australian National University, 3OPPO Research Institute, 4Shanghai AI Laboratory, 5The University of Hong Kong https://github.com/OpenNLPLab/Transnormer 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 corpora. In this paper, we examine existing kernel-based linear transformers and identify two key is- sues that lead to such performance gaps: 1) unbounded gradients in the attention computa- tion adversely impact the convergence of lin- ear 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 gra- dients and improved attention, our new linear transformer model, TRANSNORMER, demon- strates superior performance on text classifica- tion and language modeling tasks, as well as on the challenging Long-Range Arena benchmark, surpassing vanilla transformer and existing lin- ear variants by a clear margin while being sig- nificantly more space-time efficient. The code is available at TRANSNORMER. 1 Introduction Transformer models show great performance on a wide range of natural language processing and computer vision tasks (Qin et al., 2022; Sun et al., 2022b; Cheng et al., 2022a,b; Zhou et al., 2022). 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. 2 4 6 8 10 12 14 16 18 Speed (steps per sec) 54 56 58 60 62 64 Long-Range Arena Score Transformer FLASH_quad FLASH LS Performer cosFormer Linformer Nystorm Reformer TransNormer T1 TransNormer T2 a Figure 1: TRANSNORMER has smaller memory foot- prints (circle sizes) and produces clearly favorable speed (x-axis) and overall scores (y-axis), when evaluated on the challenging Long-Range Arena benchmark than the vanilla transformer and other competing methods. to the input length. Various prior works attempt to alleviate this inefficiency (Zaheer et al., 2020; Beltagy et al., 2020; Tay et al., 2020a; Kitaev et al., 2020; Child et al., 2019; Liu et al., 2022; Sun et al., 2022b). In this work, we focus on a particular sub- set of these methods, known as kernel-based lin- ear transformers (Choromanski et al., 2020; Wang et al., 2020; Katharopoulos et al., 2020; Peng et al., 2020; Qin et al., 2022) considering their desirable linear …

Content truncated to 500 words