用内部表示变化评估记忆重要性,实现无注意力矩阵的高效缓存淘汰
Epiphany-Aware KV Cache Eviction Without the Attention Matrix
- 基于前向传播中模型内部表征的变化计算令牌重要性
- 4096词缓存下在MATH-500上达72%,8192词时AIME-2024达37%
- 无需训练或定制内核,兼容FlashAttention,速度最高提升2.8倍
随着推理模型生成长达数万词的思维链,键值缓存逐渐成为部署瓶颈。现有缓存淘汰方法依赖注意力权重排序,但在长推理轨迹中该指标噪声大,且因需显式计算注意力矩阵而阻碍融合内核在生产环境中的使用。本文提出一种新度量——顿悟分数(epiphany score),直接从前向传播中读取模型内部表征的变化,无需注意力矩阵且额外状态可忽略。由此提出的EpiKV缓存淘汰方法无需训练、分类器或自定义内核,可直接接入FlashAttention推理栈,实现16倍更长的有效上下文。在上中层特征中消除位置趋势并采用因果滚动z-score归一化后,4096词缓存下在MATH-500上达到72%准确率,与最强基线ThinKV(71%)相当;8192词时,在AIME-2024上达37%准确率,超越最佳基线(33%),速度最高提升2.8倍。
原文摘要 · Abstract (English)
As reasoning models emit chains of thought tens of thousands of tokens long, KV cache increasingly becomes a deployment bottleneck. Existing cache eviction methods rank tokens by attention weight, which is a noisy importance proxy in long reasoning traces, and prohibits the use of fused kernels in production inference by forcing the model to materialize the attention matrix. In this work, we instead score tokens with a metric we term the epiphany score: the change in the model's internal representation, read directly from the forward pass with no attention matrix and negligible extra state. Our resulting cache eviction method, EpiKV, requires no training, classifier, or custom kernel, and can be used directly in FlashAttention inference stacks unchanged -- scaling to a 16x longer feasible context than attention-based scoring. upper-mid layers negatively) and remove a positional trend with a causal rolling z-score. At a 4096-token cache EpiKV reaches 72% on MATH-500, matching the strongest attention-based baseline (ThinKV 71%, H2O 67%); a lag-normalized KV variant reaches 37% on AIME-2024 at 8192 tokens against the best of them (33%), at up to 2.8x the speed.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。