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 用于后续处理，从而避免显存的过度扩展。