PaperlayerView exact revision on arXiv

Source-linked AI summary

Quality Recovery for Quantized KV Caches via Low-Rank Attention Adaptation

Seifeldin Abdellatif

arXiv:2609.04263v1cs.ARcs.CL

TL;DR

Low-bit KV caches reduce memory and bandwidth costs, but their remaining quality gap varies with the model and quantizer. With base weights and quantizer fixed, the method distills floating-cache behavior into low-rank Q/K/V updates while the student executes the packed cache path. Selected QKV placement recovers 54.24% ± 2.47% of TinyLlama’s induced PPL gap, and Gemma’s recovery is 75.96% ± 4.04%; at 2-bit scaling restores 11–12 retrieval targets but leaves most behavior absent.

  • Problem

    Low-bit KV caches reduce memory and bandwidth costs, but their remaining quality gap varies with the model and quantizer.

    Source: §1 · p. 1

  • Method

    With base weights and quantizer fixed, the method distills floating-cache behavior into low-rank Q/K/V updates while the student executes the packed cache path.

    Source: §1 · p. 1

  • Results

    Selected QKV placement recovers 54.24% ± 2.47% of TinyLlama’s induced PPL gap, and Gemma’s recovery is 75.96% ± 4.04%; at 2-bit scaling restores 11–12 retrieval targets but leaves most behavior absent.

    Source: §5.1 · p. 5, §5.5 · p. 9

  • Takeaways & Limitations

    Across tested models and fixed cache formats, low-rank attention adaptation recovers quality, but near-floating 2-bit perplexity does not imply restored retrieval.

    Source: §5.5 · p. 9, §7 · p. 10

    Placement, KIVI, and KVarN results use one seed, so they do not quantify variance or support comparisons between quantizers.

    Source: §6 · p. 9

Abstract

from arXiv · show

Low-bit key--value (KV) caches reduce the memory required for autoregressive decoding, but the resulting quality loss depends on the model and quantizer. We keep the quantizer fixed and distill the floating-cache model's behavior into low-rank Q/K/V projection updates while the student executes a physically packed incremental cache. Across three seeds, 4-bit affine-cache adapters recover $54.24\%\pm2.47\%$ of the held-out perplexity gap on TinyLlama-1.1B and $75.96\%\pm4.04\%$ on Gemma-4-12B. On the same frozen NF4 Llama-3.1-8B base, one validation-selected run per quantizer recovers $60.42\%$ under KIVI K2V2 and $37.61\%$ under KVarN K4V2, while preserving 180-case associative retrieval. Gemma's score on an official 4K/8K RULER subset rises from 42.80 with the unadapted 4-bit cache to 48.33 after adaptation (46.15 floating), with substantial task heterogeneity. Finally, a 2-bit rank--token sweep reduces TinyLlama's 2-bit PPL from 576.10 to $11.4000\pm0.0059$ across three seeds, versus 10.3988 floating, but restores only 11--12 of 180 retrieval cases. These results show that low-rank projection adaptation can recover held-out quality across fixed cache formats, while perplexity recovery need not restore long-context retrieval.

1 Introduction

The introduction asks whether low-rank attention updates can recover quality after cache formats are fixed, while the student uses the deployed packed-cache path. It frames evaluation around held-out perplexity and long-context behavior across models, bit widths, and quantizers.

1 Introduction

  • The study asks whether low-rank Q/K/V updates can recover loss under a fixed cache format by distilling floating-cache behavior into a student using the packed incremental path.The base weights and quantizer stay fixed, and calibration uses unlabeled text.Source: §1 · p. 1
  • The study separates cache storage from adapter recovery, reporting absolute perplexity, normalized gap recovery, and long-context behavior.The cache representation determines storage; held-out evaluation measures adapter recovery.Source: §1 · p. 1
  • Across tested formats, projection adaptation improves held-out quality, while attainable gains and calibration requirements vary by operating point.The evaluation spans three model scales, 2- and 4-bit affine and symmetric caches, KIVI, KVarN, and long-context NIAH and RULER tests.Source: §1 · p. 1

2 Related Work

Prior work reduces cache costs through representation design, cache-error correction, or low-rank adaptation. This paper instead leaves the cache format and attention operator unchanged and adapts ordinary Q/K/V projections.

Cache quantization

  • KIVI, KVQuant, SKVQ, KVTuner, RotateKV, and KVarN change cache representation, whereas this study adapts the model after the representation is fixed.The experiments cover grouped affine and symmetric coding, KIVI, and KVarN; KIVI’s format is used with framework attention rather than specialized decoding kernels.Source: §2 · pp. 1–2

