MHA2MLA – 復旦、上海AI Lab等推出優化LLM推理效率的方法

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

MHA2MLA是什么

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

MHA2MLA的主要功能

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

MHA2MLA的技術原理

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

MHA2MLA的項目地址

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

MHA2MLA的應用場景

• 邊緣設備部署：通過降低模型的內存占用，使其適應資源受限的智能終端和物聯網設備。
• 大規模模型推理：減少KV緩存，提高推理效率，降低硬件需求和能耗。
• 結合量化技術：與量化技術結合，進一步優化推理性能，適用于實時對話和在線翻譯等應用。
• 長文本處理：緩解長文本任務中的內存瓶頸，高效處理長文檔摘要和長篇生成。
• 快速模型遷移：僅需少量數據微調，迅速將MHA模型轉變為MLA架構，降低遷移成本。

常見問題

• MHA2MLA適合哪些場景？ MHA2MLA特別適合資源受限的環境，如邊緣設備和物聯網應用，同時也適用于大規模模型推理和長文本處理。
• 微調時需要多少數據？ 只需使用原始數據的0.3%到0.6%進行微調，即可實現有效的性能保持。
• 與其他技術兼容嗎？ 是的，MHA2MLA可以與量化技術等其他優化方法結合使用，進一步提升推理效率。