TransMLA: Multi-Head Latent Attention Is All You Need

现代大型语言模型（LLMs）在当前硬件上经常遇到通信瓶颈，而不仅仅是计算限制。多头潜在注意力机制（MLA）通过在键值（KV）层使用低秩矩阵来应对这一挑战，从而可以缓存压缩的潜在KV状态。这种方法显著减少了KV缓存的大小，相较于传统的多头注意力机制，使得推理速度更快。此外，MLA使用了一个上投影矩阵来增加表达能力，以额外的计算换取减少通信开销。尽管MLA在Deepseek V2/V3/R1中展示了效率和有效性，但许多主要模型提供商仍然依赖于组查询注意力机制（GQA），并且尚未宣布任何采用MLA的计划。在本文中，我们证明了GQA可以在保持相同KV缓存开销的情况下由MLA表示，但反之则不成立。为了鼓励更广泛地使用MLA，我们引入了TransMLA，这是一种训练后的方法，可以将广泛使用的基于GQA的预训练模型（例如LLaMA、Qwen、Mixtral）转换为基于MLA的模型。转换后，模型可以进行额外的训练以增强表达能力，而不增加KV缓存的大小。此外，我们计划开发特定于MLA的推理加速技术，以在转换后的模型中保持低延迟，从而实现更高效的Deepseek R1蒸馏。

该论文旨在解决现代大型语言模型（LLMs）在当前硬件上遇到的通信瓶颈问题，而不是纯粹的计算约束问题。这是一个现有问题，但通过引入新的方法来优化KV缓存使用，提供了新的解决方案。

关键思路是通过多头潜在注意力（MLA）机制，利用低秩矩阵在键值（KV）层中实现压缩的潜在KV状态缓存，从而显著减少KV缓存大小，提高推理速度。此外，MLA通过引入上投影矩阵增加表达能力，以额外的计算换取降低通信开销。相比传统的多头注意力机制，MLA展示了更高的效率和效果。

论文展示了GQA可以完全由MLA表示，并且保持相同的KV缓存开销，但反之则不成立。为了促进MLA的广泛应用，作者提出了一种名为TransMLA的后训练方法，可以将基于GQA的预训练模型转换为MLA模型。转换后的模型可以通过进一步训练提升表达能力而不增加KV缓存大小。此外，作者计划开发MLA特定的推理加速技术，以保持低延迟并更高效地蒸馏Deepseek R1。实验设计涉及多个广泛使用的GQA模型如LLaMA、Qwen和Mixtral，并可能包括开源代码供后续研究。

最近在这个领域中，相关的研究包括Group Query Attention (GQA) 的应用，例如《Scaling Up Language Models: Opportunities and Challenges》和《Efficient Transformers: A Survey》等论文探讨了不同类型的注意力机制及其优化方法。其他相关工作还包括对Transformer架构的各种改进，如稀疏注意力、局部注意力和其他降低计算和通信成本的方法。