Value-aware Approximate Attention

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 9567–9574 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 9567 Value-aware Approximate Attention Ankit Gupta Tel Aviv University ankitgupta.iitkanpur@gmail.com Jonathan Berant Tel Aviv University, Allen Institute for AI joberant@cs.tau.ac.il Abstract Following the success of dot-product atten- tion in Transformers, numerous approxima- tions have been recently proposed to address its quadratic complexity with respect to the input length. However, all approximations thus far have ignored the contribution of the value vectors to the quality of approximation. In this work, we argue that research efforts should be directed towards approximating the true output of the attention sub-layer, which in- cludes the value vectors. We propose a value- aware objective, and show theoretically and empirically that an optimal approximation of a value-aware objective substantially outper- forms an optimal approximation that ignores values, in the context of language modeling. Moreover, we show that the choice of kernel function for computing attention similarity can substantially affect the quality of sparse ap- proximations, where kernel functions that are less skewed are more affected by the value vec- tors. 1 Introduction The Transformer architecture (Vaswani et al., 2017) has been widely successful in a wide range of natural language processing tasks, including ma- chine translation (Edunov et al., 2018), language modeling (Roy et al., 2020), question-answering (Karpukhin et al., 2020), and many more. Trans- formers pre-trained on large amounts of text with a language modeling (LM) objective have become the standard in NLP, exhibiting surprising amounts of linguistic and world knowledge (Peters et al., 2018; Devlin et al., 2019; Petroni et al., 2019; He- witt and Manning, 2019; Roberts et al., 2020). The contextualizing component of the Trans- former is the attention layer where all positions in an input sequence of length L aggregate in- formation from the entire sequence in parallel. At its core, given L query, key and value vec- tors, the dot-product attention function outputs1 softmax(QK⊤)V where the softmax function is applied row-wise on the matrix QK⊤ ∈ RL×L, consisting of similarity scores of the query-key pairs. Unfortunately, computing Ω(L · L) similar- ity scores is prohibitive for long sequences. To alleviate this, past work proposed to com- pute an approximation of softmax(QK⊤). One major line of research focused on sparse attention variants, where only a few similarity scores are computed per position, and the rest are ignored. Methods differ by which query-key pairs are se- lected (Child et al., 2019; Ye et al., 2019; Qiu et al., 2019; Roy et al., 2020; Kitaev et al., 2020; Beltagy et al., 2020; Gupta and Berant, 2020; Vyas et al., 2020). A second line of research explored dense variants (Katharopoulos et al., 2020; Wang et al., 2020; Bello, 2021; Tay et al., 2020a) (cf. (Tay et al., 2020b) for a survey). E.g., instead of computing the attention scores exactly for only a few query- key pairs, (Choromanski et al., 2020) compute an approximation of scores for all pairs. In this work, we point …

Content truncated to 500 words