PaperlayerView exact revision on arXiv

Source-linked AI summary

KV-Kaizen: Learning Context-Adaptive Cache Compression Choices

Joao Monteiro, Louis Béthune, Anastasiia Filippova, Sonia Laguna, David Grangier, Marco Cuturi

arXiv:2609.37988v1cs.AI

TL;DR

KV-cache memory grows with sequence length and must be read at every decoding step, while query-independent eviction can discard tokens useful for other queries. KV-Kaizen trains a selector with the model to choose per-layer depth, rank, and precision under a cache budget, selecting once before pre-fill. Combined approaches achieve significant KV cache size reductions at little to no accuracy degradation; at 16k, 4× representation compression plus eviction yields a 32× smaller decode-time cache, with significant accuracy difference from the uncompressed control for a 14B model.

  • Problem

    KV-cache memory grows with sequence length and must be read at every decoding step, while query-independent eviction can discard tokens useful for other queries.

    Source: §1 · p. 1

  • Method

    KV-Kaizen trains a selector with the model to choose per-layer depth, rank, and precision under a cache budget, selecting once before pre-fill.

    Source: §1 · p. 2

  • Results

    Combined approaches achieve significant KV cache size reductions at little to no accuracy degradation; at 16k, 4× representation compression plus eviction yields a 32× smaller decode-time cache, with significant accuracy difference from the uncompressed control for a 14B model.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The findings favor training larger models and compressing them before deployment, and the selector can serve several compression targets at a larger training cost.

    Source: §6 · p. 10

    Fully realizing memory and speed benefits requires specialized kernels, and rank compression requires MLA conversion whose accuracy cost varies across evaluated models.

    Source: §6 · p. 10

Abstract

from arXiv · show

As the context size of text processed with an LLM grows, the size of KV caches can outstrip the memory allocated for the original model weights. This impacts LLM throughput negatively, since decoding is memory-bound and decode cost grows with cache size. Recent work alleviates this bottleneck by discarding the least relevant tokens. Eviction introduces a tension, since a one-off decision to discard content may prove detrimental later. Instead, we focus on alternative choices that can lead to cache compression without evicting tokens. We achieve this by learning a selector that is able to produce, based on context, a per-layer cache configuration towards an overall compression budget. The selector operates along three axes: sharing one cache across layers (depth), caching at fewer bits (precision), or truncating the low-rank latent cache representations (rank). We call the resulting method KV-Kaizen, for the many small per-layer choices it compounds. We observe that these interventions taken independently and uniformly over all layers limit achievable compression because they degrade accuracy. Crucially, composing them locally and adaptively to the context can instead preserve accuracy while achieving large memory savings. At inference, the selector runs once, before pre-fill. In evaluations on instruction following and reasoning tasks, our selectors reach the Pareto frontier of accuracy against cache size, against learning-free and post-hoc baselines. On long-context tasks, KV-Kaizen improves on eviction and can be composed with it, reaching a 32x smaller decode-time cache on a 14B model while preserving accuracy. A 4x cache size reduction incurs no accuracy degradation from 7B parameters up, and a compressed model is more accurate than a smaller uncompressed one with the same cache size. Together, these findings support pre-training large models and compressing them only afterwards.

1 Introduction

KV-Kaizen learns context-adaptive, per-layer cache configurations that combine depth, rank, and precision under a shared budget. The introduction reports accuracy-preserving cache reductions, including composition with token eviction, and device measurements showing nearly flat peak memory as sequences grow.

