假如我们输入的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选择
[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