PaperlayerView exact revision on arXiv

Source-linked AI summary

DCKV: Learned Transform Coding of the Prompt KV Cache in Hybrid Language Models

Egor Maximov, Aleksei Merkelov, Vasiliy Alekseev, Ivan Krivorotov, Valentin Beliaev, Nikita Arsenin, Dmitry Bocharov, Aleksei Goncharov

arXiv:2610.36950v1cs.LGcs.AI

TL;DR

In hybrid language models, the KV cache is the prompt-dependent state that grows and must be stored or moved to reuse a prompt. DCKV jointly trains a cache codec against the frozen model, and uses fewer bytes than KVTC at equal score on RULER 4K and LongBench.

  • Problem

    In hybrid models, the KV cache grows with prompt length and must be stored or moved to reuse a prompt, even though other state can be fixed-size.

    Source: §1 · p. 1

  • Method

    DCKV jointly trains the transform, quantizer, and conditional entropy models against the frozen model using code length and output change.

    Source: §1 · p. 1, §3.5 · p. 4

  • Results

    DCKV needs 72% fewer bytes than KVTC at equal score on RULER 4K and 62% fewer on LongBench over three rates.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The ablations show that entropy coding, the joint basis, training against model output, and conditional entropy models each add to the gain.

    Source: §6 · p. 8

    The codec was developed on Qwen3.5-0.8B, but results are reported only for Qwen3.5-2B, with no tests on other hybrid families.

    Source: Limitations · p. 8

Abstract

from arXiv · show

Hybrid language models keep a key-value (KV) cache only in their few full-attention layers, but this cache is still the part of a prompt's state that grows with the prompt and that must be stored or moved to reuse it. We present DCKV, a differentiable codec for the KV cache and the first KV-cache codec that is trained end to end. Its transform, its quantizer and its entropy models, which are conditioned on the prompt tokens and on already decoded neighbouring positions, are optimized together against the frozen language model, both for the code length that the entropy models assign and for the change of the model's output. On Qwen3.5-2B, DCKV stores a compressed token in more than 55 times fewer bytes than bf16 at a loss of at most 1.8 points on three long-context benchmarks, and in more than 100 times fewer bytes at a loss of at most 3.6 points.

1 Introduction

Hybrid models limit the KV cache to a few full-attention layers, but it remains the prompt-dependent state that must be stored or moved for reuse. DCKV jointly trains a codec against the frozen model, reducing bytes at equal benchmark score while accounting for recurrent state in request memory.

1 Introduction

  • The KV cache is 12,288 bytes per token in bf16 and grows to 1.5 GiB at 128K tokens, while Qwen3.5-2B’s recurrent state is fixed at 9.84 MiB.Six of 24 layers use full attention; the other 18 are Gated DeltaNet layers.Source: §1 · p. 1
  • KVTC uses PCA, dynamic-programming bit allocation, and a general-purpose lossless coder, but its transform and allocation do not use model output.DCKV instead optimizes the output change as part of codec training.Source: §1 · p. 1
  • DCKV jointly trains a transform, quantizer, and entropy models against the frozen model using code length and output change.The transform is shared across keys, values, and attention layers, with a small nonlinear correction and learned quantization steps.Source: §1 · p. 1
  • Conditional entropy models use token embeddings and already decoded neighbours; with the correction, they remove 12–16% of bytes at the same quality.Fitting the models in closed form to an already trained codec gives a similar saving on RULER 4K without retraining.Source: §1 · p. 1
  • 72% fewer bytes on RULER 4K and 62% fewer on LongBench than KVTC at equal score over three rates.The evaluation separates request state from codec weights.Source: §1 · pp. 1–2
  • At 32K tokens, DCKV scores 6.5 points higher than KVTC at about the same request memory.At 4K–11K tokens, recurrent state takes about three quarters of stored-request memory.Source: §1 · p. 2

2 Related Work

Prior cache-compression methods target either decoding-time caches or stored caches, while learned codecs for other representations generally train together with the model. DCKV is presented as the first KV-cache codec to train its transform, quantizer, and entropy model through both code length and model output.

Hybrid models

  • Hybrid-model reuse requires both the attention cache and a recurrent-state snapshot; DCKV compresses the former and keeps the latter exact.The recurrent state is fixed-size, whereas the attention cache grows with the prompt.Source: §2 · p. 2

