English

Residual-Mass Accounting for Partial-KV Decoding

Machine Learning 2026-05-08 v2 Computation and Language

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)(S,u) from learned positive feature maps ϕ\phi, 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-ϕ\phi approximation of the unretrieved residual mass.

Keywords

Cite

@article{arxiv.2604.05438,
  title  = {Residual-Mass Accounting for Partial-KV Decoding},
  author = {Yasuto Hoshi and Daisuke Miyashita and Jun Deguchi},
  journal= {arXiv preprint arXiv:2604.05438},
  year   = {2026}
}
R2 v1 2026-07-01T11:56:39.815Z