arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2412.12706cs.CL

更多 Token,更低精度:迈向 KV Cache 压缩中的最优 Token-精度权衡

More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression

Peking University(北京大学);Huawei Technologies(华为技术有限公司)

Jiebin Zhang, Dawei Zhu, Yifan Song, Wenhao Wu, Chuqiao Kuang, Xiaoguang Li, Lifeng Shang, Qun Liu, Sujian Li

AI总结:

本文研究 KV cache 压缩中 token 数量与量化精度的正交权衡,提出 quantized pruning,证明以低精度保留更多 token 可稳定提升长上下文与检索任务表现。

摘要:

随着大型语言模型(LLM)处理不断增大的上下文窗口,KV cache 的内存使用已成为推理过程中的关键瓶颈。主流 KV 压缩方法(包括 KV 剪枝和 KV 量化)主要分别关注 token 或精度维度。然而,这些工作在很大程度上尚未探索这两个正交维度之间的权衡。本文全面研究了 KV cache 压缩中的 token-精度权衡。实验表明,以更低精度在 KV cache 中存储更多 token(我们将该策略称为 quantized pruning,即量化剪枝)可以显著提升 LLM 的长上下文性能。对关键方面的 token-精度权衡进行深入分析表明,量化剪枝在检索相关任务上取得了显著改进,并在不同输入长度下始终表现良好。此外,量化剪枝在不同 KV 剪枝方法、量化策略和模型规模下均表现出显著的稳定性和有效性。这些发现为通过均衡的 token-精度权衡策略优化 KV cache 压缩提供了宝贵见解。我们的代码可在 https://github.com/zhzihao/QPruningKV 获取。

英文摘要:

As large language models (LLMs) process increasing context windows, the memory usage of KV cache has become a critical bottleneck during inference. The mainstream KV compression methods, including KV pruning and KV quantization, primarily focus on either token or precision dimension separately. However, these works leaving the trade-off between these two orthogonal dimensions largely under-explored. In this paper, we comprehensively investigate the token-precision trade-off in KV cache compression.Experiments demonstrate that storing more tokens in the KV cache with lower precision,a strategy we term quantized pruning, can significantly enhance the long-context performance of LLMs. In-depth analysis of the token-precision trade-off across key aspects demonstrates that, quantized pruning achieves substantial improvements in retrieval-related tasks and consistently performs well across varying input lengths. Furthermore, quantized pruning demonstrates notable stability and effectiveness across different KV pruning methods, quantization strategies, and model scales. These findings offer valuable insights into optimizing KV cache compression through balanced token-precision trade-off strategies. Our code is available at https://github.com/zhzihao/QPruningKV.

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

↑