Cache-error correction

  • GEAR, STAR-KV, and KVLinC allocate representation or operator state to correction; this method instead uses ordinary Q/K/V projection updates to absorb part of cache error.The updates can be folded when the base projection is dense.Source: §2 · p. 2

Low-rank adaptation and evaluation

  • The paper uses LoRA and QLoRA mechanisms for self-distillation rather than task tuning, and supplements needle retrieval with RULER’s broader task types.RULER spans retrieval, variable tracking, and aggregation because simple needle retrieval can conceal broader long-context failures.Source: §2 · p. 2

3 Method

The method trains low-rank Q/K/V projection updates against a fixed cache path, so calibration and evaluation use the packed-cache behavior the student will encounter. Its deployment cost depends on whether the adapter can be merged or remains explicit.

3 Method

  • Low-rank adapters modify selected query, key, and value projections while keeping pretrained weights frozen; the study compares Q-only, K-only, QK, and QKV placements.For one attention layer, the method uses hidden states H and low-rank matrices A and B with inner dimension r.Source: §3 · p. 2

3.1 Physical incremental cache

  • The deployed cache stores packed codes and metadata, appends newly packed states, and retains no unquantized copy; evaluation reads and updates it in 128-token chunks.Affine coding groups the final head dimension in groups of 64 and stores an FP16 scale and offset; the symmetric control stores a maximum-magnitude scale without an offset.Source: §3.1 · p. 2

3.2 Quantizer interface and compatibility

  • The fixed cache interface Rϕ includes coding, metadata, residual or transformed state, and reconstruction presented to attention; compatibility requires running that deployed path during calibration.An identity straight-through estimator carries gradients through hard coding to the Q/K/V adapters; stateful KVLinC must instead be treated as a complete attention operator.Source: §3.2 · p. 3

3.3 Calibration objective

  • Calibration minimizes KL divergence between floating-cache teacher and packed-cache student behavior, using straight-through gradients so attention sees reconstructed states while K/V adapters receive a backward path.The objective searches a restricted correction space rather than inverting quantization; query and key changes can be complementary.Source: §3.3 · p. 3

3.4 Deployment overhead

  • An explicit rank-r adapter adds r(din + dout) parameters and two low-rank matrix multiplications, while dense-base updates can be merged and the reported NF4 runs retain explicit adapters.The Gemma QK rank-32 adapter has 21.889M parameters and adds 21.889M MACs per processed token when run explicitly; static counts do not determine wall-clock latency.Source: §3.4 · p. 3

4 Experimental Protocol

The protocol compares fixed cache formats and adapters on held-out perplexity and synthetic long-context tasks, using repeated seeds for selected conditions and validation-selected single runs for others. Evaluation reports absolute perplexity alongside gap recovery and includes boundaries on what the comparisons establish.

Established quantizers

  • KIVI and KVarN studies use the same frozen NF4 Llama-3.1-8B base, rank-32 QKV adapters, seed 0, validation-selected checkpoints, and one test evaluation per selected checkpoint.Checkpoints are saved after 1.25M, 5M, 10M, and 20M input tokens.Source: §4 · p. 4

Two-bit capacity study

  • The TinyLlama 2-bit capacity study tests ranks 6, 12, 24, and 48 at four calibration budgets, selects by full WikiText-103 validation PPL, and evaluates selected ranks on test and shared NIAH cases.The selected rank is retrained with two additional seeds and evaluated on a fresh 180-case NIAH suite.Source: §4 · p. 4

Evaluation

  • Held-out incremental-cache perplexity uses complete windows with identical tokens across floating, unadapted packed, and adapted packed conditions; three-seed means and sample standard deviations are descriptive.TinyLlama scores 340,326 WikiText-2 tokens and Gemma scores 291,463 WikiText-103 tokens.Source: §4 · p. 4
  • The recovery metric normalizes different gaps but is unstable when Pq − Pf is near zero and can look favorable when every absolute PPL is poor, so tables retain absolute PPL.Source: §4 · p. 4

Synthetic retrieval and RULER

  • NIAH uses greedy exact-match retrieval on specified prompt suites, while Gemma RULER covers four of 13 tasks at 4K and 8K with 200 examples and string-match-all scoring.The RULER subset includes single-needle retrieval, multi-value retrieval, variable tracking, and common-word extraction.Source: §4 · pp. 4–5

5 Results

Across 4-bit affine caches and established quantizers, low-rank adaptation improves held-out perplexity, while retrieval recovery varies by operating point and task.

