XC -Cache: Cross-Attending to Cached Context for Efficient LLM Inference (2024)

Abstract: Prompts are often employed to condition decoder-only language model generation on reference information. Just-in-time processing of a context is inefficient due to the quadratic cost of self-attention operations, and caching is desirable. However, caching transformer states can easily require almost as much space as the model parameters. When the right context is not known in advance, caching the prompt can be challenging. This work addresses these limitations by introducing models that, inspired by the encoder-decoder architecture, use cross-attention to condition generation on reference text without the prompt. More precisely, we leverage pre-trained decoder-only models and only train a small number of added layers. We use Question-Answering (QA) as a testbed to evaluate the ability of our models to perform conditional generation and observe that they outperform prompt-based inference methods, are comparable to fine-tuned prompted LLMs, and drastically reduce the space footprint relative to standard KV caching by two orders of magnitude. Specifically, we introduced XC-Llama which converts a pre-trained Llama 2 into an encoder-decoder architecture by integrating cross-attention layers interleaved in between existing self-attention layers.

## introduction large language models (llms) have propelled advances in language modeling and enabled automatic production of almost human-like text. despite impressive progress, challenges persist in applying llms in practical settings such as the risk of hallucinations (or rather confabulatations @xcite ) and of nonfactual @xcite or toxic content @xcite in their generated text. moreover, without fine-tuning, it is surprisingly difficult to adapt these models to incorporate new information not included in this work focuses on grounding llm generation on contextual information provided at inference time (figure 2(a) ). the most common approach for conditioning model generation is arguably in-context learning foot_0 (icl) @xcite : one prepends the relevant context to the prompt to generate an answer conditioned on the combined query and context. this technique is a core component of popular frameworks, such as retrieval-augmented generation (rag) @xcite , with the specificity that the relevant context is not known a priori, but has to be retrieved from a corpus @xcite . while somewhat effective and straightforward, icl, as typically performed with decoder-only architectures, has flaws. on the one hand, icl-based generation is known to present high variance with respect to the prompt template so that equivalent valid-looking prompt formats produce drastically different results @xcite . on the other hand, icl is costly in terms of time and space. just-in-time processing of the context suffers from quadratic complexity on the length due to selfattention operations @xcite . the alternative is pre-processing and caching the context internal states (the so-called key-value or kv states) to speed up inference (figure 2(c )). however, this can require the same order of space as the model parameters themselves (we give details in section 2). recent work has reduced the space requirements of kv caching by sub-sampling states @xcite , although at the cost of ignoring relevant content. to overcome these limitations, we propose alternatives to icl that perform conditional generation without injecting the relevant information in the prompt (figure 2(a) ), and seek to implement lightweight cache methods as illustrated in figure 2(d). our approach is reminiscent of the, now arguably legacy, encoder-decoder architectures, as it relies on cross-attention layers to condition generation on pre-computed context encodings. more precisely, we propose cross-context-cache (xc-cache), which stores only the outputs of the encoder hidden states and relies on cross-attention to ingest the cache at inference time. we instantiate xc-cache via two parameter-efficient approaches that leverage pre-trained decoder-only models and extend them with a separate encoder to process the context: one approach uses the frozen decoder as an encoder (called xc-llama), and the other uses a small bi-directional encoder (called xc-llamaenc). crucially, our encoder-decoder architectures are more amenable to caching the context states, requiring orders of magnitude less space than their icl counterparts. when context caching is enabled, fine-tuned models result in higher accuracy performance, but demand a large memory footprint (and consequently, higher latency and cost). in contrast, our xc-cache approach substantially reduces cache memory requirements by nearly 98%; and as other efficiency-improving techniques like quantization @xcite , this reduction comes at a minor cost in accuracy, as illustrated in figure 1 . nevertheless, our method consistently outperforms icl alternatives based on llama 2 or gpt-3.5, as detailed in section 5. overall, we advocate for a conceptual shift in architecture design for conditional generation, which should recenter on caching and make it integral to a model's operation rather than an afterthought. our contributions are summarized as follows: ## caching representations let context, query and answer denote sequences of tokens from a shared vocabulary v. we write |context|, |query| and |answer| the respective length of these sequences. figure 2 (a) illustrates an llm which, conditioned on context, produces an answer to a user-specified query. ## xc: cross attending to efficiently cached context to reduce the memory footprint of caching, our approach draws inspiration from the encoder-decoder architecture, which until recently was the go-to design for conditional generation. recently, decoderonly models have become more popular in part due to their data efficiency: the entire parameter budget is allocated to the decoder and not "wasted" on an encoder, with the additional advantage that all parameters are trained against all data tokens. even though maintaining an external encoder may seem wasteful, we favor an encoder-decoder architecture: it better lends itself to pre-computing and caching context representations. indeed, only the encoder output vectors need to be stored, as opposed to intermediate states across all of the decoder's selfattention layers (kv and jit-kv cache). in what follows, we refer to a model as composed of an encoder e : v |context| → r d|context| , which takes in a context and outputs token-level representations of size d, and of a decoder d : v m × r d|context| → ∆ |v| , which takes as input the query and the context encodings and outputs an answer in the simplex of size |v|. more precisely, the decoder d autoregressively outputs the parameters of a categorical distribution over the vocabulary. for the sake of parameter efficiency and to leverage state-of-the-art pre-trained llms, we start from an existing decoder-only model and augment it with new cross-attention layers interleaved between existing self-attention layers, as illustrated in figure 4 . we consider two strategies to implement the encoder e. the first one trains a small encoder, i.e., e is orders of magnitude smaller than d. the second one uses a decoder as encoder, i.e., the frozen decoder-only model is used out-of-the-box as the encoder (e := d). more precisely, we use as encodings the representations extracted from the pre-trained d at its last layer before the language modeling head. choosing one approach over the other depends on practical considerations. if caching is possible and context representations can be computed offline, then using the decoder as an encoder is preferable for overall simplicity and parameter efficiency. otherwise, a small encoder would make just-intime processing of contexts significantly less costly. both approaches inherit the advantages of trained decoder-only models while benefiting from using an encoder during inference. in particular, contextual information can be efficiently cached since only the output at e's top layer must be stored instead of the entire set of d's intermediate states. finally, to enable context conditioning, we train exclusively the newly added modules: the crossattention layers in both settings and the small encoder in the setting that requires it. the base decoder is kept frozen in both settings (even when it acts as an encoder). consequently, our training procedure does not affect the original parameters of d, which can still be used as a general-purpose language model when the additional layers are removed. ## experimental setup we focus on the question-answering task and train on a combination of datasets where context, query, and answer triplets are available. qa is an ideal testbed for our approaches as it requires efficiently looking up external information and accounting for it at generation time. in other words, qa almost isolates the ability of the model to "read" the context, which is what we needed to test for. ## related work decoders as encoders. repurposing pre-trained decoders is becoming a popular approach to leverage powerful language models for applications other than generative modeling. for example, grit @xcite converts a pretrained causal decoder into a bi-directional encoder, yielding sentence-level embeddings while maintaining its ability to perform autoregressive generation of text. however, unlike the models we consider, this conversion requires fine-tuning all model parameters instead of additional ones. parameterefficient approaches to turn decoders into encoders were also proposed, such as in @xcite and @xcite , where a pretrained mistral decoder @xcite is fine-tuned in a contrastive setting using lora adapters to yield sentence level representations for retrieval. closer to our method is codet5+ @xcite , which also defines encoder-decoder architectures rather than turning decoders into sentence encoders. similar to one of our variants (xc-llama), it is assembled from two pre-trained decoders, one used as an encoder and the other as a decoder and linked with a few cross-attention operations. however, codet5+ requires fine-tuning the entire and relatively large encoder. we show that a good pre-trained decoder has good enough ## conclusion we introduced xc-llama as an approach to transform a pre-trained decoder-only language model into an encoder-decoder architecture that can condition generation on both the encoder inputs and the decoder query. this is achieved by integrating cross-attention layers interleaved in between existing self-attention layers of the pre-trained decoder. we describe two approaches for defining the encoder: using a copy of the decoder or introducing a trainable but small bi-directional encoder. the proposed architecture allows for a reduction in caching space by a factor exceeding 300. when evaluated in the qa setting, we observe a higher prediction accuracy than what standard icl approaches achieve, either with llama 2 or gpt-3.5 turbo. additionally, we observe accuracy levels nearly on par with caching-intensive finetuned prompted models, providing a more cachingfriendly alternative to prompted language models that proves highly practical. we note that the encoder-decoder strategy for handling contexts separately offers the extra advantage of freeing prompt space, which is now exclusively used for user interaction. identified limitations are discussed in section 8.