1 Introduction

  • KV caches grow with sequence length and every decoded token reads the cache, making decoding memory-bandwidth-bound and motivating compression without information loss.For Qwen2.5-14B, 32 sequences of 8k tokens are reported to use 48 GiB of cache.Source: §1 · p. 1
  • The introduction contrasts token eviction, which can discard content useful to other queries, with lower-precision, lower-dimensional, and cross-layer cache compression, whose aggressive use can degrade performance.It motivates composing cache-compression choices rather than relying on a single structural change.Source: §1 · pp. 1–2
  • KV-Kaizen trains a selector to choose per-layer depth, rank, and precision from the prompt under a shared cache budget; its choice is fixed after one inference-time selection before pre-fill.The model and selector are trained together, and the configuration remains in effect for generated tokens.Source: §1 · p. 2
  • At 16k context, 4× representation compression plus eviction to one position in eight yields a 32× smaller decode-time cache on Qwen2.5-14B with no significant accuracy difference.The paper also reports that a 4× cache reduction gives accuracy comparable to the corresponding backbone in evaluated models of at least 7B parameters.Source: §1 · p. 2
  • On device, composing depth, rank, and precision keeps peak memory nearly flat as sequences grow, and most combinations decode faster than an uncompressed cache.The introduction reports these outcomes for the three approaches benchmarked on device.Source: §1 · p. 3

2 Composing compression axes under one budget

KV-Kaizen combines cache sharing, latent-rank truncation, and lower precision as per-layer choices under a shared compression budget. A learned selector maps context to layer configurations while training with a cost target.

2.1 Background and notation

  • Depth shares caches across layers, rank truncates MLA latents, and precision quantizes cached tensors without dropping input tokens.Rank truncation keeps the first d latent coordinates; precision uses symmetric uniform quantization.Source: §2.1 · p. 3

2.2 The selector action space and the cache configuration cost

  • Compression ratios use the unmodified bf16 non-MLA cache as reference, so MLA conversion contributes savings before selector choices.At full width, MLA configurations start between 1.60× and 1.88×, depending on the underlying model.Source: §2.2 · p. 3
  • The 17-action space combines cache inheritance with kept-width and bit-width choices, assigning each layer a cost in bits per token.A produced cache costs (d + dR) · b; an inherited cache costs zero.Source: §2.2 · pp. 3–4

2.3 Learning the cache configuration against a budget

  • The selector maps prompt tokens to per-layer configurations and learns jointly with the backbone, using a loss term that pushes realized cost toward the target compression.Its shared prompt feature comes from a small sequence model; the selector runs once at pre-fill and the chosen plan remains fixed during generation.Source: §2.3 · p. 4
  • β rises when the target compression is missed and relaxes once met; evaluations compare learned selectors with matched-budget static and random configurations.All compared models are fine-tuned with their cache configuration already in place.Source: §2.3 · pp. 4–5

3 Evaluation

Across the evaluation, learned cache configurations retain accuracy at 4× compression on Qwen2.5-7B, while long-context QA favors the selector over matched eviction baselines. Combining compression with eviction reaches 32× cache reduction, with results varying by model size and compression axis.

3.1 Setup

  • The evaluation fine-tunes the backbone with the selector, measuring IFEval and GSM8K accuracy and 16k RULER performance against cache-size-matched baselines where available.The shared fine-tuning recipe uses the Dolci instruction mixture at length 4096 for 270,000 steps.Source: §3.1 · p. 5

3.2 Reducing the cache size by more than about 5× needs training

  • At 4× cache reduction, KV-Kaizen matches the uncompressed Qwen2.5-7B control; beyond about 5×, frontier configurations are learned.The comparison uses mean IFEval and GSM8K accuracy against cache compression.Source: §3.2 · p. 5, Figure 3 · p. 6

3.3 Scaling the base model size

  • At 3B every configuration trails its size-matched control, while at 7B and 14B depth alone is the only configuration not level with its control or ahead.Controls are matched by attention class; the comparison also extends to 32B in Appendix F.6.Source: §3.3 · p. 6
  • At 14B and 4×, precision combined with depth lands above the control, whereas depth alone loses accuracy; beyond 8×, settings without rank lose reasoning accuracy first.Depth alone at 16× scores 0.444 on IFEval and 0.056 on GSM8K.Source: §3.3 · p. 7, Figure 4 · p. 7

3.4 Long context evaluations and composition with eviction methods

  • At 4× cache reduction, RULER question-answering accuracy is 0.614 versus 0.501 for SnapKV and 0.485 for PyramidKV; combining eviction reaches 32× at 0.531.The evaluation is on Qwen2.5-14B at 16k context; single-needle retrieval is 100% across configurations.Source: §3.4 · pp. 7–8

