OBCache: Optimal Brain KV Cache Pruning for Efficient Long-Context LLM Inference

TL;DR

OBCache addresses the memory bottleneck of long-context LLM inference by rethinking KV cache eviction as a structured pruning problem. It uses a second-order Taylor expansion under the Optimal Brain Damage framework to derive closed-form, output-aware token saliency scores, including $S_p^{value}$, $S_p^{key}$, and $S_p^{joint}$, which quantify the eviction impact on attention outputs. By incorporating value states, pre-softmax logits, and attention outputs, OBCache provides richer signals than prior attention-weight heuristics and encompasses them as special cases. Experiments on LLaMA-3.1 and Qwen-2.5 demonstrate consistent improvements in long-context accuracy and perplexity when OBCache scores are integrated into existing KV cache eviction pipelines.

Abstract

Large language models (LLMs) with extended context windows enable powerful downstream applications but impose significant memory overhead, as caching all key-value (KV) states scales linearly with sequence length and batch size. Existing cache eviction methods address this by exploiting attention sparsity, yet they typically rank tokens heuristically using accumulated attention weights without considering their true impact on attention outputs. We propose Optimal Brain Cache (OBCache), a principled framework that formulates cache eviction as a layer-wise structured pruning problem. Building upon the Optimal Brain Damage (OBD) theory, OBCache quantifies token saliency by measuring the perturbation in attention outputs induced by pruning tokens, with closed-form scores derived for isolated keys, isolated values, and joint key-value pairs. Our scores account not only for attention weights but also for information from value states and attention outputs, thereby enhancing existing eviction strategies with output-aware signals. Experiments on LLaMA and Qwen models demonstrate that replacing the heuristic scores in existing works, which estimate token saliency across different query positions, with OBCache's output-aware scores consistently improves long-context accuracy.

Figures (4)

• Figure 1: Overview of the OBCache scoring mechanism. The diagram shows the eviction process using value-pruning (left) and key-pruning scores (right). Unlike prior methods based solely on attention statistics (gray region), OBCache further incorporates value states, attention logits, and outputs to estimate token saliency, explicitly targeting the minimization of eviction-induced errors.
• Figure 2: Recall rate of the top-40 salient tokens identified by the oracle eviction error. Results are collected from a 4K-context passkey retrieval task using LLaMA-3.2-3B-Instruct. The left plot (a) shows recall of top-40 tokens selected by different scoring methods. The right plot (b) demonstrates how allocating a fixed recent window improves oracle recall when perturbation windows are large, mitigating the structural bias disproportionately favoring earlier tokens.
• Figure 3: Overall long-context performance evaluation of LLaMA and Qwen. When integrated with OBCache scores, existing cache eviction baselines achieve superior compression and performance trade-off on the Needle-In-A-Haystack (a), LongBench (b), and perplexity (c) benchmarks.
• Figure 4: Language modeling perplexity evaluation on PG19 test set. We prompt Llama-3.1-8B-Instruct with 1 to 32K tokens and measure the perplexity of output tokens at varying context lengths. The KV cache budget for all methods is fixed at 1024 number of tokens.