提前确定重要键值对,让大模型长文本推理快10倍还不丢精度
Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference
- 不靠事后判断,而是提前筛选关键键值对,避免信息扭曲
- 在保持准确率几乎不变前提下,计算量减少90%以上,速度提升9.9倍
- 适合追求高吞吐、低延迟的长文本生成和问答系统
大语言模型推理的核心瓶颈在于不断增长的键值(KV)缓存注意力开销。尽管近似最优的top-k KV选择能在保持密集注意力质量的同时大幅降低计算与带宽,但现有稀疏方法普遍依赖事后启发式策略,即基于已观察到的注意力分数或代理得分进行选择。这种依赖导致事后偏差:容易扭曲真实词元重要性并遗漏关键信息,损害长程推理能力。为此,本文提出预前稀疏(Pre-hoc Sparsity, PrHS),在注意力评分前即完成KV选择,并提供明确的精度控制。通过边际到互信息分析,我们推导出互信息损失的上界仅依赖于被丢弃条目的注意力质量(记为delta)。该关系解释了事后方法的失效机理,并可通过预先控制丢弃质量实现可验证保证。在PrHS中,我们沿时间、深度、层三个维度实例化三种正交的预前选择器。在LLaMA与Mistral系列模型上的实验表明:在GSM8K和CoQA任务上,PrHS将检索开销降低超90%,相较HShare实现3倍更高的检索稀疏度,且精度相当或更优;在LongBench上平均性能下降不足1%,注意力浮点运算量相比先前稀疏基线降低约15%;在NVIDIA A100-80GB GPU上,注意力算子延迟降低9.9倍,吞吐量提升2.8倍。
原文摘要 · Abstract (English)
A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。