Inference-Time Hyper-Scaling with KV Cache Compression

TL;DR

DMS (Dynamic Memory Sparsification) is a retrofitting recipe that teaches a pretrained Transformer LLM to evict tokens from its KV cache at an 8× compression ratio with only ~1K training steps, by learning a per-head Gumbel-sigmoid eviction policy combined with delayed eviction over a sliding window. Compressing the KV cache turns into a free inference-time-scaling lever: at a fixed memory-read or peak-memory budget the compressed model can generate longer or more parallel reasoning chains, and the paper shows this trade is favorable on AIME 24, GPQA Diamond, MATH 500, and LiveCodeBench across Qwen-R1 1.5B/7B/32B. On Qwen-R1 32B the gains average +9.1 / +7.6 / +9.6 points on AIME 24 / GPQA / LiveCodeBench across compute budgets. The released NVIDIA Qwen3-8B-DMS-8x checkpoint is the recipe applied to Qwen3-8B.

Key claims

Inference-time scaling in Transformer LLMs is bottlenecked by KV-cache size (memory reads dominate decode latency at long contexts), not by raw token count — so compressing the cache lets you generate more tokens at fixed runtime/memory, defining “inference-time hyper-scaling” [§1, §2.1].
DMS retrofits an LLM with a per-head binary eviction decision $\alpha_t$ trained via Gumbel-sigmoid stochastic reparametrization, repurposing an existing K/Q neuron so no new parameters are added [§3.2].
Delayed eviction (a token marked for eviction at step $t$ remains visible until $t+w$ , $w=256$ default) is the load-bearing design choice: immediate eviction “drastically deteriorates scores for every tested sliding window length”, whereas delayed eviction is stable across windows down to $w=16$ at 4× compression [§3.2, §5.3, Fig. 5].
DMS reaches 8× compression in ~1K training steps via logit distillation against the original LLM as teacher, vs. ~44K steps for DMC at the same target — and 4× / 8× compression are reachable within 300 / 700 steps respectively [§4 Training Schedule, §5.3].
On reasoning tasks at matched KV-cache memory-read budget, DMS dominates the Pareto frontier vs. vanilla LLMs and vs. training-free baselines (TOVA, H2O, Quest); averaged across budgets the gain over the original LLM is +11.5 on AIME 24, +2.3 on MATH 500, +5.5 on GPQA Diamond, +8.3 on LiveCodeBench [Fig. 1, Fig. 3, Appendix E Table 5].
Pareto-optimal configurations combine sequential (longer chains) and parallel (more chains) scaling — neither axis alone is enough [§5.1].
On Llama 3.2 1B Instruct at CR=2/3/4, DMS is the most robust method across GSM8K, MMLU, HellaSwag, NIAH, VT; it surpasses the uncompressed vanilla LLM on long-context tasks (NIAH 96.4 → 97.8 at CR=2, VT 55.8 → 63.2 at CR=2), conjectured to fight over-squashing / lost-in-the-middle [Table 1, §5.2].
DMS exhibits non-uniform compression both along sequence length (early tokens compressed less than the target, beyond 10K tokens more aggressively than target) and along depth (early layers compressed less than later layers) [§5.3, Fig. 6].
The recipe is implementation-friendly: the eviction mask never needs to be materialized, the sparsified cache is compatible with PagedAttention, and the method does not introduce new read/write operations on the KV cache [§3.2 Performance Considerations, §3.3].

Method

DMS treats KV-cache eviction as a learned, per-head binary decision $\alpha_t \in \{0,1\}$ trained with Gumbel-sigmoid reparametrization (low temperature + logit offset $\gamma$ to start training near $\alpha_t=0$ and avoid early loss spikes). The decision predicted at step $t$ does not take effect until step $t+w$ — until then the token remains attendable inside a sliding window of size $w$ . Implementation-wise, instead of adding parameters, one neuron is repurposed from the K or Q projection per head to produce the eviction logit; the eviction decisions are passed as a vector and used to build an additive attention mask that is never materialized (compatible with FlashAttention-style kernels and PagedAttention).

Training is logit distillation from the un-retrofitted LLM as teacher, plus a one-sided L1 loss that pushes the average predicted compression rate toward a linearly-annealed target up to the final CR. Each head can land at a different compression ratio. The training schedule is linear in target CR: 100 steps per unit of CR (so 800 steps for 8×).

Results

Reasoning hyper-scaling (Qwen-R1 32B, distilled from DeepSeek-R1): averaged across compute budgets, DMS adds +9.1 on AIME 24, +7.6 on GPQA Diamond, +9.6 on LiveCodeBench over the un-retrofitted model at the same KV-cache memory-read budget [Abstract, §5.1, Fig. 1].
Pareto frontier: DMS dominates Quest (best training-free runtime baseline) and TOVA (best training-free memory baseline) on both memory-reads and peak-tokens axes across 1.5B/7B/32B [§5.1, Figs. 3–4].
General-purpose retention (Llama 3.2 1B Instruct): at CR=4 DMS still scores 39.9 on GSM8K vs. vanilla 47.0; preserves NIAH at 95.8 vs. vanilla 96.4; vs. H2O collapsing to 14.7 on GSM8K and 13.4 on NIAH at CR=4 [Table 1].
Long-context tasks: at CR=2, DMS exceeds vanilla on NIAH (97.8 vs. 96.4) and VT (63.2 vs. 55.8), interpreted as relief from dense-attention pathologies on long sequences [Table 1, §5.2].
Sample efficiency: reaches CR=4 within 300 steps and CR=8 within 700 steps, vs. DMC’s reported 44K steps at CR=8 [§5.3].

Why it’s interesting

For any model serving stack the team runs — text or otherwise — DMS is a concrete recipe that turns reasoning-time KV-cache pressure into a free knob: ~1K distillation steps to get 8× compression with no new parameters, no inference-time op changes, and on Qwen-R1 32B actually better AIME / GPQA / LiveCodeBench scores at fixed memory. The “delayed eviction” design choice is the interesting transferable idea — separating when an eviction is decided from when it executes gives the model time to integrate information from soon-to-be-evicted tokens, and is the empirical difference between this working and not. The released NVIDIA Qwen3-8B-DMS-8x checkpoint (the artifact that triggered the Slack message) is the recipe applied at 8B; the underlying paper is what is reproducible.