[论文解读] Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

Zichang Liu, Aditya Desai+6 more|arXiv (Cornell University)|May 26, 2023

Topic Modeling被引用 10

一句话总结

本文提出 Scissorhands，一种用于 LLM 推理的 KV 缓存压缩方法，利用 重要性持久性假设 的持久性仅保留关键令牌，从而在不进行微调的情况下将 KV 缓存内存压缩高达 5×。它还显示与 4-bit 量化的兼容性并保持模型质量。

ABSTRACT

Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose Scissorhands, a system that maintains the memory usage of the KV cache at a fixed budget without finetuning the model. In essence, Scissorhands manages the KV cache by storing the pivotal tokens with a higher probability. We validate that Scissorhands reduces the inference memory usage of the KV cache by up to 5X without compromising model quality. We further demonstrate that Scissorhands can be combined with 4-bit quantization, traditionally used to compress model weights, to achieve up to 20X compression.

研究动机与目标

• 动机与解决在 LLM 推理中 KV 缓存内存瓶颈的问题，其中 KV 缓存内存可能超过模型大小并限制批量吞吐量。
• 引入 Persistence of Importance Hypothesis 来识别在各步中仍然具有影响力的令牌。
• 开发一款考虑预算的 KV 缓存压缩方法，在不重新训练的前提下保持准确性。
• 提供理论依据，表明压缩注意力可以近似原始注意力输出。
• 证明与量化的兼容性，以进一步提升压缩效果。

提出的方法

• 通过经验观察重复的注意力模式并定义 重要性持久性假设。
• 开发 Algorithm 1 和 Algorithm 2（Scissorhands），在固定预算下通过保留高注意力令牌并丢弃非影响令牌来维持 KV 缓存。
• 使用历史窗口来估计令牌重要性，并采用受预算约束的 KV 缓存替换策略，灵感来自水库抽样和最近最少使用（LRU）。
• 给出一个理论界限，显示在幂律注意力分数分布下原始 KV 缓存输出与预算后输出之间的期望近似误差（Theorem 4.1）。
• 在模型层和注意头之间分配内存预算，优先考虑后层以对抗持久性下降。

实验结果

研究问题

• RQ1 在不同的 LLM 和不同层之间，重要性持久性假设是否成立？
• RQ2 在不重新训练的情况下，是否可以在固定预算下减少 KV 缓存内存，同时保持生成质量？
• RQ3 预算化 KV 缓存对注意力输出及下游任务的理论与经验影响是什么？
• RQ4 Scissorhands 是否与后训练量化（如 4-bit 量化）兼容？

主要发现

• Scissorhands 在 OPT 模型上实现 KV 缓存内存最多降低 5×，且不降低语言建模及下游任务的模型质量。
• 持久性比率在各层通常超过 95%，表明早期的关键令牌在很大程度上决定了后续的重要性。
• 基于注意力分数的令牌丢弃使输出与原模型相近，在 OPT-66B 上的 5× 压缩下精度损失可以忽略不计。
• 量化（4-bit）仍然与 Scissorhands 兼容，在评估任务上维持性能。
• 理论分析（Theorem 3.1 与 4.1）将注意力权重、MLP 跳跃连接与预算化令牌保留联系起来，解释近似行为及误差界。