arXiv:2412.12706cs.CL2024-12EMNLP被引 9

通过降低精度存储更多上下文token,提升大模型长文本处理能力。

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

  • 将低精度压缩与剪枝结合,同时优化存储的token数量和精度。
  • 在长文本任务中显著提升性能,支持不同输入长度和模型规模。
  • 适用于追求高效推理的开发者,尤其适合长上下文应用。

随着大语言模型(LLMs)上下文窗口不断增大,KV缓存的内存开销成为推理阶段的关键瓶颈。主流的KV压缩方法,如KV剪枝和KV量化,主要分别聚焦于token数量或精度维度,但对这两个正交维度之间的权衡关系研究不足。本文全面探讨了KV缓存压缩中的token-precision权衡问题。实验表明,以较低精度存储更多token(即量化剪枝)能显著提升LLM的长上下文表现。深入分析显示,该策略在检索类任务中表现优异,且在不同输入长度、剪枝方法、量化策略和模型规模下均保持稳定有效。这些发现为通过平衡token与精度的权衡优化KV缓存压缩提供了重要启示。代码已开源:https://github.com/zhzihao/QPruningKV。

原文摘要 · Abstract (English)

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.

KV缓存量化剪枝长文本生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。