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

TL;DR

This work uses Question-Answering as a testbed to evaluate the ability of the models to perform conditional generation and observe that they outperform ICL, are comparable to fine-tuned prompted LLMs, and drastically reduce the space footprint relative to standard KV caching by two orders of magnitude.

Abstract

In-context learning (ICL) approaches typically leverage prompting 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 isn't known in advance, caching ICL 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 ICL, are comparable to fine-tuned prompted LLMs, and drastically reduce the space footprint relative to standard KV caching by two orders of magnitude.

Figures (6)

• Figure 1: Average QA performance vs. caching memory footprint per context token. The closer to the top-right corner, the better. XC-Llama variants drastically reduce cache size at a small cost in accuracy.
• Figure 2: Faster inference in context-conditional language modeling. (a) A use case where the user's $\text{\smallquery}$ must be interpreted within some $\text{\smallcontext}$ to generate an $\text{\smallanswer}$. This work focuses on cases where the $\text{\smallquery}$ and the $\text{\smallanswer}$ are both small (light feather), but the $\text{\smallcontext}$ is large (heavy rock). The resulting LLM time complexity is thus $\mathcal{O}\bigl({\lvert\text{\smallcontext}\rvert^2}\bigr)$ (slow snail). (b) In-context learning (ICL) and retrieval augmented generation (RAG) are two examples where the $\text{\smallquery}$ is used to look up the $\text{\smallcontext}$from a finite corpus. (c) In many cases, the $\text{\smallcontext}$ can be processed in advance to a $\text{\smallcache}$ enabling $\mathcal{O}\bigl({\lvert\text{\smallcontext}\rvert (\lvert\text{\smallquery}\rvert + \lvert\text{\smallanswer}\rvert)}\bigr)$ (fast rocket) inference on a given $\text{\smallquery}$. (d) Finite $\text{\smallcontext}$ corpus may be processed offline, enabling fast execution at inference. Since $\text{\smallcache}$ size affects storage and communication costs, we search for models requiring smaller $\text{\smallcache}$.
• Figure 3: Stylized representation of attention execution time (area) for context-conditional language modeling. Dashed lines in (b) shows the savings when the $\text{\smallcontext}$'s (past) keys and values $\text{\smallcache}$ are provided. For causal models, the area below the diagonal represents execution time that could be saved by other means.
• Figure 4: XC-Llama's architectures. A decoder-only model implements encoder-decoder architectures. Fine-tuning out in a parameter-efficient fashion via training only a small number of cross-attention layers.
• Figure 5: Multitask training strategy. Within an epoch, each example is presented to the model twice. In the first round, the model predicts the answer conditionally based on the context and the query. In the second appearance of an example within the epoch, we train the model to repeat or infill the context.