从《你所需要的就是注意力》到《你所需要的就是多头潜在注意力》，TransMLA开启AI技术新篇章

自2017年谷歌提出了Transformer架构，以及那篇著名的论文《Attention Is All You Need》后，注意力机制迅速成为自然语言处理领域的核心技术。大型语言模型（LLMs）借助Transformer的自注意力机制，实现了对复杂语言模式的捕捉，在机器翻译、文本生成、对话系统等领域取得了革命性的突破。它们不仅改变了学术研究的方向，更深刻地影响了生产力工具的发展，提高了人们的工作效率和生活质量。

随着模型规模和数据量的不断增长，LLMs面临的挑战也日益凸显。尽管计算能力的提升使得训练更大规模的模型成为可能，但硬件的通信瓶颈却成为制约模型进一步发展的主要障碍。这种瓶颈并非源于计算资源的匮乏，而是由于自注意力机制在处理长序列时，内存和通信开销呈指数级增长。

在自注意力机制中，每个Token都需要与序列中的所有其他Token进行互动，这导致了计算量和存储需求的急剧上升，特别是键值（Key-Value，KV）缓存的大小随着序列长度的增加而迅速增长。当处理超长文本时，KV缓存的内存需求可能超过现有硬件的承载能力，导致模型无法高效地进行推理和生成。

面对这一挑战，研究人员尝试了多种解决方案。其中，组查询注意力（Group Query Attention，GQA）作为一种优化方法，通过减少键和值的头数，降低了KV缓存的尺寸。然而GQA在降低内存开销的同时，也不可避免地限制了模型的表达能力，导致性能下降。如何在不增加KV缓存大小的情况下，提高模型的表达能力，成为亟待解决的问题。

基于此，来自北京大学和小米公司的研究团队提出了多头潜在注意力（Multi-head Latent Attention，MLA）。他们认为，通过引入低秩矩阵，可以在键值层实现KV缓存的压缩，同时通过上投影矩阵，增强模型的表达能力，从而在不增加KV缓存开销的情况下，提升模型性能。

他们进一步提出了TransMLA，这是一种将现有的GQA模型转换为MLA模型的后训练方法。通过理论证明，他们证明了在相同的KV缓存开销下，MLA的表达能力始终优于GQA。实验结果也验证了这一观点，转换后的MLA模型在下游任务中表现出色，显著优于原始的GQA模型。他们开源了TransMLA架构：https://github.com/fxmeng/TransMLA

这项研究的重要性在于，它为大型语言模型的优化提供了一条新的路径。在硬件限制难以突破的情况下，如何通过模型结构的调整来提升性能，就显得尤为关键。TransMLA的提出，不仅为模型开发者提供了一种低成本、高效率的优化方案，也为LLMs的未来发展指明了方向。

值得一提的是，这个研究团队背景深厚。主要作者Fanxu Meng, Zengwei Yao, Muhan Zhang来自北京大学人工智能研究院和通用人工智能国家重点实验室，拥有扎实的理论基础和丰富的研究经验。同时，团队成员也包括来自小米公司的工程师，具备将先进技术应用于实际产品的能力。这种产学研结合的团队构成，使得他们的研究既有前瞻性的理论价值，又具备实际应用的可行性。

相关工作

自注意力机制的内存瓶颈

自注意力机制是Transformer模型的核心组件，它允许模型在序列中的任意位置建立关联。然而，其计算和内存需求与序列长度呈二次关系。具体来说，对于长度为 TT 的序列，自注意力机制需要计算并存储一个 T×TT \times T 的注意力矩阵。此外，还需要存储每个Token的键（Key）和值（Value）向量，这些键值对的数量随着序列长度的增加而迅速增加，导致KV缓存的内存占用也迅速增长。

当处理长序列时，每个Token的KV对都需要重新计算，并存储在缓存中，这种缓存的内存需求会随着序列长度呈指数级增长。这种增长超出了现有硬件的内存容量，成为系统的瓶颈，限制了大模型在处理长文本任务中的应用。例如，在长文档摘要或长篇文章生成任务中，KV缓存的内存需求可能达到甚至超过现有硬件的内存上限，导致模型无法有效运行。

