arXiv 2406.020692025-05-15cs.CLcs.AI

PyramidKV：基于金字塔式信息汇聚的动态 KV 缓存压缩

PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

University of Wisconsin - Madison（威斯康星大学麦迪逊分校）；Peking University（北京大学）；Nanjing University（南京大学）；University of Surrey（Surrey 大学）；Qwen（通义实验室）；University of California - Riverside（加州大学河滨分校）；Microsoft（微软公司）

Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Junjie Hu, Wen Xiao

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

该研究发现 LLM 存在金字塔式信息汇聚现象，并据此提出按层动态分配缓存的 PyramidKV，以极低 KV 缓存实现优异长上下文性能。

摘要：

在本研究中，我们考察大语言模型（LLM）内部基于注意力的信息流在处理长上下文时，是否通过明显的模式进行聚合。我们的观察表明，LLM 通过“金字塔式信息汇聚”聚合信息：注意力在较低层广泛分散，随后逐步巩固到特定上下文中，并最终在较高层聚焦于关键 token（也称为大规模激活或注意力汇聚）。受这些洞见启发，我们开发了 PyramidKV，这是一种新颖且有效的 KV 缓存压缩方法。该方法跨不同层动态调整 KV 缓存大小，在较低层分配更多缓存、在较高层分配更少缓存，这不同于保持统一 KV 缓存大小的传统方法。我们利用 LongBench 基准进行的实验评估表明，PyramidKV 仅保留 12% 的 KV 缓存时，性能即可与使用完整 KV 缓存的模型相匹配，从而显著降低内存占用。在强调内存效率、仅保留 0.7% KV 缓存的场景中，PyramidKV 优于其他 KV 缓存压缩技术，在 TREC 数据集上实现最高 20.5 个百分点的绝对准确率提升。在 Needle-in-a-Haystack 实验中，PyramidKV 在维持 LLM 长上下文理解能力方面优于竞争方法；值得注意的是，仅保留 128 个 KV 缓存条目即可使 LLAMA-3-70B 模型达到 100.0 Acc. 性能。

英文摘要：

In this study, we investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers, progressively consolidating within specific contexts, and ultimately focusing on critical tokens (a.k.a massive activation or attention sink) in higher layers. Motivated by these insights, we developed PyramidKV, a novel and effective KV cache compression method. This approach dynamically adjusts the KV cache size across different layers, allocating more cache in lower layers and less in higher ones, diverging from traditional methods that maintain a uniform KV cache size. Our experimental evaluations, utilizing the LongBench benchmark, show that PyramidKV matches the performance of models with a full KV cache while retaining only 12% of the KV cache, thus significantly reducing memory usage. In scenarios emphasizing memory efficiency, where only 0.7% of the KV cache is maintained, PyramidKV surpasses other KV cache compression techniques, achieving up to a 20.5 absolute accuracy improvement on TREC dataset. In the Needle-in-a-Haystack experiment, PyramidKV outperforms competing methods in maintaining long-context comprehension in LLMs; notably, retaining just 128 KV cache entries enables the LLAMA-3-70B model to achieve 100.0 Acc. performance.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。