DeepSeek的MLA实现，轻松迁移任意大模型！

复旦大学自然语言处理实验室的博士后研究员纪焘是本文的第一作者，他的研究领域包括大模型高效推理和多模态大模型。近期发表的代表作包括首个NoPE外推HeadScale、注意力分块外推LongHeads及多视觉专家大模型MouSi，他在ACL、ICLR、EMNLP等顶级会议和期刊上发表论文超过20篇。

作为人工智能产业颠覆性创新的代表，DeepSeek-R1引起了行业的广泛关注，尤其是其训练与推理的成本仅为同等性能大模型的十分之一。多头潜在注意力网络（Multi-head Latent Attention, MLA）构成了其经济推理架构的核心，通过对键值缓存实施低秩压缩，显著降低了推理成本 [1]。

然而，当前主流的大模型普遍依赖于标准注意力架构及其变种（例如，MHA、GQA、MQA），与MLA相比，其推理成本明显处于劣势。如何使预训练的任意大型语言模型（LLMs）能够快速迁移至MLA架构，而无需从头开始预训练，这一过程既具重要学术意义，也面临诸多挑战。

复旦大学自然语言处理实验室、华东师范大学、上海人工智能实验室以及海康威视联合提出了MHA2MLA框架。该框架通过实施部分RoPE保留（Partial-RoPE）与键值联合表示低秩近似（Low-rank Approximation）这两个关键步骤，成功地将任意MHA/GQA架构迁移至MLA。

DeepSeek的MLA实现，轻松迁移任意大模型！目前，MHA2MLA已在🚀alphaXiv热度榜上名列前茅🔥
复旦大学自然语言处理实验室的博士后纪焘为第一作者，副研究员桂韬为通讯作者。

论文标题：Towards Economical Inference: Enabling DeepSeek’s Multi-Head Latent Attention in Any Transformer-based LLMs
论文链接：https://arxiv.org/abs/2502.14837
开源代码：https://github.com/JT-Ushio/MHA2MLA

论文概述
本文聚焦于如何将基于MHA/GQA架构的预训练大型语言模型高效迁移至DeepSeek提出的经济推理架构——多头潜在注意力（MLA）。
MHA与MLA存在多个方面的差异，使得MHA2MLA具有相当大的挑战性：

位置编码的不同：MHA使用全维度位置编码（PE），而MLA仅在少数维度上采用PE，其余维度则与PE无关。
缓存对象的不同：MHA分别缓存键向量与值向量，而MLA则缓存带PE的键向量及与PE无关的键值联合低维表示向量。
参数矩阵的不同：MHA包含三个线性变换矩阵（即查询、键、值），而MLA则更为复杂，涉及多达七个功能不同的线性变换矩阵。
运算形式的不同：MHA的运算受到访存瓶颈的限制，而MLA则通过矩阵合并等优化手段实现更高的访存效率。

