DeepSeek-R1背后關(guān)鍵——多頭潛在注意力機(jī)制（MLA），現(xiàn)在也能輕松移植到其他模型了！

而且只需原始數(shù)據(jù)的0.3%~0.6%。

這項(xiàng)研究由復(fù)旦大學(xué)、華東師范大學(xué)、上海AI Lab等聯(lián)合提出，復(fù)旦教授邱錫鵬（Moss大模型項(xiàng)目負(fù)責(zé)人）也在作者名單之列。

他們提出了MHA2MLA這種數(shù)據(jù)高效的微調(diào)方法，使基于MHA（多頭注意力）的大語言模型（LLMs）能夠順利轉(zhuǎn)換到MLA架構(gòu)。

以Llama2-7B為例，MHA2MLA在降低推理成本（如減少KV緩存大小92.19%）的同時，能將性能損失控制在較小范圍（如LongBench性能僅下降0.5%）。

具體咋回事，下面我們接著看。

掌握DeepSeek核心秘訣

多頭注意力MHA（Multi-Head Attention）是Transformer架構(gòu)中的一個核心組件，允許模型同時關(guān)注輸入的不同部分，每個注意力頭都獨(dú)立地學(xué)習(xí)輸入序列中的不同特征。

然而，隨著序列長度的增長，鍵值（Key-Value，KV）緩存的大小也會線性增加，這給模型帶來了顯著的內(nèi)存負(fù)擔(dān)。

為了解決MHA在高計(jì)算成本和KV緩存方面的局限性，DeepSeek突破性地引入了多頭潛在注意力機(jī)制MLA。

簡單說，MLA最大創(chuàng)新之處在于：

利用低秩聯(lián)合壓縮鍵值技術(shù)，減少了推理時的KV緩存，從而在保持性能的同時顯著降低內(nèi)存占用。

這一技術(shù)也被視為DeepSeek-V3、DeepSeek-R1等當(dāng)紅炸子雞模型背后的關(guān)鍵。

而現(xiàn)在，為了進(jìn)一步降低其他LLMs的推理成本，研究人員開發(fā)了一種能將采用MHA的模型快速適配MLA架構(gòu)的方法——MHA2MLA。

這一數(shù)據(jù)微調(diào)方法包含兩個關(guān)鍵部分：

• partial-RoPE，即從對注意力分?jǐn)?shù)貢獻(xiàn)較小的查詢和鍵的維度中移除旋轉(zhuǎn)位置嵌入（RoPE）；
• 低秩近似，基于預(yù)訓(xùn)練的鍵和值參數(shù)引入聯(lián)合奇異值分解（SVD）近似。

先說第一個。Transformer架構(gòu)中，RoPE（旋轉(zhuǎn)位置編碼，Rotary Position Embedding) 通過旋轉(zhuǎn)操作將位置信息融入查詢向量Q和鍵向量K ，幫助模型捕捉序列位置關(guān)系。

但研究發(fā)現(xiàn)，在計(jì)算注意力分?jǐn)?shù)時，并非所有維度的RoPE對結(jié)果貢獻(xiàn)相同。

換句話說，即使去除那些對注意力分?jǐn)?shù)影響較小的部分維度的RoPE，理論上不會對模型理解上下文的能力造成關(guān)鍵影響。

基于此，研究人員通過計(jì)算敏感度指標(biāo)來確定哪些維度的RoPE貢獻(xiàn)較小。

具體而言，對于每個維度，計(jì)算RoPE變化時注意力分?jǐn)?shù)的變化程度。一旦變化程度低于特定閾值的維度，即被判定為對注意力分?jǐn)?shù)貢獻(xiàn)小。在后續(xù)計(jì)算中，這些維度將不再應(yīng)用RoPE。

最終實(shí)驗(yàn)證明，partial-RoPE這一策略在不顯著影響模型性能的前提下，減少了計(jì)算量。

再說低秩近似策略。

