Summary of ABRIDGEKV: Compressing Long Reasoning into the KV Cache of Learned Pause Tokens

Summary (Overview)

• Core contribution: ABRIDGEKV trains a reasoning model to compress its own chain-of-thought (CoT) into the KV cache entries of learned "pause tokens" — special input tokens with trained embeddings that produce no output — achieving a 4× reduction in persistent KV cache size without accuracy loss.

• Key innovation: Instead of evicting cache entries (which degrades performance) or writing textual summaries (which requires expensive annotation pipelines), ABRIDGEKV uses the next-token loss to teach pause tokens what information to preserve, enabling online, end-to-end learned compression during decoding.

• Three critical design choices: (1) code rows take positions inside the block they replace (not after it), (2) cache budget is allocated across layers according to measured value of each layer's entries, and (3) compression operates on relatively large blocks (512 tokens).

• Results: At a 32K-token cap, ABRIDGEKV matches full-KV accuracy (43.0 vs 43.8 avg@8 on 446 competition problems) while using 3.54× less resident KV memory per sequence and finishing problems 3.2× faster on a saturated H100. The strongest eviction baseline (R-KV) generates 40–69% more tokens, and the closest trained compressor (LightThinker) loses 10 accuracy points.

• Information flow verified: Controlled experiments show that specific facts planted in a block are stored in that block's code rows and retrieved from them up to 8 compressions later (73% recall vs 98% for full KV), confirming that later tokens attend to code rows as a direct substitute for raw content.

Introduction and Theoretical Foundation

Background and Motivation

Test-time scaling improves reasoning accuracy by generating longer chains of thought, but decoding long CoTs is memory-bound: the KV cache grows linearly with generated tokens, and its size (not model weights) bounds batch size and throughput. Unlike long-prompt settings, this cache is produced by the model's own decoding and grows with exactly the test-time compute being spent.

Limitations of KV Eviction

KV eviction keeps a fixed number of cache entries and discards the rest, but has two fundamental limitations:

1. The model is never trained on a cache with entries removed, so it cannot gracefully handle missing information.
2. Selection is made from past attention, whereas a CoT reuses intermediate results at unpredictable later points; an evicted entry leaves no representation behind.

Empirically, eviction failure appears as longer generations: keeping a quarter of the cache, simple eviction policies let 38–58% of samples run to the 32K limit, and R-KV generates 40–69% more tokens than the uncompressed model.

Learned Compression for CoTs

Learned compression (e.g., Mu et al., 2023; Ge et al., 2024) was developed for fixed inputs compressed once before decoding. Compressing a CoT differs fundamentally:

"The model must compress its own output online, dozens of times per response, and condition on the compressed states thousands of tokens later, so each summary has to be written without knowing which of its contents the later reasoning will need."

Prior textual-summary approaches (e.g., MEMENTO) require sophisticated annotation pipelines where a frontier LLM scores block boundaries, writes summaries, and grades them. ABRIDGEKV avoids this entirely by using a fixed set of tokens as summaries, letting their KV entries learn from the next-token loss what information to preserve.

Theoretical Foundation

The key insight is that a small, fixed set of KV entries can replace the information carried by an explicit textual summary if three conditions are met:

1. Positioning: Code rows must take positions inside the block they replace, preserving relative geometry to later queries.
2. Layer allocation: The cache budget must be allocated across layers according to measured value, not uniformly.
3. Block size: Compression must operate on relatively large blocks (512 tokens), not paragraph-level granularity.

Methodology

Block Layout and Attention Mask

The prompt (P tokens) is kept raw and never compressed. The response is generated in blocks of $C = 512$ tokens; block $b$ has content $x_{b,0}, \ldots, x_{b,C-1}$ and $K_{\max} = 256$ code slots $p_{b,0}, \ldots, p_{b,K_{\max}-1}$ (pause tokens with learned input embeddings).

Visibility rules: A content query $x_{b,t}$ attends to the prompt, codes of every earlier block $b' < b$, and its own block causally. A code query $p_{b,k}$ attends to the prompt, codes of earlier blocks, all of its block, and other codes of its block.

Position Assignment

Content keeps its uncompressed numbering, and code rows sit inside the block they summarize using bit-reversed ordering:

$$\mathrm{pos}(x_{b,t}) = P + bC + t, \qquad \mathrm{pos}(p_{b,k}) = P + bC + \mathrm{bitrev}_{K_{\max}}(k) \frac{C}{K_{\max}} \tag{1}$$

where $\mathrm{bitrev}_{K_{\max}}(k)$ reverses the $\log_2 K_{\max}$ bits of $k$ (for $K_{\max}=8$: $0, 4, 2, 6, 1, 5, 3, 7$). Three properties follow:

• (i) No position inflation: A response of $T$ tokens occupies exactly $T$ positions, staying within the base model's trained range.
• (ii) Nested, K-independent quotas: Keeping the first $q$ slots in bit-reversed order gives an even grid of stride $C/q$; quota $q$ is a prefix of quota $2q$.
• (iii) Stable relative geometry: Later queries see each code at roughly the distance the content it replaces would have had.

Training Objective