为了最大限度地利用MHA的预训练参数矩阵，同时对齐MLA的缓存对象和运算形式，本文提出的MHA2MLA框架首先通过部分RoPE保留（Partial-RoPE）方法，将PE相关表示（少量维度，如1/8）与PE无关表示（大量维度）进行分离，其中PE相关的键向量对齐于MLA。其次，将值的变换矩阵（W_v）与PE无关的键的变换矩阵（W_{k, nope}）进行拼接，并通过奇异值分解（SVD）获取降维变换矩阵和升维变换矩阵，从而完成键值联合低秩表示的对齐，进而实现缓存对象与运算形式的对齐。 DeepSeek的MLA实现，轻松迁移任意大模型！在135M到7B参数量级上的实验表明，仅需使用预训练数据的0.3%至0.6%进行高效微调，即可基本恢复架构迁移所造成的性能损失。同时，MHA2MLA能够与其他高效推理技术相结合，例如通过结合4-bit KV缓存量化，Llama2-7B减小了92.19%的KV缓存，而在LongBench上的性能仅下降0.5%。
部分RoPE保留（Partial-RoPE）
为实现从标准的MHA（多头注意力机制）向MLA（多头潜在注意力机制）的迁移，作者提出了部分RoPE微调（partial-RoPE finetuning）策略，该策略通过从大量维度中移除RoPE（旋转位置编码），并将其转换为NoPE（无位置编码），以解决MLA与RoPE之间的冲突。 DeepSeek的MLA实现，轻松迁移任意大模型！作者主要尝试了四种移除RoPE的策略：1）保留高频位置信息S_high，该方法简单直接，能够保留与局部语义特征相关的高频特征 [2]；2）保留低频位置信息S_low，旨在检验低频成分在语义理解任务中的作用；3）均匀采样策略S_uniform，通过等间隔均匀采样频率，保留相关位置频率；4）使用查询、键向量范数乘积（2-norm）来近似注意力贡献值的筛选策略S_{2-norm}，针对每个注意力头计算所有频率的平均2-norm分数，选择得分较高的频率以保留位置信息。该策略能够自适应地识别对模型性能关键的特征频率。 DeepSeek的MLA实现，轻松迁移任意大模型！经对Partial-RoPE的消融实验表明：1）保留低频位置信息的S_low导致了最大性能损失，而保留高频位置信息的S_high则造成的损失明显较小，显示出高频维度的重要性；2）S_uniform与S_{2-norm}均表现出更为优越的性能，分别在135M模型和1.7B模型上实现了最小的性能损失。最终，作者选择S_{2-norm}作为默认配置，原因在于在结合低秩近似时，注意力贡献分数较低的维度通常呈现出更小的损失。秩近似，从而显著减少存储需求。为最大限度地保留预训练知识，本文提出了两种基于奇异值分解（SVD）的投影矩阵初始化策略：一是 SVD_split，这种方法分别对矩阵进行低秩分解，以保留各自的表征特性；二是 SVD_joint，它考虑到键值矩阵间的关联性，在将参数矩阵进行拼接后整体进行低秩分解。
消融实验的结果显示：无论是在 GQA 基座还是 MHA 基座中，SVD_joint 方法始终优于 SVD_split 方法。
实验结果 DeepSeek的MLA实现，轻松迁移任意大模型！作者在多规模语言模型（包括 SmolLM-135M/360M/1B7 和 Llama2-7B）上以及不同压缩比例配置下，对所提出的方法进行了评估。实验结果表明：1）在相同微调设置下，压缩比例越高，性能损失越显著，尤其是对于两个 GQA 模型而言；2）在相同压缩比例下，原始模型参数越多，性能损失则相对较小，这揭示了 MHA2MLA 的潜在规模法则；3）MHA2MLA 的微调所需数据量仅为预训练数据的 0.3% 至 0.6%，这大大降低了从头预训练 MLA 模型的高成本。 DeepSeek的MLA实现，轻松迁移任意大模型！在 LongBench 长文本生成任务中，作者对结构迁移后的 Llama2-7B 模型进行了评估，并将 KV 缓存量化作为基准对比方案。实验结果显示，在 d_{kv}=16 的情况下，MHA2MLA 能够实现与 2-bit 量化相同的压缩率（87.5%），并仅损失一半的性能（-3.0% 对比 -6.2%）；此外，结合 4-bit 量化后，不仅实现的压缩比例超过 2-bit 量化，而且性能损失均优于所有 2-bit 的基线方法，例如 92.19% 压缩比例仅损失 0.5%，96.87% 压缩比例仅损失 3.2%，证明了 MHA2MLA 显著减轻了推理期间的内存瓶颈问题。
总结与展望
本研究主要集中于如何将基于多头自注意力（MHA）的预训练大型语言模型（LLMs）及其变体调整为具有更高 KV 缓存效率的多头注意力（MLA）架构，从而显著降低推理时的内存瓶颈。通过精心设计的架构，MHA2MLA 仅需 0.3% 至 0.6% 的预训练数据。这一框架展现了与现有压缩技术的良好兼容性，同时保留了常识推理及长上下文处理能力，为资源高效部署 LLMs 提供了切实可行的路径。
作者指出本研究受限于硬件条件，当前实验未能涵盖需进行 128K 长上下文微调的 Llama3 模型，对 7B 参数规模的验证也存在瓶颈。未来工作将扩展至更多基座模型，同时，作者还计划结合参数高效的微调策略，进一步降低架构迁移过程中参数更新的规模。
^{参考文献：}^{[1] DeepSeek-AI, Aixin Liu, Bei Feng et al. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. ArXiv preprint.}^{[2] Federico Barbero, Alex Vitvitskyi, Christos Perivolaropoulos, Razvan Pascanu, Petar Veličković. Round and Round We Go! What makes Rotary Positional Encodings useful? CoRR 2024}