用Sigmoid注意力可让缓存淘汰更平滑,训练时软删除也能在推理中高效硬删除。
Sigmoid Attention as a Better Substrate for Learned KV Cache Eviction

- 采用Sigmoid注意力配合学习门控,实现训练与推理间更一致的缓存淘汰机制
- 在不显著增加困惑度的情况下,模型能安全删除大量键值缓存条目
- 适合关注大模型推理效率优化的研究者和工程团队
学习型键值缓存淘汰常面临软-硬不匹配问题:训练时可微门控仅减弱令牌贡献,而推理阶段仅当键值条目被物理删除才能节省内存。我们探究注意力机制是否影响这一软到硬的转换。在OpenWebText上训练的GPT-2规模Transformer中,进行2×2×2的受控对比,考察注意力类型、学习门控与位置编码的影响。尽管Sigmoid注意力作为密集语言模型表现较差,但结合学习型硬淘汰后,其在删除键值条目时几乎无困惑度损失。在相同稠密主干模型上,使用匹配的活跃缓存协议,学习型Sigmoid门控的困惑度低于H₂O与KeyDiff等后处理方法,而Softmax门控并未普遍优于这些方法。结果表明,注意力归一化方式显著影响训练时的软门控能否顺利转化为硬删除。
原文摘要 · Abstract (English)
Learned KV-cache eviction often faces a soft-to-hard mismatch: during training, differentiable gates typically attenuate token contributions, whereas inference saves memory only when KV entries are physically removed. We ask whether the attention substrate affects this soft-to-hard transition. Using GPT-2-scale Transformers trained on OpenWebText, we run a controlled $2\times2\times2$ comparison over attention type, learned gating, and positional encoding. Although sigmoid attention is worse as a dense language model, learned hard eviction changes the useful operating points: sigmoid-gated models delete KV entries with negligible PPL change relative to their own no-eviction references. Under a matched live-cache protocol on the same dense backbones, learned sigmoid gates obtain lower PPL than our H$_2$O and KeyDiff implementations, whereas softmax gates do not uniformly beat these post-hoc methods. The results suggest that attention normalization can substantially affect whether a training-time soft gate transfers cleanly to hard KV deletion.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。