【摘要】大型语言模型（LLM）在处理大量上下文方面取得了显著进展，键值缓存在提高其性能方面发挥了至关重要的作用。然而，KV缓存因输入长度增加而增长，这对内存和时间效率提出了挑战。为了解决这一问题，本文介绍了SnapKV，这是一种无需微调的创新方法，可以有效地最小化KV缓存大小，同时在实际应用中仍能提供相当的性能。 我们发现模型中的每个注意头在生成过程中始终专注于特定的即时注意特征。同时，可以从位于提示末尾的“观察”窗口中获得这种稳健模式。利用这一见解，SnapKV通过为每个注意力头选择群集的重要KV位置来自动压缩KV缓存。我们的方法显著降低了处理长输入序列时不断增长的计算开销和内存占用。具体而言，在处理16K令牌的输入时，SnapKV实现了一致的解码速度，与基线相比，生成速度提高了3.6倍，内存效率提高了8.2倍。同时，它在16个长序列数据集上保持了与基线模型相当的性能。此外，SnapKV可以使用HuggingFace实现在单个A100-80GB GPU上处理多达38万个上下文令牌，只需稍加更改，在大海捞针测试中仅表现出微不足道的准确性下降。进一步的综合研究表明SnapKV具有实际应用的潜力。

原文：SnapKV : LLM Knows What You are Looking for Before Generation
地址：https://arxiv.org/abs/2404.14469v1
代码：https://github.com/FasterDecoding/SnapKV
出版：未知
机构: University of Illinois Urbana-Champaign, Cohere \

1 研究问题

本文研究的核心问题是: 如何减少大语言模型(LLMs)在处理长上下文时的key-value(KV)缓存开销,提高推理效率。

当前LLMs在聊天机器人等实际应用中,经常需要处理非常长的上下文,如多轮对话历史、冗长的文章等。随着上下文长度的增加,LLMs在注意力计算时使用的KV缓存也会线性增长,导致存储和计算开销剧增。这就好比一个图书管理员要管理一个rapidement expanding的图书馆,如何在有限的时间和空间内高效地检索每本书的信息,成为一个棘手的问题。

本文研究问题的特点和现有方法面临的挑战主要体现在以下几个方面:

• 现有的KV缓存压缩方法主要关注生成阶段的优化,忽略了输入阶段KV缓存的压缩。但在实践中,输入的tokens数量往往远大于生成的tokens,因此输入阶段的KV缓存是内存和计算瓶颈所在。

• 压缩输入阶段的KV缓存面临信息丢失的风险,尤其是在存在大量噪声上下文的情况下,如何保证压缩后的KV缓存仍能准确捕捉到关键信息是一个挑战。

• 不同的用户指令会关注输入上下文中的不同部分,如何根据指令动态调整KV缓存的压缩策略是现有方法欠缺考虑的。

针对这些挑战,本文提出了一种灵活高效的"SnapKV"方法:

SnapKV的核心思想是:LLMs在生成过程中对输入tokens的关注模式具有很强的一致性和稳定性,而这种模式可以在输入序列末尾的一个小窗口内就观察到。SnapKV利用这个洞见,只保留那些被注意力heads持续关注的输入tokens对应的KV,大幅压缩KV缓存的尺寸。同时,SnapKV还引入了池化等聚类技术,既能捕捉局部的完整信息,又能过滤掉大部分无关tokens。这就像图书管理员先根据书的类别和借阅频率把书分门别类存放,再只在每个类别中选取一些代表性的书籍摆放在显眼位置,既提高了检索效率,又不影响读者找到想要的书。SnapKV的设计完全基于对LLMs注意力模式的观察,不需要对预训练模型做任何修改,因此使用起来非常简单灵活。实验表明,在多个长文本数据集上,SnapKV在大幅降低推理开销的同时,性能与使用完整KV相当。

2 研究方法

2.1 SnapKV整体流程

本文提出了SnapKV,一种高效压缩大型语言模型(LLMs)中Key-Value(KV)缓存的方法。在处理长上下文时,KV缓存的急剧增长会带来显著的内存和计算开销。SnapKV旨在以极小的性能损失换取可观的效率提升。它的核心思想源自一个观察:在生成过程中,注意力头会一致地关注Prompt中的特定位置。SnapKV主要包含两个步骤:

