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

What it claims

DCKV is a learned codec for the stored prompt KV cache of Qwen3.5-2B. Its transform, per-component quantization steps and entropy models are trained jointly against the frozen model, using code length plus output KL divergence. Over three rates it needs 72% fewer bytes than KVTC at equal score on RULER 4K and 62% fewer on LongBench. The large ratios over bf16 shrink much less at the request level. At 4K to 11K tokens, the fixed 9.84 MiB Gated DeltaNet recurrent state, which DCKV keeps exact, is about three quarters of a stored request's memory. At 32K tokens, where DCKV's payload is over a third of request memory, DCKV scores 6.5 points above KVTC at about the same request memory.

How it works

Qwen3.5-2B has 6 full-attention layers out of 24. Its KV cache is 12,288 bytes per token in bf16, which is 48 MiB at 4K tokens and 1.5 GiB at 128K. The first 4 and last 128 positions stay in bf16 (1.55 MiB). For each other token, keys (taken before RoPE, which is reapplied after decoding) and values of all six layers are stacked into one 6,144-dim row.

• Transform. One shared basis covers keys, values and all layers, unlike KVTC's separate key and value bases. It projects each row to 2,048 coefficients. The basis starts from PCA and adds a small nonlinear correction: one hidden layer of 256 GELU units, initialised to zero. Components whose scale falls below 0.05 quantization steps are dropped.
• Coding. Codes are rANS-coded on GPU via the torch-ans library under Laplace distributions. Codes beyond ±127 are escaped at 64 bits each.
• Conditional entropy models. Even positions predict their code means from the prompt token's input embedding. Odd positions also use the decoded coefficients of both neighbours, a 1-D checkerboard order. Any cache length therefore decodes in two parallel passes. Token ids are assumed available from the reused prompt and are not counted; they would add 1 to 5%.
• Training. Each sample is a 4,096-token prompt and a 256-token continuation. The loss is bits plus λ times the mean KL between continuations over the exact and reconstructed caches. Rate is computed with additive noise, and reconstruction uses rounded values with a straight-through gradient. The authors report that pure noise overestimates entropy by 1.47× to 5.66× at low scales. One λ gives one rate. Rescaling steps or training one codec over a λ range costs bytes or quality.

Evidence

The conditional models plus the correction remove 12 to 16% of bytes at equal quality. Fitting the conditional models in closed form to an already-trained codec gives a similar saving on RULER 4K without retraining.

The stored payload is 5% (linear codec) and 17% (DCKV) above the code length the models predict.

The KVTC baseline is a re-implementation, since no official code was found. It builds on a public proof of concept with the paper's dynamic-programming bit allocation added. At 16×, greedy allocation lost 13.0 points on RULER 4K against 0.53 for dynamic programming. KVTC at 16×, 24× and 32× gives 627, 410 and 304 bytes per token.

Benchmarks are RULER 4K and 32K (13 tasks each), LongBench (18 of 21 tasks) and 5-shot MMLU. Three LongBench tasks were excluded. With all 21, losses are 0.3 to 1.0 points smaller and the method order is unchanged.

How to use it

Codec weights are 176 MiB fp32, about 5% of the 2B model, against 96 MiB for the linear codec and 10 to 20 MiB for KVTC. Peak GPU memory rises by 182 MiB. Training ran on FineWeb-Edu for 1B tokens with AdamW on two A100s, taking 33.9 hours for a loop of four codecs. Calibration used 48 documents (FineWeb and OpenR1-Math). The fetched text links torch-ans but does not mention a DCKV code or weights release.

Limits

All evaluations encode and decode in the same process. Table selection requires encoder and decoder to compute identical predictions, which the authors flag in their Limitations. Results cover one model, bf16, greedy decoding, one request at a time and Hugging Face Transformers on one A100 40GB. Each rate needs its own trained codec. The fetched text ends partway through the experimental setup, so per-rate result tables are not shown.