KV-Kaizen: Learning Context-Adaptive Cache Compression Choices

Summary (Overview)

• KV-Kaizen is a novel method for KV cache compression in LLMs that learns per-layer, context-adaptive choices along three orthogonal compression axes: depth (cross-layer cache sharing), rank (low-rank latent truncation via MLA), and precision (per-layer bit-width quantization).
• The method trains a lightweight selector (2.6–3.7M parameters, <0.05% of a 7B model) that reads the prompt once at pre-fill time and outputs a per-layer cache configuration satisfying a target compression budget, expressed in bits per token.
• Key finding: composing mild compression along multiple axes preserves accuracy far better than pushing any single axis aggressively. At 4× compression, KV-Kaizen matches uncompressed accuracy on models ≥7B parameters.
• KV-Kaizen composes multiplicatively with token eviction methods: combining 4× representation compression with SnapKV (keep 1/8) achieves 32× smaller decode-time cache on Qwen2.5-14B with negligible accuracy loss.
• The method supports budget-conditioned selectors (one model serving multiple compression targets) and demonstrates that compressed 7B models outperform smaller uncompressed models at matched cache sizes.

Introduction and Theoretical Foundation

Background: The KV Cache Bottleneck

Attention-based LLMs trade memory for compute: the KV cache stores key/value vectors for every prefilled or generated token, avoiding recomputation at each decoding step. However, this cache grows with sequence length and batch size, creating a memory bottleneck:

• Qwen2.5-14B uses 192 KiB per token of cache memory; serving 32 sequences of 8k tokens uses 48 GiB—more than the model weights in bf16.
• Decoding is memory-bandwidth-bound: every decoded token reads the entire KV cache, so cache size directly limits throughput.

Three Non-Eviction Compression Axes

The paper focuses on compression that does not drop tokens, distinguishing it from eviction-based approaches:

1. Depth (cross-layer attention): Layers share KV caches—an "anchor" layer produces a cache; subsequent layers "inherit" from the nearest anchor. The number of anchor layers determines effective depth.

2. Rank (MLA truncation): Multi-head latent attention (MLA) stores a shared latent representation $c^{KV}_l \in \mathbb{R}^r$ plus a decoupled RoPE key $k^R_l \in \mathbb{R}^{d_R}$. Truncating the latent to its first $d \leq r$ coordinates controls cache size. Non-MLA models are converted via CARE.

3. Precision (quantization): Cached tensors are stored at $b \leq 16$ bits via symmetric uniform quantization (one scale per token), with straight-through gradient estimation during fine-tuning.

Key Insight

"Applying these structural changes aggressively quickly results in performance degradation."

The central hypothesis: individually mild compression along each axis, composed per-layer and adapted to context, achieves compression ratios no single axis reaches alone without accuracy loss.

Methodology

Action Space and Cost Model

Each layer $l$ chooses one action $a_l$ from the space:

$$\mathcal{A} = \{\text{inherit}\} \cup \{(d, b) : d \in \mathcal{D}, b \in \mathcal{B}\}$$

where $\mathcal{B} = \{2, 4, 8, 16\}$ (bit-widths) and $\mathcal{D} = \{r/8, r/4, r/2, r\}$ (kept latent widths). This yields 17 actions per layer.

The per-token cost of action $a$ is:

$$c(a) = \begin{cases} 0 & a = \text{inherit} \\ (d + d_R) \cdot b & a = (d, b) \end{cases} \quad \text{(bits per token)}$$

The compression ratio for a plan $a_{1:L}$ is:

$$\rho(a_{1:L}) = \frac{C_0}{\sum_{l=1}^L c(a_l)}, \quad C_0 = L \cdot 2 n_{kv} d_h \cdot 16$$

where $C_0$ is the uncompressed bf16 non-MLA cache cost.

Selector Architecture and Training

• Selector: Two bidirectional transformer blocks ingest the backbone's token embeddings; representations are mean-pooled over time. One linear head per layer maps this shared feature to logits over $\mathcal{A}$.
• Selection: Gumbel-softmax with straight-through estimator (hard argmax forward, softmax backward), temperature annealed from 2.0 to 0.1.
• Initialization: Heads initialized so every layer's argmax is the least-compressed action (start from uncompressed model).
• Inference: Plain argmax; selector runs once before pre-fill (~3 ms, 0.20–3.3% of pre-fill time); configuration fixed for all generated tokens.

Training Objective

Standard language modeling loss plus a rate penalty:

$$\mathcal{L}_{\text{rate}} = \beta \cdot \mathbb{E}_i \left| \frac{1}{\rho_i} - \frac{1}{\rho^\star} \right|, \quad \frac{1}{\rho_i} = \frac{1}{C_0} \sum_l \sum_{a \in \mathcal{A}} w^{(i)}_{l,a} c(a)$$

where $w^{(i)}_l$ is layer $l$'s straight-through one-hot selection vector. The hyperparameter $\beta$ is adaptively adjusted during training (raised when compression target is missed, relaxed when met), enabling one setting to reach all targets.

Budget-Conditioned Selectors

