Sketching as a Tool for Understanding and Accelerating Self-attention for Long Sequences

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5187 - 5199 July 10-15, 2022 ©2022 Association for Computational Linguistics Sketching as a Tool for Understanding and Accelerating Self-attention for Long Sequences Yifan Chen1∗, Qi Zeng1∗, Dilek Hakkani-Tur2, Di Jin2, Heng Ji1, Yun Yang1 1University of Illinois Urbana-Champaign 2Amazon Alexa AI 1{yifanc10, qizeng2, hengji, yy84}@illinois.edu 2{hakkanit, djinamzn}@amazon.com Abstract Transformer-based models are not efficient in processing long sequences due to the quadratic space and time complexity of the self-attention modules. To address this limitation, Linformer and Informer reduce the quadratic complex- ity to linear (modulo logarithmic factors) via low-dimensional projection and row selection, respectively. These two models are intrinsically connected, and to understand their connection we introduce a theoretical framework of ma- trix sketching. Based on the theoretical analy- sis, we propose Skeinformer to accelerate self- attention and further improve the accuracy of matrix approximation to self-attention with col- umn sampling, adaptive row normalization and pilot sampling reutilization. Experiments on the Long Range Arena benchmark demonstrate that our methods outperform alternatives with a consistently smaller time/space footprint1. 1 Introduction Transformer (Vaswani et al., 2017) utilizes softmax self-attention modules to capture the dependency between tokens in a sequence and has been widely used in various Natural Language Processing tasks. The time and space complexity of the dot-product self-attention is quadratic in the input sequence length, which restricts the largest sequence length and batch size. To adapt transformers to long se- quences, documents have to be truncated, chunked using a sliding window, or processed in parallel on multiple GPUs. These additional operations usu- ally cause the loss of long-range dependency and introduce additional computational costs. In this paper, we focus on efficient self-attention methods (Xiong et al., 2021; Qiu et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020; Kitaev ∗Equal contribution. The majority of this work was done prior to the first author’s internship at Amazon Alexa AI. 1Our code is released at https://github.com/ pkuzengqi/Skeinformer et al., 2020a; Roy et al., 2021), among which Lin- former (Wang et al., 2020b) and Informer (Zhou et al., 2020) are two representative approaches to reducing the O(n2) self-attention to an � O(n) oper- ation ( � O(·) means O(·) modulo poly-log terms and n is the sequence length) in both space and time complexity. Linformer forms a low-rank factoriza- tion of the original attention by decomposing it into smaller attentions, while Informer allows each key to only attend to a certain number of queries. To better understand self-attention, we intro- duce a theoretical framework, sketching (Woodruff, 2014), to help explain the key ideas in Informer and Linformer from the perspective of matrix approxi- mation. Specifically, sketching methods replace the original matrix B with its random sketch BS to reduce computations. In Section 3.3 we introduce some concrete instances of commonly used distri- butions for constructing the random sketching ma- trix S. Furthermore, taking matrix approximation as a general guideline, we recognize the deficiency …

Content truncated to 500 words