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

Summary (Overview)

• DCKV is the first end-to-end trainable codec for KV caches in hybrid language models, jointly optimizing the transform, quantizer, and conditional entropy models against a frozen language model.
• On Qwen3.5-2B, DCKV achieves 55-60× compression (205-221 bytes/token vs. 12,288 in bf16) with a loss of at most 1.8 points on long-context benchmarks, and >100× compression with at most 3.6 points loss.
• DCKV requires 72% fewer bytes than KVTC at equal score on RULER 4K and 62% fewer on LongBench (BD-rate across three rates).
• The codec uses a joint transform for keys and values across all attention layers, with a small nonlinear correction, learned per-component quantization steps, and GPU rANS entropy coding.
• Conditional entropy models conditioned on prompt tokens and neighboring decoded positions remove an additional 12-16% of bytes at the same quality.

Introduction and Theoretical Foundation

Background and Motivation

Hybrid language models interleave a few full-attention layers with linear-attention or state-space layers. In Qwen3.5-2B, only 6 of 24 layers use full attention, while 18 are Gated DeltaNet layers with a fixed-size recurrent state (9.84 MiB per request). Although the KV cache is small (12,288 bytes/token in bf16), it is the only state component that grows with the prompt:

• 48 MiB at 4K tokens
• 1.5 GiB at 128K tokens

Reusing prompt caches requires storing them outside GPU memory and moving them between devices. Since such caches are written once and read one or more times, a codec can afford entropy coding—the key metric being stored bytes at a given answer quality.

Theoretical Foundation

The work builds on two research threads:

1. KV Cache Transform Coding (KVTC) [43]: Decorrelates keys and values with PCA, allocates bits by dynamic programming, and applies a general-purpose lossless coder. Its transform minimizes reconstruction error but does not see the model's output.

2. Learned Image/Video Codecs [3, 4, 5, 18]: Train transforms, quantizers, and entropy models on rate–distortion objectives, conditioning entropy models on information available at the decoder.

The key insight: for a language model, the natural distortion metric is the change in output distribution, not reconstruction error. DCKV is the first KV-cache codec to train all components end-to-end against the frozen model, using both code length and KL divergence of model outputs.

Methodology

What is Compressed

DCKV compresses the attention cache of a prefilled prompt once and restores it completely before decoding. The first 4 and last 128 cache positions remain in bf16; other positions form the compressed region $T$. For token $t \in T$, keys and values from all $L=6$ attention layers are stacked into a single row $x_t \in \mathbb{R}^d$ with $d = 2LH d_h = 6144$ (for $H=2$ key-value heads of dimension $d_h = 256$). Keys are taken before rotary position embedding (RoPE) and re-applied after decompression to maintain position independence.

Transform and Quantization

The encoder maps a row to $n = 2048$ coefficients, and the decoder reconstructs:

$$y_t = (x_t - m)W_e + g_e(x_t - m), \quad q_{t,k} = \text{round}(y_{t,k}/\Delta_k), \quad \tilde{y}_t = q_t \odot \Delta$$ $$\hat{x}_t = m + \tilde{y}_t W_d + g_d(\tilde{y}_t)$$

where:

• $m$ is the mean row of calibration data
• $W_e \in \mathbb{R}^{d \times n}$ and $W_d \in \mathbb{R}^{n \times d}$ are transform matrices
• $\Delta \in \mathbb{R}^n_{>0}$ holds one quantization step per component
• $g_e, g_d$ are nonlinear corrections (one hidden layer, 256 units, GELU) whose last layers start at zero

The matrices initialize from the leading $n$ principal directions $V$ ($W_e = V$, $W_d = V^\top$), making the untrained codec equivalent to PCA with uniform quantization.

Entropy Coding

Each code $q_{t,k}$ is written with rANS using a static zero-mean Laplace distribution per component:

$$P(q_{t,k}) = F_{\beta_k}\left(q_{t,k} + \frac{1}{2}\right) - F_{\beta_k}\left(q_{t,k} - \frac{1}{2}\right)$$

where $F_\beta$ is the Laplace CDF with scale $\beta_k = b_k/\Delta_k$ (initialized from mean absolute coefficients). Codes beyond ±127 are escaped and stored at 64 bits each.

Conditional Entropy Models

Token-conditioned model: Predicts coefficients from the input embedding $e(w_t)$ of the token:

$$\mu_t = e(w_t)AB + c$$