To serve multiple compression targets with one selector, the target $\rho^\star$ is made an input via a zero-initialized FiLM layer on the pooled feature; training draws a random target per sequence.

Empirical Validation / Results

Main Results (Qwen2.5-7B)

Figure 3 shows mean IFEval + GSM8K accuracy vs. cache compression:

• 4× compression with KV-Kaizen matches uncompressed accuracy (0.679).
• Static (uniform per-layer) and random configurations degrade accuracy at matched budgets.
• KIVI (post-hoc 2-bit quantization) stays on the frontier up to ~5×; beyond that, only learned configurations work.
• Applying learned configurations to a frozen backbone collapses accuracy (must fine-tune under the configuration).
• Compressed 7B beats smaller uncompressed models: past 2×, a compressed 7B cache outperforms full-cache 3B, 1.5B, and 0.5B models.

Scaling with Model Size

Table 1 (mean IFEval + GSM8K accuracy differences at 4× compression):

Bold marks differences ≥ −0.053 (largest seed range).

Key patterns:

• Larger models tolerate compression better: gaps close from 3B to 14B.
• Depth alone is the weakest axis (must sacrifice 3/4 of layers to reach 4×).
• Combining axes doesn't add costs: precision+depth ≈ precision alone, while depth alone is far worse.
• Depth costs more on reasoning (GSM8K) than instruction following (IFEval).

Long-Context Evaluation and Eviction Composition (Table 2)

RULER at 16k on Qwen2.5-14B (QA accuracy):

Key result: KV-Kaizen at 4× beats eviction-only methods at matched cache size (0.614 vs 0.501 for SnapKV). Combining KV-Kaizen with eviction multiplies compression: 32× total with only −0.012 accuracy difference. KIVI + SnapKV is worse on both axes (29.5×, 0.464).

Budget-Conditioned Selectors (Table 3)

• Trained on {1×, 2×, 4×} with precision+depth: one selector covers all three, ahead of uncompressed control at 1×, only 0.067 behind single-target at 4×.
• Wider range [2×, 8×] with all three axes requires more training: 540k steps gets within 0.038/0.023/0.039 of single-target selectors at 2×/4×/8×.
• Too-wide target ranges degrade performance (selector stops following requests).

Hardware Benchmarking (Apple M3 Ultra)

Qwen2.5-14B with 4-bit weights, batch size 1:

• Composed compression (4-bit cache + MLA rank 128 + 12 groups of 4 layers): 152× total compression.
• Peak decode memory: 12.3 GB (bf16) → 0.17 GB (composed).
• Depth always helps decode throughput; MLA pays off at long sequences; 4-bit cache adds de-quantization compute cost but wins at long horizons.

Theoretical and Practical Implications

Why Composing Mild Compression Works

The paper's central theoretical contribution is demonstrating that compression axes are complementary rather than additive in their costs. Precision and rank degrade a layer "by degrees" (coarser values, narrower latent), while depth is a discrete change (layer computes from foreign keys/values). The selector naturally uses depth least because it's the most disruptive, but depth enables compression factors unreachable by precision alone (which floors at 8× due to 2-bit minimum).

Training Under the Configuration is Mandatory

"Applying a configuration to a model that never trained under it degrades it systematically."

This is a critical practical insight: post-hoc compression (even with good quantizers like KIVI) hits a ceiling around 5×; beyond that, the model must be fine-tuned under the target configuration.

Larger Models are More Compression-Tolerant

The consistent finding that larger models tolerate compression better supports a clear practical recommendation: pre-train large models, compress them afterward. Compressed 7B models outperform smaller uncompressed models at matched cache sizes, favoring the "train big, serve compressed" paradigm.

Selector Diversity and Deployment Considerations

• Some selectors return a single global configuration for all prompts (could be replaced by a static config); others adapt to context. This must be verified per selector.
• Budget-conditioned selectors enable deployment flexibility (one model, multiple cache budgets) at some accuracy cost.

Conclusion

Main Takeaways

1. KV-Kaizen learns per-layer, context-adaptive cache configurations along depth, rank, and precision axes under a unified bit-per-token budget.
2. Composing mild compression along multiple axes preserves accuracy at compression ratios (4×+) that no single axis achieves alone.
3. 4× compression is essentially free for models ≥7B; compressed larger models beat smaller uncompressed ones.
4. Compression composes multiplicatively with eviction (32× total on 14B at 16k context).
5. Training under the serving configuration is essential; post-hoc methods plateau at ~5×.

Limitations and Future Work

• Specialized kernels needed to fully realize speed/memory benefits of mixed per-layer configurations.
• MLA conversion cost varies across model families; rank axis unavailable without it.
• Budget conditioning degrades with wider target ranges; finer control could enable per-deployment adaptation.
• Depth compression disproportionately hurts reasoning tasks; mechanisms to mitigate this are unexplored.
• The 32× result combines with eviction but was only demonstrated at 16k context; longer contexts remain untested.

Future Directions

• Finer-grained budget control for hardware-adaptive deployment.
• Specialized inference kernels for mixed-width/rank/depth configurations.
• Mitigating depth compression's impact on reasoning.
• Extending budget-conditioned selectors to wider target ranges without accuracy loss.