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

> DCKV is the first end-to-end trainable KV-cache codec for hybrid language models, achieving 55-60x compression with minimal quality loss via joint transform, quantizer, and entropy model training.

- **Source:** [arXiv](https://arxiv.org/abs/2610.36950)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/JdWAnp
- **Whiteboard:** https://picx.dev/p/JdWAnp/image

## Summary

# 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)

| Method | Payload (B/tok) | RULER 4K | LongBench | RULER 32K | MMLU |
|--------|-----------------|-----------|-----------|------------|------|
| **Uncompressed** | 12,288 | 94.54 | 42.66 | 94.25 | 59.34 |
| KVTC 16× | 627 | -0.53 | -0.31 | -0.02 | -0.62 |
| KVTC 24× | 410 | -1.87 | -1.97 | -2.45 | -0.99 |
| KVTC 32× | 304 | -3.88 | -4.17 | -6.94 | -0.93 |
| **DCKV λ=3·10⁵** | 214 | **-0.56** | -1.75 | **-0.45** | -1.10 |
| **DCKV λ=10⁵** | 113 | **-1.69** | **-3.55** | **-2.15** | -0.87 |
| DCKV λ=5·10⁴ | 74 | -4.60 | -5.51 | -6.45 | -0.65 |

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:

| Configuration | RULER 4K (B/tok) | RULER 4K (Δscore) |
|---------------|-------------------|---------------------|
| KVTC 24× | 410 | -1.87 |
| + entropy coding (sep. bases) | 392 | -0.71 |
| + joint basis | 364 | -1.30 |
| + trained transform (linear) | 282 | -0.60 |
| + nonlinear correction | 232 | -0.71 |
| + trained cond. models | 224 | -0.59 |
| **DCKV (all)** | **214** | **-0.56** |

### 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

---

_Markdown view of https://picx.dev/p/JdWAnp, served by PicX — AI-generated visual whiteboard summaries of research papers._