5.1 Three-seed quality recovery

  • Across three seeds, QKV recovers 54.24% ± 2.47% of TinyLlama’s 4-bit PPL gap and 75.96% ± 4.04% for Gemma.The corresponding adapted PPLs are 10.5471 ± 0.0080 for TinyLlama and 26.7745 ± 0.0450 for Gemma.Source: Table 1 · p. 5, §5.1 · p. 5

5.2 Where should the adapter be placed?

  • QKV gives the strongest matched-budget recovery: 56.95% on TinyLlama and 63.21% on Llama-3.1, versus 44.01% for Llama-3.1 QK.On TinyLlama, Q and QK are nearly tied, K alone is weaker, and QKV reaches 10.5383 PPL.Source: §5.2 · pp. 5–6, Table 2 · p. 6

5.3 Recovery transfers to KIVI and KVarN

  • KIVI K2V2 recovers 60.42% of its test PPL gap and preserves 179/180 NIAH cases; KVarN K4V2 recovers 37.61% and preserves 180/180.Each quantizer uses one validation-selected run, so these results do not estimate variance or compare quantizer quality.Source: §5.3 · p. 6, Figure 2 · p. 7

5.4 Bit width and adapter capacity

  • At 2 bits, rank-48 QKV reaches test PPL 11.3935 versus 576.1030 unadapted and 10.3988 floating; three seeds average 11.4000 ± 0.0059.Across the tested rank and calibration-token grid, validation PPL improves with both greater rank and more calibration.Source: Table 3 · p. 6, §5.4 · p. 7

5.5 Behavioral retrieval beyond perplexity

  • On a fresh 180-case set, rank-48 2-bit adapters retrieve 11, 12, and 11 targets versus 0 unadapted and 170/180 floating.The corresponding accuracy remains 6.1–6.7%, far below the 94.4% floating reference.Source: §5.5 · pp. 8–9
  • Gemma’s three QK adapters raise retrieval from 112/120 packed cases to 115–117, while TinyLlama’s three QKV adapters raise it from 163/180 to 167–169.For TinyLlama, adapters fix 10–11 packed-cache failures while breaking 5–7 successes; no individual paired test reaches p < .05.Source: §5.5 · pp. 7–8, Table 4 · p. 8
  • Gemma gains 20 percentage points for shallow needles at 8,192 tokens, while TinyLlama’s largest cited gain is from 40% to 55% at middle depth and 1,024 tokens.These cell patterns are descriptive across three independently trained adapters.Source: Figure 4 · p. 9, §5.5 · p. 9
  • On Gemma’s official RULER subset, the eight-cell macro-average rises from 42.80 packed to 48.33 adapted, above 46.15 floating.Task changes are heterogeneous: multi-value retrieval and common-word extraction improve, single-needle retrieval declines, and variable tracking remains essentially unsolved.Source: §5.5 · p. 8, Table 5 · p. 9

6 Discussion and Limitations

Low-rank adaptation reproducibly recovers quality under tested fixed-cache formats, but long-context recovery is uneven and evaluation and deployment scope remain limited.

6 Discussion and Limitations

  • In the 2-bit stress test, increased capacity and calibration recover 11–12/180 NIAH cases, far below the 170/180 floating result.The paper reports that NIAH and the four-task RULER subset measure synthetic behavior rather than broad downstream performance.Source: §6 · p. 9
  • Placement, KIVI, and KVarN results use one seed, and three model families do not establish universality.Gemma has no placement sweep, and its optimizer skips a small fraction of updates despite close seed results.Source: §6 · p. 9
  • The reported NF4 path retains an explicit Gemma adapter, adding 21.889M MACs per processed token; fused-kernel benchmarks are needed to determine deployment costs.Compatibility with untested outlier, layer-mixed, and learned-state methods remains theoretical.Source: §6 · p. 10

7 Conclusion

Across tested models, low-rank attention adapters recover quality lost to fixed cache formats, but perplexity recovery does not ensure retrieval recovery.

7 Conclusion

  • Low-rank attention adapters recover quality lost to physically packed affine, symmetric, KIVI, and KVarN caches across tested models.The conclusion positions this as a post-training mechanism for recovering quality under fixed cache formats.Source: §7 · p. 10
  • Matched budgets favor QKV, while rank and calibration nearly restore 2-bit perplexity but not retrieval.Four-bit NIAH and RULER show useful but nonuniform recovery.Source: §7 · p. 10
  • All experiments and results are available at github.com/saifmb0/kvlora.Source: §7 · p. 10
Loading 2609.04263v1…