Residual-Mass Accounting for Partial-KV Decoding

TL;DR

This paper introduces a residual-mass accounting method for partial-KV decoding in language models, improving retrieval accuracy with minimal support budgets across multiple benchmarks.

Contribution

It proposes a novel residual-mass accounting rule that enhances partial-KV decoding performance without altering the core language model architecture.

Findings

Improves over Top-K baseline at 1% exact-support budget on RULER and BABILong.

Maintains favorable summarization results on LongBench.

Residual subtraction reduces main error source to learned feature approximation.

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 $\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.