现有解决方案及局限性

为了缓解自注意力机制带来的内存瓶颈，研究人员提出了多种方法。这些方法在降低内存和计算需求的同时，也带来了各自的局限性。

图1：将GQA等效转换为MLA的过程。（a） 在GQA中，在计算注意力之前，重复K以匹配Q头的数量。（b） 相反，WK被重复；与X相乘后，直接产生K′，相当于GQA。（c） 分解W′K产生一个r维表示，它与GQA中的KV缓存维度相匹配，同时保持等价性。

线性注意力：例如Linear Transformer、RWKV、Mamba等方法，通过线性化注意力机制，使计算和内存需求与序列长度线性相关。这些方法通过重新设计注意力机制，将复杂的矩阵乘法转化为向量内积，从而避免了大量的矩阵计算。然而，这种方法可能降低模型对长距离依赖的捕捉能力，导致在处理需要长距离依赖的任务时，模型性能下降。

动态Token剪枝：如LazyLLM、A2SF和SnapKV，通过动态地移除不重要的Token，以减少KV缓存的大小。这些方法选择性地保留对后续预测最重要的Token，从而减少内存使用。然而，动态剪枝可能错失关键的上下文信息，尤其是在需要全局理解的任务中。此外，为了确定哪些Token可以被剪枝，需要额外的计算和复杂的策略设计，这增加了系统的复杂性。

剪枝头部维度：方法如SliceGPT、Sheared和Simple Pruning，通过减少注意力头的数量或降低每个头的维度，来减少计算和内存需求。这些方法通过剪去模型中冗余或不重要的头部来优化内存使用。然而，过度的剪枝可能削弱模型的表达能力，因为注意力头的多样性对于捕捉复杂的语言模式至关重要，过少的头部可能导致模型无法充分表示复杂的关系，进而影响性能。

跨层共享KV表示：如YONO、MiniCache和MLKV，通过在不同Transformer层之间共享KV表示，减少需要存储的KV缓存的总量。这种方法通过让每层使用相同的KV表示来节省内存。然而，由于不同层可能需要不同的表示来捕获不同层次的特征，共享KV表示可能会限制模型的学习能力，并引入不同层间的干扰，影响性能的稳定性。

KV量化：如GPTQ、Kivi和KVQuant，通过对KV向量进行量化，以较低的比特数表示，从而减少内存占用。例如，将浮点数表示的KV向量量化为8位或更低的精度。这种方法可以显著减少内存使用和计算开销。然而，量化可能引入误差，导致表示精度下降，影响模型的准确性。过度量化还可能导致梯度消失或爆炸，影响模型的训练和稳定性。

组查询注意力（GQA）的应用与局限

组查询注意力（Group Query Attention，GQA）是一种优化策略，通过将多个查询头分组，使每个组共享相同的键和值表示，从而减少了KV缓存的大小。在这种机制下，键和值的头数减少了，从而降低了内存需求。

GQA的应用优势

降低内存开销：通过共享键和值，显著减少KV缓存的大小，适用于内存受限的环境。

提升计算效率：减少计算量，加快模型的推理速度。

然而GQA也存在一定的局限性，在减少内存开销的同时，也限制了模型的表达能力。由于多个查询共享相同的键和值，模型的表达多样性受到限制，这使得模型难以捕捉复杂的模式和关系。在一些需要细粒度注意力的任务中，GQA的性能可能无法达到理想效果。

尽管GQA在降低KV缓存方面表现出色，但其对模型表达能力的限制使得在一些应用场景中性能受损。为此，寻找一种既能降低KV缓存大小，又不牺牲模型表达能力的方法，成为了亟待解决的问题。这正是多头潜在注意力（MLA）所要解决的核心挑战。

TransMLA方法详解

MLA的核心思想

