Linearizing Transformer with Key-Value Memory (2022)

Abstract: Efficient transformer variants with linear time complexity have been developed to mitigate the quadratic computational overhead of the vanilla transformer. Among them are low-rank projection methods such as Linformer and kernel-based Transformers. Despite their unique merits, they usually suffer from a performance drop comparing with the vanilla transformer on many sequence generation tasks, and often fail to obtain computation gain when the generation is short. We propose Memsizer, an approach towards closing the performance gap while improving the efficiency even with short generation. It projects the source sequences into lower dimension representations like Linformer, while enjoying efficient recurrent-style incremental computation similar to kernel-based transformers. This yields linear computation time and constant memory complexity at inference time. Memsizer also employs a lightweight multi-head mechanism which renders the computation as light as a single-head model. We demonstrate that Memsizer provides an improved balance between efficiency and accuracy over the vanilla transformer and other efficient transformer variants in three typical sequence generation tasks, including machine translation, abstractive text summarization, and language modeling.

## introduction transformer @xcite has become the de facto standard for almost all nlp tasks across the board. at the core of the vanilla transformer is the attention mechanism that captures the interactions between feature vectors at different positions in a sequence. despite its great success, the vanilla transformer models are typically computationally expensive as the computation of the * equal contribution. a number of transformer variants have been proposed to reduce the computational overhead @xcite . one family of methods leverages low-rank projections to reduce the number of pair-wise interactions (i.e., the size of attention matrices) @xcite @xcite . these methods first project the input sequence into a low-resolution representation. for example, @xcite project the length dimension to a fixed feature dimension. nevertheless, these methods have difficulties modeling variable-length sequences and autoregressive (causal) attention, impeding their applications in sequence generation tasks. recent works propose to approximate the softmax attention through kernelization @xcite @xcite . for sequence generation tasks, these works can cache computation in a recurrent manner, leading to constant memory complexity in sequence length during inference. despite the improved efficiency in long-form generation, the computation gain of these kernel-based approaches vanishes when the generation is as short as a typical sentence length. additionally, they usually suffer from a performance loss when training from scratch @xcite . in this work, we propose an approach called memsizer, an efficient transformer variant which follows the paradigm of low-rank projections while enjoying memory-efficient recurrent-style generation as the kernel-based transformers. concretely, we develop a key-value memory layer @xcite to substitute the multi-head attention layer in the vanilla transformer. we pack the information in the source sequence into a fixed-sized set of memory values in a length-dynamic manner, and use input-independent parametric matrices as the memory keys. in this way, we emphasize more on modeling the values and significantly simplify the design of keys. this unbalanced design of keys and values further enables us to suppress the multi-head computation to be as fast as with single head. memsizer is conceptually simple yet can handle variable-length sequences and causal attention for generation thanks to the length-dynamic projection. with the unbalanced memory layer and dynamic projection, memsizer enjoys linear time complexity and constant memory complexity. our experiments in three typical sequence generation tasks (machine translation, abstractive text summarization, and language modeling) show that the proposed method achieves comparable or better performance to state-of-the-art linear recurrent transformers, with more substantial reductions in inference latency, memory consumption, and model size. the advantages are more prominent with longer input lengths. in some tasks, the proposed memsizer can maintain or even improve the performance of the vanilla transformer, offering an appealing alternative for sequence generation tasks. ## memsizer: a different perspective of attention mechanism as discussed in section 2.2, the sa in the vanilla transformer can be perceived as an instantiation of the key-value memory network in section 2.1, where the memory key k and value v are pointwise projections of the source x s . in this work, we replace the sa module with a different memory mechanism which achieves recurrent inference computation thus linear complexity. our memory mechanism comes with a different specification of query, key and value in sa. specifically, following eq. (1-2), we specify the key-value memory layer as ## experiments we present extensive experiments on three typical sequence generation tasks in nlp, including machine translation, abstractive text summarization, and language modeling. ## conclusion we present memsizer, a method that leverages a novel key-value memory network specification to accelerate the original self-attention module. mem-sizer compresses source information to a set of global memory entries and uses an unbalanced keyvalue mechanism which further leads to lightweight multi-head computation. memsizer advances recent recurrent transformers with kernel approximation with lower time, memory, and storage cost during generation. our experiments in three standard generation tasks demonstrate that our model achieves an improved balance between efficiency and accuracy. the proposed method can be stacked with other computation reduction techniques to further advance the efficiency of transformers.