Memory-efficient Transformers via Top-k Attention (2021)

Abstract: Following the success of dot-product attention in Transformers, numerous approximations have been recently proposed to address its quadratic complexity with respect to the input length. While these variants are memory and compute efficient, it is not possible to directly use them with popular pre-trained language models trained using vanilla attention, without an expensive corrective pre-training stage. In this work, we propose a simple yet highly accurate approximation for vanilla attention. We process the queries in chunks, and for each query, compute the top-*k* scores with respect to the keys. Our approach offers several advantages: (a) its memory usage is linear in the input size, similar to linear attention variants, such as Performer and RFA (b) it is a drop-in replacement for vanilla attention that does not require any corrective pre-training, and (c) it can also lead to significant memory savings in the feed-forward layers after casting them into the familiar query-key-value framework. We evaluate the quality of top-*k* approximation for multi-head attention layers on the Long Range Arena Benchmark, and for feed-forward layers of T5 and UnifiedQA on multiple QA datasets. We show our approach leads to accuracy that is nearly-identical to vanilla attention in multiple setups including training from scratch, fine-tuning, and zero-shot inference.

## introduction the transformer architecture @xcite has been successful in a wide range of natural language processing tasks, including machine translation @xcite , language modeling @xcite , question-answering @xcite , and many more. transformers 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 @xcite @xcite @xcite . the contextualizing component of the transformer is the attention layer where all positions in an input sequence of length l aggregate information from the entire sequence in parallel. at its core, given l query, key, and value vectors q, k, v respectively, the dot-product attention function outputs softmax(qk )v where the softmax function is applied row-wise on the matrix @xmath0, leading to an expensive ω(l 2 ) memory requirement. to alleviate this, past work proposed approximation methods for the computation of softmax(qk ). one major line of research focused on sparse attention variants, where only a few similarity scores are computed per query, and the rest are ignored. methods differ by which query-key pairs are selected @xcite @xcite @xcite @xcite . a second line of research explored dense variants @xcite @xcite (cf. @xcite for a survey). for example, instead of computing the attention scores exactly for only a small number of querykey pairs, @xcite compute an approximation of scores for all pairs. in this work, we adopt the sparse attention approach, but rather than approximating the k most similar key vectors per query vector, we compute this quantity exactly. specifically, we propose topk attention where, for each query vector, we only keep its k largest similarity scores with respect to the l keys, where k l. we show that top-k attention can be implemented in a memory-efficient manner by (a) chunking the query vectors when 2k 4k 6k 8k 10k 0 10 20 30 memory (gib) 0k 10k 20k 30k 40k 50k 60k @xcite . this allows us, e.g., to train a typical 12-layer transformer decoder over 32k-long inputs on a 30gib gpu (figure 3a ). we extensively evaluate top-k attention on a wide range of tasks and demonstrate its mentioned advantages. training from scratch, we show top-k attention performs as well as vanilla self-attention on long range arena, a benchmark dedicated to evaluating the ability of transformers to handle long sequences, and in a language modeling task (wikitext-103). second, we show top-k attention can be used as a drop-in replacement for vanilla attention at inference time without any additional training at the feed-forward layer of the unifiedqa model @xcite on 12 different question answering (qa) datasets, reducing the number of keys used per query by more than 99%. last, we show top-k attention obtains similar performance to vanilla attention on a wide range of qa tasks when fine-tuning t5 @xcite , without the need for any corrective pre-training. overall, our results demonstrate that top-k attention is a simple and effective method for dramatically reducing the memory footprint of transformers without loss of accuracy that can allow resource-constrained researchers enjoy the benefits of large pre-trained transformer-based models. our code is available at https://github. com/ag1988/top_k_attention . ## efficient transformer through top-k attention in this section, we briefly review the transformer architecture, its sparse approximations, and show how to cast the feed-forward layer into the querykey-value framework ( §2.1). we then describe topk attention and our memory-efficient implementation for it ( §2.2). ## benchmarking in this section, we benchmark top-k attention in terms of time and memory, and compare it to vanilla attention, query-chunking without the top-k operation, and to performer @xcite , as a representative of state-of-the-art linear attention variants. we separately benchmark (a) a single self-attention layer over long sequences, (b) a single feed-forward layer with a large feedforward dimension, and (c) a 12-layer transformer decoder with same architecture as bert-base @xcite . experimental details for all models, we benchmark by running a forward and backward pass over random inputs. each measurement is an average over 3 runs on an nvidia a100 gpu and is discarded if memory usage exceeds 30gib. we use causal masking for self-attention layers to highlight the simplicity of our approach that can seamlessly handle arbitrary attention masks, unlike other methods @xcite @xcite , where implementing causal masking requires customized cuda implementations. for performer, we use 256 random features, and the cuda implementation from @xcite . ## experimental evaluation of top-k attention having established top-k attention as a memory efficient alternative to vanilla attention, we now show that, even for small values of k, top-k attention provides a high-quality approximation of vanilla attention, both at the multi-head attention and feedforward layers. we empirically show this in a wide range of setups including (a) training from scratch on tasks that require handling long-range dependencies ( §4.1) and on language modeling ( §4.2), (b) fine-tuning pre-trained language models (t5) on multiple qa datasets ( §4.5), and (c) performing zero-shot inference using pre-trained language models (unifiedqa) without any training ( §4.3).