3.5 One selector covers a set of budgets

  • A selector trained on {1, 2, 4} serves all three budgets and is 0.067 behind a single-target selector at 4×; wider target ranges require more training.The wider-range selector trails per-target selectors by 0.038, 0.023, and 0.039 at 2×, 4×, and 8×, respectively.Source: §3.5 · p. 8

3.6 Cache configuration diversity

  • Configurations vary across datasets, so whether a selector can be replaced by one global configuration should be checked on representative test data.The selector costs about 3 ms, negligible compared with pre-fill.Source: §3.6 · p. 8
  • Depth is used least because inheriting another layer’s cache means a layer works from keys and values it did not produce; non-MLA configurations above 8× must give up at least half their layers.Precision bottoms out at two bits, while rank is unavailable on non-MLA models.Source: §3.6 · pp. 8–9

4 Benchmarking cache compression approaches

The on-device benchmark measures pre-fill time, decode throughput, and peak decode memory for cache-compression settings on Qwen2.5-14B. Composed compression sharply reduces peak memory, while throughput and pre-fill effects vary by method and sequence length.

4 Benchmarking cache compression approaches

  • 12.3 GB of peak memory for bf16 falls to 0.17 GB for the composed cache on the Apple M3 Ultra benchmark.The benchmark uses Qwen2.5-14B with 4-bit weights at batch size 1.Source: Figure 5 · p. 9, §4 · p. 9
  • A 4-bit cache, rank-128 MLA latent, and 12 groups of four layers sharing caches compose to a reported 152× reduction.The individual settings are reported as 3.56×, 10.7×, and 4×, respectively.Source: §4 · p. 9
  • Depth always helps throughput, MLA mostly pays off at long sequence lengths, and the bf16 cache is fastest at pre-filling.The reported comparison sweeps context lengths for pre-fill and completion lengths for decode throughput.Source: §4 · p. 9
  • A 4-bit KV cache adds de-quantization compute, increasing pre-fill and MLA costs, but pays off at long horizons.The figure reports pre-fill time relative to the uncompressed bf16 cache.Source: §4 · p. 9, Figure 5 · p. 9

5 Related work

Prior cache-compression work often fixes a structural choice or bit-width at pre-training or conversion, while other approaches allocate selectively. KV-Kaizen combines a wider set of actions under a shared budget.

5 Related work

  • Non-evicting approaches include sharing KV states within or across layers, low-rank latents, and fewer-bit storage, often fixed across layers at pre-training.Prior work also explores non-uniform allocation and learned allocation, including rank selection and per-layer rates.Source: §5 · p. 9
  • Compared with prior learned allocators, KV-Kaizen's wider action space combines multiple mechanisms, which the authors claim can reach a given reduction with milder settings.The cited comparison discusses Bhatnagar et al. (2026) and Lin et al. (2025).Source: §5 · p. 10

6 Conclusions, limitations, and future work

KV-Kaizen learns per-layer depth, rank, and precision choices under one cache budget; combining these approaches can reduce cache size with little to no accuracy degradation. The related work contrasts this composed, budgeted allocation with token eviction and prior architectural or adaptive compression strategies.

6 Conclusions, limitations, and future work

  • At 16k context, 4× representation compression combined with token eviction gives a 32× smaller decode-time cache, with significant accuracy difference from the uncompressed control on a 14B model.The conclusion says larger models generally tolerate compression better and notes that budget-conditioned selectors can serve multiple targets at greater training cost.Source: §6 · p. 10
  • Realizing the full memory and speed benefits of configurations mixing widths and bit-widths across layers requires specialized kernels.The accuracy cost of MLA conversion varies across evaluated models, and widening the training-target range can reduce accuracy even at less-compressed settings.Source: §6 · p. 10

