今天小編分享的科學經驗:DeepSeek-R1秘籍輕松遷移,最低只需原始數據0.3%,歡迎閱讀。
DeepSeek-R1 背後關鍵——多頭潛在注意力機制(MLA),現在也能輕松移植到其他模型了!
而且只需原始數據的 0.3%~0.6%。
這項研究由復旦大學、華東師範大學、上海 AI Lab 等聯合提出,復旦教授邱錫鵬(Moss 大模型項目負責人)也在作者名單之列。
他們提出了MHA2MLA這種數據高效的微調方法,使基于 MHA(多頭注意力)的大語言模型(LLMs)能夠順利轉換到 MLA 架構。
以 Llama2-7B 為例,MHA2MLA 在降低推理成本(如減少 KV 緩存大小 92.19%)的同時,能将性能損失控制在較小範圍(如 LongBench 性能僅下降 0.5%)。
具體咋回事,下面我們接着看。
掌握 DeepSeek 核心秘訣
多頭注意力 MHA(Multi-Head Attention)是 Transformer 架構中的一個核心組件,允許模型同時關注輸入的不同部分,每個注意力頭都獨立地學習輸入序列中的不同特征。
然而,随着序列長度的增長,鍵值(Key-Value,KV)緩存的大小也會線性增加,這給模型帶來了顯著的内存負擔。
為了解決 MHA 在高計算成本和 KV 緩存方面的局限性,DeepSeek 突破性地引入了多頭潛在注意力機制 MLA。
簡單說,MLA 最大創新之處在于:
利用低秩聯合壓縮鍵值技術,減少了推理時的 KV 緩存,從而在保持性能的同時顯著降低内存占用。
這一技術也被視為 DeepSeek-V3、DeepSeek-R1 等當紅炸子雞模型背後的關鍵。
而現在,為了進一步降低其他 LLMs 的推理成本,研究人員開發了一種能将采用 MHA 的模型快速适配 MLA 架構的方法——MHA2MLA。
這一數據微調方法包含兩個關鍵部分:
partial-RoPE,即從對注意力分數貢獻較小的查詢和鍵的維度中移除旋轉位置嵌入(RoPE);
低秩近似,基于預訓練的鍵和值參數引入聯合奇異值分解(SVD)近似。
先說第一個。Transformer 架構中,RoPE(旋轉位置編碼,Rotary Position Embedding ) 通過旋轉操作将位置信息融入查詢向量 Q 和鍵向量 K ,幫助模型捕捉序列位置關系。
但研究發現,在計算注意力分數時,并非所有維度的 RoPE 對結果貢獻相同。
換句話說,即使去除那些對注意力分數影響較小的部分維度的 RoPE,理論上不會對模型理解上下文的能力造成關鍵影響。
基于此,研究人員通過計算敏感度指标來确定哪些維度的 RoPE 貢獻較小。
具體而言,對于每個維度,計算 RoPE 變化時注意力分數的變化程度。一旦變化程度低于特定阈值的維度,即被判定為對注意力分數貢獻小。在後續計算中,這些維度将不再應用 RoPE。
最終實驗證明,partial-RoPE 這一策略在不顯著影響模型性能的前提下,減少了計算量。
再說低秩近似策略。
該方法基于預訓練的鍵和值參數,引入聯合奇異值分解(SVD)近似。
SVD 是一種矩陣分解技術,通過對鍵值矩陣進行 SVD 分解,可以用低秩矩陣近似原始矩陣,從而減少參數數量。
具體實現中,研究人員首先提取預訓練模型中的鍵和值參數矩陣,對這些矩陣進行聯合 SVD 分解;然後根據模型的性能和壓縮需求,構建低秩近似矩陣,用這些低秩近似矩陣替代原始的鍵值矩陣參與後續計算。
最終結果顯示,此舉有效降低了模型推理時的計算量和内存占用。
性能幾乎不變,将 Llama2 KV 緩存減少 90% 以上
實驗環節也驗證了 MHA2MLA 方法的有效性。
能在顯著降低推理成本的同時,保持甚至提升模型性能。
研究人員選取了用 MHA 或 GQA 預先訓練的不同規模(135M-7B)的 LLMs,然後設定了對照組。
一組是基于傳統 MHA 的原始模型,用于直接對比 MHA2MLA 方法在相同任務和數據集上的性能表現;另一組是采用分組查詢注意力(GQA)的模型,GQA 作為 MHA 的變體,在一定程度上優化了計算成本,将其與 MHA2MLA 對比,能更清晰地展現 MHA2MLA 的優勢。
在評估其常識性推理能力的六個基準測試中,研究發現:
與原始 LLMs 性能相比,四個基礎模型的性能變化極小,135M 模型性能下降 0.25%,360M、1B7 和 7B 模型分别有 0.03% 、0.03% 和 0.37% 的性能提升或保持。
這表明微調數據未顯著影響原模型性能,MHA2MLA 能有效實現架構遷移,而且微調數據僅需預訓練數據的 0.3%-0.6%。
甚至,較大模型在轉換到 MLA 架構時性能下降更少,這說明這一方法對規模更大的模型更有效。
此外,在長文本生成能力評估中,以 LongBench 為基準,MHA2MLA 相比訓練後量化方法,在壓縮率和精度平衡上表現出色。
當 dkv=16 時,MHA2MLA 可實現 87.5% 的壓縮率,精度損失僅 3%;與 4-bit 量化結合後,壓縮率可達 92.19%(dkv=64 + Int4HQQ)和 96.87%(dkv=16 + Int4HQQ),精度損失分别為 -0.5% 和 -3.2%,優于所有 2-bit 量化的基線模型。
這也反映了 MHA2MLA 方法能夠與量化技術良好兼容。
綜合以上實驗,可以看到以 Llama2-7B 為例,MHA2MLA 在降低推理成本(如減少 KV 緩存大小 92.19%)的同時,能将性能損失控制在較小範圍(如 LongBench 性能僅下降 0.5%)。
不過,論文也提到了研究局限性。
受計算資源限制,未在更大、更多樣化的開源大語言模型上驗證 MHA2MLA;且由于 Deepseek 未開源 MLA 的張量并行推理框架,難以探索大于 7B 的模型。
下一步,研究人員計劃在更多模型上進行驗證。
感興趣的童鞋可以查看原論文 ~
論文:
https://arxiv.org/abs/2502.14837
代碼:
https://github.com/JT-Ushio/MHA2MLA