ThinK: Thinner Key Cache by Query-Driven Pruning

发表时间: 2025-04 · arXiv:2407.21018 (ICLR 2025)

原文: https://arxiv.org/abs/2407.21018

Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo (Salesforce AI Research, The Chinese University of Hong Kong)

主要贡献

大型语言模型（LLMs）在处理长序列时面临着计算和内存需求的显著挑战。在长上下文推理场景中，键值（KV）缓存的内存消耗会随着序列长度和批处理大小线性增长，从而造成严重的内存瓶颈。现有的KV缓存优化方法主要集中在减少序列长度（$S$）维度或降低数据精度（量化），而很大程度上忽略了通道维度（$D$）的优化潜力。

本文的研究目标是针对长上下文场景，通过挖掘KV缓存通道维度的冗余性，来降低推理过程中的内存消耗。作者的核心创新点包括：
1. 首次揭示了KV缓存通道维度存在显著冗余，具体表现为Key缓存通道幅度的极度不平衡分布，以及注意力权重矩阵呈现出的低秩结构。
2. 提出了一种名为THINK的查询驱动（query-dependent）的KV缓存剪枝方法。该方法将寻找最不重要通道的问题转化为最小化注意力权重损失的优化任务，通过贪婪策略选择性地剪枝对当前查询最不重要的通道。
3. THINK作为一种即插即用的技术，与现有的KV缓存驱逐（Eviction）和量化（Quantization）方法完全正交。实验表明，与纯粹的KV缓存驱逐或量化方法相比，THINK在保持甚至提升模型精度的同时，能额外减少20%以上的KV缓存内存开销。例如，与KIVI结合使用时，THINK实现了2.8倍的峰值内存缩减，使单GPU上的批处理大小从4倍提升至5倍。
4. 探索了将THINK扩展至Value缓存剪枝（THINKV）的潜力，展示了该方法的广泛适用性。

背景知识/关键Observation/设计原则

识别关键观察结果。作者确定了几个关键的观察结果来激发对KV cache通道进行剪枝的方法。具体来说，作者可视化了KV cache的幅度，并对LLaMA模型的注意力机制执行了奇异值分解（SVD）。

KV cache通道的幅度分布。附录A中的图4可视化了每个通道中跨token的KV cache绝对值。与先前的发现（【30，Awq: Activation-aware weight quantization for on-device llm compression and acceleration+2024】；【50，Smoothquant: Accurate and efficient post-training quantization for large language models+2023】；【32，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024】）一致，作者观察到在key cache中只有特定通道具有显著的幅度，而value cache缺乏明显的模式。例如，在第14层中，所有token在第50个通道附近的key cache幅度显著更高。在第20层第一个head的第50个和第150个通道中观察到类似的模式。基于此观察，Liu等人（【32，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024】）提出在key cache通道上执行量化。除了量化，作者的发现表明，对注意力机制贡献较小的特定key cache通道可以被剪枝。此外，通道量化和剪枝是正交的技术，这意味着它们可以同时应用以进一步提高模型效率。

奇异值分析。作者对指定层的注意力权重进行了奇异值分解（SVD）（【10，Applied numerical linear algebra+1997】），以研究它们的主成分。从SVD导出的奇异值捕获了注意力矩阵的有效秩，表明信息如何跨不同组件分布。

奇异值的能量分布。附录A中的图5(a)说明了奇异值的能量分布，在对数尺度上绘制以增强差异的可见性。值得注意的是，在所有head和层中，只有少数奇异值表现出超过$0.01$的高能量水平，突出了它们的相对重要性。这一观察结果与先前的发现（【3，Eigen analysis of self-attention and its reconstruction from partial computation+2021】）一致，即一小部分奇异值通常捕获注意力机制中的大部分信息。此外，能量的快速衰减表明低秩近似可以有效地捕获key cache中的基本信息。

注意力矩阵的低秩特性。附录A中的图5(b)中，归一化的累积能量总和表明前50个奇异值占总能量的$90\%$以上。这些发现表明注意力矩阵本质上是低秩的（【48，Linformer: Self-attention with linear complexity+2020】；【7，Scatterbrain: Unifying sparse and low-rank attention+2021】；【14，Get more with less: Synthesizing recurrence with kv cache compression for efficient llm inference+2024】），表明key cache可以使用低维向量来近似（【41，Loki: Low-rank keys for efficient sparse attention+2024】）。

