用残差质量会计法提升部分键值解码效率,显著改善长文本生成效果。
Residual-Mass Accounting for Partial-KV Decoding

- 基于提取的锚点精确计算,其余部分用学习的残差估计补全。
- 在1%精确支持预算下,长文本任务上性能超越传统Top-K基线。
- 适合需要高效处理长上下文的部署场景,尤其对冻结模型有效。
我们研究了一种受控的部分键值解码设置:对尾部锚点和检索到的标记集精确计算未归一化的softmax贡献,其余预填充标记则由残差估计表示。重点在于查询相关精确支持选定后的会计规则,仅将完全的Top-K作为理想选择器,而非可部署的检索系统。所提规则保持主干语言模型和精确分支的键值张量不变,从学习到的正特征图ϕ构建固定大小的摘要状态(S,u),减去已检索标记的特征贡献以保证精确与残差集合无重叠,并将估计的残差分子和分母合并到单一归一化中。在1%精确支持预算下,我们的残差补全方法在RULER和BABILong数据集上,对所有报告的上下文长度均优于仅使用选择的Top-K基线,在1%-4%预算范围内趋势基本持续。在LongBench上,摘要任务结果总体有利,多文档问答则表现不一。注意力输出诊断表明,检索标记减法是符合划分一致性的会计规则,但主要剩余误差源于学习到的ϕ对未检索残差质量的近似不完善。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。