with $A \in \mathbb{R}^{2048 \times r}$, $B \in \mathbb{R}^{r \times n}$. The code is written as the residual $q_{t,k} - \text{round}(\mu_{t,k}/\Delta_k)$ using one of 576 precomputed tables (0.84 MiB).

Checkerboard context model: Positions split into even/odd (1D checkerboard order). Even positions coded first with token-conditioned model alone; odd positions coded second using a linear map of rank 1024 that predicts means from both neighbors' decoded coefficients and the token embedding. Each pass is parallel over tokens.

Rate–Distortion Training

The loss is:

$$\mathcal{L} = \frac{R}{|T|} + \lambda \frac{1}{C} \sum_{i=1}^{C} \text{KL}(p_i \| \hat{p}_i)$$

where $R$ is the rate in bits under conditional models, KL divergence is in nats, and the model reads a continuation of $C=256$ tokens twice (exact vs. reconstructed cache).

Mixed quantization is used: rate computed on noisy coefficients, reconstruction from rounded coefficients with straight-through gradient. This is important because components with $\beta_k \leq 0.2$ have >90% zero coefficients, and noise-based rate overestimates entropy by factors of 1.47–5.66 at low scales.

Empirical Validation / Results

Main Results (Table 1)

Key findings:

• At λ=3·10⁵: 56-60× fewer bytes than bf16 with ≤1.75 points loss
• DCKV at λ=10⁵ scores higher than KVTC at 32× while storing 2.2–2.7× fewer bytes
• BD-rate: DCKV needs 72% fewer bytes than KVTC on RULER 4K, 62% fewer on LongBench

Ablation Study (Table 2)

Progressive gains from KVTC to DCKV:

Request Memory at Scale (Figure 3)

The recurrent state dominates request memory at moderate lengths:

• 81% of request memory on RULER 4K
• 73% on LongBench (10.8K tokens)
• 54% on RULER 32K

At 32K tokens, DCKV scores 6.5 points higher than KVTC at approximately the same request memory.

Theoretical and Practical Implications

Significance

1. First end-to-end trained KV-cache codec: DCKV demonstrates that jointly training transform, quantizer, and entropy models against a frozen language model is feasible and dramatically outperforms hand-designed codecs.

2. Output-aware distortion: Using KL divergence of model outputs as distortion (rather than reconstruction error) proves crucial—it enables 3+ points better scores at similar payloads, and this advantage grows at stronger compression.

3. Conditional entropy models for KV caches: The paper shows that conditioning on prompt tokens (available at decode time) and neighboring positions yields substantial rate savings (12-16%) that are complementary to transform improvements.

Practical Implications

• Cache serving systems: DCKV enables storing prompts in roughly 1/55th the space of bf16 while maintaining answer quality—critical for long-context serving where cache transfer dominates latency.
• Hybrid models: The method targets the only growing state component, complementing recurrent state compression methods.
• Deployment: Codec weights (176 MiB fp32, ~5% of the 2B model) are loaded once and shared across requests; peak GPU memory grows by 182 MiB.
• Latency: Time-to-first-token with stored cache is 115 ms at 4K tokens and 287 ms at 32K, versus 862 and 4,601 ms for prefill of the uncompressed model.

Conclusion

Main Takeaways

DCKV trains transform, quantizer, and conditional entropy models end-to-end against a frozen language model, storing the KV cache at ~210 bytes/token (2.9× fewer than KVTC at equal loss on RULER 4K). Each component contributes: entropy coding over bit allocation, joint basis over separate bases, rate–distortion training over untrained transforms, and conditional models over static entropy models.

Future Directions

1. Generalization: Whether results carry over to other hybrid model families
2. Multi-rate codecs: Training one codec for a range of rates without quality loss at the upper end
3. Cross-device determinism: Making decoding produce identical results on different devices (e.g., integer inference as in DCVC-RT)
4. Co-training with pre-training: A model pre-trained together with the codec, with cache code length in the loss, may compress even better

Key Limitations

• Evaluated on only one model (Qwen3.5-2B) with 13–18 tasks, limiting statistical power
• Each rate requires its own separately trained codec
• DCKV restores full bf16 cache before decoding—does not reduce peak GPU memory or per-token latency
• All evaluations encode and decode in the same process; cross-device behavior unverified
• Some tasks (variable tracking, passage retrieval) lose disproportionately more than others, and the cause is not yet understood