方法细节

符号定义。作者使用大写字母（例如，$X, Y$）表示标量值，使用粗体大写字母（例如，$\mathbf{Q}, \mathbf{K}$）表示矩阵和张量。符号$\|\cdot\|_p$表示向量的$l_p$范数。除非另有说明，$\|\cdot\|$表示$l_2$范数。Frobenius范数由$\|\cdot\|_F$表示。向下取整函数由$\lfloor\cdot\rfloor$表示，向上取整函数由$\lceil\cdot\rceil$表示。

KV cache优化的初步研究。在具有扩展上下文或批处理的场景中，内存和速度方面的主要限制是由于处理KV cache大小引起的。考虑对大型语言模型（LLM）服务的一批请求，该服务提供由token $[x_{B1}, ..., x_{BS}]$组成的长输入提示，总KV cache大小可以计算如下：$2 \times B \times S \times L \times N \times D$，其中$L$是层数，$N$是head数，$D$是head维度。KV cache大小随着batch size $B$和序列长度$S$线性增长。对于具有多头注意力（MHA）（【46，Attention is all you need+2017】）的模型，例如LLaMA2-7B（【45，Llama 2: Open foundation and fine-tuned chat models+2023】），$2048$的上下文长度和$13$的batch size需要存储13 GB的KV cache，这等同于模型参数的大小。对于生成的每个token，必须将KV cache从片外内存（HBM）（【22，Dissecting the nvidia volta gpu architecture via microbenchmarking+2018】）传输到片内内存（cache），从而导致内存瓶颈。这种巨大的内存需求凸显了管理大规模模型的挑战以及对高效内存利用策略的需求。当前的方法基于序列长度$S$（【51，Efficient streaming language models with attention sinks+2023】；【57，H2o: Heavy-hitter oracle for efficient generative inference of large language models+2024】；【28，Snapkv: Llm knows what you are looking for before generation+2024】）和精度（【20，Kvquant: Towards 10 million context length llm inference with kv cache quantization+2024】；【32，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024】）来优化KV cache。作者将引入一种新方法THINK，从head维度数量$D$的角度对其进行优化。

基于幅度的剪枝。基于图4中的观察，该图描绘了不同通道之间幅度的显著变化，一个直接的标准是使用幅度的范数来衡量key cache中不同通道的重要性。其计算公式为：
$M_{n,d} = \|\mathbf{K}[n, :, d]\|_p$。给定剪枝率$\lambda$，作者仅在每个head的$D$个通道中保留$T = \lfloor(1 - \lambda)D\rfloor$个最重要的通道：$I = \mathbf{Top}_T(M, T)$，其中$\|\cdot\|_p$是每个通道的$l_p$范数。$n \in [1, N]$和$d \in [1, D]$是key cache中head和通道的指示符。$I \in (\mathbb{Z}^+)^{N \times T}$存储每个head张量$M$中前$T$个值的指示符。

基于幅度剪枝的性能分析。在表1中，作者展示了将具有各种剪枝率的key cache剪枝应用于LLaMA-3-8B模型的结果。作者利用$l_1$和$l_2$范数作为评估标准，并使用LongBench基准（【2，Longbench: A bilingual, multitask benchmark for long context understanding+2023】）验证性能。与基线方法H2O（【57，H2o: Heavy-hitter oracle for efficient generative inference of large language models+2024】）和SnapKV（【28，Snapkv: Llm knows what you are looking for before generation+2024】）（两者的KV长度均为512）相比，作者进一步修剪了key cache的通道。$30\%$的剪枝率可以保持准确性；然而，将其增加到$40\%$会导致显著的性能下降，特别是对于基于$l_1$范数的剪枝。基于幅度的剪枝结果支持了作者的假设，即key cache在通道维度上是冗余的。这些结果还表明需要更好的剪枝指标来有效地实现更高的剪枝率。

