PAGE: 分区感知的门控KV缓存驱逐

中文摘要 AI 辅助

KV缓存驱逐方法决定保留哪些令牌，但不决定是否完全驱逐，因此基准平均值可能掩盖一类输入，在这些输入上压缩将准确率从99%降至0%。我们将驱逐重新定义为逐输入准入决策，并表明输入分为两类：容量受限类，在任何预算下驱逐都是灾难性的；以及易稀释类，驱逐是安全或有益的。从预填充注意力中计算出的单个无标签标量，即成对top-$k$头一致性的早期到晚期下降，可在任何解码之前预测此类别。PAGE对该下降进行阈值化：当下降较大时应用任何基础驱逐器，否则保留完整缓存，无需训练且无需准确率标签。该下降在四个架构系列中一致地按驱逐安全性对输入排序，每个模型约100个输入的无标签试点可为新系列重新校准阈值。作为保护措施，PAGE在四个驱逐器、四个模型和两个基准上将容量受限区域的危害率从0.75降至0.026，即减少29倍，将99%至0%的崩溃转变为平坦的89%，而无需重新训练驱逐器。该门在驱逐已安全之处不生效，其保护的容量受限类是少数可识别的输入，因此收益是针对性的安全增益而非平均增益。PAGE是逐输入保护措施，而非压缩器：在名义16倍预算下，实际压缩为1.8至3.4倍（平均2.9倍），在静态配置下批次16时衰减至接近1，且经过训练的驱逐器在匹配内存下表现更优。

英文摘要

KV-cache eviction can do more than compress. In long-context LLMs, keeping only some cached tokens sometimes matches or exceeds full-cache accuracy, because many redundant prefill tokens otherwise dilute attention away from the tokens that carry the answer. This benefit is not uniform, and evicting the wrong tokens can drop accuracy to zero on tasks that require precise retrieval, so the useful question is not only which tokens to keep but also whether to evict this input at all. We show that one label-free number computed from the prefill attention, the drop between early and late layers in how much attention heads agree on which tokens to read, predicts per input, before any decoding, which of the two cases an input falls under. We build this into PAGE (Partition-Aware Gated Eviction), a wrapper that runs any SnapKV-style evictor when the drop is large and keeps the full cache when it is small, with no training, labels, or fine-tuning. PAGE is a safety mechanism rather than a compressor, so we measure it by the failures it prevents. It cuts the harm rate on capacity-bound inputs from 0.75 to 0.026, and on multi-key retrieval with Mistral-7B plain SnapKV falls from 99\% to 0\% as the budget shrinks, while PAGE holds it at 89\%. Elsewhere, it passes the base evictor through unchanged, which is the intended behaviour and is what we observe in 8 of 16 cells. Code is available at https://anonymous.4open.science/r/PAGE-018239.

发表机构

• National Institute of Science Education and Research（国家科学教育与研究学院）
• Homi Bhabha National Institute（霍米·巴巴国家研究所）

机构由 AI 辅助整理，请以论文原文为准。