一种 KV Cache 友好的大模型 Attention 机制。

原标题：阶跃公开了自家新型注意力机制：KV缓存消耗直降93.7%，性能不减反增
文章来源：机器之心
内容字数：4981字

高效大语言模型推理：多矩阵分解注意力机制MFA的突破

机器之心AIxiv专栏报道了阶跃星辰、清华大学等机构的研究成果：一篇关于新型注意力机制——多矩阵分解注意力（MFA）及其变体MFA-Key-Reuse的论文。该研究旨在解决大语言模型推理阶段的内存瓶颈问题，即传统注意力机制的键值缓存（KV Cache）随批处理大小和序列长度线性增长的问题。

1. 问题与挑战

现有解决KV Cache内存占用问题的注意力机制变体，例如MQA和MLA，存在性能与复杂度之间的权衡难题。MQA牺牲了模型表达能力以换取极低的内存使用；而MLA虽然结构灵活，但其性能上限受限于最小维度，增加中间维度并不能有效提升性能。

2. MFA机制的创新

MFA的提出旨在最大限度地节省资源并接近理论性能上限。研究团队通过三个关键创新实现了这一目标：

1. 显著增加注意力头数量和维度：突破传统设计局限，极大提升模型容量。

2. 创新的低秩分解策略：在扩展模型注意力头数量和维度时保持高参数效率。

3. 单键值头设计：即使增加模型复杂度，也能保持最低水平的内存使用。

3. 理论分析与比较

研究团队提出了广义多头注意力（GMHA）的概念框架，并引入“总有效秩（TER）”和“共享隐空间维度（SLSD）”两个指标来衡量GMHA系列模型的容量。通过此框架，MFA与MQA、MLA、MHA的比较分析表明，MFA在参数效率和性能之间取得了最佳平衡。MFA在更高的TER和SLSD下，实现了更低的KV Cache占用。

4. 实验结果与结论

实验结果显示，MFA和MFA-KR在不同规模模型（1B到7B参数）和训练数据量（10B到1T）下均表现出色。MFA与传统MHA具有相当的性能扩展能力，并在内存节省方面展现出显著优势，最大规模模型上实现了高达87.5%的内存节省。消融实验也验证了MFA和MFA-KR设计的有效性。

5. 展望

MFA以简洁的设计解决了LLM高效推理的显存瓶颈问题，并能无缝集成到现有Transformer生态中。这项研究为大语言模型的广泛应用提供了重要的技术支撑。

联系作者

文章来源：机器之心
作者微信：
作者简介：专业的人工智能媒体和产业服务平台