多头潜在注意力（MLA）的核心思想是通过在键值层使用低秩矩阵来压缩KV缓存，同时引入上投影矩阵以增强模型的表达能力。MLA通过低秩矩阵降低了KV缓存的尺寸，在减少内存占用的同时，保持了必要的注意力信息。而上投影矩阵则提供了更高的表达能力，使得模型可以在不增加KV缓存大小的情况下，提高注意力机制的灵活性和多样性。

理论证明：MLA表达能力优于GQA

定理1：在相同的KV缓存开销下，MLA的表达能力始终大于GQA。

证明思路：

GQA通过复制键值头，使得组内的查询共享相同的键和值，这在一定程度上限制了模型的表达多样性。而MLA则通过低秩分解，使得键和值的表示更加丰富，允许更多的通道间多样性。这种差异使得MLA在保持相同KV缓存开销的情况下，能够提供更强的表达能力。

具体证明步骤：

首先，我们展示如何在GQA中通过复制键来匹配查询头的数量。设X为输入序列，WK为键的投影矩阵，则键矩阵KK可以表示为：

在GQA中，为了匹配查询头的数量，键矩阵需要被复制。设复制因子为s，则复制后的键矩阵可以表示为：

接下来，我们将复制过程移到参数侧，即先复制投影矩阵，然后再进行计算。设WK′为复制后的投影矩阵，则有：

这样做可以实现相同的效果，同时不增加KV缓存的大小。

为了进一步增强表示能力，我们引入低秩分解。使用奇异值分解（SVD），将复制的投影矩阵分解成低秩矩阵：

其中，UK和VK是正交矩阵，SK是对角矩阵。我们可以截断SVD，仅保留前r个奇异值，从而实现低秩表示：

设低秩矩阵为WKa和WKb，则有：

最终，键矩阵表示为：

这样，MLA通过低秩分解，不仅压缩了KV缓存，还增强了模型的表达能力。通过构造一个WKbW^b_K中向量正交的例子，可以证明MLA在某些情况下的表达能力是GQA无法达到的。

TransMLA的模型转换方法

TransMLA提供了一种将基于GQA的预训练模型转换为MLA模型的方法。具体步骤如下。

调整权重矩阵：将原始GQA模型的键和值投影矩阵拆分，并引入额外的低秩矩阵。通过低秩分解，我们得到两个新的投影矩阵WKaW^a_K和WKbW^b_K，用以替换原始的投影矩阵。

保持KV缓存大小不变：在转换过程中，通过低秩分解和矩阵调整，确保转换后的MLA模型的KV缓存大小与原始GQA模型相同。这意味着，我们在增强模型表达能力的同时，不会增加内存开销。

增强模型表达能力：转换后的MLA模型可以在不增加内存开销的情况下，提升注意力机制的多样性和灵活性。具体地，通过引入上投影矩阵WKbW^b_K，模型可以更好地捕捉复杂的模式和关系，提升整体性能。

通过上述方法，TransMLA不仅成功地将GQA模型转换为MLA模型，还显著提升了模型的表达能力和实际性能，为大模型的优化提供了一种高效的解决方案。

实验结果与分析

实验设置

为了验证TransMLA的有效性，研究团队设计了一系列实验，选择了两个基于GQA的预训练模型：Qwen2.5-7B和Qwen2.5-14B。通过将这些模型转换为MLA模型，并调整相应的权重矩阵维度，研究团队希望评估MLA在实际应用中的表现。

图2:TransMLA和基于GQA的Qwen模型之间训练损失和准确性的比较。

在实验中，使用了包含数学和编程任务的指令微调数据集SmolTalk，具体包括MetaMathQA和SelfOSS-Starcoder2-Instruct数据集。实验设置的训练配置为：批次大小为16，学习率为2e-5，训练2个周期。为了减少对原始模型的影响，实验中仅对键值层进行训练。这样的配置旨在保持实验的公平性和一致性，使得结果具有较高的可信度。

微调性能比较

在训练过程中，TransMLA模型的损失下降速度明显快于原始GQA模型。这表明，TransMLA在数据拟合方面具有明显优势。在相同的训练条件下，TransMLA能够更快速地适应数据，并在较短的时间内达到更低的损失值。这一结果验证了TransMLA在训练效率上的改进，表明其优化方法在实际应用中具有较高的效能。