查询驱动的剪枝。对于每个head，使用查询（queries）和键（keys）计算注意力分数，然后应用于值（values）。head $i$的注意力公式为：$\mathrm{Attention}(\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i) = \mathrm{softmax}(\frac{\mathbf{Q}_i\mathbf{K}_i^T}{\sqrt{D}})\mathbf{V}_i$，其中$\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i \in \mathbb{R}^{S \times D}$。当$\mathbf{K}_i$的一个通道被剪枝时，$\mathbf{Q}_i$中对应的通道也将被移除。作者旨在找到要剪枝的最佳通道子集，由选择矩阵$\mathbf{S} \in \{0, 1\}^{D \times D}$表示，其中$\mathbf{S}$是一个具有二进制条目的对角矩阵（$1$表示保留通道，$0$表示剪枝通道）。为了更好地保持剪枝通道后的性能，作者最小化了原始和剪枝注意力权重之间差异的Frobenius范数：$\min_{\mathbf{S}} \|\mathbf{Q}_i\mathbf{K}_i^T - \mathbf{Q}_i\mathbf{S}(\mathbf{K}_i\mathbf{S})^T\|_F$。给定剪枝率$\lambda$，它可以进一步展开为：
$\min_{\mathbf{S}} \|\mathbf{Q}_i\mathbf{K}_i^T - \mathbf{Q}_i\mathbf{S}\mathbf{K}_i^T\|_F$
受限于 $\mathrm{trace}(\mathbf{S}) = \lfloor(1 - \lambda)D\rfloor$，且 $\mathbf{S} = \mathrm{diag}(s_1, s_2, \ldots, s_D)$，其中 $s_j \in \{0, 1\}$。

优化选择矩阵的贪婪算法。为了简单起见，作者使用贪婪算法来优化$\mathbf{S}$。为了实现剪枝目标，作者定义了一个评估每个通道重要性的标准，并贪婪地选择得分最大的通道：$\mathrm{Score}_i[j] = \|\mathbf{Q}_i[:, j]\mathbf{K}_i[:, j]^T\|_F, \quad I_i = \mathbf{Top}_T(\mathrm{Score}_i, T)$。以下是它为何优化选择矩阵的详细解释。$\mathrm{Score}_i[j]$测量了每个head $i$中通道$j$的查询和键向量之间相互作用的幅度。通过选择具有最高相互作用幅度的通道，作者旨在保留对注意力机制最显著的贡献。该标准确保所选通道保留了注意力计算中的主要信息流，从而最大限度地减少重要信息的丢失。

观察窗口的设定。遵循SnapKV（【28，Snapkv: Llm knows what you are looking for before generation+2024】），为了降低计算成本，作者仅使用最后的$S_{obs}$窗口来计算分数，因为输入序列的最后一个窗口识别出与生成高度相似的注意力模式：$\|\mathbf{Q}_i[-S_{\mathrm{obs}}:, j]\mathbf{K}_i[:, j]^T\|_F$。

THINK的实现过程。作者选择不对最近的token和新生成的键进行剪枝。因此，作者的键值（KV）cache有两个类别：一个子集由通道大小减小的剪枝键组成，而另一个子集保持其原始大小。此外，作者维护一个二进制掩码（其内存开销可以忽略不计），以指示哪些通道已被剪枝。图2说明了THINK在解码阶段的一种实现。首先，作者使用存储的掩码对查询进行剪枝，确保查询维度和剪枝后的key cache保持一致。接着，将剪枝后的查询乘以剪枝后的键，同时将未剪枝的查询应用于未剪枝的键。最后，将两个输出连接起来。作者的方法去除了不重要的Key cache通道以获得真正的内存减少。

实验环境

• 数据集：

  • LongBench：用于综合评估长上下文理解能力，包含17个数据集，覆盖6种任务（单文档问答、多文档问答、摘要、少样本学习、合成任务、代码补全），平均输入长度为6,711词。
  • Needle-in-a-Haystack (大海捞针)：测试模型在长文档中准确定位关键信息的能力，评估KV缓存压缩方法是否会在压缩过程中丢失关键信息。

• 模型架构：LLaMA-2-7B-chat, LLaMA-3-8B-Instruct, LLaMA-3-70B-Instruct, Mistral-7B-Instruct-v0.2。

• 硬件配置：NVIDIA A100 GPUs。
• 软件配置：基于HuggingFace实现。基线方法包括H2O、SnapKV（基于Token驱逐）和KIVI（基于KV缓存量化）。

