KeyFormer：使用注意力分数压缩KV缓存

Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

202403，发表在Mlsys

Introduction

优化KV cache的策略，主要是集中在系统级别的优化上，比如FlashAttention、PagedAttention，它们能够优化KV缓存的访问，但是没有减少缓存的读取。即使有MQA、GQA这样，能够真正减少缓存访问量的技术，也是需要在预训练阶段就用上。

有些其他丢弃不重要KV的方式，比如滑动窗口（下图b），它会限制模型捕获全面语义的能力

基于某些token比其他token更有意义的观察，Keyformer提出了一种丢弃不重要token的KV的方式。下图是这种方法和其他方法的对比。

具体来说，观察到90%的注意力权重会集中到关键token的小子集上。这些token对于LLM理解上下文至关重要，但可能超出了滑动窗口注意力的范围。

Keyformer引入了一种混合注意力方法，如上图(d)所示，它在生成下一个token时结合了最近的token和前面的关键token。

实验表明，Keyformer在与H2O等方法相比，有显著改进。这是因为Keyformer在识别关键标记时考虑了被丢弃标记的重要性，而不同于H2O仅基于注意力分数来识别“重头标记”。这句话有点难懂，还看不出来和H2O相比有什么优势

具体步骤

• 在生成过程中，Keyformer 会在prefill阶段丢弃掉 n - k 个 token，保持 k 个 token 用于后续处理，从而避免显存的过度扩展。

• 为了适应丢弃 token 后可能出现的挑战，Keyformer 引入了logits 正则化，即通过向未归一化的 logits 添加噪声来增强模型的鲁棒性。这可以帮助模型在推理时识别关键的 token，即使某些上下文信息缺失。

• Keyformer 选择使用 Gumbel 分布 来进行 logits 调整。Gumbel 分布有一个很好的特性，即它倾向于描述一组样本中的最大值，因此非常适合用于识别关键 token。也就是给归一化前的logits添加Gumbel分布的噪声。

• Keyformer有一个大小为w的滑动窗口，用于保留最近的w个token。因此每次decode用到的KV就是k-w个关键token的KV和w个最近token的KV

• 针对每一个token，有一个分数函数：

  其中是温度，和解码采样一样，温度越高越平均，温度越低越不平均；是Gumbel噪声；就是token的注意力分数。这个分数函数用来衡量token对当前解码的重要程度

• 这个注意力分数每次decode都累积

这个方法需要在显存中一直保存所有的KV缓存，只不过是在计算时只取部分缓存。没有减少显存占用，但减少了访存和计算。

实验部分

baseline：以full attention作为golden标准，目标是达到它的99%-99.9%；以滑动窗口注意力和H2O这两种方法进行比较。

实验从两个方面进行比较，一个是使用了多少比例的KV缓存，能够达到可接受的正确率；另一个是和baseline相比，吞吐量上有多少提升

正确率的实验结果：

横坐标百分之多少，表示保留了百分之多少的KV缓存。显示使用了70%的缓存，就能达到使用全部缓存的99%的效果。

在吞吐量提升的实验上，由于前面实验表明H2O如果只使用50%缓存的情况下，准确率是明显不如keyformer的；因此为了公平起见，就将keyformer设置为保留50% cache，而H2O保持90% cache。

那么显然keyformer的吞吐量会高于H2O

还做了消融实验，主要是这几点：

1. 计算每个token得分时，是对每个层进行累积分数还是每个层各保留一个分数：结果表明，每个层各保留一个分数的准确率更高

2. 位置编码使用原始文本中的位置进行编码，还是使用丢弃非关键token后的位置进行编码：前者更好

3. 不同正则化方式的影响，比较了无正则化、恒定数正则化、高斯正则化和Gumbel正则化，发现高斯正则化还不如不使用正则化，而显然是Gumbel正则化分数最高

4. 滑动窗口大小在比例为20%-30%时效果最好