A Extended related work

  • The method leaves the time axis untouched, so representation compression can be combined with token eviction rather than replacing it.The cited discussion describes eviction's difficulty: token importance can depend on the query, so a discarded position may be needed later.Source: A Extended related work · p. 15
  • Prior architectural reductions and quantization typically fix one structural choice or bit-width at conversion or pre-training, whereas this work trains under a budget across three composed axes.The passage contrasts this approach with fine-tuning for retrieval and measuring compression robustness only as a downstream property.Source: A Extended related work · p. 15
  • Cross-layer cache-sharing methods include post-hoc layer merging and alignment approaches, while xKV fixes its group and rank before training under them.Filippova et al. (2026) instead randomize the source layer during fine-tuning so one checkpoint tolerates test-time sharing patterns.Source: A Extended related work · p. 15
  • Earlier adaptive methods allocate token budgets, ranks, or bit-widths non-uniformly; the authors distinguish their approach by jointly choosing depth, rank, and precision under one budget.They further state that the budget can be an input at runtime, allowing one selector to answer a range of requests.Source: A Extended related work · pp. 15–16

B Experimental details

The appendices specify the backbones, training and evaluation setup, feasible compression ranges, and MLA cache representation used in the experiments. They describe how model geometry and depth choices constrain those ranges and how rank truncation is represented during training.

B Experimental details

  • OLMo-3 runs use full attention in every layer, and rank configurations convert non-MLA backbones to MLA; the reported backbone table gives model and cache dimensions.Converted models use CARE with balanced key-value scaling from TransMLA; Qwen2.5 projection biases are retained.Source: B Experimental details · p. 16, Table 4 · p. 16
  • Cache quantization uses symmetric uniform quantization with one scale per stored vector, and training quantizes then de-quantizes in the forward pass with straight-through gradients.Non-MLA backbones quantize keys and values; MLA backbones quantize the retained latent and positional key at the same bit-width.Source: B Experimental details · p. 16
  • Fine-tuning uses Dolci-Instruct-SFT1 for 270,000 steps, about one epoch, with sequences truncated to 4096 tokens and assistant turns supervised.Training uses eight sequences per step across eight GPUs, with at most 32,768 tokens per step.Source: B Experimental details · pp. 16–17
  • A shared token feature feeds independent per-layer action heads, while a FiLM MLP conditions on the normalized target compression factor.Each training sequence samples a target; the Gumbel temperature decreases from 2.0 to 0.1, and constant-multiplier runs warm up β linearly.Source: B Experimental details · p. 17
  • Evaluation uses greedy decoding on fixed sets; timing uses batch size one, four warmups, and the median of 12 repetitions.Pre-fill and selector latency are measured on an H100 at 4× compression for precision and depth.Source: B Experimental details · p. 17

C.1 Compression ranges

  • The feasible compression range depends on available actions and whether depth is enabled; at the precision-only maximum of 8×, every layer must use two bits.The maximum uses the lowest-cost action per layer without depth, or at one anchor with depth; the action table specifies bit-widths and latent widths.Source: §C.1 · pp. 17–18, Table 5 · p. 18

C.2 MLA cache representation

  • An MLA cache stores a latent and a positional key per token; rank truncation zeroes coordinates after the chosen width during training.The training tensor retains full shape for differing widths in a batch, while a deployed kernel would write only the retained coordinates.Source: §C.2 · p. 18

C.3 Quantization and storage overhead

The section describes quantization choices and storage overhead, then adapts the rate multiplier to the gap between requested and realized compression.

C.3 Quantization and storage overhead

  • The symmetric quantizer uses one scale per stored vector, while grouped asymmetric quantization uses a scale and zero-point per channel group.At two bits, the symmetric quantizer has three levels; the grouped asymmetric quantizer uses all 2^b levels.Source: §C.3 · p. 18
  • At two bits and group size 128, KIVI reaches 5.15× compression for 2048 tokens, or 7.11× without its residual window.At 16k context, four-bit KIVI with group size 128 uses 4.34 bits per element, corresponding to 3.7× compression.Source: §C.3 · p. 18, Table 6 · p. 19

D.1 Update rule

  • The rate multiplier adapts to the requested-versus-realized compression gap, with its update parameters specified and the multiplier all-reduced across data-parallel ranks.The update uses λ = 0.9, η = 0.05, δ = 0.02, βmax = 104, and initializes β = 0.Source: §D.1 · p. 19