該方法基于預(yù)訓(xùn)練的鍵和值參數(shù)，引入聯(lián)合奇異值分解（SVD）近似。

SVD是一種矩陣分解技術(shù)，通過對鍵值矩陣進(jìn)行SVD分解，可以用低秩矩陣近似原始矩陣，從而減少參數(shù)數(shù)量。

具體實(shí)現(xiàn)中，研究人員首先提取預(yù)訓(xùn)練模型中的鍵和值參數(shù)矩陣，對這些矩陣進(jìn)行聯(lián)合SVD分解；然后根據(jù)模型的性能和壓縮需求，構(gòu)建低秩近似矩陣，用這些低秩近似矩陣替代原始的鍵值矩陣參與后續(xù)計(jì)算。

最終結(jié)果顯示，此舉有效降低了模型推理時的計(jì)算量和內(nèi)存占用。

性能幾乎不變，將Llama2 KV緩存減少90%以上

實(shí)驗(yàn)環(huán)節(jié)也驗(yàn)證了MHA2MLA方法的有效性。

能在顯著降低推理成本的同時，保持甚至提升模型性能。

研究人員選取了用MHA或GQA預(yù)先訓(xùn)練的不同規(guī)模（135M-7B）的LLMs，然后設(shè)置了對照組。

一組是基于傳統(tǒng)MHA的原始模型，用于直接對比MHA2MLA方法在相同任務(wù)和數(shù)據(jù)集上的性能表現(xiàn)；另一組是采用分組查詢注意力（GQA）的模型，GQA作為MHA的變體，在一定程度上優(yōu)化了計(jì)算成本，將其與MHA2MLA對比，能更清晰地展現(xiàn)MHA2MLA的優(yōu)勢。

在評估其常識性推理能力的六個基準(zhǔn)測試中，研究發(fā)現(xiàn)：

與原始LLMs性能相比，四個基礎(chǔ)模型的性能變化極小，135M模型性能下降0.25%，360M、1B7和7B模型分別有0.03% 、0.03%和0.37%的性能提升或保持。

這表明微調(diào)數(shù)據(jù)未顯著影響原模型性能，MHA2MLA能有效實(shí)現(xiàn)架構(gòu)遷移，而且微調(diào)數(shù)據(jù)僅需預(yù)訓(xùn)練數(shù)據(jù)的0.3%-0.6%。

甚至，較大模型在轉(zhuǎn)換到MLA架構(gòu)時性能下降更少，這說明這一方法對規(guī)模更大的模型更有效。

此外，在長文本生成能力評估中，以LongBench為基準(zhǔn)，MHA2MLA相比訓(xùn)練后量化方法，在壓縮率和精度平衡上表現(xiàn)出色。

當(dāng)d_kv=16時，MHA2MLA可實(shí)現(xiàn)87.5%的壓縮率，精度損失僅3%；與4-bit量化結(jié)合后，壓縮率可達(dá)92.19%（d_kv=64 + Int4_HQQ）和96.87%（d_kv=16 + Int4_HQQ），精度損失分別為-0.5%和-3.2%，優(yōu)于所有2-bit量化的基線模型。

這也反映了MHA2MLA方法能夠與量化技術(shù)良好兼容。

綜合以上實(shí)驗(yàn)，可以看到以Llama2-7B為例，MHA2MLA在降低推理成本（如減少KV緩存大小92.19%）的同時，能將性能損失控制在較小范圍（如LongBench性能僅下降0.5%）。

不過，論文也提到了研究局限性。

受計(jì)算資源限制，未在更大、更多樣化的開源大語言模型上驗(yàn)證MHA2MLA；且由于Deepseek未開源MLA的張量并行推理框架，難以探索大于7B的模型。

下一步，研究人員計(jì)劃在更多模型上進(jìn)行驗(yàn)證。

感興趣的童鞋可以查看原論文~

論文：https://arxiv.org/abs/2502.14837
代碼：https://github.com/JT-Ushio/MHA2MLA