KV cache compression for decoding

  • Decoding-time methods evict, quantize, project, or share cache entries, but retain fixed-width caches that attention can read directly.That fixed-width requirement rules out variable-length codes, leaving statistical redundancy within the cache.Source: §2 · p. 2

Learned compression

  • Learned image and video codecs train transforms, quantizers, and entropy models together, while DCKV works with an already-trained language model.Earlier language-model codec work described here trains a small model together with a codec for one hidden state.Source: §2 · p. 2

Entropy-coded storage of KV caches

  • DCKV is presented as the first KV-cache codec to train its transform, quantizer, and entropy model through code length and model output.Its entropy models also condition on prompt tokens and already decoded neighbours.Source: §2 · p. 2
  • Stored-cache codecs use entropy coding after quantization, rotations, prediction, PCA, or rate–distortion bit allocation, but their probability models are fixed or profiled.Codec-Gauge trains channel transforms around an existing codec with a smooth proxy for rate.Source: §2 · p. 2

3 Method

DCKV compresses the prompt’s attention-cache region once after prefill and restores it before decoding, while keeping recurrent state exact. It combines a joint transform, decoder-conditioned entropy models, and rate–distortion training against a frozen model.

3.1 What is compressed

  • Each token’s keys and values across six attention layers form one row; the first four and last 128 cache positions remain in bf16.Keys are taken before RoPE and receive RoPE again after decompression.Source: §3.1 · pp. 2–3

3.2 Transform and quantization

  • A joint transform maps each cache row to 2,048 coefficients, with one quantization step per component and an optional nonlinear correction.The transform starts from a joint PCA basis; components with scale below 0.05 step units are neither stored nor reconstructed.Source: §3.2 · p. 3

3.3 Entropy coding

  • Each code is written with rANS, while the static entropy model assigns component-wise zero-mean Laplace probability over each quantization cell.DCKV replaces the static model with conditional models that shift each code’s Laplace distribution by a predicted mean.Source: §3.3 · p. 3, §3.4 · p. 3

3.4 Conditional entropy models

  • Conditional entropy models predict code means from the prompt token and, for odd positions, the decoded codes of both neighboring tokens.Even positions are coded first and odd positions second, with each pass parallel over tokens.Source: §3.4 · pp. 3–4

3.5 Rate–distortion training

  • Training minimizes conditional code length plus output-distribution KL divergence against the frozen model.The prompt is prefetched once; the model’s next-token distributions are compared using the exact and reconstructed caches. Mixed quantization computes rate on noisy coefficients but reconstructs from rounded coefficients.Source: §3.5 · p. 4

3.6 Memory accounting

  • Request memory counts the compressed payload, 132 uncompressed positions, and recurrent state; DCKV’s codec weights are counted separately as model weights.The recurrent state is 9.84 MiB, and DCKV adds 176 MiB of fp32 codec parameters.Source: §3.6 · p. 4

4 Experimental Setup

The experiments compare DCKV with cache-compression baselines on Qwen3.5-2B using long-context and short-prompt benchmarks. The setup specifies paired task-level comparisons, codec training, and request-memory accounting.

Benchmarks

  • Evaluation uses RULER at 4K and 32K, 18 LongBench tasks, and 5-shot MMLU, with benchmark-specific metrics.The LongBench subset excludes LCC, RepoBench-P, and SAMSum; including all 21 tasks makes reported losses 0.3–1.0 points smaller without changing method order.Source: §4 · pp. 4–5

Baselines

  • Baselines include KVTC with PCA bases, dynamic bit allocation, and DEFLATE, plus xKV with prompt-specific low-rank bases.For xKV, the prompt-specific basis is counted in request memory and payload; KVTC is implemented from a public proof of concept with dynamic programming bit allocation.Source: §4 · p. 5

Training

  • Codecs train on FineWeb-Edu for 1B tokens, using calibration rows from 48 documents of web text and mathematical reasoning.Training uses AdamW on two A100 GPUs; the reported DCKV runs take 25.9 or 33.9 hours, depending on configuration.Source: §4 · p. 5

Comparisons

  • Comparisons report paired differences across shared tasks, W/T/L counts and sign tests, while BD-rate estimates payload differences over overlapping score ranges.Evaluations are deterministic across repeated launches; BD-rate uses linear interpolation from three to five points per curve.Source: §4 · p. 5

5 Results

