通过智能合并与注意力补偿,大幅压缩大模型缓存并提升推理速度。
SelKV: Selective KV Cache Merging with Per-Token Merge-or-Drop and Attention Compensation

- 按向量相似性动态决定是否合并令牌,保留语义信息。
- 在解码时用注意力比例补偿修正合并导致的权重失衡。
- 适用于长文本问答,可提速3.3倍且保持高生成质量。
大型语言模型自回归生成文本依赖于随上下文长度线性增长的关键值(KV)缓存,造成显著内存瓶颈。现有压缩方法通过合并令牌缓解开销,但常采用无差别聚合,损害表征质量并引入注意力塌陷——合并后的令牌在软最大值中获得与单个令牌相同的权重,尽管其编码了多个输入。我们提出一种无需训练的双组件框架解决上述问题:首先,基于值向量相似性的软余弦门动态调节合并决策,抑制或丢弃不相似的令牌以保持语义保真度;其次,引入注意力比例补偿机制,在解码时应用基于预填充阶段注意力统计的对数偏置,纠正合并引发的软最大值不平衡。在包含16个英文数据集的LongBench上评估,仅保留25%的KV缓存,性能优于代表性一次性基线。尤其在分组查询注意力(GQA)模型上表现稳健,近乎无损生成。此外,该方法在复杂多文档问答任务中超越全缓存基线,并在10万令牌条件下实现3.3倍解码加速。
原文摘要 · Abstract (English)
Large Language Models (LLMs) generate text autoregressively, relying on a key-value (KV) cache whose memory footprint grows linearly with context length, creating a major bottleneck. Recent compression methods mitigate this cost via token merging; however, these approaches often rely on indiscriminate aggregation, which degrades representations and introduces attention sag, a mismatch where merged tokens receive the same softmax mass as individual tokens despite encoding multiple inputs. We propose a training-free, dual-component framework for KV cache compression that addresses these limitations. First, a soft cosine gate adaptively modulates merging decisions based on value-vector similarity, suppressing or discarding dissimilar tokens to preserve semantic fidelity. Second, we introduce an attention-ratio compensation mechanism that applies a decoding-time logit bias derived from prefill attention statistics, correcting the softmax imbalance induced by merging. Evaluated on LongBench (16 English datasets) while retaining only 25% of the KV cache, our framework achieves strong compressed performance against representative one-shot baselines. It is especially robust on the evaluated grouped-query attention (GQA) models, maintaining nearlossless generation quality. Furthermore, the method outperforms the full-cache baseline on complex multi-document QA tasks and delivers a 3.3x decoding speedup at 100k tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。