D.2 Budget accuracy and task performance

The adaptive rule reliably reaches requested cache budgets, though meeting a budget does not itself ensure the best accuracy. Capping the multiplier improves several high-compression results, and representative runs approach their targets during training.

D.2 Budget accuracy and task performance

  • The adaptive rule reaches the target in all 51 evaluated fixed-target configurations.Constant multipliers are less reliable near maximum feasible compression; with precision only on Qwen2.5-7B, tested settings reached 6.32× against an 8× target.Source: §D.2 · p. 19
  • At matched realized compression, swept constant multipliers outperform adaptive ones in four of ten pairs, while six pairs differ by at most 0.053 in accuracy.The four wins include all three precision-and-depth pairs; reaching the target budget therefore does not guarantee best task accuracy.Source: §D.2 · p. 19, Table 7 · p. 20
  • Capping β at 50 raises Qwen2.5-7B precision-and-depth accuracy from 0.255 to 0.496 at 8× and from 0.162 to 0.327 at 16×.At 14B and 16× with all three approaches, the cap changes accuracy by +0.005, with almost no layer inheriting.Source: §D.2 · p. 19
  • In representative all-axis runs, compression stabilizes near its target by about 30% of training, as the rate multiplier first rises and then falls.Figure 6 covers ten Qwen2.5 runs at 1.5B, 3B, 7B, and 14B with 4×, 8×, and 16× targets; 1.5B is evaluated only at 4×.Source: §D.2 · p. 19, Figure 6 · p. 20

E.1 Alternatives to the learned configuration

The appendix tests whether learned configurations outperform alternatives at matched cache costs, and compares training-dependent compression with post-hoc and axis-specific baselines. Its findings emphasize that configuration choice and training context matter alongside cache size.

E.1 Alternatives to the learned configuration

  • With trained 7B weights fixed, no tested alternative beats the selector in either action space, and choosing configurations per prompt group gains almost nothing.With all three approaches, alternative losses spread by 0.30 and the selector's margin is 0.014; the test does not establish which configuration would be best if each were trained independently.Source: §E.1 · pp. 20–21, Table 8 · p. 21

E.2 Offline search

  • At equal cache size, the searched configuration barely changes loss, random configurations raise it clearly, and the uniform configuration raises it far more.Per-layer estimates predict the searched configuration well but miss most of the uniform configuration's loss when every layer is compressed too much.Source: §E.2 · p. 21, Table 9 · p. 21

F.1 Quantization and smaller-model baselines

  • Post-hoc plain quantization destroys accuracy, whereas training under the same quantizer preserves control-level accuracy; KIVI holds up at four bits but loses accuracy at two bits.The smaller uncompressed Qwen2.5-1.5B model saves only 2× and performs far below a 7B model compressed 4×.Source: §F.1 · p. 21

F.2 Accuracy by compression axis

  • At 4× on Qwen2.5-7B, configurations using rank or precision stay within the seed range against the relevant converted or unconverted control, while depth alone is the exception.Depth alone loses much more on reasoning than on instruction following.Source: §F.2 · p. 22, Table 11 · p. 22

F.3 Long-context evaluation

Long-context tests compare cache compression with eviction and show that splitting large reductions between them can preserve question-answering performance. The appendix also reports fine-tuning effects and tests learned versus fixed sharing patterns.

F.3 Long-context evaluation

  • At 7B, the 4× configuration matches query-informed SnapKV and PyramidKV, while combining it with SnapKV at 16× matches four-bit KIVI under the same evictor.Policies selecting tokens before the query arrives fall far behind, as do depth or rank alone at this context length.Source: §F.3 · p. 23, Table 12 · p. 23
  • At 14B and 16k context, splitting a 32× cache reduction between 4× compression and eviction that keeps one position in eight stays level with the uncompressed control.At the same retained fractions, KVK beats four-bit KIVI by 0.098, 0.074, and 0.067; the panel uses a non-MLA backbone because conversion alone costs most question-answering accuracy at this length.Source: §F.3 · p. 23, Table 13 · p. 24

