MHA2MLA – 復(fù)旦、上海AI Lab等推出優(yōu)化LLM推理效率的方法

MHA2MLA是一種由復(fù)旦大學(xué)、華東師范大學(xué)及上海AI Lab等多家機(jī)構(gòu)共同開發(fā)的數(shù)據(jù)高效微調(diào)方法。該方法采用了DeepSeek的多頭潛在注意力機(jī)制（MLA），旨在提升任何基于Transformer的大型語言模型（LLM）的推理效率，并顯著降低推理成本。MHA2MLA通過兩個核心策略實現(xiàn)其目標(biāo)：首先是部分旋轉(zhuǎn)位置編碼（partial-RoPE），去除了對注意力分?jǐn)?shù)貢獻(xiàn)較小的查詢和鍵的旋轉(zhuǎn)位置編碼（RoPE）維度；其次是低秩近似，通過聯(lián)合奇異值分解（SVD）對鍵和值進(jìn)行壓縮，從而減少KV緩存的內(nèi)存占用。該方法僅需使用原始數(shù)據(jù)的0.3%至0.6%進(jìn)行微調(diào)，即可在極大降低KV緩存（如高達(dá)92.19%）的同時，將性能損失控制在微小范圍內(nèi)（例如LongBench性能僅下降0.5%）。

MHA2MLA是什么

MHA2MLA是由復(fù)旦大學(xué)、華東師范大學(xué)和上海AI Lab等機(jī)構(gòu)聯(lián)合研發(fā)的一種高效微調(diào)方法，旨在優(yōu)化基于Transformer的LLM的推理過程。通過引入DeepSeek的多頭潛在注意力機(jī)制（MLA），MHA2MLA能夠降低推理成本并提升推理效率。該方法依賴于兩個重要策略：部分旋轉(zhuǎn)位置編碼（partial-RoPE）和低秩近似（Low-Rank Approximation），實現(xiàn)了在數(shù)據(jù)量極小的情況下，依然能夠保持模型的高性能。

MHA2MLA的主要功能

• 顯著降低KV緩存：基于低秩壓縮技術(shù)，最大限度地減少KV緩存的大小（最高可達(dá)96.87%），從而降低推理過程中的內(nèi)存占用。
• 保持模型性能：在極少量數(shù)據(jù)（0.3%至0.6%原始訓(xùn)練數(shù)據(jù)）下進(jìn)行微調(diào)，確保性能損失極小（如LongBench性能僅下降0.5%）。
• 兼容現(xiàn)有技術(shù)：可與量化技術(shù)（如4-bit量化）結(jié)合使用，進(jìn)一步提升推理效率。
• 數(shù)據(jù)高效性：僅需少量數(shù)據(jù)即可完成從MHA到MLA的架構(gòu)轉(zhuǎn)變，適合在資源受限的環(huán)境中快速實施。

MHA2MLA的技術(shù)原理

• 部分旋轉(zhuǎn)位置編碼（Partial-RoPE）：在多頭自注意力機(jī)制中，旋轉(zhuǎn)位置編碼（RoPE）通過旋轉(zhuǎn)操作將位置信息融入查詢向量（Q）和鍵向量（K），幫助模型捕捉序列中的位置信息。MHA2MLA根據(jù)每個維度對注意力分?jǐn)?shù)的貢獻(xiàn)，移除貢獻(xiàn)較小的RoPE維度，減少計算量和內(nèi)存占用，這一過程稱為部分RoPE，從而為低秩壓縮騰出空間。
• 低秩近似（Low-Rank Approximation）：MLA利用低秩聯(lián)合壓縮鍵值（KV）來減少內(nèi)存占用。MHA2MLA借鑒這一思想，對MHA中的鍵和值參數(shù)矩陣進(jìn)行奇異值分解（SVD），將其分解為低秩矩陣的乘積，以更少的參數(shù)近似原始矩陣。為了更好地保留鍵和值之間的交互信息，MHA2MLA采用聯(lián)合SVD（SVDjoint）策略，對鍵和值矩陣進(jìn)行聯(lián)合分解，而非各自處理。

MHA2MLA的項目地址

• GitHub倉庫：https://github.com/JT-Ushio/MHA2MLA
• arXiv技術(shù)論文：https://arxiv.org/pdf/2502.14837

MHA2MLA的應(yīng)用場景

• 邊緣設(shè)備部署：通過降低模型的內(nèi)存占用，使其適應(yīng)資源受限的智能終端和物聯(lián)網(wǎng)設(shè)備。
• 大規(guī)模模型推理：減少KV緩存，提高推理效率，降低硬件需求和能耗。
• 結(jié)合量化技術(shù)：與量化技術(shù)結(jié)合，進(jìn)一步優(yōu)化推理性能，適用于實時對話和在線翻譯等應(yīng)用。
• 長文本處理：緩解長文本任務(wù)中的內(nèi)存瓶頸，高效處理長文檔摘要和長篇生成。
• 快速模型遷移：僅需少量數(shù)據(jù)微調(diào)，迅速將MHA模型轉(zhuǎn)變?yōu)镸LA架構(gòu)，降低遷移成本。

常見問題

• MHA2MLA適合哪些場景？ MHA2MLA特別適合資源受限的環(huán)境，如邊緣設(shè)備和物聯(lián)網(wǎng)應(yīng)用，同時也適用于大規(guī)模模型推理和長文本處理。
• 微調(diào)時需要多少數(shù)據(jù)？ 只需使用原始數(shù)據(jù)的0.3%到0.6%進(jìn)行微調(diào)，即可實現(xiàn)有效的性能保持。
• 與其他技術(shù)兼容嗎？ 是的，MHA2MLA可以與量化技術(shù)等其他優(yōu)化方法結(jié)合使用，進(jìn)一步提升推理效率。