Residual-Mass Accounting for Partial-KV Decoding

Yasuto Hoshi; Daisuke Miyashita; Jun Deguchi

cs.LGApr 7, 2026

Residual-Mass Accounting for Partial-KV Decoding

Authors: Yasuto Hoshi, Daisuke Miyashita, Jun Deguchi

Organizations: Kioxia Corporation

Abstract

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
	$\mathrm{Top}\text{-}\mathrm{K}$	0.266	0.271	0.275	0.276	0.323	0.324	0.320	0.321
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	0.284	0.283	0.285	0.281	0.332	0.333	0.329	0.327
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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

\mathrm{Top}\text{-}\mathrm{K}{+}\phi

rows use the 64k-trained

\phi

. Bold marks the best partial-KV method within each task/backbone/budget block, including ties. “Full exact” is the unmodified backbone LM. “

\mathrm{Top}\text{-}\mathrm{K}{+}\phi

(

\mathrm{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
	$\mathrm{Top}\text{-}\mathrm{K}$	0.732	0.606	0.541	0.878	0.787	0.688
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	0.753	0.617	0.548	0.889	0.814	0.711
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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

e_{\mathrm{Top}\text{-}\mathrm{K}}-e_{\mathrm{Top}\text{-}\mathrm{K}{+}\phi}

, measured by relative

\ell_{1}

attention-output error against full attention. Positive values indicate that subtractive residual completion reduces error relative to selection-only

\mathrm{Top}\text{-}\mathrm{K}

Corpus	Length	$K$	$\mathrm{Top}\text{-}\mathrm{K}$ $\downarrow$	$\mathrm{NoSub}$ $\downarrow$	$\mathrm{Sub}\text{-}\phi$ $\downarrow$	$\mathrm{ExactNoSub}$ $\downarrow$	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

\ell_{1}

attention-output error and residual log-partition calibration.

\mathrm{ExactNoSub}

uses the true whole-mid contribution without subtraction, double-counts

\mathcal{K}(q)

, and is diagnostic rather than a valid partition. The final column reports mean residual log-partition error

\log\hat{Z}_{\mathcal{R}}-\log Z_{\mathcal{R}}

; negative values indicate underestimation. Bold marks the lowest approximate method among

\mathrm{Top}\text{-}\mathrm{K}

\mathrm{NoSub}

, and

\mathrm{Sub}\text{-}\phi

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Exact softmax kernels and learned $\phi$ 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 $\phi$ scores are used only in the residual estimator.

Backbone	Steps / hardware	Wall-clock
Llama-3.2-1B-Instruct	200k steps on 8 $\times$ H200	19 h
Llama-3.2-3B-Instruct	200k steps on 8 $\times$ H200	35 h

Appendix

Table 4: Training cost for the 64k-trained

\phi

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
	$\mathrm{Top}\text{-}\mathrm{K}$	0.547	0.732	0.762	0.784	0.698	0.878	0.898	0.912
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	0.603	0.753	0.771	0.787	0.750	0.889	0.904	0.914
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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

\phi

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
	$\mathrm{Top}\text{-}\mathrm{K}$	0.263	0.308	0.306	0.310	0.401	0.488	0.498	0.496
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	0.298	0.324	0.321	0.319	0.431	0.497	0.494	0.495
	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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

\phi

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
$\mathrm{Top}\text{-}\mathrm{K}$	0.732	0.692	0.606
$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ (4k-trained)	0.758	0.017	0.014
$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ (8k-trained)	0.748	0.718	0.016
$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ (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

\phi

; evaluation is performed at 4k, 8k, and 16k.

Figure 3: Residual denominator calibration for Llama-3.2-1B-Instruct. Histograms of

\log\hat{Z}_{\mathcal{R}}-\log Z_{\mathcal{R}}

for

\mathrm{Sub}\text{-}\phi

on Wikitext and CWE at 4k and 64k. Negative values indicate residual-mass underestimation.

Figure 4: Head-wise

\mathrm{Sub}\text{-}\phi

–

\mathrm{NoSub}

difference for Llama-3.2-1B-Instruct. Histograms of

\Delta=e_{\mathrm{Sub}\text{-}\phi}-e_{\mathrm{NoSub}}

across heads. Negative values favor

\mathrm{Sub}\text{-}\phi

Method	RULER 4k	BABILong 4k
Full exact	0.803	0.331
$\mathrm{Top}\text{-}\mathrm{K}$	0.732	0.308
$d_{\phi}=32$	0.758	0.335
$d_{\phi}=64$	0.758	0.339
$d_{\phi}=128$	0.753	0.334

Appendix

Table 8: Feature-dimension ablation with 4k-trained

\phi

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
$\mathrm{Top}\text{-}\mathrm{K}$	0.732	0.308
$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	0.758	0.339
Hedgehog $\phi$ + our loss	0.016	0.079
Our $\phi$ + LoLCATs loss	0.018	0.079
Our $\phi$ + KL only	0.082	0.087

Appendix

Table 9: Loss/module ablation with 4k-trained

\phi

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 $\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ 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 $\mathcal{L}_{\mathrm{ours}}=\lambda_{\mathrm{KL}}\mathcal{L}_{\mathrm{KL}}+(1-\lambda_{\mathrm{KL}})\mathcal{L}_{\mathrm{aux}}$ ( Equation 31 ); panels (b)–(d) keep our feature map fixed and vary the loss ( $\mathcal{L}_{\mathrm{lolcats}}$ , KL only, or aux only). Points closer to the diagonal indicate more faithful approximation of unretrieved mid-region kernels.

Context	$K$	Full exact	$\mathrm{Top}\text{-}\mathrm{K}$	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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	$\mathrm{Top}\text{-}\mathrm{K}$	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$	$\mathrm{Top}\text{-}\mathrm{K}{+}\phi$ ( $\mathrm{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_{\phi}$	Per head	q heads	KV heads	$\phi_{\mathrm{q}}$ total	$\phi_{\mathrm{k}}$ total	Total $\phi$ 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

\phi

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 $H_{\mathrm{mid}}$ . Panel (b) plots $H_{\mathrm{mid}}$ against the estimated residual denominator fraction $\rho_{\mathrm{res}}=\hat{Z}_{\mathcal{R}}/(Z_{\mathcal{E}}+\hat{Z}_{\mathcal{R}})$ . The fraction $\rho_{\mathrm{res}}$ depends on the learned $\phi$ residual estimate. It is not the teacher-side mass-recovery quantity $C_{\mathrm{mid}}(K)$ .

Figure 7: Teacher-side mid-region mass-recovery curves

C_{\mathrm{mid}}(K)

for representative heads on Wikitext-64k. Low mid-entropy heads have larger

C_{\mathrm{mid}}(K)

at small

K

, meaning that a small number of keys recover most teacher mid-region mass. High mid-entropy heads have smaller

C_{\mathrm{mid}}(K)

at the same

K

. Selection-only

\mathrm{Top}\text{-}\mathrm{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

H_{\mathrm{mid}}

. Each panel compares selection-only

\mathrm{Top}\text{-}\mathrm{K}

and subtractive

\mathrm{Top}\text{-}\mathrm{K}{+}\phi

completion against full attention using relative

\ell_{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

\mathrm{Top}\text{-}\mathrm{K}

and

\mathrm{Top}\text{-}\mathrm{K}{+}\phi

; therefore, differences isolate the learned residual completion mechanism. Panel (c) plots completion gain

e_{\mathrm{Top}\text{-}\mathrm{K}}-e_{\mathrm{Top}\text{-}\mathrm{K}{+}\phi}

versus

H_{\mathrm{mid}}

, 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 .

Softmax Attention

Softmax attention weights token representations by query–key similarity and underpins Transformer retrieval, sequence modeling, and long-context computation.

See all 75 papers

Counterexamples to Local Reconstruction Gain as a Proxy for Final Fidelity in Residual Completion
Sep 28, 2026Yasuto Hoshi, Daisuke Miyashita, Jun DeguchiSparse Attention
Kioxia Corporation
ResKV: Reconstructing Omitted Attention Contributions for Fixed-Budget KV Cache Compression
Jul 31, 2026Yuhang Zhan, Lisi Chen, Shuo ShangSoftmax Attention KV Caching
University of Electronic Science and Technology of China
KV-Rescue: Recovering Reasoning Language Model KV Eviction Loss via Stepwise Interleaving
Aug 16, 2026Minsoo Cheong, Woosang Lim, Vincent-Daniel Yun +1KV-Cache Eviction Efficient Language Model Inference
Seoul National University · University of Southern California

Sep 28, 2026J/K move · Enter open · S save

Counterexamples to Local Reconstruction Gain as a Proxy for Final Fidelity in Residual Completion

Yasuto Hoshi, Daisuke Miyashita, Jun Deguchi

Kioxia Corporation

Sparse Attention

Open paper AskPDF arXiv

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.

Abstract

Softmax Attention

Explore similar work

Counterexamples to Local Reconstruction Gain as a Proxy for Final Fidelity in Residual Completion

ResKV: Reconstructing Omitted Attention Contributions for Fixed-Budget KV Cache Compression

KV-Rescue: Recovering Reasoning Language Model KV Eviction Loss via Stepwise Interleaving