专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

KV cache压缩时如何权衡Token和Precison？北大、华为工作发现保留更多的Token远比更多的精度重要！

PaperWeekly · 公众号 · 科研 · 2025-01-10 18:18

正文

©PaperWeekly 原创 · 作者 | 张介宾

单位 | 北京大学计算语言所

研究方向 | 推理加速

论文标题：

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

论文作者：

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

论文链接：

https://arxiv.org/pdf/2412.12706

代码链接：

https://github.com/zhzihao/QPruningKV

简介

KV cache 技术通过缓存已经计算好的 key、value 向量（KV），大幅减少了推理的时间。然而随着大语言模型（LLM）的窗口长度增长，KV cache 的内存占用成为了推理时候的新瓶颈。

很多工作开始探讨如何压缩 KV cache 的内存。压缩 KV cache 的内存不仅能减少推理时候的内存开销，还可以减少解码时候的时间消耗，因为解码步骤一般是带宽瓶颈的。

现有的 KV cache 压缩工作主要包括 KV Pruning 和 KV Quantization。KV Pruning 方法通过一些特征筛选掉 KV cache 中不重要的 Token，从 Token 的角度进行压缩。KV Quantizaiton 方法则通过使用更低的精度保留 KV cache，从 Precision 的角度进行压缩。

这些方法都取得了很好的压缩效果，然而却忽视了在实际使用时候，一般是在给定的内存 Budget 下进行推理。此时一个很重要的研究问题应该是如何在给定的内存 Budget 下权衡 Token 或者 Precsion 以达到最佳的推理效果，而不仅仅依赖于对单一的维度进行压缩。

本文着重分析了这个问题，即在给定 KV cache 的内存 Budget 下如何权衡 Token 和 Precision 以达到更好的效果。具体来说，本文首先论证了 Token 和 Precision 联合压缩的可行性，并发现不同的内存 Budget 下使用更少的精度来保留更多的 Token 都可显著提高推理的性能。

文章进而细致地分析了 Token 和 Precision 的 Trade-off 在不同任务类型和输入长度上的效果，以及受模型 Scale，量化策略，层间策略等多种因素的影响情况。

主要结果

上表为对 KV Pruning 后进一步量化的结果，表中汇报了 Llama-3-8B 在 LongBench 和大海捞针上的效果。

KV Pruning 方法使用了比较有代表性的 StreamingLLM，H2O，SnapKV，PyramidKV；KV Quantization 方法则使用了 KIVI。可以看出 KV Pruning 下保存不同的 Token 数量时继续量化至 4bit 都只有很小的性能损失。而 Token 数量的变化相比 Precision 对结果的影响则更显著一些。

上图进而比较了不同的 KV Cache Budget 下三种设置的性能变化。这三种设置分别是：

1 倍的 Token 数目，用 16bit 精度保存（即单纯的 KV Pruning）；
2 倍的 Token 数目，用 8bit 精度保存；
4 倍的 Token 数目，用 4bit 精度保存。

这三种设置有着接近的内存占用。可以看出更多 Token 更少精度的设置在各种情况下都比单纯的 KV Pruning 有可观的性能提升，尤其是在极低内存 Budget 的情况下。

深入分析

本文继续对 Token 和 Precision 的 Trade-Off 进行了深入的分析，主要包括这个 Trade-off 对任务类型和输入长度的影响。同时还分析了其他重要的因素，例如模型 Scale、量化策略的影响，以及层粒度的 Token 和精度的权衡。

上表列出了三种设置在 LongBench 上的 6 种子任务类型以及 RULER 这个检索任务上的效果。

可以看出更多的 Token 显著的提高了检索任务 RULER 的效果，在比较需要检索的任务上例如单文档 QA 或者摘要任务上都取得了不错的性能提升。而对于不那么需要检索而更多需要推理的合成任务或者多文档 QA 上，性能提升则比较有限。

本文进一步探索了不同设置在不同层上的影响，上图用 4 倍 Token，4bit 作为 baseline，并调整横轴中包括的层的设置，并汇报相对于 baseline 的性能变化。可以看出中间层给更多 Token 比较重要，尤其对检索任务。而在开头和结尾的那些层改变 Token 和精度则影响较小。

对于更多其他因素的影响请参见原文，总体来说更多的 Token 更少精度的策略受各种因素的影响都比较小，展现了较好的性能和稳定性。

总结与意义

本文的贡献主要包括三部分：

首先分析了 KV Pruning 进一步量化的可行性，并发现大多数 KV Pruning 方法都可以继续量化至 4bit 而仅有很少的性能损失。
发现在 KV cache 的 Budget 一定的情况下，在 KV cache 中使用更低的 Precision 保留更多的 Token 可以显著提高模型的效果，尤其在 Budget 极低的情况下，提升更为显著。
进一步分析了不同情况下 Token-Precision 的 Trade-Off，包括任务类型，输入长度，量化策略，模型 Scale 以及层间策略的影响。

更多阅读