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

本文介绍了“Keyformer”，一种创新的推理时间方法，以缓解与KV缓存大小和内存带宽利用相关的挑战。Keyformer利用了这样一个观察结果，即生成推理中约90%的注意力权重集中在一组特定的标记上，称为“关键”标记。Keyformer通过使用一种新的评分函数识别这些关键标记，仅保留KV缓存中的关键标记。这种方法有效地减少了KV缓存大小和内存带宽使用，同时不影响模型的准确性。我们评估了Keyformer在三个基础模型上的性能：GPT-J、Cerebras-GPT和MPT，这些模型采用了各种位置嵌入算法。我们的评估涵盖了各种任务，特别是涉及扩展上下文的摘要和对话任务。Keyformer减少KV缓存可以将推理延迟降低2.1倍，将标记生成吞吐量提高2.4倍，同时保持模型的准确性。

解决Key-Value Cache大小和内存带宽利用方面的挑战，提高大型语言模型的生成推理效率。

Keyformer通过识别关键令牌并仅保留关键令牌来减少KV缓存大小和内存带宽使用，从而提高生成推理效率。

实验表明，Keyformer可以将推理延迟降低2.1倍，提高令牌生成吞吐量2.4倍，同时保持模型的准确性。该方法在GPT-J、Cerebras-GPT和MPT等模型上进行了评估，并涵盖了多种任务，特别是涉及扩展上下文的摘要和对话任务。

最近的相关研究包括使用不同的技术来提高大型语言模型的效率和性能，如压缩、量化、分布式训练等。相关论文包括《Compressing Transformer-based Language Models by Sparse Factorization》、《Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference》、《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》等。