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

> LKV achieves state-of-the-art KV cache compression via end-to-end learned budgeting and token selection, recovering 98.4% of full-cache performance at 15% retention with zero inference overhead.

- **Source:** [arXiv](https://arxiv.org/abs/2605.06676)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/NTtnRC
- **Whiteboard:** https://picx.dev/p/NTtnRC/image

## Summary

# 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)

| Method | Avg. | Single-Doc QA | Multi-Doc QA | Summarization | Few-Shot | Synthetic | Code |
|--------|------|---------------|--------------|---------------|----------|-----------|------|
| Full Cache | 47.48 | 43.57 | 45.41 | 28.89 | 65.43 | 53.48 | 51.42 |
| SnapKV | 38.05 | 27.27 | 34.08 | 22.93 | 58.11 | 40.73 | 50.05 |
| PyramidKV | 32.12 | 25.17 | 26.75 | 22.57 | 49.29 | 24.99 | 46.29 |
| DuoAttention | 31.03 | 21.76 | 25.94 | 21.66 | 50.71 | 17.75 | 50.40 |
| Ada-SnapKV | 39.06 | 27.38 | 35.46 | 23.29 | 59.28 | 43.77 | 50.56 |
| **LKV (Ours)** | **46.73** | **41.66** | **42.79** | **28.76** | **66.30** | **53.38** | **51.25** |

### Ablation Study: Budgeting vs. Selection

| Selection | Budgeting | Avg. | ΔBase | Multi-Doc QA | Synthetic |
|-----------|-----------|------|-------|--------------|-----------|
| SnapKV | Uniform | 38.05 | - | 34.08 | 40.73 |
| SnapKV | AdaKV | 39.06 | +1.01 | 35.46 | 43.77 |
| SnapKV | LKV-H | 44.06 | +6.01 | 41.75 | 52.87 |
| ExpAttn | Uniform | 37.12 | - | 37.39 | 17.77 |
| ExpAttn | AdaKV | 38.16 | +1.04 | 37.33 | 23.26 |
| ExpAttn | LKV-H | 45.72 | +8.60 | 43.27 | 45.00 |
| LKV-T | Uniform | 39.74 | - | 40.27 | 22.41 |
| LKV-T | AdaKV | 43.45 | +3.71 | 38.83 | 47.76 |
| **LKV-T** | **LKV-H** | **46.73** | **+6.99** | **42.79** | **53.38** |

### 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

| Method | Complexity | Time (s) | Overhead |
|--------|------------|----------|----------|
| Full Cache | - | 82.40 | - |
| DuoAttention | O(1) | 87.21 | +5.8% |
| Ada-SnapKV | O(W·T) | 88.37 | +7.2% |
| **LKV (Ours)** | **O(T)** | **88.76** | **+7.7%** |

- **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

---

_Markdown view of https://picx.dev/p/NTtnRC, served by PicX — AI-generated visual whiteboard summaries of research papers._