ABC: Attention with Bounded-memory Control

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 7469 - 7483 May 22-27, 2022 c⃝2022 Association for Computational Linguistics ABC: Attention with Bounded-Memory Control Hao Peng♠ Jungo Kasai♠ Nikolaos Pappas⋆∗ Dani Yogatama♣ Zhaofeng Wu♢∗ Lingpeng Kong♦ Roy Schwartz♡ Noah A. Smith♠♢ ♠Paul G. Allen School of Computer Science & Engineering, University of Washington ⋆Amazon Web Services ♣DeepMind ♢Allen Institute for Artificial Intelligence ♡School of Computer Science & Engineering, Hebrew University of Jerusalem ♦Department of Computer Science, The University of Hong Kong {hapeng,jkasai,npappas,zfw7,nasmith}@cs.washington.edu dyogatama@deepmind.com, lpk@cs.hku.hk roy.schwartz1@mail.huji.ac.il Abstract Transformer architectures have achieved state- of-the-art results on a variety of natural lan- guage processing (NLP) tasks. However, their attention mechanism comes with a quadratic complexity in sequence lengths, making the computational overhead prohibitive, especially for long sequences. Attention context can be seen as a random-access memory with each to- ken taking a slot. Under this perspective, the memory size grows linearly with the sequence length, and so does the overhead of reading from it. One way to improve the efficiency is to bound the memory size. We show that dis- parate approaches can be subsumed into one abstraction, attention with bounded-memory control (ABC), and they vary in their organi- zation of the memory. ABC reveals new, unex- plored possibilities. First, it connects several efficient attention variants that would otherwise seem distinct. Second, this abstraction gives new insights—an established approach (Wang et al., 2020b) previously thought to not be appli- cable in causal attention, actually is. Last, we present a new instance of ABC, which draws in- spiration from existing ABC approaches, but re- places their heuristic memory-organizing func- tions with a learned, contextualized one. Our experiments on language modeling, machine translation, and masked language model fine- tuning show that our approach outperforms pre- vious efficient attention models; compared to strong transformer baselines, it significantly im- proves the inference time and space efficiency with no or negligible accuracy loss. 1 Introduction Transformer architectures are now central in natural language processing (Vaswani et al., 2017). They rely on the attention mechanism (Bahdanau et al., 2015) to contextualize the input. The context can be seen as a random access memory whose size lin- early grows with the sequence length; each query ∗This work was done while Zhaofeng Wu and Nikolaos Pappas were at the University of Washington. reads from it using a softmax-normalized linear combination, with overhead linear in the memory size. This amounts to a quadratic complexity over- all, making transformers’ computational overhead prohibitive, especially for long sequences. One way to improve attention’s efficiency is to bound its memory size. Imposing a constant- sized constraint over the memory ensures that read- ing from it has constant time and space overhead, yielding a linear overall complexity in sequence lengths. This is in fact a common strategy adopted by several recent works. In this work, we show that some of these works are closely connected in ways that, to date, have gone unremarked. We propose attention with bounded-memory …

Content truncated to 500 words