MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

自回归推理中的Transformer模型非常受Key-Value（KV）缓存的好处，但随着模型大小、批量大小和序列长度的增长，可能会导致主要的内存瓶颈。我们引入了多层Key-Value（MLKV）共享，这是一种新颖的方法，将KV共享扩展到Transformer层，以减少内存使用，超越了多查询注意（MQA）和分组查询注意（GQA）所能实现的范围。在各种NLP基准测试和推理指标上进行评估，使用未经训练的Pythia-160M变量，证明MLKV显著减少了内存使用，同时最小化性能损失，将KV缓存大小降低了6倍，相对于MQA。这些结果突出了MLKV在大规模有效部署Transformer模型方面的潜力。我们在https://github.com/zaydzuhri/pythia-mlkv提供了代码。

本论文旨在解决Auto-regressive inference of transformers中Key-Value (KV) caching所带来的内存瓶颈问题，探讨一种新的方法以减少内存使用。

本论文提出一种新的方法——Multi-Layer Key-Value (MLKV) sharing，将KV共享扩展到transformer层，以进一步减少内存使用。

通过在各种NLP基准测试和推理指标上进行评估，本论文证明了MLKV可以显著减少内存使用，而性能损失很小。相比Multi-Query Attention (MQA)，MLKV将KV缓存大小降低了6倍。此外，本论文提供了开源代码。

最近在这个领域中，还有一些相关的研究，如Multi-Query Attention (MQA)和Grouped-Query Attention (GQA)。