[论文解读] Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

Suyu Ge, Yunan Zhang+4 more|arXiv (Cornell University)|Oct 3, 2023

Topic Modeling被引用 5

一句话总结

FastGen 通过利用轻量级注意力分析，为 LLM 实现自适应 KV 缓存压缩，动态应用定制化压缩策略（如逐出长距离上下文、仅保留特殊标记或丢弃低频标记），在无需微调或重新训练的情况下，实现高达 36% 的内存减少，且生成质量损失可忽略不计。

ABSTRACT

In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.

研究动机与目标

• 在不需模型微调或重新训练的情况下，减少自回归 LLM 推理中 KV 缓存的内存占用。
• 识别并利用注意力模块中的结构模式，实现有针对性的、上下文感知的 KV 缓存压缩。
• 开发一种即插即用的方法，根据注意力头的功能角色自适应调整压缩策略。
• 通过智能、结构感知的剪枝实现显著的内存节省，同时保持生成质量。
• 通过最小化 GPU 内存压力，实现大语言模型在设备上的高效部署。

提出的方法

• 该方法采用轻量级注意力分析步骤，分析注意力图并将每个头分类为五种结构类别之一。
• 基于分析结果，FastGen 应用五种压缩策略之一：局部上下文逐出、特殊标记保留、基于频率的剪枝、列级稀疏化或标准缓存。
• 系统通过仅选择每个头的相关标记构建自适应 KV 缓存，其策略集包括频率、特殊标记、标点符号和局部注意力模式。
• 分析的开销极低——仅占 KV 缓存内存的 0.78%，使其适用于实时推理。
• 该方法采用贪心策略排序，按优先级顺序依次应用压缩策略，优先处理高影响力策略（如特殊标记和高频标记）。
• 该方法完全兼容现有 LLM 推理流水线，无需模型重新训练或架构修改。

Figure 1: Different attention heads usually have different structures. Left: Four common attention structures (more details are elaborated in Section 3 and Section 4 ). Right: Attention map compositions of three attention heads that are in the same layer.

实验结果

研究问题

• RQ1 LLM 中的注意力模块能否被分类为不同的结构模式，从而指导高效、有针对性的 KV 缓存压缩？
• RQ2 如何在不降低生成质量的前提下，对每个注意力头实现自适应的 KV 缓存压缩？
• RQ3 为最大化内存节省并保持模型性能，压缩策略的最佳应用顺序与组合是什么？
• RQ4 即插即用、无需微调的方法能否在 KV 缓存中实现显著的内存减少，且质量损失可忽略？
• RQ5 压缩性能对超参数（如频率和局部注意力阈值）的选择有多敏感？

主要发现

• 在 Llama 1-65B 上，FastGen 实现了高达 36.04% 的 KV 缓存内存使用减少，同时在 AlpacaEval 上保持 49.75% 的胜率，几乎与完整缓存基线持平。
• 移除特殊标记策略 (${\bm{C}}_{\text{special}}$) 导致性能最大降幅——胜率下降 2.11%，凸显其关键作用。
• 基于频率的策略 (${\bm{C}}_{\text{frequent}}$) 对压缩贡献最大，当与其他策略结合时，可将缓存减少 36.04%。
• 策略应用顺序显著影响性能：默认贪心顺序（特殊 → 标点 → 频繁 → 局部）实现了最高的胜率和压缩比。
• 该方法对超参数变化具有鲁棒性，在所有测试配置中均保持超过 45% 的胜率，且压缩比因设置不同而显著变化。
• 分析开销可忽略不计——仅占完整 KV 缓存大小的 0.78%，使其适用于实时部署。

Figure 2: Performance of Adaptive KV Cache (FastGen) and Fixed KV Cache (Frequency, Local, and Frequency+Local; Zhang et al. , 2023 and Liu et al. , 2023a ) on AlpacaEval.