Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

KV缓存在加速基于transformer的自回归大型语言模型（LLMs）解码中发挥着至关重要的作用。然而，在长序列长度和大批量大小时，存储KV缓存所需的内存量可能变得禁止。自transformer发明以来，用于减小KV缓存大小的最有效的两种干预措施是多查询注意力（MQA）及其概括，分组查询注意力（GQA）。MQA和GQA都修改了注意块的设计，以便多个查询头可以共享单个键/值头，从而大大减少了不同键/值头的数量，同时仅最小程度地降低了准确性。本文表明，通过在相邻层之间共享键和值头，可以进一步推进多查询注意力，从而产生一种新的注意力设计，称为跨层注意力（CLA）。使用CLA，我们发现可以将KV缓存的大小再减小2倍，同时保持几乎与未修改的MQA相同的准确性。在从头开始训练1B和3B参数模型的实验中，我们证明CLA提供了一种帕累托改进，可以实现比传统MQA更好的内存/准确性权衡，使推理序列长度和批量大小比以前更大。

论文旨在解决KV缓存大小随着序列长度和批量大小增加而变得过大的问题，提出一种新的多头注意力设计Cross-Layer Attention（CLA）来进一步减小KV缓存的大小。

CLA通过共享相邻层之间的键和值头来减少KV缓存的大小，进一步优化了Multi-Query Attention（MQA）的设计。

论文在1B和3B参数模型的训练中进行了实验，证明CLA相比MQA在内存/精度权衡方面提供了Pareto改进，可以实现比传统MQA更长的序列长度和更大的批量大小的推理。

与MQA和GQA等先前的多头注意力设计相关的研究。