基于前置稀疏的近 oracle 键值选择用于长上下文推理
机器学习
2026-02-10 v1 人工智能
信息论
math.IT
摘要
大型语言模型 (LLM) 推理的核心瓶颈之一是对不断增长的键值 (KV) 缓存进行注意力的成本。虽然近 oracle top-k KV 选择可以在保持稠密注意力质量的同时显著降低计算和带宽开销，但现有稀疏方法通常依赖后验启发式，即基于观察到的注意力或代理得分的选择器。此类条件引入后验偏差：倾向于扭曲真实 token 重要性并遗漏关键 token，从而损害长程推理。为解决此问题，我们提出 Pre-hoc Sparsity (PrHS)，其在注意力评分之前选择 KV 条目，可提供显式的准确度控制。通过边际到互信息分析，我们推导出互信息损失的上界，该上界仅取决于被丢弃的质量。这一关系解释了后验启发式的失效模式，并通过控制被丢弃的质量在 advance 中实现可验证的保证。在 PrHS 框架内，我们在时间、深度和层面的三个正交前置选择器上实现原型。针对 LLaMA 和 Mistral 系列模型进行大量实验验证 PrHS。实验表明，PrHS 在 GSM8K 和 CoQA 上将检索开销降低超过 90%，在保持或改善准确度的同时，将检索稀疏性提高 3 倍于 HShare。它在 LongBench 上平均降低约 1% 的准确度，注意力 FLOPs 约降低 15%，注意力运算延迟提升 9.9 倍，NVIDIA A100-80GB GPU 上吞吐量提升 2.8 倍。
引用
@article{arxiv.2602.08329,
title = {Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference},
author = {Yifei Gao and Lei Wang and Rong-Cheng Tu and Qixin Zhang and Jun Cheng and Dacheng Tao},
journal= {arXiv preprint arXiv:2602.08329},
year = {2026}
}
备注
An effective method for accelerating LLM's inference via selective KV processing