LKV: End-to-End Learning of Head-wise Budgets and Token Selection for LLM KV Cache Eviction

Summary (Overview)

• Core contribution: LKV (Learned KV Eviction) reformulates KV cache compression as an end-to-end differentiable optimization problem, replacing heuristic-based approaches with learned policies for both budget allocation and token selection.
• Two key components: LKV-H learns global budget allocation across all attention heads via a flattened 1D competition, while LKV-T performs matrix-free token selection using lightweight MLPs that predict token importance from intrinsic KV features.
• State-of-the-art results: Achieves 46.73 average score on LongBench with Llama-3.1-8B-Instruct at 15% KV retention, recovering 98.4% of Full Cache performance (47.48) and outperforming the strongest baseline (Ada-SnapKV) by 7.67 points.
• Key insight: Learned budgeting (LKV-H) is identified as the dominant driver of compression quality, contributing significantly more to performance than token selection improvements.
• Practical efficiency: LKV scales to 262k+ context length (vs. 225k for Full Cache), achieves 6.6× KV storage reduction, and incurs only ~7.7% prefill latency overhead.

Introduction and Theoretical Foundation

Problem Context

Long-context inference in LLMs faces a memory bottleneck due to the linear growth of KV cache. The KV cache stores key-value states for all tokens across all layers and heads, consuming memory proportional to sequence length. This limits maximum context length and reduces inference throughput.

Limitations of Existing Approaches

The paper identifies three fundamental limitations of current KV eviction methods:

1. Fragmented budget allocation: Methods like H2O and SnapKV use uniform budgets across layers/heads, ignoring heterogeneity in attention patterns. PyramidKV enforces rigid layer-wise decay; Ada-KV only adjusts within-layer head budgets.

2. Costly token selection: Methods like H2O and SnapKV require computing query-key attention scores to evaluate token importance, creating a "chicken-and-egg" problem—the attention matrix computation is exactly what compression aims to avoid.

3. Heuristic proxies: Many methods optimize proxy metrics (attention concentration, attention weights) rather than the actual task objective, and rely on manual safeguards like minimum budget thresholds or attention sinks.

Theoretical Formulation

The paper formulates KV compression as minimizing the discrepancy between full-cache and compressed models:

$$\min_{\phi} \mathbb{E}_{x \sim \mathcal{D}} \left[ \mathcal{L}\left(\text{LLM}(\mathcal{C}_{\text{full}}), \text{LLM}(\mathbf{M} \odot \mathcal{C}_{\text{full}}; \phi)\right) \right]$$

subject to the budget constraint $\|\mathbf{M}\|_0 \leq B_{\text{total}}$, where $\mathbf{M} \in \{0,1\}^{L \times H \times t}$ is a binary retention mask operating at the KV-head level.

The Differentiable Soft-TopK Operator

A critical enabler is the Soft-TopK operator $\mathcal{ST}_k(\cdot)$, which provides:

1. Normalization: Output sums exactly to $k$, enforcing budget constraints
2. Order preservation: Higher scores yield higher mask values
3. Translation invariance: $\mathcal{ST}_k(\mathbf{x} + c) = \mathcal{ST}_k(\mathbf{x})$, decoupling selection from absolute magnitude

The operator has a closed-form solution for the threshold:

$$\lambda(\mathbf{x}) = \log\left(\frac{\sqrt{(k-m)^2 + S_1 S_2} - (k-m)}{S_1}\right)$$

where $S_1 = \sum_{i=1}^{m} e^{-x_i}$ and $S_2 = \sum_{i=m+1}^{n} e^{x_i}$.

Methodology

LKV-H: Learning Global Budget Allocation

Mechanism: Each KV head $(l,h)$ has a learnable embedding $\mathbf{e}^{(l,h)} \in \mathbb{R}^{d_e}$, mapped to importance scores via a shared MLP:

$$s^{(l,h)} = \text{MLP}_{\text{head}}(\mathbf{e}^{(l,h)})$$

Global 1D Competition: The score tensor is flattened into a 1D vector, enabling cross-layer budget transfer. The Soft-TopK operator generates retention ratios:

$$\mathbf{r}_{\text{flat}} = \mathcal{ST}_{R_{\text{agg}}}(\mathbf{s}_{\text{flat}}), \quad \sum_{i=1}^{L \cdot H} (\mathbf{r}_{\text{flat}})_i = R_{\text{agg}}$$