The model is trained with the ordinary next-token loss on every response token, computed while history is held as codes:

$$\mathcal{L}_{\mathrm{CE}} = -\frac{1}{|R|} \sum_{t \in R} \log p_\theta\left(x_t \mid \text{prompt}, \text{codes}_{< b(t)}, x_{b(t), < t}\right) \tag{2}$$

The whole response (code rows included) is one forward pass under the block mask, so the loss of later blocks flows back into the pause tokens that wrote the codes they read. No auxiliary terms are added.

Per-Layer Budget Allocation

The budget is a per-layer quota $q_\ell \leq K_{\max}$: layer $\ell$ keeps $q_\ell$ of the 256 code rows of every completed block. The allocation procedure:

1. Train a model with uniform-128 schedule.
2. Measure each layer's value: halve one layer's quota (128→64) at evaluation time, record the increase in held-out cross-entropy, repeat for each of the 36 layers.
3. Allocate: the 6 most expensive layers get 256 rows, the 12 cheapest get 64, the other 18 keep 128 (preserving the mean: $6 \cdot 256 + 18 \cdot 128 + 12 \cdot 64 = 36 \cdot 128$), and train a new model on this schedule.

Empirical Validation / Results

Main Results (Table 1, 32K cap)

†Eviction baselines run on a 60-problem subset.

Key Findings

Accuracy and generation length: ABRIDGEKV's length distribution is statistically indistinguishable from full KV's (Kolmogorov–Smirnov distance 0.024 vs 0.017 for a full-KV re-run; significance threshold 0.032). R-KV at budget 4096 has KS distance 0.33 from full KV, with 32% of samples hitting the cap.

Serving efficiency: ABRIDGEKV's resident KV per sequence is 3.54× smaller than full KV's. It decodes at 1.05×, 1.38×, and 2.43× full KV's throughput at batch 1, 8, and 32. A running ABRIDGEKV sequence holds 369 MiB of KV vs 1,277 MiB for full KV. Saturated, ABRIDGEKV finishes a problem in 0.31× full KV's time.

Ablation results:

• Appended positions (instead of in-block): −6.2 accuracy points, the largest measured effect. The loss is not from position inflation (restricting to problems within trained range leaves gap unchanged at 6.2).
• Uniform allocation (instead of per-layer): −2.3 points; per-layer recovers 1.6 points.
• Block size at fixed 4:1 ratio: accuracy rises with block size (40.7, 41.5, 42.5 avg@8 for C=256, 512, 1024), while time to finish stays constant.

Comparison with eviction: R-KV at budget 4096 matches ABRIDGEKV's accuracy but generates 40% more tokens and takes longer than full KV to finish a problem. At budget 2048, it falls 3.8 points below ABRIDGEKV's 8:1 arm.

Information Flow Analysis

A controlled experiment planted a random four-digit value in block $b$ and asked for it $k$ blocks later. Over 93 edited solutions (372 cases), the value was recalled in 73% of cases (vs 98% for full KV), with no decline from 1 to 8 compressions, and zero recall when block $b$'s code rows are removed. The recalling tokens attend to block $b$'s code rows far more than to any other block's.

Theoretical and Practical Implications

Layer Allocation Insights

The measured cost of halving code rows concentrates in a few layers (the first layer and a band in the middle of the network), while most early and late layers can lose half their rows at almost no cost. Critically:

"A proxy from the uncompressed model points the wrong way" — schedules chosen from full-KV attention raise exactly the early layers the measurement finds cheap.

Single-layer costs are not additive: halving every layer at once raises held-out CE by twice the sum of single-layer increases, motivating the use of only three quota levels and cutting only a third of the layers.

Practical Deployment

ABRIDGEKV integrates with standard paged serving engines (vLLM) with a simple patch: a per-layer schedule is stored as one KV cache group per quota level, so rows a layer does not keep occupy no memory. The served model matches the trained one (per-token CE agreement to $4 \times 10^{-4}$).

Broader Applications

1. Coding agents: ABRIDGEKV could complement existing compaction mechanisms (which trigger at 250K–1M tokens) by reducing KV memory growth throughout the long interval before compaction.

2. Multimodal reasoning: Learned KV compression may be especially promising for modalities difficult to summarize faithfully in text (e.g., computer-use agents reasoning over screenshots), where KV space provides a more natural compression medium.

3. Recursive compression: Compressing earlier code rows into new code rows could eventually replace textual compaction altogether.

Conclusion

ABRIDGEKV demonstrates that a reasoning model can learn to compress its own chain of thought into the KV cache of learned pause tokens, achieving a 4× reduction in persistent cache size with no accuracy loss and 3.2× faster problem completion. The three key design choices — in-block code positions, measured per-layer budget allocation, and large blocks — are each essential, with in-block positioning being the largest single factor.

Limitations: Results start from a base model fine-tuned with the authors' own SFT (not directly on RL-tuned checkpoints); compression ratio is fixed per block regardless of information content; and the quota schedule assumes single-layer measurements are independent.

Future directions: Larger blocks (accuracy still rising at C=1024), recursive ABRIDGEKV for bounded cache without textual rewriting, multimodal and computer-use agents, and interaction-aware quota allocation methods.