F.4 Effect of fine-tuning

  • SFT changes accuracy by −0.052 to +0.014 for Qwen models at 7B and above, while Mistral loses 0.152.These changes motivate reporting compression relative to a control trained with the same recipe.Source: §F.4 · p. 24, Table 14 · p. 24

F.5 Learned and static cache configurations

  • At 7B, evenly spaced anchors outperform learned placements, while the results suggest that the number of inheriting layers matters more than their placement.Simple sharing patterns can place anchors when sharing is required, but whether to share is worth learning because precision and rank can often meet the budget without it.Source: §F.5 · p. 24, Table 15 · p. 25

F.6 Training seeds and model size

The appendix examines selector behavior across model sizes, budget conditioning, cross-layer sharing, configuration diversity, and selector architecture. It finds that target-range training and compression choices involve trade-offs in accuracy, training cost, and when sharing is used.

F.6 Training seeds and model size

  • At 32B, MLA conversion lowers accuracy by 0.131 while language-modeling loss improves from 0.448 to 0.353; accuracy changes are non-monotonic across model sizes.The largest accuracy range across training seeds is 0.053 for the uncompressed 7B control.Source: §F.6 · p. 25, Table 16 · p. 25, Table 17 · p. 26

F.7 Discrete and continuous budget conditioning

  • Selectors trained on {2, 4, 8} and [2, 8] perform similarly wherever both meet requests; at two-thirds of their total training steps, the range selector approaches the single-target selectors it replaces.Wider sets can compress more than requested, and their rate multiplier does not settle, ending training up to ten times higher than for {1, 2, 4}.Source: §F.7 · p. 26

F.8 Use of cross-layer sharing

  • Depth is used past 8× on non-MLA backbones, but not at fixed targets when precision and rank suffice; imposing unnecessary sharing costs accuracy.Past 8×, all trained non-MLA configurations use at least half their layers to inherit another layer’s cache.Source: §F.8 · pp. 26–27, Table 19 · p. 27

F.9 Cache configuration diversity

  • Across model sizes, six selectors return several per-layer configurations and six reuse a single configuration across inputs, measured over 224 prompts.The prompts cover IFEval, GSM8K, and RULER at 4k, 8k, and 16k tokens.Source: §F.9 · p. 27, Table 20 · p. 27

F.10 Selector architecture

  • The embedding transformer is retained despite higher accuracy from the hidden-state variant with all three approaches, because it runs before pre-fill; training takes 1.0–3.1× ordinary fine-tuning depending on actions.The hidden-state variant requires a full backbone pass to select a configuration and then a second pre-fill.Source: §F.10 · pp. 27–28, Table 21 · p. 28

G On Device benchmarks

The on-device benchmarks examine how weight precision, KV-cache precision, and MLA rank relate to throughput and memory. They report that lower-precision attention can raise long-sequence attention cost in this setup, while reducing MLA rank increases decode throughput.

G On Device benchmarks

  • Benchmarks run on an M3 Ultra with 512GB of RAM, with temperatures and frequencies monitored to avoid thermal throttling.The measurements are reported in the on-device benchmark appendix.Source: G On Device benchmarks · p. 28

G.1 The Role of Precision

  • For quantized matrices, decomposed SDPA matmuls do not benefit from Flash, increasing low-precision self-attention cost on long sequences.Figure 8 compares three KV-cache precisions on Qwen2.5-14B with 4-bit weights.Source: §G.1 · p. 28, Figure 8 · p. 28
  • At fixed KV-cache precision and sequence length around 30K, weight precision has a large impact on throughput and total memory footprint, orthogonal to KV-cache precision.Figure 7 reports Qwen2.5-14B at three weight precisions with an uncompressed bf16 KV cache.Source: §G.1 · p. 28, Figure 7 · p. 28
  • Reducing MLA rank increases decode throughput.Figure 9 shows Qwen2.5-14B with a bf16 KV cache.Source: §G.1 · p. 28, Figure 9 · p. 29
Loading 2609.37988v1…