TransMLA：多头隐式注意力机制足矣

TransMLA: Multi-head latent attention is all you need

Source | HN Comments

文章介绍了 TransMLA，一种基于 Multi-head Latent Attention (MLA) 的后训练方法，旨在解决大型语言模型（LLMs）的通信瓶颈。MLA 通过压缩的隐式 KV 状态减少 KV 缓存大小，从而加速推理。研究证明了 Group Query Attention (GQA) 可以用 MLA 表示，但反之则不行。TransMLA 将基于 GQA 的预训练模型（如 LLaMA, Qwen, Mixtral）转换为 MLA 模型，转换后可进行额外训练以提高表达性，且不增加 KV 缓存大小。作者计划开发 MLA 专属的推理加速技术，以保持低延迟，实现更高效的 Deepseek R1 蒸馏。

arXiv:2502.07864 (cs) [2025年2月11日提交 (v1), 最近修订于2025年2月13日 (此版本, v2)]

Title: TransMLA: Multi-Head Latent Attention Is All You Need

作者：Fanxu Meng, Zengwei Yao, Muhan Zhang

查看由 Fanxu Meng 和另外 2 位作者撰写的题为 TransMLA: Multi-Head Latent Attention Is All You Need 的论文 PDF 版本

View PDF HTML (实验性)

摘要：现代大型语言模型（LLMs）通常在当前的硬件上遇到通信瓶颈，而不仅仅是计算限制。Multi-head Latent Attention (MLA) 通过在 key-value (KV) 层中使用低秩矩阵来解决这一挑战，从而允许压缩的隐式 KV 状态被缓存。相对于传统的多头注意力机制，这种方法显著减少了 KV 缓存大小，从而加快了推理速度。此外，MLA 采用上投影矩阵来增加表达性，用额外的计算来换取减少的通信开销。尽管 MLA 在 Deepseek V2/V3/R1 中已经展示出了效率和有效性，但许多主要的模型提供商仍然依赖于 Group Query Attention (GQA)，并且尚未宣布任何采用 MLA 的计划。在本文中，我们证明了 GQA 总是可以用 MLA 来表示，同时保持相同的 KV 缓存开销，但反之则不成立。为了鼓励更广泛地使用 MLA，我们介绍了 TransMLA，这是一种后训练方法，可以将广泛使用的基于 GQA 的预训练模型（例如 LLaMA, Qwen, Mixtral）转换为基于 MLA 的模型。转换后，模型可以进行额外的训练以提高表达性，而无需增加 KV 缓存大小。此外，我们计划开发 MLA 专属的推理加速技术，以保持转换后模型的低延迟，从而实现更高效的 Deepseek R1 蒸馏。

注释：| this https URL ---|--- 主题：| 机器学习 (cs.LG); 人工智能 (cs.AI) 引用方式：| arXiv:2502.07864 [cs.LG] (或 arXiv:2502.07864v2 [cs.LG] 对于此版本) https://doi.org/10.48550/arXiv.2502.07864 Focus to learn more arXiv-issued DOI via DataCite

提交历史

来自: Meng Fanxu [查看电子邮件] [v1] Tue, 11 Feb 2025 18:20:18 UTC (326 KB) [v2] Thu, 13 Feb 2025 18:07:04 UTC (327 KB) 全文链接：

访问论文:

查看由 Fanxu Meng 和另外 2 位作者撰写的题为 TransMLA: Multi-Head Latent Attention Is All You Need 的论文 PDF 版本