1. 基于Observation Window对重要的Prompt信息进行Voting筛选;

2. 将筛选出的KV与Observation Window进行拼接,得到压缩后的KV缓存用于后续生成。

举个简单例子,假设Prompt序列长度为1000,Observation Window大小为16,我们希望将KV缓存压缩至256。SnapKV首先基于最后16个Token的注意力分布,通过Voting算法选出最重要的240个位置。然后将这240个位置对应的Key和Value与Observation Window拼接,形成大小为256的新KV缓存。这样,后续的生成过程就只需在显著缩减的KV缓存上进行注意力计算,从而大幅提升了效率。

2.2 基于Observation Window的Voting算法

SnapKV的第一步是利用Observation Window进行重要信息的筛选。Observation Window定义为Prompt的最后一个片段,实验表明该窗口能够很好地指示对后续生成至关重要的注意力特征。给定形状为的注意力张量,其中为注意力头数,为Observation Window长度,为Window之前的Prompt长度,Voting过程可以形式化为:

其中表示在张量的每个注意力头上选出前大的值对应的索引,定义为,超参数控制压缩率。选定后,即可从完整KV缓存中提取相应的Key和Value用于后续步骤。Observation Window的大小对算法性能有一定影响。太小可能无法捕捉足够的重要信息,而太大又会引入较多无关噪声,实验表明16-32之间的取值较为合适。

2.3 基于Pooling的重要信息Clustering

在初步实验中,论文作者发现仅根据注意力值选取离散位置,无法很好地保留完整的上下文信息,导致生成质量下降。为解决这一问题,SnapKV在Voting之前引入了基于Pooling的聚类操作。具体而言,在Sum注意力权重后,使用带Padding的1维Pooling层进行滑动窗口聚合:

其中为Pooling的Kernel大小。Pooling的作用是在选中某个位置的同时,也能捕捉其周围的上下文,从而使筛选更加鲁棒完整。的大小控制了聚类的程度,论文选取了3-7之间的值。之后的Voting过程就在Pooling后的注意力权重上进行。实验证实,引入Pooling不仅提升了下游任务性能,也使SnapKV对超参数更加鲁棒。

2.4 SnapKV算法的伪代码与实现

论文给出了SnapKV的PyTorch风格伪代码,如下所示:

其中的关键步骤包括:

1. 基于Observation Window计算Prompt各个位置的注意力权重之和行

2. 使用1维Pooling进行权重聚类,选出Top-k的位置索引行

3. 利用gather操作提取选中位置的Key和Value行

4. 将压缩的KV与Observation Window对应的KV拼接,得到最终结果行

可以看到,SnapKV的实现简洁高效,只需对现有LLM的KV缓存管理模块进行少量修改即可无缝集成。综上所述,SnapKV提供了一种实用的KV缓存压缩方案,在显著提升推理效率的同时保持了生成性能,为大型语言模型的实际应用铺平了道路。 第四步、实验部分详细撰写:

5 实验

5.1 实验场景介绍

本文探索了SnapKV在多个极长上下文语言模型中的表现,涉及任务包括Needle-in-a-Haystack、基准测试LongBench、搜索增强生成等,考察了SnapKV在推理速度、内存占用、生成质量等方面的影响。

5.2 实验设置

• Datasets: Needle-in-a-Haystack、LongBench (涵盖16个不同任务)、bioasq (搜索增强生成任务)

• Baseline: LWM-Text-Chat-1M、LongChat-7b-v1.5-32k、Mistral-7B等开源长上下文模型

• Implementation details:

  • SnapKV最大Prompt KV Cache设为1024、2048、4096等不同大小

  • Pooling操作使用Max Pooling,卷积核大小为5-13不等

• metric:

  • Needle-in-a-Haystack准确率

  • LongBench各子任务对应指标

  • 生成速度(ms/token)、内存占用

5.3 实验结果

实验一、LWM-Text-Chat-1M上的Needle-in-a-Haystack及速度内存测试