DCKV reduces KV-cache payload substantially while retaining benchmark quality, and its advantage over KVTC grows at stronger compression. Ablations attribute gains to entropy coding, a joint basis, rate–distortion training, and conditional models; request-memory benefits depend on prompt length, and losses vary by task.

5.1 Quality and memory

  • 205–221 bytes per token gives DCKV 56–60 times fewer bytes than bf16, with score changes of −0.56, −1.75 and −0.45 points on RULER 4K, LongBench and RULER 32K.KVTC at 16× stores 2.5–3.1 times as many bytes and loses at most 0.53 points; on LongBench it is 1.44 points better than DCKV (p=0.03).Source: §5.1 · p. 5, Table 1 · p. 6
  • At λ = 105, DCKV stores 113–122 bytes per token and scores higher than KVTC at 32× on all three benchmarks while using 2.2–2.7 times fewer bytes.Across three rates, DCKV needs 72% fewer bytes at equal score on RULER 4K and 62% fewer on LongBench; at λ = 5·104 it uses about a quarter to a third of KVTC’s bytes without a significant score difference.Source: §5.1 · p. 6

5.2 Where the gains come from

  • Replacing dynamic bit allocation and DEFLATE with one quantization step, a Laplace model and rANS gains about one point on RULER 4K and LongBench at slightly lower payload.A joint key–value basis uses 7–10% fewer bytes than separate bases at the same step; at step 4.0 it scores 4.3 points higher on RULER 4K, but the difference is not significant.Source: §5.2 · p. 6
  • At λ = 105, rate–distortion training scores 3.0 points higher on RULER 4K than the untrained joint basis at a similar payload.Mixed quantization scores 0.7 and 1.7 points higher than the noise proxy on RULER 4K and LongBench and is better on 20 of 31 tasks, worse on 6; at λ = 3·105 the estimators do not differ significantly.Source: §5.2 · p. 7
  • At λ = 3·105, adding the correction and conditional models reduces bytes by 12–14% versus the codec without additions, without a significant score difference.On RULER 4K, the correction alone saves 5.5% and conditional models alone 8.7%; fitted conditional models save 12.5%, while jointly training additions saves 22% at equal score under stronger compression.Source: §5.2 · p. 7

5.3 Similar request memory

  • The recurrent state accounts for 81%, 73% and 54% of request memory on RULER 4K, LongBench and RULER 32K; at 32K, DCKV scores 6.5 points higher than KVTC at about the same request memory.At λ = 105, DCKV needs 5–22% less request memory than KVTC at 32× and scores higher on all three benchmarks.Source: §5.3 · pp. 7–8, Figure 3 · p. 7

5.4 Behaviour by task

  • Task-level losses are uneven: at λ = 3·105, passage retrieval loses 10.0 points in English and 6.0 in Chinese on LongBench, while the largest RULER 32K losses are 4.2 and 3.7 points.At λ = 105, LongBench passage retrieval loses 27.0 points in English and 12.5 in Chinese; variable tracking loses 5.8 points on RULER 4K and 11.0 on RULER 32K.Source: §5.4 · p. 8

6 Conclusion

DCKV jointly trains a transform, quantizer and conditional entropy models against a frozen language model to compress stored prompt caches. On Qwen3.5-2B, its byte savings over KVTC grow at stronger compression, while recurrent-state memory limits request-memory savings at moderate prompt lengths.

6 Conclusion

  • DCKV trains the transform, quantizer and conditional entropy models end to end against the frozen language model, combining code length with output-distribution change.The codec is designed for the stored prompt cache of a hybrid language model.Source: §6 · p. 8
  • At about 210 bytes per token, DCKV uses 2.9 times fewer bytes than KVTC at the same loss on RULER 4K; across three rates it needs about a third of KVTC’s bytes at equal score on RULER 4K and LongBench.The conclusion reports that its advantage grows at stronger compression.Source: §6 · p. 8
  • At moderate prompt lengths, the recurrent state takes most stored-request memory, while at 32K tokens DCKV scores higher than KVTC at about the same request memory.Because DCKV is trained against a frozen model, it can serve a released model without changing its weights.Source: §6 · p. 8
  • Whether these results carry over to other hybrid models, how to train one codec across rates without upper-end loss, and how to make decoding identical across devices remain open questions.Joint pre-training of a model and codec may yield a cache that compresses better.Source: §6 · p. 8

Limitations