实验结果

1. LongBench长上下文理解测试（表2、表3、表4、表9）：

   • 实验内容：在不同KV-size（128至2048）下，比较H2O、SnapKV以及它们与THINK结合（不同剪枝率$\lambda$）的性能。
   • 结果与结论：
     • THINK在H2O和SnapKV压缩的基础上，成功对Key缓存通道进一步剪枝。在LLaMA-3-8B上，THINK不仅降低了内存，甚至在H2O和SnapKV上略微提升了性能；在Mistral-7B上，内存减少的同时仅有极小的性能下降；在LLaMA-3-70B上，剪枝40%通道后性能与SnapKV基线持平或更优。
     • 当剪枝率$\lambda=0.4$（即总KV内存减少20%）时，查询驱动的通道剪枝方法显著优于基于$l_1$或$l_2$范数的方法。
     • 随着KV-size从128增加到2048，THINK的性能进一步提升。特别是在KV-size 2048且$\lambda=0.4$时，甚至超越了全量KV缓存的LLaMA-3-8B基线性能。
     • 与KIVI（量化到2-bit）结合时，THINK(0.4)在进一步减少20%内存的情况下，几乎没有性能下降（表4）。
     • 将THINK直接应用于未经其他压缩的Vanilla模型（表9），在剪枝40%的情况下依然保持了卓越的性能。

2. Needle-in-a-Haystack测试（表5、附录图6）：

   • 实验内容：在不同KV-size下，评估SnapKV结合THINK (剪枝率0.4, 0.5, 0.6) 的精确信息检索能力。
   • 结果与结论：当$\lambda=0.4$时，THINK在所有模型和KV-size下均达到或超过了原始SnapKV的准确率，证明该方法能有效保留信息通道剔除噪声。当$\lambda \ge 0.5$且KV-size较小（128, 512）时，准确率有所下降；但在较大的KV-size（1024, 2048）下，THINK依然保持了与SnapKV相当的鲁棒性能。图6的可视化进一步表明，THINK甚至成功找回了一些SnapKV遗漏的“针”。
     

3. 消融实验：Recent Size的影响（表6）：

   • 实验内容：评估保留最近KV embedding数量（0, 32, 128）对Mistral-7B模型性能的影响。
   • 结果与结论：Recent-size为32的表现显著优于0，证明了保留最近KV的重要性。而32与128之间的性能差异可忽略不计，表明保留最近32个KV即可维持最佳性能。

4. 同等内存消耗下的性能对比（表7、图3a）：

   • 实验内容：调整H2O或SnapKV的KV-size，使其与结合了THINK(0.4)的方法占用完全相同的内存，进行公平对比。
   • 结果与结论：在完全相同的内存限制下，结合THINK的H2O或SnapKV始终优于未结合THINK的基线方法，证明了查询驱动通道剪枝与KV压缩方法集成的有效性。

5. 内存使用与吞吐量对比（图3b、附录表10、表11）：

   • 实验内容：评估KIVI结合THINK(0.4)在LLaMA-2-7B上不同Batch Size下的峰值内存、首字延迟（TTFT）、单字生成时间（TPOT）和吞吐量。
   • 结果与结论：THINK极大缓解了内存瓶颈。在相同内存下，KIVI+THINK能将KIVI原本4倍的Batch Size提升至5倍。TPOT从0.27 ms/token降至0.25 ms/token（40%剪枝），吞吐量从5168 tokens/s提升至5518 tokens/s。由于在线计算通道重要性，TTFT略有增加（从7.0ms增加至10.4ms）。
     

6. Value Cache剪枝探索（附录表8）：

   • 实验内容：同时对Key和Value缓存的通道应用不同比例的剪枝。
   • 结果与结论：在LLaMA-3-8B上，同时剪枝Key和Value通道的性能与未剪枝相当甚至更好。但在Mistral-7B上，剪枝Value通道会导致轻微的性能下降。这印证了图4的观察：Value缓存通道幅度更均匀，稀疏性较低，因此识别冗余通道更困难，但仍具备进一步压缩内存的潜力。

结论

