We study a controlled partial-KV decoding setting in which exact unnormalized softmax contributions are computed for sink/tail anchors and a retrieved token set, while the remaining prefill tokens are represented by a residual estimate. We focus on the accounting rule after the query-dependent exact support has been selected, and use exhaustive Top-K only as an oracle selector, not as a deployable retrieval system. The proposed rule leaves the backbone language model and the exact-branch KV tensors unchanged. It builds fixed-size summary states (S,u) from learned positive feature maps φ, subtracts retrieved-token feature contributions to keep the exact and residual sets non-overlapping, and merges the estimated residual numerator and denominator with the exact branch under one normalization. At a 1% exact-support budget, our residual-completion method improves over the selection-only Top-K baseline on RULER and BABILong across frozen 1B and 3B Llama-3.2-Instruct backbones at all reported context lengths. In the 0.5-4% exact-support budget sweeps, this trend largely persists. On LongBench, summarization results are mostly favorable, while multi-document QA is mixed. Attention-output diagnostics support retrieved-token subtraction as the partition-consistent accounting rule, while indicating that the main remaining error is imperfect learned-φ approximation of the unretrieved residual mass.
Figures & tables
Llama-3.2-1B-Instruct
Llama-3.2-3B-Instruct
Task
Method
0.5%
1%
2%
4%
0.5%
1%
2%
4%
GovReport
Full exact
0.280
0.331
Top-K
0.266
0.271
0.275
0.276
0.323
0.324
0.320
0.321
Top-K+ϕ
0.284
0.283
0.285
0.281
0.332
0.333
0.329
0.327
Top-K+ϕ ( NoSub )
0.288
0.285
0.282
0.278
0.336
0.330
0.331
0.325
QMSum
Full exact
0.227
0.242
Table 1: LongBench exact-support budget-fraction sweep: summarization (GovReport and QMSum; ROUGE-L) and multi-document QA (MuSiQue and HotpotQA; F1). Sink/tail anchors count toward each percentage budget. All Top-K+ϕ rows use the 64k-trained ϕ . Bold marks the best partial-KV method within each task/backbone/budget block, including ties. “Full exact” is the unmodified backbone LM. “ Top-K+ϕ ( NoSub )” intentionally omits retrieved-token subtraction.
Llama-3.2-1B-Instruct
Llama-3.2-3B-Instruct
Benchmark
Method
4k
16k
64k
4k
16k
64k
RULER
Full exact
0.803
0.665
0.581
0.918
0.816
0.721
Top-K
0.732
0.606
0.541
0.878
0.787
0.688
Top-K+ϕ
0.753
0.617
0.548
0.889
0.814
0.711
Top-K+ϕ ( NoSub )
0.754
0.617
0.543
0.886
0.809
0.708
BABILong
Full exact
0.331
0.244
0.140
0.474
0.357
0.215
Table 2: RULER and BABILong results at a 1% total exact-support budget. RULER reports the official aggregate score. BABILong reports average accuracy over QA1–QA5. Bold marks the best partial-KV method within each benchmark/length/backbone block, including ties.
Figure 1: Head-wise completion gain versus mid-normalized attention entropy on Wikitext-64k. Each point corresponds to one layer–query-head. Gain is eTop-K−eTop-K+ϕ , measured by relative ℓ1 attention-output error against full attention. Positive values indicate that subtractive residual completion reduces error relative to selection-only Top-K .
Corpus
Length
K
Top-K↓
NoSub↓
Sub-ϕ↓
ExactNoSub↓
mean log- Z err.
Wikitext
4k
21
0.267
0.238
0.191
0.156
−1.615
CWE
4k
21
0.210
0.290
0.173
0.188
−2.378
Wikitext
64k
636
0.160
0.147
0.103
0.180
−1.843
CWE
64k
636
0.184
0.188
0.151
0.183
0.485
Table 3: Residual-accounting diagnostics: mean relative ℓ1 attention-output error and residual log-partition calibration. ExactNoSub uses the true whole-mid contribution without subtraction, double-counts K(q) , and is diagnostic rather than a valid partition. The final column reports mean residual log-partition error logZ^R−logZR ; negative values indicate underestimation. Bold marks the lowest approximate method among Top-K , NoSub , and Sub-ϕ .
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2: Exact softmax kernels and learned ϕ scores for a representative Llama-3.2-1B-Instruct head (layer 0, head 31) on Wikitext-64k using the 64k-trained maps. The diagnostic illustrates that learned ϕ scores are used only in the residual estimator.
Backbone
Steps / hardware
Wall-clock
Llama-3.2-1B-Instruct
200k steps on 8 × H200
19 h
Llama-3.2-3B-Instruct
200k steps on 8 × H200
35 h
Appendix
Table 4: Training cost for the 64k-trained ϕ modules used in the experiments.
Llama-3.2-1B-Instruct
Llama-3.2-3B-Instruct
Length
Method
0.5%
1%
2%
4%
0.5%
1%
2%
4%
4k
Full exact
0.803
0.918
Top-K
0.547
0.732
0.762
0.784
0.698
0.878
0.898
0.912
Top-K+ϕ
0.603
0.753
0.771
0.787
0.750
0.889
0.904
0.914
Top-K+ϕ ( NoSub )
0.605
0.754
0.769
0.784
0.745
0.886
0.902
0.913
8k
Full exact
0.743
0.870
Appendix
Table 5: RULER exact-support budget fraction sweep with one 64k-trained ϕ set per backbone LM. Sink/tail anchors count toward each percentage budget.
Llama-3.2-1B-Instruct
Llama-3.2-3B-Instruct
Length
Method
0.5%
1%
2%
4%
0.5%
1%
2%
4%
4k
Full exact
0.331
0.474
Top-K
0.263
0.308
0.306
0.310
0.401
0.488
0.498
0.496
Top-K+ϕ
0.298
0.324
0.321
0.319
0.431
0.497
0.494
0.495
Top-K+ϕ ( NoSub )
0.293
0.322
0.322
0.320
0.434
0.493
0.496
0.497
8k
Full exact
0.289
0.406
Appendix
Table 6: BABILong exact-support budget fraction sweep with one 64k-trained ϕ set per backbone LM. Sink/tail anchors count toward each percentage budget.
RULER @1%
4k
8k
16k
Full exact
0.803
0.743
0.665
Top-K
0.732
0.692
0.606
Top-K+ϕ (4k-trained)
0.758
0.017
0.014
Top-K+ϕ (8k-trained)
0.748
0.718
0.016
Top-K+ϕ (16k-trained)
0.759
0.704
0.634
Appendix
Table 7: Cross-length mismatch diagnostic on RULER at 1% total exact-support budget for Llama-3.2-1B-Instruct. Rows differ only in the maximum length used to train ϕ ; evaluation is performed at 4k, 8k, and 16k.
Figure 3: Residual denominator calibration for Llama-3.2-1B-Instruct. Histograms of logZ^R−logZR for Sub-ϕ on Wikitext and CWE at 4k and 64k. Negative values indicate residual-mass underestimation.
Figure 4: Head-wise Sub-ϕ – NoSub difference for Llama-3.2-1B-Instruct. Histograms of Δ=eSub-ϕ−eNoSub across heads. Negative values favor Sub-ϕ .
Method
RULER 4k
BABILong 4k
Full exact
0.803
0.331
Top-K
0.732
0.308
dϕ=32
0.758
0.335
dϕ=64
0.758
0.339
dϕ=128
0.753
0.334
Appendix
Table 8: Feature-dimension ablation with 4k-trained ϕ on Llama-3.2-1B-Instruct at a 1% total exact-support budget. Sink/tail anchors count toward the 1% budget.
Method
RULER 4k
BABILong 4k
Full exact
0.803
0.331
Top-K
0.732
0.308
Top-K+ϕ
0.758
0.339
Hedgehog ϕ + our loss
0.016
0.079
Our ϕ + LoLCATs loss
0.018
0.079
Our ϕ + KL only
0.082
0.087
Appendix
Table 9: Loss/module ablation with 4k-trained ϕ on Llama-3.2-1B-Instruct at a 1% total exact-support budget. Sink/tail anchors count toward the 1% budget. Hedgehog and LoLCATs refer to diagnostic swaps inside our residual-completion pipeline, not evaluations of those methods in their native settings.
Figure 5: Mid-region kernel fidelity for Top-K+ϕ completion on Wikitext-4k using Llama-3.2-1B-Instruct (layer 0, head 31). Panel (a) changes the feature map parameterization while keeping the loss fixed to Lours=λKLLKL+(1−λKL)Laux ( Equation 31 ); panels (b)–(d) keep our feature map fixed and vary the loss ( Llolcats , KL only, or aux only). Points closer to the diagonal indicate more faithful approximation of unretrieved mid-region kernels.
Context
K
Full exact
Top-K
Top-K+ϕ
Top-K+ϕ ( NoSub )
4k
21
6.69 / 235.91
50.51 / 1594.37
76.85 / 2421.98
67.83 / 2142.50
16k
144
6.41 / 340.92
49.38 / 1673.12
84.41 / 2798.91
67.87 / 2288.69
64k
636
6.84 / 1430.63
52.04 / 2834.26
84.58 / 3996.32
66.12 / 3418.48
Appendix
Table 10: Prototype latency disclosure for Llama-3.2-1B-Instruct. Each method column reports latency as per-step / end-to-end in milliseconds. The per-step value is averaged over generated tokens 2–32; end-to-end latency includes the full prefill pass plus 31 generation steps. K is the number of retrieved mid-region tokens induced by the 1% total exact-support budget after reserving sink/tail anchors.
Context
K
Full exact
Top-K
Top-K+ϕ
Top-K+ϕ ( NoSub )
4k
21
11.68 / 431.06
87.27 / 2774.16
141.85 / 4486.39
129.18 / 4093.23
16k
144
12.48 / 728.27
89.82 / 3125.23
148.23 / 5006.89
123.59 / 4242.39
64k
636
17.09 / 3380.97
88.54 / 5606.19
144.03 / 7571.53
122.12 / 6897.49
Appendix
Table 11: Prototype latency disclosure for Llama-3.2-3B-Instruct, in the same format as Table 10 .
Backbone / dϕ
Per head
q heads
KV heads
ϕq total
ϕk total
Total ϕ params
1B / 32
575,010
512
128
294.4M
73.6M
368.0M
1B / 64
591,426
512
128
302.8M
75.7M
378.5M
1B / 128
624,258
512
128
319.6M
79.9M
399.5M
3B / 64
624,194
672
224
419.5M
139.8M
559.3M
Appendix
Table 12: Head-specific ϕ module parameter counts. “Per head” is the parameter count for one query-head or KV-head feature map. The total counts include all query-head maps and all KV-head maps.
Figure 6: Diffuse-attention and residual-denominator diagnostics for Llama-3.2-1B-Instruct on Wikitext-64k. Panel (a) shows the head/layer map of mid-normalized entropy Hmid . Panel (b) plots Hmid against the estimated residual denominator fraction ρres=Z^R/(ZE+Z^R) . The fraction ρres depends on the learned ϕ residual estimate. It is not the teacher-side mass-recovery quantity Cmid(K) .
Figure 7: Teacher-side mid-region mass-recovery curves Cmid(K) for representative heads on Wikitext-64k. Low mid-entropy heads have larger Cmid(K) at small K , meaning that a small number of keys recover most teacher mid-region mass. High mid-entropy heads have smaller Cmid(K) at the same K . Selection-only Top-K therefore leaves more teacher mid-region mass unrecovered in those heads.
Figure 8: Same- K attention-output error by mid-entropy quartile for Llama-3.2-1B-Instruct on Wikitext-64k. Panels (a)–(d) are ordered from low to high Hmid . Each panel compares selection-only Top-K and subtractive Top-K+ϕ completion against full attention using relative ℓ1 error, with the same retrieved set at each K .
Figure 9: Fixed- K breakdown for Llama-3.2-1B-Instruct on Wikitext-64k. The same retrieved set is used for Top-K and Top-K+ϕ ; therefore, differences isolate the learned residual completion mechanism. Panel (c) plots completion gain eTop-K−eTop-K+ϕ versus Hmid , corresponding to the main-text gain diagnostic in Figure 1 .
Figure 10: Fixed- K breakdown for Llama-3.2-3B-Instruct on Wikitext-64k, in the same format as Figure 9 .
Residual completion augments query-aware sparse attention by estimating the contribution of tokens omitted from the exact sparse computation. We ask whether improving a layer's attention-output reconstruction on the same incoming Q/K/V and selected support necessarily improves the fidelity of the final model output. We study training-free RESA and learned Top-K+φ with frozen backbone language models. A prespecified single-layer screen yields two Qwen3-0.6B/Multi-LexSum interventions for which direct-runtime measurements show positive prespecified request-aggregate local reconstruction gain but worse final KL fidelity than the corresponding all-abstain Exact Top-K baseline on both discovery and prompt-token-disjoint holdout requests. Exact restoration at the same layer instead improves final fidelity, showing that the reversal is specific to approximate completion in these cases. In complementary multi-layer experiments, a task-independent local diagnostic often repairs the tested completion estimators, although the repaired models do not consistently outperform Exact Top-K. Together, these results show that better local reconstruction need not translate into better final-model fidelity.
KV cache compression is essential for efficient long-context inference. Existing eviction methods permanently discard unselected tokens and consequently remove their aggregate contribution to attention. Merging-based alternatives preserve more information but can perturb retained keys and values that should remain exact. We observe that the information omitted by cache eviction can be formulated as residual statistics in both the numerator and denominator of softmax attention. Based on this observation, we propose ResKV, which divides a fixed KV budget into an exact main cache and a compact residual cache that reconstructs the contribution of omitted tokens. ResKV lets main-cache tokens and residual entries participate in the same softmax normalization, so residual entries restore both attention numerator and denominator mass rather than acting as a post-hoc correction. A construction-time validation proxy determines residual allocation for each layer and KV head, while a decode-time dynamic gate adjusts residual contributions for individual queries. Comprehensive evaluations on LongBench and RULER, covering query-aware and query-agnostic settings, multiple backbones, cache budgets, and representative compression baselines, demonstrate broad improvements under the same retained KV budget while preserving the practical efficiency of compressed decoding, including peak memory usage and long-context decode throughput.
Yuhang Zhan, Lisi Chen, Shuo Shang
University of Electronic Science and Technology of China
KV-cache eviction caps the memory cost of long reasoning traces but is inherently lossy because the model decodes from a partial view of its history. Under aggressive budgets, this not only lowers accuracy but can also cause runaway degeneration, where the model produces incoherent or repetitive tokens until reaching the length limit. We characterize much of this loss as an information gapf caused by missing context, rather than a capability gap caused by limited model capacity. An evicted 7B model and a full-context 1.5B model make complementary errors, and an oracle choice between their answers recovers 79% of the accuracy gap to the full-KV 7B model. Based on this observation, we propose KV-Rescue, a training-free inference framework that bridges the information gap introduced by KV eviction using a lightweight full-context helper. KV-Rescue interleaves reasoning steps from the two models into a shared trajectory. An online detector uses entropy and compressibility to terminate the generation of incoherent or repetitive base-model candidates early. Across five math benchmarks with Qwen2.5-Math 7B and 72B, KV-Rescue recovers an average of 87% of the accuracy lost to eviction at eviction budget B=64. A decode-cost analysis further shows that preventing runaway degeneration cuts base-model token generation by 43% on average.
Minsoo Cheong, Woosang Lim, Vincent-Daniel Yun +1
Seoul National University · University of Southern California