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

📄 arXiv: 2405.12981v1 📥 PDF

作者: William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly

分类: cs.LG, cs.CL

发布日期: 2024-05-21

💡 一句话要点

提出跨层注意力（CLA）机制，在保证精度前提下进一步压缩Transformer KV缓存

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer KV缓存 注意力机制 模型压缩 大语言模型 推理加速 跨层注意力

📋 核心要点

  1. 现有Transformer模型KV缓存占用大量内存，限制了长序列和高并发场景的应用。
  2. 论文提出跨层注意力（CLA）机制，通过在相邻层间共享Key/Value头，进一步压缩KV缓存。
  3. 实验表明，CLA在1B和3B参数模型上实现了与MQA相当的精度，同时显著减少了内存占用。

📝 摘要（中文）

Key-value (KV) 缓存对于加速基于Transformer的自回归大语言模型（LLM）的解码至关重要。然而，存储KV缓存所需的内存量在长序列长度和大批量大小时会变得非常巨大。在Transformer发明之后，减少KV缓存大小最有效的两种干预措施是多查询注意力（MQA）及其推广形式分组查询注意力（GQA）。MQA和GQA都修改了注意力块的设计，使得多个查询头可以共享单个键/值头，从而大大减少了不同键/值头的数量，同时仅略微降低了准确性。在本文中，我们展示了通过在相邻层之间共享键和值头，可以进一步推进多查询注意力，从而产生一种新的注意力设计，我们称之为跨层注意力（CLA）。通过CLA，我们发现可以将KV缓存的大小再减少2倍，同时保持与未修改的MQA几乎相同的准确性。在从头开始训练10亿和30亿参数模型的实验中，我们证明CLA提供了优于传统MQA的内存/准确性权衡的帕累托改进，从而能够以比其他方式更长的序列长度和更大的批量大小进行推理。

🔬 方法详解

问题定义：Transformer模型在自回归解码过程中，需要缓存每一层的Key和Value向量（KV缓存）以加速后续token的生成。随着序列长度和batch size的增加，KV缓存的内存占用迅速增长，成为部署大型语言模型的瓶颈。现有的MQA和GQA虽然能减少KV缓存大小，但仍有进一步优化的空间。

核心思路：论文的核心思路是在MQA的基础上，进一步利用层与层之间的冗余性，允许相邻Transformer层共享Key和Value头。这样可以在不显著损失模型性能的前提下，更大幅度地减少KV缓存的大小。

技术框架：CLA的核心在于修改了Transformer的注意力机制。传统的Transformer每一层都有独立的Key和Value头。CLA则允许相邻层共享这些头。具体来说，假设有L层Transformer，CLA将每两层（或其他数量的层）的Key和Value头进行共享。Query头的数量保持不变。这样，KV缓存的大小就减少了。

关键创新：CLA的关键创新在于打破了Transformer层间Key和Value头独立性的假设，通过跨层共享实现了更极致的KV缓存压缩。与MQA/GQA相比，CLA在减少KV缓存大小的同时，尽可能地保持了模型的表达能力。

关键设计：CLA的关键设计在于确定哪些层之间共享Key和Value头。论文中主要考虑相邻层共享。此外，需要仔细调整训练超参数，以补偿由于共享Key和Value头带来的模型容量损失。损失函数与标准Transformer相同，但可能需要调整学习率等参数。

📊 实验亮点

在1B和3B参数模型的实验中，CLA在保持与MQA几乎相同准确率的情况下，将KV缓存大小减少了2倍。这表明CLA在内存/精度权衡方面优于传统MQA，使得模型能够处理更长的序列长度和更大的batch size。实验结果证明了CLA在实际应用中的有效性。

🎯 应用场景

CLA技术可应用于各种基于Transformer的大语言模型推理加速场景，尤其是在资源受限的边缘设备或需要处理超长序列的场景下。通过减少KV缓存大小，CLA可以显著降低内存需求，提高推理吞吐量，并支持更大的batch size，从而降低部署成本，加速AI应用落地。

📄 摘要（原文）

Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another 2x while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, enabling inference with longer sequence lengths and larger batch sizes than would otherwise be possible