目的: 评估SnapKV在极长上下文模型中处理Needle-in-a-Haystack任务的表现,并对比推理速度和内存占用

涉及图表: 图5、图6

实验细节概述：Needle-in-a-Haystack考察模型从380k token长文本中准确找出单个目标句子的能力。使用SnapKV将Prompt KV Cache压缩至1024,只用单个A100-80G GPU即可处理如此长的文本,且准确率下降很小。速度和内存测试则考察不同批大小下SnapKV的加速和节省内存效果。

结果:

• SnapKV可将380k token长文本压缩380倍,且准确率仅有微小下降

• 在16k长度、批大小2时,SnapKV将推理速度提升3.6倍、最大可处理长度提升8.2倍

实验二、SnapKV中Pooling操作的消融实验

目的: 探究SnapKV中Pooling操作对模型表现的影响

涉及图表: 图7

实验细节概述：在更具挑战性的LongEval-Lines任务中,对比使用Pooling和不使用时模型表现。该任务需要在含噪上下文中识别特定格式的键值对。

结果: 使用Pooling可显著提升检索准确率。推测是因为Pooling有助于保留token间的上下文完整性,避免仅保留零散token而丢失信息。

实验三、在LongBench上评估SnapKV

目的: 全面评估SnapKV在多个长上下文开源模型和任务上的表现

涉及图表: 表1

实验细节概述：对比SnapKV在不同压缩率下与原始模型在LongBench 16个任务上的表现,涉及单文档问答、多文档问答、摘要、小样本学习等。平均输入token长度在1.2-1.3万。

结果:

• 即使压缩到1024 token,SnapKV在大多数任务上的表现与原模型相当,最高达到92%的压缩率

• SnapKV在所有任务上均显著优于另一压缩方法H2O

实验四、在Command-R上评估SnapKV用于搜索增强生成

目的: 评估SnapKV在搜索增强生成任务中的表现

涉及图表: 表2、表3、表4

实验细节概述：先在Command-R上重复Needle-in-a-Haystack实验,以更鲁棒的随机上下文设置测试。然后评估SnapKV在搜索相关文档、端到端问答表现。平均输入长度1.6万token,压缩4倍。

结果:

• 在修正的Needle-in-a-Haystack测试中,即使压缩32倍,SnapKV仍能保持与原模型相当的结果

• 在搜索相关文档任务中,SnapKV保留了原模型98.8%的F1

• 在端到端搜索问答中,SnapKV的回答质量下降2.1%,证明其鲁棒性

实验五、SnapKV与并行解码的兼容性分析

目的: 分析SnapKV与并行解码策略Medusa的兼容性

涉及图表: 图8

实验细节概述：将SnapKV集成到并行解码框架Medusa中,在不同长度输入下测试生成速度。

结果:

• 在1万token长度下,SnapKV+Medusa比单独Medusa提速1.3倍,比原生解码提速2.2倍

• 表明SnapKV能与并行解码协同工作,进一步提升长上下文处理效率

4 总结后记

本论文针对大语言模型(LLM)在处理长序列时存在计算和存储效率低下的问题,提出了一种名为SnapKV的压缩方法。通过对注意力机制中关键信息的挖掘和压缩,SnapKV能够在保持模型性能的同时,显著降低计算开销和内存占用。在多个长文本任务上的实验结果表明,SnapKV取得了与原模型相当的效果,同时大幅提升了推理速度和内存效率。

疑惑和想法：

1. SnapKV目前主要关注输入序列的压缩,是否可以进一步优化生成序列的KV cache管理?

2. 除了pooling之外,是否存在其他更高效的聚类算法可以用于压缩注意力矩阵?

3. SnapKV在压缩过程中是否会损失一些关键信息?如何在压缩率和信息损失之间取得更好的平衡?

4. 能否将SnapKV与知识蒸馏等模型压缩技术相结合,进一步提升LLM的效率?

可借鉴的方法点：

1. 挖掘注意力模式并据此压缩KV cache的思想可以推广到其他Transformer类模型,用于提升效率。

2. 感觉可以和prompt压缩进行结合。