在大型语言模子（LLM）中，KV Cache（键值缓存）的压缩方法及其紧张性。
为什么要压缩KV Cache？
- 盘算效率：在生成文本的过程中，每个生成的token都需要与之前全部的token的键值（K和V）进行交互。假如不使用KV Cache，则每次生成新token时都需要重新盘算，这将导致巨大的盘算量。
- 显存斲丧：KV Cache虽然镌汰了盘算量，但随着token数量的增长，它占用的显存也会增长。显存是有限的，因此需要对KV Cache进行压缩。
KV Cache的分类及方法
- 生成阶段压缩：大多数方法集中在生成阶段的KV Cache压缩，以支持长文本的生成或加快盘算。这些方法通常接纳驱逐策略，保留紧张的KV，舍弃不紧张的。
- ClusterKV：通过聚类KV Cache中的token，并使用聚类中心来判断紧张性，镌汰盘算量同时保持精度。
- LeanKV：根据每个attention head的需要保留不同的KV，分为高精度和低精度保留，动态调整以压缩KV Cache。
输入prompt的KV Cache压缩
- SnapKV：针对输入prompt太长导致的题目，通过识别输入中的模式来压缩KV Cache。它假设这些模式在生成过程中是一致的。
- ZigZagKV：在SnapKV的基础上进行优化，提出不应该给每个attention head分配相同大小的显存，而应根据每个head所需的最少token数来分配。
ZigZagKV的空间分配
- 题目：给模子的每一层分配相同大小的显存来存储KV Cache是不公道的，由于不同层的LMBA（维持attention score所需的最小KV数）差别很大。
- 解决方案：提出了一种按比例分配显存的方法，并为每个层保留一个最小大小，以避免某些层分配到的空间过小。
结果
- 文章最后提到，ZigZagKV在各种测试中表现最佳，特别是在大海捞针测试中。
总结
在大型语言模子中，如何通过压缩KV Cache来进步盘算效率和镌汰显存斲丧。不同的压缩方法针对不同的场景和需求，而ZigZagKV方法在保持模子性能的同时，有效地解决了显存分配的题目。
source：ZigZagKV，KVCache80%的压缩率可以达到近乎无损的结果
为什么要压缩？
输入Prompt和decoding阶段生成的每个token都会产生对应的K和V。每次生成一个新的token的时间，当前的需要和汗青上全部的K和V进行交互得到结果，以是我们一样平常把已经盘算出来的K和V都生存起来供背面的token使用，生存的结果就叫做KV Cache。
假如不使用KV Cache来生存汗青结果，那么每一个token生成的时间都需要重新盘算prompt和之前生成的tokens的K和V。显然，KV Cache减小了盘算量。但是，它增长了显存的斲丧。随着token的数量增多，需要越来越多的显存。而我们的显存显然不是无穷大的，这就要求我们需要做KV Cache的压缩。别的，随着KV的增多，每生成一个token的盘算量也越来越大(由于需要和全部的K和V都进行盘算)。这进一步增长了KV Cache压缩的需求。
KV Cache分类以及方法
KV Cache的内容来源于两个方面：
之前的大多数文章内里的KV Cache方法都集中在处理生成阶段压缩KV Cache，以此来支持生成长文本或者在生成过程中加快盘算。这一类KV Cache压缩的方法大多采取驱逐策略，保留体系认为紧张的KV，舍弃那些不紧张的。此中紧张性的判断非常紧张，大部门的体系都是对于每一层按照attention score的大小，保留top-B的结果。
前面的文章：
- ClusterKV，靠近无损的KV Cache压缩方法
- LeanKV, 靠近无损的5倍的KV Cache压缩以及2.5倍的吞吐提升
介绍了两种结果都不错的KVCache压缩方法。此中ClusterKV的紧张观察点是对于一个q紧张的全部token对应的K之间的cos距离都比较小。别的就是KV Cache内里某个token的紧张性是动态变化的。基于这两个观察，ClusterKV把KV Cache内里的token进行聚类，使用聚类的中心来做紧张性判断，假如一个类被判断为紧张那么就把这一个类内里全部的token都拿来做盘算，这样既镌汰了盘算量还保持了精度。而KVCache并不常驻显存，这样就可以镌汰显存使用。这篇文章的方法实在是显著不同于大多数接纳驱逐策略的方法，由于每个KV都没有完全被驱逐，都是可重入的。
LeanKV也是一篇做得非常过细的文章。核心观察是每一层的每个attention head需要保留的紧张KV不一样，而且具有稀疏性。基于这个观察，它不以layer为单位，而是以attention head为单位，去保留对应的紧张KV。保留的KV又会分为高精度保留、低精度保留。每次新生成了KV的时间，每个attention head会去盘算是否保留新的K和V，假如保留，是否需要驱逐之前的。用这种动态的方法，达到压缩KV又保持精度的结果。
别的一类文章，以SnapKV为代表，研究的是别的一个题目：输入太长，压缩输入prompt的KV Cache。这里的压缩并不是要让生成的token很长，而是不压缩的话可能输入就已经把显存撑爆。或者即便显存还可以，每一步盘算也会由于需要交互的token太多而非常慢。以是压缩输入的KV解决的是这两个题目。
ZigZagKV实在是在SnapKV上继承做优化，也就是说，它也是输入prompt的KV Cache压缩算法。它的核心改进点在于不应该给每个attention head分配同样多的显存，而应该根据每个head可以保留绝大部门信息需要的token数来进行分配(SnapKV给全部的layer的head是平均分配的)。
ZigZagKV的一些新的观察
我们起首定义MBA(MinimumBudget size required to maintain 90% of the totalAttention score)，实在就是维持90%的attention score需要的最小KV数：
此中，是的第i个元素。是当前KV Cache内里的token数量。
由于每个attention head都可以盘算出一个MBA值，我们把一层内里全部的attention head的MBA的均值叫做LMBA(LayerMinimumBudget size to maintainAttention score)：
LMBA越高也就意味着需要更多的token来维持精度，也就意味着需要分配更多的显存来存储KV Cache。
我们分别盘算了Mistral和LLaMA在WikiMQA数据集上的LMBA，结果如下：
可以看出，不同的模子的对应层的LMBA差距比较大；同一个模子不同层的LMBA差距也很大。
刚才我们分析是基于保持attention score，我们再以输出的相似性这个维度分析一下需要保留多少的KV才气让输出类似。我们定义LMBO(Layer-wiseMinimumBudget forOutput):
此中是用全部的KV得到的输出，是用一部门得到的结果。
我们再次发现，同一个模子不同层在保持输出基本一致的限定下需要的KV数量差距很大。
基于上面的两个观察，我们得出的结论就是给模子的每一层分配同样大小显存来存储KV Cache是不公道的。
ZigZagKV空间分配
为每一层分配多大的空间存储KV Cache是公道的呢？
一个直观的想法就是按照比例来分配，大层多份，小的少分。
上面公式内里的就是每一层按照LMBA盘算得到的比例，B是单层的Budget，是模子的L层总共的budget，是第层分到的大小。
这么分实在有一个题目，就是假如某一层的LMBA特别大的时间可能导致某些层分配到的空间特别小，也就意味着这些层基本无法生存任何KV。一个层基本没有KV Cache，结果自然包管不了。以是，我们给每个层保留一个最小大小，分配公式改为
这就是ZigZagKV的主要改进点。ZigZagKV和SnapKV都假设给prompt的KV Cache空间是，SanpKV给每一层均匀分配B空间。而ZigZagKV自顺应地分配空间。
解决了每一层KV Cache的空间分配题目，还有一个题目就是要生存那些K和V。这里需要先介绍SnapKV。
SnapKV
假如我们输入的prompt有16K，那么假如存储对应的全部KV无论对于显存照旧盘算都有极大的压力。那显然，我们需要压缩prompt生成的KV。为了压缩输入KV，SnapKV有一个非常紧张的观察：
Pattern can be identified before generation and is consistent during generation
什么意思呢？实在就是说，输入内里哪些紧张可以使用输入prompt本身就可以决定。
作者把输入128个token为单位分别window，取靠近尾部的20个window来做盘算。prompt前面的部门叫做prefix, 20个窗口叫做observation window。对于prefix内里的任意一个token，假设叫做。我们取20个window内里的某一个window，假设叫。的每一个元素和都会有交互，由于在的元素前面。那么，的每个元素和都会盘算出一个attention score。这样，对于窗口，会有128个attention score。取这128个attention score的均值作为的紧张性值。
这样，对于窗口，prefix的每个token都可以盘算出紧张性值。留意，我们盘算的时间是按照attention head为单位盘算的。盘算完之后就可以得到。此中是attention head的数量；是一个窗口内token的数量，这里实在就是128；是prompt除了观察的20个窗口的token之外其他的token的总数。
对于每个attention head，prefix内里每个token和窗口内的每个token都有一个attention score，把它们取平均得到最后的score。这是一个的矩阵。然后每个head可以根据平均score选出紧张的topK个token。
也可以根据设定的K值选出topK紧张的token，此中，代表压缩率。
在生成一个token的时间，对于当前step的q，它和prefix内里的全部token都会有交互。对于模子的每一层，可以盘算出的attention score。我们设定一个阈值，我们认为大于这个阈值的attention score对应的token对于当前生成的token来说都是紧张的。以此来选择出生成阶段prefix中的紧张token。
我们比较窗口选择出来的紧张token和生成阶段选择出来的紧张token的重叠率，也就是论文所说的Hit Rate。
用数学公式表达就是：
此中。公式稍微有点复杂，更多照旧看我的形貌吧，笔墨容易懂点。
通过比较20个窗口选择的紧张token和生成阶段现实选择出来的紧张token，我们发现最后一个窗口选择出来的和生成阶段选择的是高度重合的。也就意味着，Pattern can be identified before generation。
在生成了若干个token之后，我们把生成的token也像输入一样分成4个窗口，每个窗口包罗128个token。然后使用这四个窗口去选择prefix中的紧张token。最后发现，选择出来的和prompt内里最后一个窗口选择的是高度吻合的。
这也就是开始所说的Pattern is consistent during generation.
SnapKV基于这两个观察，放心地用prompt内里最后一个窗口来压缩输入prompt的KV Cache。现实上，作者也观察到只使用上面直接用窗口选择的内容，会造成生成的信息不全。以是，现实上的操作是在盘算完均值矩阵之后，对每一行都使用1D Pooling操作。其头脑很简单，比如第5个token很紧张，那么它相邻的4和6我们也应该选进来，不然5的上下文不全。而要选到4和6最简单的就是pool一下，可以简单理解为跟高斯含糊一个性质。
最后选择的token就是pool之后的topK，再机上最后一个窗口内里全部的token。
ZigZagKV的token选择
SnapKV内里token的选择跟上面的SnapKV非常类似。只不过每个head的配额少了，在相同压缩率的情况下topK的K对于每个head是变化的。以是对每个head按照配额大小盘算好K，其它的都跟上面的方法一致。
结果
大海捞针测试，ZigZagKV最好：
其它一些细分测试内里，ZigZag也是最好的。
参考文献
[1].ClusterKV，靠近无损的KV Cache压缩方法
[2].LeanKV, 靠近无损的5倍的KV Cache压缩以及2.5倍的吞吐提升
[3]. ZigZagKV: Dynamic KV Cache Compression for Long-context Modeling based on Layer Uncertainty
[4]. SnapKV: LLM Knows What You are Looking for Before Generation
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
