ABC : Attention with Bounded-memory Control (2022)

Abstract: Transformer architectures have achieved state- of-the-art results on a variety of natural language 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 token 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 disparate approaches can be subsumed into one abstraction, attention with bounded-memory control (ABC), and they vary in their organization of the memory. ABC reveals new, unexplored possibilities. First, it connects several efficient attention variants that would otherwise seem apart. Second, this abstraction gives new insights—an established approach (Wang et al., 2020b) previously thought to not be applicable in causal attention, actually is. Last, we present a new instance of ABC, which draws inspiration from existing ABC approaches, but replaces their heuristic memory-organizing functions with a learned, contextualized one. Our experiments on language modeling, machine translation, and masked language model finetuning show that our approach outperforms previous efficient attention models; compared to the strong transformer baselines, it significantly improves the inference time and space efficiency with no or negligible accuracy loss.

## introduction transformer architectures are now central in natural language processing @xcite . they rely on the attention mechanism @xcite to contextualize the input. the context can be seen as a random access memory whose size linearly grows with the sequence length; each query reads from it using a softmax-normalized linear combination, with overhead linear in the memory size. this amounts to a quadratic complexity overall, making transformers' computational overhead prohibitive, especially for long sequences. one way to improve attention's efficiency is to bound its memory size. imposing a constantsized constraint over the memory ensures that reading 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 control (abc), a unified abstraction over them. in abc, constant-sized memories are organized with various control strategies, e.g., induced from heuristic patterns @xcite @xcite , inter alia), locality assumptions @xcite , or positions @xcite . these strategies, by and large, are "contextagnostic." in response to this, we propose abc mlp , a particular instance of abc that learns a contextualized control strategy from data. specifically, abc mlp uses a neural network to determine how to store each token into the memory (if at all). compared to previous bounded-memory models, it strikes a better trade-off between accuracy and efficiency: controlling for the accuracy, abc mlp can get away with much smaller memory sizes. abc models (including abc mlp ) come with a linear complexity in sequence lengths, and admit recurrent computation graphs in causal attention (self-attention over the prefix). therefore they are appealing choices in a variety of applications, including text encoding, language modeling and text generation. this leads to a surprising finding. linformer @xcite , an established efficient attention method, was previously thought not to be applicable in causal attention or autoregressive decoding @xcite . through the abc view, we show that it actually is, and achieves competitive performance in our machine translation experiments. abc connects existing models that would otherwise seem distinct, reveals new insights into established methods, and inspires new efficient attention architectures. we explore its applications in transformers, as a drop-in substitute for the canonical softmax attention. abc offers a novel lens that can help future research in the analysis of transformers, where the theoretical insights are still catching up with empirical success. experiments on language modeling, machine translation, and masked language model finetuning show that our abc mlp model outperforms previous abc approaches in accuracy with a much smaller memory size. compared to the strong transformer baseline, abc mlp achieves a significant speedup and memory savings at inference time, with no or negligible accuracy loss. the efficiency improvements are more prominent for long sequences, suggesting that the asymptotic savings are even more appealing in applications involving long sequences. we release our code at https://github.com/noahs-ark/abc . ## an outer-product view of attention this section presents our outer-product memory perspective of attention, which allows for a smooth transition to later discussion. in attention, a sequence of queries {q i } n @xmath0, each storing a key and value pair: 1 query q reads from the memory using a softmax-normalized linear combination, producing a d-dimensional vector: this takes o(n ) time and space. when the attention with n queries can be parallelized (e.g., in text encoding), it takes linear time and quadratic space; when it cannot be (e.g., in decoding), it takes quadratic time and linear space. the memory can be equivalently represented as sums of vector outer products: the n -by-d matrix on the last line has its ith row being k ⊤ i and all others zeros; in this sense, k i is stored in the ith slot by e i , not affecting others. ## attention with bounded memory a straightforward way to improve attention's efficiency is to bound its memory size. our outerproduct view of attention provides a straightforward way to devise this, by replacing {e i } with control vectors that select n ≪ n vectors to attend to. we dub this approach attention with bounded-memory control (abc). concretely, let k, @xmath1, with n set a priori. we will discuss various ways to construct {ϕ i } in the subsequent sections. reading from the memory takes a constant o(n) time and space; therefore abc's overall complexity is o(n n), linear in the sequence length. 2 eq. 3 offers an equivalent recurrent computation, which is particularly useful in causal attention where only the prefix is looked at, in what follows, we study several existing efficient attention approaches and show that they are in fact instances of the abc abstraction. ## learned memory control the abc abstraction connects several existing approaches that would otherwise seem distinct. this inspires the design of new architectures. we hypothesize that learning a contextualized strategy can achieve better performance. this section introduces abc mlp . it parameterizes ϕ with a singlelayer multi-layer perceptron (mlp) that takes as input the token's representation x i , and determines which slots to write it into and how much. a key design choice of abc mlp is that its ϕ i depends only on current input x i . this helps (1) keep the recurrent computation efficient in practice @xcite , and (2) make it applicable in not only encoder self-attention and cross attention, but also causal attention. concurrently to this work, @xcite and @xcite also proposed methods to learn contextualized control. they compute ϕ i from previous layer's memory, revealing the full sequence to the control vectors. as a result, these two approaches are unsuitable for causal attention. 6 abc mlp , as other abc models, can be used as a drop-in replacement for the canonical softmax attention, and we apply its multihead variant in transformers. with proper parameter sharing, the number of additional parameters abc mlp incurs is small: inspired by @xcite , we tie ϕ-mlp's parameters across different layers, which adds less than 1% parameters to the models. ## experiments we evaluate abc models on language modeling ( §5.1), sentence-level and document-level machine translation ( §5.2), and masked language model finetuning ( §5.3). dataset statistics and implementation details are summarized in §c. ## analysis decoding efficiency over varying sequence lengths. abc's efficiency gains can be more prominent for long sequences. we study abc mlp 's decoding overhead with varying sequence lengths. following @xcite , we consider a sequence-to-sequence generation experiment. three linear-complexity models are compared: rfa (with 256/128 cross/causal memory sizes; @xcite , t2r (32/4; @xcite , and abc mlp (32/8). the sizes are chosen to maximize efficiency without accuracy drop. t2r needs to be finetuned from a pretrained transformer to match its performance, while others don't. all linear-time models achieve consistent decoding speed for different lengths (figure 1a ), substantially outpacing the softmax attention baseline, especially for long sequences. in particular, abc mlp decodes ∼1.25 times faster than rfa, another competitive model that can match transformer's accuracy without a warm start from a pretrained model. this can be attributed to the fact that abc mlp achieves similar accuracy with a much smaller memory. t2r's memory sizes are similar to abc mlp 's, but it decodes about 20% faster. abc mlp also achieves significant savings in terms of memory overhead (figure 1b ). abc mlp , rfa, and t2r's curves are similar. ## conclusion we presented attention with bounded-memory control (abc). it provides a unified perspective of several recently-proposed models, and shows that they vary in the organization of the bounded memory. abc reveals new insights into established methods and inspires new architectures. we proposed abc mlp , a particular instance of abc that learns a contextualized memory control. on language modeling, machine translation, and masked language model finetuning, abc mlp outperforms previous abc models. compared to the strong transformer baseline, abc mlp achieves substantial efficiency improvements with no or negligible accuracy loss. figure 1b compares abc mlp 's (32-8 memory sizes) attention memory overhead with softmax attention's. following @xcite , we consider a synthetic sequence-to-sequence generation task with varying sequence lengths. a batch size of 16 and greedy decoding is used. the models are of the same size as those in §5.2.