FINCH : Prompt-guided Key-Value Cache Compression for Large Language Models (2024)

Abstract: Recent large language model applications, such as Retrieval-Augmented Generation and chatbots, have led to an increased need to process longer input contexts. However, this requirement is hampered by inherent limitations. Architecturally, models are constrained by a context window defined during training. Additionally, processing extensive texts requires substantial GPU memory. We propose a novel approach, Finch, to compress the input context by leveraging the pre-trained model weights of the self-attention. Given a prompt and a long text, Finch iteratively identifies the most relevant Key (K) and Value (V) pairs over chunks of the text conditioned on the prompt. Only such pairs are stored in the KV cache, which, within the space constrained by the context window, ultimately contains a compressed version of the long text. Our proposal enables models to consume large inputs even with high compression (up to 93x) while preserving semantic integrity without the need for fine-tuning.

## introduction large language models (llms), built upon the transformer architecture, have delivered breakthroughs in numerous applications. with their generalization and reasoning capabilities, models such as chatgpt have revolutionized fields where extensive input prompts are necessary for generating precise responses, such as retrieval-augmented generation, chain-of-thought, conversational chatbots, and in-context learning @xcite @xcite . however, the use of llms in production is limited by their increasing requests in terms of gpu memory @xcite . first, as the computational complexity grows along with the size of the models, their memory consumption increases. second, this issue becomes more pronounced when llms process larger inputs, as demanded by their ever-increasing context size. third, the key-value (kv) cache mechanism, typically employed by llms to speed up the generation process, prioritizes efficiency by retaining and reusing previously computed kv vectors during attention computation, bypassing re-calculations at each token generation step @xcite . nevertheless, this solution comes with the trade-off of increased memory consumption. 1 to offer more efficient solutions to operate these models, it has been proposed to compress input prompts, exploiting the redundancy in natural language @xcite . by preserving critical token information while compressing less crucial details, these models reduce the context in a compact description, without noticeably degrading the functional accuracy @xcite . compression also enables the llms to process large inputs that do not fit the model's context size. however, most of these models require a training/fine-tuning process or a large number of calls to an external model for the compression @xcite . we revisit the llms' generative inference mechanism to deal with the memory constraint problem and the limitations of current solutions in processing large inputs. we propose a novel approach targeting the reduction of the kv cache memory footprint while avoiding resource-intensive retraining or fine-tuning processes. drawing insights from the patterns inherent in attention modules, and guided by the understanding that not all attention modules engage with every token, our solution compresses the cached vectors, leading to a reduction in memory usage and efficient text generation. our approach, termed finch, 2 facilitates faster generative inference through adaptive kv cache at each step in the prefill stage, finch sequentially consumes a document chunk (two dashed border squares, blue background), alongside the input prompt (one dashed border square, yellow background) as depicted at the bottom. at each step, it processes the key, value pairs (solid squares in the transformer) and identifies the most relevant to the prompt. it then carries them to the cache processing the next chunk (where they appear with a violet background). in the generation stage, the model synthesizes a response that is informed by the compressed cached information from the entire document. the white square is the space reserved for producing output tokens in the generation stage. finch incrementally feeds the kv cache with the compressed context without any learning or external summarization module; it can be used in a plug-and-play fashion with any decoder-based model. the compression rate is specified by setting the target size of the kv cache as an input parameter constrained by the model context size. even with high compression ratios, our method ensures that the correctness of the model response is preserved. we test finch on two popular benchmarks covering tasks in question answering, summarization, code completion, synthetic tasks and few-shot learning. compared against the original llm (without compression) over the squad v2 benchmark @xcite , finch achieves comparable generation quality at 2.35x compression and 90% of the reference accuracy score at 3.76x compression, while being faster in terms of end-to-end execution times in most cases. when compared to the state-of-the-art compression method longllmlingua @xcite finch reports the best quality scores in most of the tasks in longbench @xcite , both with llama 2 and mistral @xcite . our method achieves a compression range of 2x to 93x across various tasks, consistently outperforming a truncation baseline in most experiments. remarkably, finch even surpasses the performance of the llms operating with the full, uncompressed context in certain cases. finally, in question answering tasks, we also include a rag baseline, and our method outperforms it in 10 out of 12 experiments. ## related work we position our work w.r.t. two main topics. first, we discuss strategies for improving computational efficiency, i.e., making llms accessible for real-time applications or use on devices with limited resources. second, we focus on attention patterns in llms, as our work shows that those contribute significantly towards optimizing the models to handle larger inputs in a limited context size. efficiency improvements in llms. methods targeting the reduction of inference and fine-tuning costs include model modification, such as quantization @xcite and model compression @xcite . other efforts enhance model efficiency for llms by eliminating redundant input words based on attention scores @xcite and compressing the input sequence by augmenting the encoding modules with pooling layers @xcite . proposed solutions also involve learning to skip layers in the transformer architecture @xcite or to select the most critical tokens for performance @xcite . other approaches pursue prompt compression, either by limiting the number of tokens that are processed in inference by learning special ''compressed'' tokens @xcite @xcite or by pruning and merging tokens @xcite , e.g., learning thresholds for pruning unimportant ones @xcite . however, some of these strategies require an additional re-training or fine-tuning phase and others have been designed for encoder models and are not well suited for autoregressive llms such as @xcite . in contrast with such solutions, our approach condenses auto-regressive llms input contexts during the prefill stage by using the caching mechanism without model re-training and even faster inference. finally, recent methods focus on optimizing the generation stage to improve efficiency @xcite @xcite @xcite . we leave to future work the study of how to use our prompt-guided token selection strategy in such approaches. the role of attention. our work relies on self-attention to make the most relevant information in a context available in a concise manner. the development of transformer models provoked studies to unravel the underlying mechanisms of self-attention, e.g., heads prominently pay attention to separator and adjacent tokens @xcite . our solution capitalizes on the attention mechanism structure to heighten inference efficiency by exploring the kv cache for the most important key, value pairs w.r.t. the given prompt. related work evaluates the informativeness of lexical units using a language model and drops less informative content for compression @xcite @xcite , for example, by regarding tokens with lower perplexity as more influential in the inference process. these techniques view llms as a compressor for world knowledge and work by further compressing information within prompts @xcite . in contrast with these solutions, our approach instead optimizes the management of the kv cache during the prefill stage without requiring a separate llm. other approaches look at how to select the most important tokens in the prefill stage, but, differently from our method that dynamically identifies the most important tokens, they rely on manually defined policies for token selection @xcite . finally, we focus on a plug-and-play solution for existing models, with an emphasis on limited computing resources. this is in contrast with other solutions that demand more devices to handle a very large input context @xcite . ## background self-attention is foundational in transformer models @xcite , enabling language understanding and generation capabilities. transformers learn the contextual relationships between words or subwords within a sentence. central to this mechanism are three types of vectors-queries (q), keys (k), and values (v)-that are learned from the input embeddings. the self-attention mechanism computes the dot product of the query with all keys to determine their similarity. a softmax function normalizes these scores, creating a distribution that determines how much attention to allocate to each token. the output is a weighted sum of the values. in several nlp tasks, transformers generate a response sequence from a given context/document and a user prompt. consider a sequence of tokens representing the context x @xmath0, which may also include instructions, the goal is to enable the model to generate a response sequence @xmath1. this process can be divided into two stages. prefill stage. as a first step, both the context and the prompt sequence are concatenated to form the input sequence @xmath2, where: this sequence is then embedded into an embedding matrix @xmath3, where d denotes the embedding dimension and processed through multiple layers of multi-head self-attention and feed-forward networks, which operate in parallel across the sequence length and attention heads. each attention layer calculates and stores the corresponding key and value @xmath4, @xmath5. in the transformer architectures, the k and v matrices encapsulate historical token information. unlike other components of the transformer (e.g., feedforward layer or layer norm), which process current inputs independently of their past tokens, the k and v matrices are the only matrices that retain information from previously encountered tokens. caching the key and value matrices for every layer for the context eliminates the necessity to recompute them for each new token generated. ## problem formulation as discussed, k and v are the only matrices that retain information about previous tokens. we can therefore formulate the problem of compression as reducing the size of these two matrices during the prefill stage and before the actual answer generation takes place. specifically, we have to find k and ṽ where k, ṽ ∈ r k×d such that two properties are satisfied: our approach aims at compressing contexts into a manageable form for llms, particularly when faced with extensive documents and the need to maintain computational efficiency. our methodology is motivated by the following observation: the softmax of self-attention distributes attention across all elements to varying degrees, effectively capturing a spectrum of contextual relationships in the data. we hypothesize that the ''smooth'' distribution of attention may include superfluous information for the given prompt at hand. ## experimental setup we evaluate finch using a variety of datasets and nlp tasks, with a focus on its application to the llama 2 7b-chat @xcite and the mistral 7b-instruct-v0.2 @xcite models. experiments are conducted with 4-bit normalfloat quantization and double quantization @xcite . method compl. per layer seq. ops initial cache size cache growth/op. vanilla o(cd) o(a) @xmath6: for an assessment of finch's ability to preserve quality when compressing according to equation 1, we use short texts that let us run the entire document as input. we use squad v2 @xcite , a benchmark which includes both questions that can and cannot be answered with the given documents. we measure how our model maintains or improves its accuracy, despite having reduced context, against two baselines. first, we report for vanilla, the standard model configuration which has access to the full context. second, a truncate strategy that reduces the input to the same size used by finch. given a budget, we truncate the input after a number of tokens equal to half the reduced context both from the start and from the end, i.e., we take the beginning and the end of the document. longbench: to assess the robustness of our method with long documents and a variety of tasks, we also evaluate on the longbench benchmark @xcite . this is a suite of tasks that involve extended contexts, including single-document question answering (qa), multi-document qa, document summarization, few-shot learning, code completion, and a synthetic task. the tasks span 16 datasets and presents a challenge due to the length of the input texts; for the size of the output, we use the original values in the dataset. for this dataset, our model is also compared against a third baseline, longllmlingua @xcite , a state-of-the-art method for compression of long input texts. for longllm-lingua, we use phi-2 @xcite as the compressor and llama 2 7b-chat (or mistral 4 finch's code and datasets are available at https:// github.com/giulio98/context-compression/ . lost in the middle: a critical challenge for llms is the ''lost in the middle'' issue @xcite , where models exhibit degraded performance if relevant information is situated in the middle of long contexts. we evaluate the robustness of our compression technique also in their dataset. ## results and discussion we discuss five questions over our results. ## conclusion and future work we have shown how attention can be used to identify and prioritize important information within the input data, effectively reducing the need for truncation. finch tackles the limitations of llms in processing large inputs, offering a balance between computational efficiency and maintaining high language model quality. our solution leverages the pre-trained model weights of the self-attention mechanism to provide an economically feasible method for operating llms. as future work, we envision a dynamic threshold mechanism to avoid that a fixed amount of kv states are selected in every chunk of the prefill stage, exploiting the fact that some chunks are not relevant and can be compressed more. another interesting research question is about the use of the proposed method to compress the generated output tokens. this extension would be especially valuable in settings where the llm is requested to generate long outputs, such as chain-of-thought reasoning. our approach could be used to identify the important tokens to preserve in the generation step -this is aligned with results showing that preserving a fraction of the original context is sufficient to obtain high quality generated outputs @xcite . finally, we are interested in studying how cache compression techniques can be extended to structured data, e.g., for replacing the current data retrieval and filtering solution in table question answering @xcite .