在具体的任务表现上，TransMLA模型也展现出了优越的性能。在数学任务（如GSM8K）和编程任务（如HumanEval）中，TransMLA模型的准确率显著高于原始GQA模型。例如，在7B模型的实验中，GSM8K的准确率从原始GQA模型的81%提升到了TransMLA模型的86%。这种显著的性能提升，不仅表明TransMLA在特定任务中的适应性强，也展示了其在提升模型表达能力方面的潜力。

TransMLA的性能提升可以归因于以下几个方面。

增强的键值维度在模型表达能力上的提升起到了关键作用，通过调整权重矩阵，TransMLA扩大了键和值的维度，使得模型能够捕捉到更多的特征和模式，从而提升了整体性能。

正交分解在优化模型结构方面也发挥了重要作用，实验结果表明，仅增加参数量并不能带来显著的性能提升，正交分解方法通过优化模型的表示，使得模型在保持较低内存开销的情况下，实现了性能的显著提升。这一结果强调了结构优化在提升模型性能中的重要性。

综上所述，TransMLA通过优化KV缓存的表示和引入先进的矩阵分解技术，实现了在不增加内存开销的情况下，显著提升模型性能的目标。这不仅为大规模语言模型的进一步发展提供了新的思路，也为未来的研究和应用指明了方向。

结论与展望

TransMLA成功地展示了多头潜在注意力（MLA）的强大优势，证明了其在表达能力上显著优于组查询注意力（GQA）。通过理论分析和实验验证，研究团队展示了MLA在相同KV缓存开销的情况下，能够提供更强的模型性能。转换后的MLA模型在多种下游任务中表现出色，为大型语言模型（LLMs）的注意力机制设计提供了全新的视角和思路。研究表明，TransMLA不仅有效地减少了KV缓存的内存需求，还大幅提升了模型的表达能力和推理速度。

TransMLA的研究不仅在技术层面上取得了重要突破，也为整个行业带来了深远的影响。

TransMLA显著提升了资源效率，在不增加硬件负担的情况下，通过优化注意力机制和键值缓存，提升了模型的性能。这为企业和研究机构在有限的计算资源下，进一步提高模型的表现提供了新的解决方案。

TransMLA对生态环境也具有友好性，通过减少模型对内存和计算资源的需求，降低了能源消耗和碳排放。这一特性使得TransMLA在绿色计算和可持续发展方面表现出色，符合当今社会对环保和节能的追求。

TransMLA提供了一种低成本的模型转换方法，使现有的基于GQA的模型可以轻松迁移到MLA架构。这样的迁移不仅能够保留原有模型的优点，还能大幅提升其性能和效率。这为模型开发者和应用者提供了极大的便利，加速了新技术的推广和应用。

研究团队计划在以下几个方面继续拓展和优化TransMLA。

团队将扩展到更大规模的模型，包括LLaMA、Qwen、Mistral等大型语言模型，验证TransMLA在更大规模下的有效性。这将进一步检验和提升MLA在处理超长序列和大规模数据上的表现。

研究团队将使用DeepSeek R1蒸馏技术，进一步提升转换后模型的性能。通过蒸馏技术，可以优化模型的参数和结构，使其在保持高性能的同时，具有更高的计算效率和稳定性。

针对MLA的特性，研究团队计划开发特定的推理加速策略。通过设计高效的推理算法，最大程度地发挥MLA的优势，保持模型的低延迟和高响应速度。这将为实时应用和高频交互场景中的LLMs提供更加优质的性能支持。

总的来说，TransMLA的成功不仅为现有的大模型优化提供了有效的解决方案，也为未来的研究和应用开辟了新的路径。随着技术的不断进步和优化，相信MLA将在更多领域发挥重要作用，推动人工智能技术的发展和创新。

参考资料：​​​https://arxiv.org/pdf/2502.07864​​

本文转载自​​独角噬元兽​​，作者： FlerkenS ​​​​