大语言模型KV缓存压缩策略详解：从Streaming LLM到PyramidKV
大语言模型KV缓存压缩策略详解：从Streaming LLM到PyramidKV
随着大语言模型的不断发展，如何处理长文本和多轮对话成为了一个重要问题。本文介绍了几种主流的KV缓存压缩策略，包括Streaming LLM、SnapKV、H2O和PyramidKV等方法，并对比了它们的优缺点。这些策略在保证模型性能的同时，有效地解决了显存不足的问题，为大语言模型的进一步发展提供了重要参考。
一、Streaming LLM(LCLR 2024)
1.问题
现有的大语言模型无法处理长度过长的文本，比如较长的QA，多轮对话等等。主要的原因有以下两个：
（1）KV会随着文本长度的变长会变得越来越大，GPU的显存是存不了太多的KV；
（2）由于LLM在预训练的时候，文本窗口是有限的（e.g. 每次只能输出4096个token），所以如果输出的内容太长，是进行不下去的。
2.重要性
类似于长文本生成、多轮对话的应用现有的LLM没有一个很好的解决方案
3.现有的解决方案以及局限性
（a）没有任何改变，当文本过长时，会出现显存爆炸的错误；
（b）只保存最后L个token的KV状态，这样虽然能够处理长文本，但是一但初始token的KV被剔除后，性能下降的很厉害；
（c）最后L个token的KV重新算，虽然保证了性能，但是计算很慢。
4.创新方法
文章观察到了attention sink的现象，就是说attention值比较大的都集中在了初始的几个token上面，所以一旦把初始的token剔除掉后，会导致softmax的时候数值的分布发生较大的变化，所以streaming llm采取以下策略：缓存初始4个tokens的KV，再加上最后L个tokens的KV，中间的KV全部剔除。
5.实验结果
与重新计算相比，streaming llm能够提高22倍的性能
6.局限性
streaming llm虽然能够处理很长的文本，但是这个文本的上下文依赖不能很强，如果文本有着很强的上下文依赖关系，很有可能把前面文本的KV剔除掉，从而导致生成文本的正确性收到较大的影响，streaming llm是不具备记忆性的长文本输出系统。
二、SnapKV(2024)
1.问题
现有的一些KV压缩策略都是针对生成部分的内容进行压缩，但prompt的长度往往更长
2.创新方法
prompt可以分为两个部分：前缀部分和观察窗口部分，文章认为，观察窗口对于KV的压缩非常关键，即使前缀部分完全相同，但是观察窗口不同，也会导致attention在prefix中的分布不同。
首先，使用观察窗口中的q和所有k作attention操作，仅对观测窗口作mask操作，然后进行softmax。将前缀部分进行聚集操作，即把他们相加起来，然后再进行一个池化操作，这是为了更好的关注上下文，然后找到前k个值最大的sum_attn所在位置。
最后，将这些位置对应的KV取出，和观测窗口的KV一起构成新的KV
3.实验结果
3.1 Needle-in-a-Haystack
3.2 Decoding Speed and Memory Bound
3.3 Ablation Study of Effectiveness of Pooling
4.局限性
仅对prompts进行了KV压缩，未考虑生成部分的，而且无法处理过长的prompts
三、H2O(2023 NeurIPS)
1.问题
KV太大，显存成了瓶颈
2.创新方法
文章中的两个观察：
Q-K的相乘矩阵是一个稀疏矩阵，有95%左右的数据都基本为0，是一个稀疏矩阵；
Q-K矩阵中有用的值集中在一个地方
H2O和snapKV基本差不多，都是用后面一部分Q-K矩阵来选择应该保留哪些KV，与snapKV不同的是没有池化这一步
H2O保留的KV主要有两个部分：最后tokens的KV，和前面一些比较重要token的KV。
3.实验结果：
文章用吞吐量来进行衡量实验结果，即每秒钟能够产生多少个tokens
四、PyramidKV(2024)
1.创新方法
文章的一个观察：
较低层的Q-K矩阵数值分布的比较均有，越往上面，数值就越集中，数据出现了聚集现象，这就说明了底层需要的KV cache比较多，顶层需要的比较少，因此每一层的KV cache出现了金字塔分布情况。
PyramidKV仍然完整保留最后一部分tokens的KV cache，对于前面的tokens也是选择最重要的tokens所对应的KV，选择方法与前面方法基本相同，唯一不一样的是每一层的KV cache大小不一样，呈现金字塔分布。
2.实验结果
目前的KV压缩方法，PyramidKV取得的效果最好
总结
前面的4中KV压缩方法，都有一个特点都是完整保留最后一部分tokens的KV值，然后对于前面tokens的KV，不同的方法去选择他们认为比较重要tokens的KV，把这两部分的KV拼接起来便压缩了KV。
最近tokens的KV所有方法都保留了，虽然每种方法都有自己的说法，说明了最近tokens KV的重要性。
共享KV权重
共享权重KV的4个工作
Multi-Query Attention（MQA 2019' CoRR）
Grouped-Query Attention（GQA 2023' EMNLP）
You Only Cache Once（YOCO 2024'CoRR）
Cross Layer Attention（CLA 2024'CoRR）
MQA与GQA采用了层内KV共享方法，在多头注意力机制中，每个头的Q都独立使用一对KV，不同头之间的KV权重不共享，MQA则只保留了一个头的KV，所有头的Q都共享一个头的KV，不同头之间的差异只通过Q来反映；GQA则采取一种折中的方法，首先对不同头的Q进行分组，然后每一组Q共享一组KV，如果分组数量为1，则对应的就是MQA，如果分组数量为注意力头数，则对应原始的transformer。
YOCO
CLA
YOCO和CLA都是采用层间共享的方法来对KV进行压缩，YOCO在前1/2层采用了self-attention，这一部分的计算是用来产生KV的，后1/2层采用了cross-attention，不再产生KV，而是使用前1/2产生的KV，CLA并不是选择固定的前几层来产生KV Cache（比如YOCO，使用的是前L/2层），而是将产生KV Cache的层交替分布在模型不同深度的层，然后邻近层复用附近层产生的KV Cache进行Attention计算。