The evaluation covers one hybrid model, with limited task counts and several untested settings; rate control, context length, memory and cross-device decoding also have defined boundaries. The uniform-noise rate estimate is biased for many trained coefficient scales.

Limitations

  • Results are limited to Qwen3.5-2B and 13 or 18 tasks; the authors report no sample intervals or multiple-comparison correction and have not analyzed why some tasks lose more.The KVTC variant uses one quantization range per component for the whole prompt, unlike the cited paper’s per-block choices, and the authors did not evaluate that variant.Source: Limitations · pp. 8–9
  • Each rate requires its own codec and training run; the tested one-codec range-rate variant falls behind single-rate codecs at the upper end, for an unexplained cause.Changing a trained codec’s step factor adjusts rate without training but costs bytes.Source: Limitations · p. 8
  • The codec is trained for a frozen model; joint pre-training might improve cache compression, but the authors lacked compute to test it.The frozen-model setup lets the codec work with a released model as is.Source: Limitations · p. 8
  • Codecs trained on 4K-token prompts were evaluated only up to 32K; beyond that, Figure 3 curves are computed, and at 32K the measured payloads make those curves overstate DCKV’s advantage.At 32K, DCKV’s measured payload per token is 3% above its RULER 4K value, while KVTC’s is 12–14% below it.Source: Limitations · p. 8
  • DCKV restores the full bf16 cache before decoding, so it does not reduce peak GPU memory or time per output token; its weights add about 5% to model size.Cross-device decoding has not been verified because conditional-model predictions must match bit for bit.Source: Limitations · pp. 8–9
  • With a stored cache, DCKV time to first token, including decompression, was 115 ms at 4K and 287 ms at 32K, versus 862 ms and 4,601 ms for uncompressed prefill.Compression alone was not timed; KVTC time on the shared node was unstable, measuring 6.8 s at 32K in this run and 0.6 s in an earlier run.Source: Limitations · p. 9

A Bias of the Uniform-Noise Rate Estimate

  • At λ = 5·104, 105 and 3·105, the training rate estimate exceeds held-out rounded-code rates by 45%, 37% and 13%; the authors have not tried training on rounded-code entropy.Many trained scales lie in the biased range: 76%, 65% and 36% of linear-codec components have βk < 0.2, as do 43–60% of DCKV components depending on the entropy model.Source: A Bias of the Uniform-Noise Rate… · p. 13

B Additional Ablations

The ablations show that DCKV’s additions reduce payload, while rate changes without retraining and a single codec for a range of rates have trade-offs. Other variants and alternative KVTC settings reveal limits in compression, quality, and model size.

Step multiplier instead of retraining

  • At 433 bytes per token on RULER 4K, halved-step DCKV uses 31% fewer bytes than KVTC at 16×, with a 0.28-point loss; on LongBench, its score does not differ significantly from the uncompressed model.At 441 bytes per token on RULER 32K, it likewise does not differ significantly from the uncompressed model and uses 19% fewer bytes than KVTC at 16×.Source: B Additional Ablations · p. 13, Table 4 · p. 14, Figure 4 · p. 14

The four codecs of the training run of DCKV

  • On RULER 4K, the four training-run codecs use 245, 232, 224, and 214 bytes per token, respectively, as the correction, trained conditional models, and both additions are introduced.Their scores are 0.96, 0.71, 0.59, and 0.56 points below the uncompressed model; the closed-form conditional fit also reaches 214 bytes with the codec’s original scores.Source: B Additional Ablations · pp. 13–14

One codec for a range of λ

  • For a codec trained over λ ∈ [5·10^4, 3·10^5], DCKV has a −31.0% BD-rate against the range codec without additions and a −69.9% BD-rate against KVTC.At λ = 3·10^5 on LongBench, range DCKV scores 1.01 points below the range codec without additions and 1.42 points below DCKV trained for that rate alone.Source: B Additional Ablations · p. 14, Table 5 · p. 15

Conditional models on a development codec

  • With a fixed development codec, token conditioning saves 6.2–6.5% of payload, and adding decoded neighbours saves 9.9–14.3%, without changing scores.Longer rANS streams save another 7%, but their decoding-time cost could not be measured reliably.Source: B Additional Ablations · p. 15, Table 6 · p. 15

Other variants

  • Offline, separate layer bases cost 13–27% more bits than a joint basis at equal KL divergence, while a 16-basis mixture uses 18–25% fewer bits but needs about 1.2 GB of weights.A learned hyperprior saved at most 1.2% of payload beyond the token-conditioned model on Qwen3.5-0.8B.Source: B Additional Ablations · p. 15

