Eigen Attention: Attention in Low-Rank Space for KV Cache Compression (2024)

Abstract: Large language models (LLMs) represent a groundbreaking advancement in the domain of natural language processing due to their impressive reasoning abilities. Recently, there has been considerable interest in increasing the context lengths for these models to enhance their applicability to complex tasks. However, at long context lengths and large batch sizes, the key-value (KV) cache, which stores the attention keys and values, emerges as the new bottleneck in memory usage during inference. To address this, we propose Eigen Attention, which performs the attention operation in a low-rank space, thereby reducing the KV cache memory overhead. Our proposed approach is orthogonal to existing KV cache compression techniques and can be used synergistically with them. Through extensive experiments over OPT, MPT, and Llama model families, we demonstrate that Eigen Attention results in up to 40% reduction in KV cache sizes and up to 60% reduction in attention operation latency with minimal drop in performance.

## introduction the recent boom in artificial intelligence applications and their widespread public adoption can be attributed to the human-like capabilities of large language models (llms). llms have demonstrated remarkable performance across a wide range of natural language processing (nlp) tasks @xcite . the maximum number of tokens these models can process simultaneously to understand and generate text is referred to as the context/sequence length, and this closely determines their performance limits. hence, there has been considerable interest in increasing the context length to enhance their capabilities @xcite @xcite . longer context lengths open up new possibilities, such as summarizing lengthy documents, retrieving information to answer questions about extensive texts, and analyzing code. to make applications enabled by llms more accessible, developing techniques to serve these models efficiently is crucial. one standard technique to accelerate llm inference on gpus is caching the intermediate attention keys and values through a kv cache to avoid expensive re-computation for every generated token. however, it is observed that at long context lengths, the kv cache becomes the new memory and latency bottleneck @xcite . furthermore, while batching multiple requests together amortizes the weight access cost of llms, it exacerbates the kv cache memory overhead. consider the weight and kv cache memory footprint for the llama-2 7b model with a batch size of 16 and a context length of 32k tokens. here, the weight memory occupies 14 gb, while the kv cache requires a significantly higher memory of 256 gb at 16-bit precision. in essence, increasing the context/sequence length of llms has transformed llm inference into a memory-bound problem. the entire kv cache must be bought on-chip from the off-chip gpu memory for each newly generated token while the computation core stays idle, waiting for data. existing methods to address the kv cache bottleneck can be broadly classified into four distinct categories. first, some approaches focus on reducing the number of kv cache heads in the multihead attention block through grouped query and multi-query attention @xcite . second, a few methods aim to alleviate the memory overhead by utilizing a low-precision quantized kv cache @xcite . third, certain strategies involve evicting kv cache values associated with unimportant tokens, thereby caching only the keys and values of important tokens determined through some metrics @xcite . finally, some approaches tackle the issue from a systems perspective by utilizing the cpu and disk memory for the kv cache or by extending virtual memory and incorporating paging techniques into the attention mechanism @xcite . in contrast to the above-mentioned techniques, we propose eigen attention, a strategy to alleviate the kv cache overhead through low-rank approximation. eigen attention leverages the observation that attention inputs in llms (i.e., key, query, and value) can be reasonably approximated using a few principal basis vectors or eigenvectors @xcite . expressing keys and values as a linear combination of these principal vectors essentially reduces their dimension, leading to a lower memory footprint of the kv cache. to achieve this, we use a very small subset of training data as a calibration dataset to generate a set of query, key, and value matrices for the trained model. subsequently, we obtain the basis vectors through singular value decomposition (svd) on these matrices and choose the most important directions through a pre-defined error threshold. eigen attention is a post-training technique that can be applied without requiring any additional fine-tuning. moreover, our approach is orthogonal to existing techniques for kv cache compression and can be used in conjunction with them. figure 1 illustrates the differences between standard attention and our proposed eigen attention. columns of the u k matrix are the principal basis vectors for keys and queries. similarly, u v contains the important basis vectors for the value matrix. we project the attention inputs into a low-rank space defined by the eigenvectors u k and u v by assimilating these projections into the corresponding weight matrices w q , w k , w v and w o offline. during inference, the low-dimensional w k u k and w v u v are multiplied with the input vector x to generate lower-dimensional k and v matrices. while the primary goal of eigen attention is to reduce the memory overhead of the kv cache, the proposed low-rank approximation also leads to a computeefficient implementation of the attention block (section 5.2). for the standard dimension d and the compressed dimension r(≪ d), ## background multi-head attention. a typical llm consists of l decoder layers, each with two components: multi-head attention (mha) and the fully connected feed-forward network (ffn). for an input token embedding @xmath0, the mha block performs attention operations in parallel across h heads. here, n is the sequence length, and d is the hidden dimensionality of the model. for each attention head @xmath1, 2, ...h}, x is transformed into key, query, and value matrices as follows: the attention operation a at each head i is computed, and the results from all heads are concatenated to obtain the final output of the mha block, as shown in equation 2 . llm inference. the inference consists of the prefill and generation phases. in the prefill phase, the keys k i and v i values are computed for an input token embedding @xmath2. the total size of kv cache (in bits) can be derived by 2 * b * n * d * h * l * p, where l corresponds to the number of decoder layers in the llm, and p corresponds to the precision of cached vectors. in the generation phase, the model uses and updates the kv cache to generate the output autoregressively, one token at a time. note that this phase is memory-bound. for an incoming token embedding @xmath3, the key k i and value v i computed through equation 1 are appended to the kv cache: the query q i obtained through equation 1 is used to compute the attention output for each head: then, similar to equation 2, the heads' outputs are concatenated and multiplied with w o to produce the final output of the attention block. ## methodology this section describes eigen attention, which achieves kv cache compression by performing the attention operation in a low-rank space determined by a few basis vectors. we first discuss generating these principal basis vectors and then describe our proposed approach detailed in algorithm 1. ## conclusion in this work, we propose eigen attention, a novel technique that reduces the memory overhead associated with kv cache in llms. eigen attention is inspired by the observation that keys, queries, and values can be approximated using a few basis vectors, enabling the possibility of performing the attention operation in a low-rank space with minimal performance loss. to achieve this, we project the attention inputs into low-rank subspaces defined by a set of principal basis vectors computed offline using a calibration dataset in a one-shot manner. these projections are integrated into the weight matrices, which are utilized during inference to generate lower dimensional key and value matrices, thereby reducing the kv cache memory footprint. our approach is orthogonal to existing kv cache compression strategies and can be used in synergy. extensive experiments across a range of llms and language tasks demonstrate that eigen attention reduces the kv cache memory footprint by 40% with minimal performance loss and achieves up to a 60% improvement in attention operation latency compared to the standard attention baseline.