论文针对大语言模型长上下文推理中的KV缓存内存瓶颈，观察到Key缓存通道幅度的极度不平衡和低秩特性，提出了一种名为THINK的查询驱动剪枝方法。该方法基于注意力分数进行优化，为每个输入查询保留最关键的信息通道。THINK作为一种即插即用的模块，能够与现有的Token级KV缓存驱逐技术（如SnapKV、H2O）和量化技术（如KIVI）无缝集成。大量实验证明，THINK在缩减40%的Key缓存大小（总KV内存减少20%）的情况下，依然保持甚至超越了基线的性能。

局限性与未来工作：由于需要在线计算通道重要性，THINK略微增加了首字生成时间（TTFT）。此外，虽然该方法在Key缓存剪枝上非常有效，但Value缓存的剪枝仍是未来探索的开放方向。

补充细节

结合量化的实现。图7说明了当与KV cache量化方法KIVI（【32，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024】）集成时作者方法的实现。在预填充（prefill）阶段，作者首先在应用量化之前剪枝$X_K$的不重要通道。在解码阶段，每个新到达的key cache $t_K$被添加到$X_{K_r}$。一旦$X_{K_r}$达到$G$个token（KIVI中的残差长度超参数），作者对数据进行剪枝和量化，然后将其与先前量化的$Q(P(X_{K_g}))$连接起来。

Value cache剪枝的策略。与用于Key cache的方法类似，Value cache中通道的剪枝可以由两个主要标准指导：基于幅度的剪枝和查询驱动的剪枝。作者发现查询驱动的剪枝仍然优于基于幅度的剪枝。其公式为：$\mathrm{Score}_{v,i}(\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i)[j] = \|\mathrm{softmax}(\frac{\mathbf{Q}_i[-S_{\mathrm{obs}}:]\mathbf{K}_i^T}{\sqrt{D}})\mathbf{V}_i[:, j]\|_F$，且$I_i = \mathbf{Top}_T(\mathrm{Score}_{v, i}, T)$，其中$\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i \in \mathbb{R}^{S \times D}$。作者定义了一个标准$\mathrm{Score}_{v, i}$来指示value cache的head $i$中每个通道的重要性。然后，仅保留前T个通道。表8报告了同时剪枝key和value通道的结果，表明剪枝Value cache通道比剪枝Key cache通道更困难。然而，在LLaMA-3-8B-Instruct上对Key cache和Value cache同时剪枝$30\%$仍然能产生可接受的性能。这证明了Value cache在通道维度上也有剪枝的潜力。如图4所示，Key cache沿通道维度的幅度高度不平衡，而Value cache沿通道维度的幅度更均匀。这表明Value cache中的通道稀疏性不如Key cache中显著，从而使得识别用于剪枝的冗余通道更具挑战性。作者将把value cache剪枝策略作为未来工作的一部分进行研究。

与基于SVD的方法的比较。在本节中，作者将THINK与基于SVD的方法（【39，Eigen attention: Attention in low-rank space for kv cache compression+2024】；【53，Asvd: Activation-aware singular value decomposition for compressing large language models+2023】；【6，Palu: Compressing kv-cache with low-rank projection+2024】）进行了比较。基于SVD的KV cache压缩方法使用一些校准数据离线分解KV cache权重，依赖于存储潜在表示并在推理期间恢复cache。作者的方法是一种即插即用的在线剪枝策略，不需要更改模型的权重或架构。这使得THINK轻量级且易于集成到现有系统中。此外，作者的方法直接动态剪枝Key cache通道，而不需要任何重建。作者在表12中提供了作者的方法与基于SVD的方法的比较。遵循Palu，作者使用Mistral-7B-Instruct-v0.2作为基线模型。作者的结果表明，当KV cache被压缩$80\%$时，作者的方法保持了性能，而Palu经历了$1.28\%$的准确率下降。同样，对于ASVD，作者在相同的压缩率下进行评估。作者的方法表现出明显更好的性能，与ASVD观察到的$5.21\%$的性能下降相比，仅有$0.12\%$的性能下降。结果证明了THINK的有效性，在相同压缩率下，与基于SVD的方法相比，性能下降明显更少。作者的剪枝方法有潜力与基于SVD的方法集成。探索这种结合可能会在KV cache压缩方面产生进一步的进步，作者旨在在未来的工作中对此进行研究。