DeepSeek的MLA，任意大模型都能輕鬆遷移了

2025-04-22 14:15 機器之心

復旦 NLP 實驗室博士後紀燾是這篇文章的第一作者，研究方向為大模型高效推理、多模態大模型，近期代表工作為首個NoPE外推HeadScale、注意力分塊外推LongHeads、多視覺專家大模型MouSi，發表ACL、ICLR、EMNLP等頂會頂刊論文 20 餘篇。
DeepSeek-R1 作為 AI 產業顛覆式創新的代表轟動了業界，特別是其訓練與推理成本僅為同等效能大模型的數十分之一。多頭潛在注意力網路（Multi-head Latent Attention, MLA）是其經濟推理架構的核心之一，透過對鍵值快取進行低秩壓縮，顯著降低推理成本 [1]。
然而，現有主流大模型仍然基於標準注意力架構及其變種（e.g., MHA, GQA, MQA），推理成本相比 MLA 呈現顯著劣勢。使預訓練的任意 LLMs 快速遷移至 MLA 架構而無需從頭預訓練，這既有重大意義又具有挑戰性。
復旦 NLP 實驗室、華東師大、上海 AI Lab、海康威視聯合提出 MHA2MLA 框架，透過部分 RoPE 保留（Partial-RoPE）和鍵值聯合表示低秩近似（Low-rank Approximation）兩個關鍵步驟，成功將任意 MHA/GQA 架構遷移到 MLA。
目前，MHA2MLA 已位列🚀alphaXiv 熱度榜🔥
復旦 NLP 實驗室博士後紀燾為第一作者，副研究員桂韜為通訊作者。
  • 論文標題：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 極具挑戰：
  1. 位置編碼不同：MHA 採用全維度位置編碼（PE），MLA 僅少量維度採用 PE，剩餘維度則 PE 無關
  2. 快取物件不同：MHA 快取分離的鍵向量及值向量，MLA 快取帶 PE 的鍵向量及 PE 無關的鍵值聯合低維表示向量
  3. 引數矩陣不同：MHA 包含查詢、鍵、值三個線性變換矩陣，MLA 則更加複雜、多達七個目的不同的線性變換矩陣
  4. 運算形式不同：MHA 的運算受限於訪存瓶頸，MLA 則能透過矩陣吸收等最佳化實現更高的訪存效率
本文提出的 MHA2MLA 為了最大化利用 MHA 預訓練引數矩陣並對齊 MLA 的快取物件和運算形式，首先透過部分 RoPE 保留（Partial-RoPE）分離出 PE 相關表示（少量維度，如 1/8）和 PE 無關表示（大量維度），其中 PE 相關的鍵向量對齊 MLA。其次拼接值的變換矩陣（W_v）和 PE 無關的鍵的變換矩陣（W_{k, nope}），並進行 SVD 分解得到降維變換矩陣和升維變化矩陣，中間的鍵值聯合低秩表示對齊 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 衝突的問題。
作者主要嘗試了四種移除 RoPE 的策略：1）保留高頻位置資訊 S_high，該方法最簡單直接，保留了區域性語義特徵相關的高頻特徵 [2]；2）保留低頻位置資訊 S_low，與保留高頻位置資訊的策略形成對比，檢驗低頻成分在語義理解任務中的潛在作用；3）均勻取樣策略 S_uniform，等間隔均勻取樣頻率保留位置頻率；4）使用查詢、鍵向量範數乘積 (2-norm) 近似注意力貢獻值 [2] 的篩選策略 S_{2-norm}，針對每個注意力頭，計算所有頻率的平均 2-norm 分數，隨後選擇得分較高的頻率保留位置資訊。該策略能自適應識別對模型效能關鍵的特徵頻率。
Partial-RoPE 的消融實驗表明：1）保留低頻位置資訊的 S_low 導致了最大的效能損失，保留高頻位置資訊的 S_high 導致的效能損失明顯小於保留低頻，說明了高頻維度的重要性；2）S_uniform 和 S_{2-norm} 均展現出更優的效能，分別在 135M 模型和 1.7B 模型上取得了最少的效能損失。最終作者選擇 S_{2-norm} 作為預設配置，是因為注意力貢獻分數較低的維度在結合低秩近似時損失更少。
鍵值聯合表示低秩近似
移除了大量維度的 RoPE 之後，MHA2MLA 就可以對值向量和 PE 無關的鍵向量進行低秩近似，從而大幅減少快取空間。為最大化保留預訓練知識，本文提出兩種基於奇異值分解 (SVD) 的投影矩陣初始化策略：1）SVD_split，分別對矩陣進行低秩分解，保持各自的表徵特性；2）SVD_joint，考慮鍵值矩陣之間的關聯性，引數矩陣拼接後整體進行低秩分解。
消融實驗表明：無論是在 GQA 基座還是 MHA 基座上，SVD_joint 方法始終優於 SVD_split 方法。
實驗結果
作者在多種規模的語言模型（SmolLM-135M/360M/1B7 和 Llama2-7B）以及不同壓縮比例的配置下評估了所提出的方法。實驗表明：1）相同微調設定下，壓縮比例越高，效能損失越大，特別是對於兩個 GQA 模型；2）相同壓縮比例下，原始模型引數越多，效能損失越小，揭示了 MHA2MLA 的潛在 scaling law。3）MHA2MLA 的微調資料量僅需預訓練資料的 0.3%~0.6%，避免了從頭預訓練 MLA 模型的高昂成本。
作者在 LongBench 長文字生成任務中評估了結構遷移後的 Llama2-7B 模型，將 KV 快取量化作為基準對比方案。實驗表明，MHA2MLA 能在 d_{kv}=16 的情況下實現與 2-bit 量化相同的壓縮比例（87.5%），同時僅損失一半的效能（-3.0% vs. -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 提供了一條實用路徑。
作者提到該研究受限於硬體條件，當前實驗未能覆蓋 Llama3 等需 128K 長上下文微調的模型，也未突破 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
© THE END 
轉載請聯絡本公眾號獲得授權
投稿或尋求報道：[email protected]


相關文章