用概率视角重看缓存淘汰,让模型推理更稳定高效
A Probabilistic Interpretation of KV Cache Eviction
- 将缓存淘汰建模为期望估计问题,通过采样逼近最优解
- 发现现有方法存在偏差,可借此在解码时修正淘汰影响
- 新方法在多种任务上更鲁棒,相同压缩率下表现更优
KV缓存淘汰的核心理念是:通过丢弃部分缓存条目以提升吞吐量,而对生成质量的影响微乎其微。尽管该策略在实践中被广泛采用,但多数方法依赖启发式规则选择淘汰对象。本文从概率推理角度正式化这一问题,并证明其计算上是难解的。我们进一步表明,将淘汰问题转化为期望估计后,可通过采样近似求解;同时揭示了此前被忽视的解码阶段补偿机制的可行性。此外,我们发现现有方法本质上是零方差有偏估计器,可轻松改进以支持解码时修正。实验显示,结合概率框架与解码时修正的新方法,在不同任务上均比现有方法更鲁棒,且在相同压缩预算下达到相当甚至更优性能。
原文摘要 · Abstract (English)
The premise and promise of KV (cache) eviction is simple: higher throughput can be achieved by evicting some entries from the KV cache, at a negligible cost to quality. This holds empirically for many existing methods, though most rely on creative heuristics for selecting which entries to drop. Despite recent advances, the problem of KV eviction has remained informal in the literature. This paper aims to properly formalize this problem through the lens of probabilistic reasoning and reveal what can be learned from this perspective. Concretely, we (1) formalize the problem of KV eviction and, unfortunately, prove that it is computationally hard, (2) show that by framing it probabilistically, KV eviction reduces to the problem of expectation estimation, which can be approximated through sampling, (3) show that through this probabilistic interpretation, correcting for evicted entries during decoding---a previously ignored problem---becomes feasible, and (4) reveal that existing methods in the literature are zero-variance biased estimators that can be easily adapted in order to enable decode time correction. In practice, we show that this probabilistic version of KV eviction coupled with decode time correction is more robust to different tasks compared to existing eviction methods and achieves competitive performance at the same compression budget.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。