只会开新坑不会填旧坑的屑Mr.Xau🕊️🕊️🕊️ 又打算分享一篇基于KV Cache优化经典工作

本人近期任务太多以后一定会填坑的(下次一定 🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺🥺

StreamLLM 基于注意力汇聚现象的KV Cache策略

论文地址,Streaming这个词很有意思,不是我们玩的steam ,它可以是喷射的意思,也可以是涓涓细流的意思;我觉得从这篇工作的内容来看,翻译为娟娟溪流可能更加合适一点。那么这个娟娟细流到底指的是LLM中的什么,但是正所谓铁打的衙门流水的官,KV Cache中有没有铁打不变的东西呢?且听后文分析。

知识补充

主要是KV Cache的介绍:什么是KV Cache,为什么只缓存KV?

什么是KV Cache?

回忆一下Transformer中的注意力机制,在经典的Transformer中我们有向量化的语料输入Tokens序列,如果batch size=1, 的形状是 ,其中是输入序列的长度 一句话单词数是给每个单词的编码向量的维度 ;经过注意力编码之后有(暂时以一个头的注意力为例子):

我们通过这个编码之后的矩阵,可以计算输出的注意力分数:

我们把展开写看看:

如果采用仅解码器的架构,由于掩码的存在,会有:

我们都知道每次LLM会将上一次输出的一个token放在下次输入的最后一个,那么下一轮的注意力分数是:

显然因为掩码的存在,注意力分数变成了一个下三角矩阵,因此当我们产生了新的tokens的时候,只需要计算最新tokens的q查询向量即可;这样的结构我们很自然就能想到,每次LLM进行推理的时候将之前都缓存起来,下次计算注意力分数的时候便无须再次重新计算了;

从上面的公式的形状,我们也不难发现,完全没有必要缓存矩阵,因为每次下三角矩阵扩展的时候都只用到了这一列,所以,之前计算的对于后面的计算是没有作用的我们可以完全舍弃以节省显存开销。

KV Cache有什么问题

随着输入序列的增加,如果缓存所有的矩阵,那么对显存的需求将会二次上升。这就要求我们设计高效的KV Cache的动态调整策略。详细的KV Cache大小估算公式可以参考附录。

这篇工作为什么叫StreamLLM,我的理解是因为即使随着对轮对话的进行,上下文越来越长,StreamLLM也能够基于注意力汇聚这一现象合理的关键的初始和一定上文窗口的cache来使得即使上下文增长也能像涓涓细流一样稳定流畅的推理。


蓑鱿剪切线


附录

to be continue