Published on

阶跃新型注意力机制：KV缓存消耗降93.7% 性能不减反增

作者
  • avatar
    姓名
    Ajax
    Twitter

引言

随着大语言模型（LLM）的广泛应用和推理时扩展的新范式的出现，高效的大规模推理成为一个关键挑战。尤其是在语言模型的推理阶段，传统的注意力机制中的键值缓存（KV Cache）会随着批处理大小和序列长度线性增长，这成为了制约大语言模型规模化应用和推理时扩展的瓶颈，被称为“内存杀手”。

虽然业界已经提出了多种注意力机制的变体，如 MQA、GQA、MLA 等，试图解决这个问题，但这些方案要么难以在严格的显存限制下保持理想性能，要么在模型结构上引入额外的复杂度，带来了巨大的工程挑战和生态兼容性问题。

MFA：多矩阵分解注意力

近期，阶跃星辰和清华大学等机构的研究人员在论文《Multi-matrix Factorization Attention》中提出了一种新型注意力机制架构——多矩阵分解注意力（MFA）及其变体 MFA-Key-Reuse。这一机制不仅显著降低了语言模型推理的成本，还实现了性能的显著提升。

核心发现：

  • MFA 和 MFA-KR 在实验中超越了 MLA 的性能。
  • 在减少高达 93.7% 的 KV Cache 使用量的情况下，MFA 和 MFA-KR 与传统的 MHA 性能相当。
  • MFA 实现简单，容易复现，对超参数敏感度低，并且兼容各种位置编码 (Pos-embedding)。

MFA 方法与分析

研究团队通过深入研究注意力机制的设计和容量分析，明确了注意力机制容量相关的两个关键维度，并提出了一系列分析方法和设计原则。

为了更好地理解注意力机制的本质，研究团队提出了广义多头注意力（GMHA）的概念框架，为理解不同的 MHA 变种提供了一个统一的视角。他们从推理角度研究了键值的计算和存储方式，并从分解角度探讨了模型的容量特征。这种创新的分析方法为理解不同策略如何在模型性能和计算效率之间取得平衡提供了全新视角。

研究人员确立了完全参数化双线性注意力（FPBA）作为理论上的性能上限标准。他们发现，目前广泛使用的 MHA 及其各种变体实际上都可以被视为 FPBA 的低秩分解版本。在分析现有解决方案时，研究团队重点考察了两个代表性的改进方案：多查询注意力（MQA）和多头潜在注意力（MLA）。

  • MQA： 采用更激进的参数共享策略，让所有注意力头共享同一组键值参数。虽然降低了内存使用，但可能影响模型的表达能力。
  • MLA： 引入共享的潜在空间进行参数压缩，节省内存。然而，模型的实际表达能力仍然受限于最小维度的约束，增加中间维度并不能真正提升模型的性能上限。

基于对现有方案的深入分析，研究团队提出了一个目标：开发一种新的注意力机制，既能最大限度地节省资源，又能尽可能接近理论性能上限。由此，他们开发出了多矩阵分解注意力（MFA）。

MFA 的设计体现了三个关键创新：

  • 突破传统设计： 显著增加了注意力头的数量和维度，极大提升了注意力模块的模型容量。
  • 矩阵分解创新： 采用激进的低秩分解策略，在扩展模型注意力头的数量和维度的同时，保持了极高的参数效率。
  • 单键值头设计： 确保即使在增加模型复杂度的情况下，内存使用仍然保持在最低水平。

MFA 与其他注意力机制的对比

为了进一步分析 MFA 和其他注意力机制的不同，研究团队引入了两个关键指标来度量 GMHA 系列模型的容量：

  • 模型总有效秩 (TER)： 注意力头数量与每个头部分解秩（FRH）的乘积。
  • 共享隐空间维度 (SLSD)： 代表所有注意力头共同使用的隐空间维度。

TER 和 SLSD 越大，模型具有更高的容量。每个头的分解秩（FRH）不超过 SLSD，而 KV Cache 的占用受制于 FRH 和 SLSD 之间的较大值，这构成了容量和效率之间的关键权衡。

通过这个框架分析：

  • 与 MQA 相比，MFA 同时实现了更高的 SLSD 和更高的 TER。
  • 与 MLA 相比，在相似的参数预算下，MFA 不仅达到了更小的 KV cache 尺寸，还实现了更高的 TER，同时保持了相当的 SLSD。
  • 与传统的 MHA 相比，虽然 MFA 的 SLSD 较小，但其 TER 更高，这解释了为什么在实验中 MFA 能够获得更好的性能。

实验结果

研究团队进行了一系列深入的扩展性实验，系统地测试了从 1B 到 7B 参数的不同规模的模型，训练数据量从 10B 扩展到 1T。

  • 性能扩展： MFA 方案展现出与传统 MHA 完全相当的扩展能力，即使在更大规模下，MFA 依然能保持其优异的性能。MFA-KR 虽然在性能上略有降低，但其扩展趋势与 MHA 保持一致。
  • 内存节省： 随着模型规模的增加，MFA 和 MFA-KR 在内存节省方面的优势不仅得到保持，还呈现出进一步扩大的趋势。在最大规模模型上，MFA 实现了 87.5% 的内存节省，而 MFA-KR 更是将内存使用降低到原来的 6.25%。

此外，研究人员还进行了一系列的消融实验，证明 MFA 和 MFA-KR 的设计的有效性，并在其他的主流的位置编码上也验证了 MFA 和 MFA-KR 的性能优势。

展望

MFA 用简洁的设计带来了显著的提升，在不增加额外工程复杂度的前提下，解决了 LLM 的高效推理显存瓶颈问题，并能无缝集成到现有的 Transformer 生态中。这种简单而强大的创新，有望加速大语言模型在更多场景中的落地应用。