Zero Inference Overhead: Ratios depend only on static embeddings, pre-computed once and frozen. At inference, the discrete budget for head $(l,h)$ is simply $b^{(l,h)} = \lfloor r^{(l,h)} \cdot t \rfloor$.

LKV-T: Matrix-Free Intrinsic Token Selection

Head-Specific Scoring: Each head has a unique lightweight MLP that predicts token importance from concatenated key-value features:

$$u_i^{(l,h)} = \text{MLP}_{\text{token}}^{(l,h)}([\mathbf{k}_i^{(l,h)}; \mathbf{v}_i^{(l,h)}])$$

This is query-agnostic—importance derives solely from KV states, avoiding dependency on the current query.

Gumbel-Based Training: During training, Gumbel noise is injected with temperature annealing:

$$\tilde{\mathbf{m}} = \mathcal{ST}_b(\mathbf{u} + \mathbf{g}; \tau), \quad \sum_{i=1}^{t} \tilde{m}_i = b$$

Temperature decays from high (exploration) to near-zero (deterministic). At inference, hard Top-K selection is used.

Training Objective

Self-distillation with a frozen teacher (full cache) and student (compressed):

$$\mathcal{L}_{\text{total}} = D_{\text{KL}}(P_T \parallel P_S) + \frac{\beta}{L} \sum_{l=1}^{L} \|\mathbf{H}_T^{(l)} - \mathbf{H}_S^{(l)}\|_2^2$$

A custom FlashAttention kernel uses multiplicative masking for numerical stability under BF16 precision.

Empirical Validation / Results

Experimental Setup

• Models: Llama-3.1-8B-Instruct, Qwen3-8B (both GQA)
• Benchmarks: LongBench (16 subtasks), RULER (synthetic)
• Training: ~23k samples, <2 hours on 8×A100 GPUs, ~0.1% trainable parameters
• Baselines: SnapKV, PyramidKV, DuoAttention, Ada-SnapKV, ExpAttn

Main Results (LongBench, 15% KV Budget)

Ablation Study: Budgeting vs. Selection

Robustness Results

• Length scalability: Stable performance from 4k to 32k contexts (trained on 16k), demonstrating policy transferability
• Extreme compression: At R=0.1, LKV maintains high performance while baselines degrade significantly
• Rigid priors fail: PyramidKV and DuoAttention drop to 24.99 and 17.75 on Synthetic tasks due to aggressive pruning of critical deep layers/heads

Efficiency Analysis

• Memory: LKV scales to 262k+ context (vs. 225k OOM for Full Cache)
• Storage: 6.6× reduction (25.0 GB → 3.75 GB at 200k length)

Theoretical and Practical Implications

Theoretical Contributions

1. Unified framework: Bridges existing selection mechanisms—the Soft-TopK operator approximates Softmax when k=1 and converges to Hard Top-K as τ→0⁺
2. End-to-end differentiability: Enables joint optimization of global budgets and token selection through gradient propagation
3. Cross-layer resource transfer: Flattened competition allows critical deep-layer heads to claim capacity from redundant shallow-layer heads

Practical Implications

1. Near-lossless compression: 98.4% of Full Cache performance at 15% retention enables extended contexts on commodity hardware
2. Query-agnostic efficiency: Matrix-free selection avoids O(t²) attention computation during eviction
3. Zero runtime overhead: Pre-computed budgets eliminate online allocation costs
4. Architecture compatibility: Fully compatible with GQA and orthogonal to other efficiency paradigms (quantization, sparse attention, prompt compression)

Key Insight

The ablation study reveals that learned budgeting is the dominant factor in compression quality. LKV-H provides +6.01 to +8.60 average gains across selectors, while LKV-T contributes an additional +2.67 over SnapKV under uniform budgets. This suggests that resource allocation—not token selection—is the primary bottleneck in KV cache compression.

Conclusion and Future Directions

LKV establishes a data-driven framework for KV cache compression, achieving state-of-the-art performance through:

• Global learned budgeting (LKV-H) that bypasses rigid structural priors
• Matrix-free token selection (LKV-T) that avoids attention computation
• End-to-end differentiable optimization via Soft-TopK with closed-form thresholds

Limitations and future work:

• Requires a lightweight training phase (though only ~2 hours on 8×A100)
• Needs specialized kernels (flattened FlashAttention) for maximum sparse access efficiency
• Future directions may include extending to other model architectures, integrating with quantization methods, and exploring even more aggressive compression regimes