通过细粒度分配维度，让KV缓存仅用不到1%的内存实现长文本精准推理。
Beyond Token Eviction: Mixed-Dimension Budget Allocation for Efficient KV Cache Compression
- 为每个令牌分配可变维度，而非整体丢弃，实现更精细的缓存压缩
- 在5万上下文长度下仅用0.26%缓存仍保持100%准确率
- 适合需要超长文本处理的场景，如文档问答、代码生成
键值（KV）缓存广泛用于加速Transformer推理，但其内存消耗随输入长度线性增长，限制了长上下文应用。现有令牌淘汰方法通过丢弃不重要令牌来降低内存，可视为粗粒度的降维，即每个令牌被赋予零或完整维度。我们提出MixedDimKV，一种细粒度分配维度的KV缓存压缩方法，并进一步结合头级别重要性信息提出MixedDimKV-H。在长上下文基准测试中，MixedDimKV优于不依赖头级别重要性分析的现有压缩方法；当使用相同头级别重要性信息时，MixedDimKV-H始终优于HeadKV。值得注意的是，该方法在LongBench上仅需6.25%的KV缓存即可达到全注意力性能；在Needle-in-a-Haystack测试中，于5万上下文长度下仅用0.26%缓存仍保持100%准确率。
原文摘要 · Abstract (English)
Key-value (KV) caching is widely used to accelerate transformer inference, but its memory cost grows linearly with input length, limiting long-context deployment. Existing token eviction methods reduce memory by discarding less important tokens, which can be viewed as a coarse form of dimensionality reduction that assigns each token either zero or full dimension. We propose MixedDimKV, a mixed-dimension KV cache compression method that allocates dimensions to tokens at a more granular level, and MixedDimKV-H, which further integrates head-level importance information. Experiments on long-context benchmarks show that MixedDimKV outperforms prior KV cache compression methods that do not rely on head-level importance profiling. When equipped with the same head-level importance information, MixedDimKV-H consistently outperforms HeadKV. Notably, our approach achieves comparable performance to full attention on LongBench with only 6.25% of the KV cache. Furthermore, in the Needle-in-a-Haystack test, our solution maintains 100% accuracy at a 50K context length while using as little as 0.26% of the cache.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。