KVTC closer to the original

  • With exact dynamic programming and DEFLATE, the original-style KVTC used 30–53% fewer bits at equal KL divergence than greedy-allocation KVTC.It was not compared with the dynamic-programming KVTC used in the benchmarks.Source: B Additional Ablations · p. 16

C Training and Evaluation Details

The appendix reports the gap between entropy-model code-length estimates and stored payload, then explains how paired benchmark comparisons are summarized and qualified.

C Training and Evaluation Details

  • On RULER 4K, stored payload is 16% above expected code length for the linear codec and 35% above for DCKV.The expected lengths are 1,944 and 1,262 bits per token, versus stored payloads of 282 and 214 bytes per token, respectively.Source: C Training and Evaluation Details · p. 16

D Paired Comparisons

  • Paired comparisons report the A-minus-B point difference, task-level wins/ties/losses, and a two-sided sign test over differing tasks.With 13–18 tasks, 1–2-point differences between strong codecs are often not significant, and p-values were not corrected for multiple comparisons.Source: D Paired Comparisons · pp. 16–17, Table 8 · p. 17

E Benchmarks with Short Prompts

Short-prompt benchmarks provide limited evidence about cache compression: GPQA scores show no relation to payload, while IFEval largely does not test the codec and shows an unexplained positive offset.

E Benchmarks with Short Prompts

  • On GPQA Diamond, codec-run scores range from −6.1 to +1.5 points versus the uncompressed model, with no relation to payload.GPQA has 198 zero-shot questions and a mean prompt length of 255 tokens.Source: E Benchmarks with Short Prompts · p. 17, Table 9 · p. 18
  • IFEval does not test the codec because 99.3% of prompts fit within 132 uncompressed positions, yet all 29 codec and KVTC runs score 0.9–1.5 points above the reference.The authors attribute the offset to an unexplained part of the evaluation harness that runs a compressed cache; reference reruns matched the old scores.Source: E Benchmarks with Short Prompts · pp. 17–18, Figure 5 · p. 18

F Memory and Efficiency

Memory accounting includes recurrent state and codec weights, so payload savings do not translate uniformly into request-memory savings. Training curves also show distinct rate–KL trade-offs between mixed quantization and the noise proxy.

F Memory and Efficiency

  • Compression does not reduce peak GPU memory because decoding uses the restored bf16 cache; instead, peak memory grows by 182 MiB for DCKV and 102 MiB for the linear codec.These increases are about the size of the codec weights.Source: F Memory and Efficiency · p. 18
  • If codec weights are charged to one request, DCKV’s 176 MiB of weights mean it stores less than the uncompressed model only from about 15.4K prompt tokens onward.Its 156 MiB of extra weights over KVTC at 16× pay off after about 110 stored requests on RULER 4K or 36 on LongBench.Source: F Memory and Efficiency · p. 18, Table 12 · p. 19

G Additional Figures

  • At λ = 3·10^5, the mixed-quantization run keeps about the untrained codec’s KL divergence while reaching 44% of its rate.At 1B tokens, the noise-proxy run has a 3–5% lower rate than the mixed run at the same λ, but higher KL divergence.Source: G Additional Figures · p. 19, Figure 6 · p. 20

H Per-Task Results

This section reports per-task score changes against the uncompressed model for RULER 4K and LongBench. Figure 7 visualizes RULER 4K changes across codec configurations and orders tasks by mean change.

H Per-Task Results

  • Figure 7 compares DCKV at trained and halved steps with KVTC at 16× and 32× on RULER 4K payloads.Tasks are ordered by mean change across the four configurations, with the hardest at the bottom; UUID-valued needles are included.Source: Figure 7 · p. 20
  • Tables 13 and 14 report each task’s score against the uncompressed model, with paired mean changes for the included tasks.Table 13 covers 13 RULER 4K tasks; Table 14 reports the paired mean over 18 LongBench tasks.Source: H Per-Task Results · p. 20, Table 13 · p. 21, Table 14 · p. 21
  • Table 14 excludes LCC and RepoBench-P from its means because both score near zero for every configuration in the harness.On SAMSum and RepoBench-P, every codec and KVTC run scores above the uncompressed reference.Source: Table 14 · p. 21

Related papers

Loading 2610.36950v1…