自2017年谷歌提出了Transformer架構，以及那篇著名的論文《Attention Is All You Need》后，注意力機制迅速成為自然語言處理領域的核心技術。大型語言模型（LLMs）借助Transformer的自注意力機制，實現了對復雜語言模式的捕捉，在機器翻譯、文本生成、對話系統等領域取得了革命性的突破。它們不僅改變了學術研究的方向，更深刻地影響了生產力工具的發展，提高了人們的工作效率和生活質量。

隨著模型規模和數據量的不斷增長，LLMs面臨的挑戰也日益凸顯。盡管計算能力的提升使得訓練更大規模的模型成為可能，但硬件的通信瓶頸卻成為制約模型進一步發展的主要障礙。這種瓶頸并非源于計算資源的匱乏，而是由于自注意力機制在處理長序列時，內存和通信開銷呈指數級增長。

在自注意力機制中，每個Token都需要與序列中的所有其他Token進行互動，這導致了計算量和存儲需求的急劇上升，特別是鍵值（Key-Value，KV）緩存的大小隨著序列長度的增加而迅速增長。當處理超長文本時，KV緩存的內存需求可能超過現有硬件的承載能力，導致模型無法高效地進行推理和生成。

面對這一挑戰，研究人員嘗試了多種解決方案。其中，組查詢注意力（Group Query Attention，GQA）作為一種優化方法，通過減少鍵和值的頭數，降低了KV緩存的尺寸。然而GQA在降低內存開銷的同時，也不可避免地限制了模型的表達能力，導致性能下降。如何在不增加KV緩存大小的情況下，提高模型的表達能力，成為亟待解決的問題。

基于此，來自北京大學和小米公司的研究團隊提出了多頭潛在注意力（Multi-head Latent Attention，MLA）。他們認為，通過引入低秩矩陣，可以在鍵值層實現KV緩存的壓縮，同時通過上投影矩陣，增強模型的表達能力，從而在不增加KV緩存開銷的情況下，提升模型性能。

他們進一步提出了TransMLA，這是一種將現有的GQA模型轉換為MLA模型的后訓練方法。通過理論證明，他們證明了在相同的KV緩存開銷下，MLA的表達能力始終優于GQA。實驗結果也驗證了這一觀點，轉換后的MLA模型在下游任務中表現出色，顯著優于原始的GQA模型。他們開源了TransMLA架構：https://github.com/fxmeng/TransMLA

這項研究的重要性在于，它為大型語言模型的優化提供了一條新的路徑。在硬件限制難以突破的情況下，如何通過模型結構的調整來提升性能，就顯得尤為關鍵。TransMLA的提出，不僅為模型開發者提供了一種低成本、高效率的優化方案，也為LLMs的未來發展指明了方向。

值得一提的是，這個研究團隊背景深厚。主要作者Fanxu Meng, Zengwei Yao, Muhan Zhang來自北京大學人工智能研究院和通用人工智能國家重點實驗室，擁有扎實的理論基礎和豐富的研究經驗。同時，團隊成員也包括來自小米公司的工程師，具備將先進技術應用于實際產品的能力。這種產學研結合的團隊構成，使得他們的研究既有前瞻性的理論價值，又具備實際應用的可行性。

相關工作

自注意力機制的內存瓶頸

自注意力機制是Transformer模型的核心組件，它允許模型在序列中的任意位置建立關聯。然而，其計算和內存需求與序列長度呈二次關系。具體來說，對于長度為 TT 的序列，自注意力機制需要計算并存儲一個 T×TT \times T 的注意力矩陣。此外，還需要存儲每個Token的鍵（Key）和值（Value）向量，這些鍵值對的數量隨著序列長度的增加而迅速增加，導致KV緩存的內存占用也迅速增長。

當處理長序列時，每個Token的KV對都需要重新計算，并存儲在緩存中，這種緩存的內存需求會隨著序列長度呈指數級增長。這種增長超出了現有硬件的內存容量，成為系統的瓶頸，限制了大模型在處理長文本任務中的應用。例如，在長文檔摘要或長篇文章生成任務中，KV緩存的內存需求可能達到甚至超過現有硬件的內存上限，導致模型無法有效運行。

現有解決方案及局限性

為了緩解自注意力機制帶來的內存瓶頸，研究人員提出了多種方法。這些方法在降低內存和計算需求的同時，也帶來了各自的局限性。

圖1：將GQA等效轉換為MLA的過程。（a） 在GQA中，在計算注意力之前，重復K以匹配Q頭的數量。（b） 相反，WK被重復；與X相乘后，直接產生K′，相當于GQA。（c） 分解W′K產生一個r維表示，它與GQA中的KV緩存維度相匹配，同時保持等價性。

線性注意力：例如Linear Transformer、RWKV、Mamba等方法，通過線性化注意力機制，使計算和內存需求與序列長度線性相關。這些方法通過重新設計注意力機制，將復雜的矩陣乘法轉化為向量內積，從而避免了大量的矩陣計算。然而，這種方法可能降低模型對長距離依賴的捕捉能力，導致在處理需要長距離依賴的任務時，模型性能下降。

動態Token剪枝：如LazyLLM、A2SF和SnapKV，通過動態地移除不重要的Token，以減少KV緩存的大小。這些方法選擇性地保留對后續預測最重要的Token，從而減少內存使用。然而，動態剪枝可能錯失關鍵的上下文信息，尤其是在需要全局理解的任務中。此外，為了確定哪些Token可以被剪枝，需要額外的計算和復雜的策略設計，這增加了系統的復雜性。

剪枝頭部維度：方法如SliceGPT、Sheared和Simple Pruning，通過減少注意力頭的數量或降低每個頭的維度，來減少計算和內存需求。這些方法通過剪去模型中冗余或不重要的頭部來優化內存使用。然而，過度的剪枝可能削弱模型的表達能力，因為注意力頭的多樣性對于捕捉復雜的語言模式至關重要，過少的頭部可能導致模型無法充分表示復雜的關系，進而影響性能。

跨層共享KV表示：如YONO、MiniCache和MLKV，通過在不同Transformer層之間共享KV表示，減少需要存儲的KV緩存的總量。這種方法通過讓每層使用相同的KV表示來節省內存。然而，由于不同層可能需要不同的表示來捕獲不同層次的特征，共享KV表示可能會限制模型的學習能力，并引入不同層間的干擾，影響性能的穩定性。

KV量化：如GPTQ、Kivi和KVQuant，通過對KV向量進行量化，以較低的比特數表示，從而減少內存占用。例如，將浮點數表示的KV向量量化為8位或更低的精度。這種方法可以顯著減少內存使用和計算開銷。然而，量化可能引入誤差，導致表示精度下降，影響模型的準確性。過度量化還可能導致梯度消失或爆炸，影響模型的訓練和穩定性。

組查詢注意力（GQA）的應用與局限

組查詢注意力（Group Query Attention，GQA）是一種優化策略，通過將多個查詢頭分組，使每個組共享相同的鍵和值表示，從而減少了KV緩存的大小。在這種機制下，鍵和值的頭數減少了，從而降低了內存需求。

GQA的應用優勢

降低內存開銷：通過共享鍵和值，顯著減少KV緩存的大小，適用于內存受限的環境。

提升計算效率：減少計算量，加快模型的推理速度。

然而GQA也存在一定的局限性，在減少內存開銷的同時，也限制了模型的表達能力。由于多個查詢共享相同的鍵和值，模型的表達多樣性受到限制，這使得模型難以捕捉復雜的模式和關系。在一些需要細粒度注意力的任務中，GQA的性能可能無法達到理想效果。

盡管GQA在降低KV緩存方面表現出色，但其對模型表達能力的限制使得在一些應用場景中性能受損。為此，尋找一種既能降低KV緩存大小，又不犧牲模型表達能力的方法，成為了亟待解決的問題。這正是多頭潛在注意力（MLA）所要解決的核心挑戰。

TransMLA方法詳解

MLA的核心思想

多頭潛在注意力（MLA）的核心思想是通過在鍵值層使用低秩矩陣來壓縮KV緩存，同時引入上投影矩陣以增強模型的表達能力。MLA通過低秩矩陣降低了KV緩存的尺寸，在減少內存占用的同時，保持了必要的注意力信息。而上投影矩陣則提供了更高的表達能力，使得模型可以在不增加KV緩存大小的情況下，提高注意力機制的靈活性和多樣性。

理論證明：MLA表達能力優于GQA

定理1：在相同的KV緩存開銷下，MLA的表達能力始終大于GQA。

證明思路：

GQA通過復制鍵值頭，使得組內的查詢共享相同的鍵和值，這在一定程度上限制了模型的表達多樣性。而MLA則通過低秩分解，使得鍵和值的表示更加豐富，允許更多的通道間多樣性。這種差異使得MLA在保持相同KV緩存開銷的情況下，能夠提供更強的表達能力。

具體證明步驟：

首先，我們展示如何在GQA中通過復制鍵來匹配查詢頭的數量。設X為輸入序列，WK為鍵的投影矩陣，則鍵矩陣KK可以表示為：

在GQA中，為了匹配查詢頭的數量，鍵矩陣需要被復制。設復制因子為s，則復制后的鍵矩陣可以表示為：

接下來，我們將復制過程移到參數側，即先復制投影矩陣，然后再進行計算。設WK′為復制后的投影矩陣，則有：

這樣做可以實現相同的效果，同時不增加KV緩存的大小。

為了進一步增強表示能力，我們引入低秩分解。使用奇異值分解（SVD），將復制的投影矩陣分解成低秩矩陣：

其中，UK和VK是正交矩陣，SK是對角矩陣。我們可以截斷SVD，僅保留前r個奇異值，從而實現低秩表示：

設低秩矩陣為WKa和WKb，則有：

最終，鍵矩陣表示為：

這樣，MLA通過低秩分解，不僅壓縮了KV緩存，還增強了模型的表達能力。通過構造一個WKbW^b_K中向量正交的例子，可以證明MLA在某些情況下的表達能力是GQA無法達到的。

TransMLA的模型轉換方法

TransMLA提供了一種將基于GQA的預訓練模型轉換為MLA模型的方法。具體步驟如下。

調整權重矩陣：將原始GQA模型的鍵和值投影矩陣拆分，并引入額外的低秩矩陣。通過低秩分解，我們得到兩個新的投影矩陣WKaW^a_K和WKbW^b_K，用以替換原始的投影矩陣。

保持KV緩存大小不變：在轉換過程中，通過低秩分解和矩陣調整，確保轉換后的MLA模型的KV緩存大小與原始GQA模型相同。這意味著，我們在增強模型表達能力的同時，不會增加內存開銷。

增強模型表達能力：轉換后的MLA模型可以在不增加內存開銷的情況下，提升注意力機制的多樣性和靈活性。具體地，通過引入上投影矩陣WKbW^b_K，模型可以更好地捕捉復雜的模式和關系，提升整體性能。

通過上述方法，TransMLA不僅成功地將GQA模型轉換為MLA模型，還顯著提升了模型的表達能力和實際性能，為大模型的優化提供了一種高效的解決方案。

實驗結果與分析

實驗設置

為了驗證TransMLA的有效性，研究團隊設計了一系列實驗，選擇了兩個基于GQA的預訓練模型：Qwen2.5-7B和Qwen2.5-14B。通過將這些模型轉換為MLA模型，并調整相應的權重矩陣維度，研究團隊希望評估MLA在實際應用中的表現。

圖2:TransMLA和基于GQA的Qwen模型之間訓練損失和準確性的比較。

在實驗中，使用了包含數學和編程任務的指令微調數據集SmolTalk，具體包括MetaMathQA和SelfOSS-Starcoder2-Instruct數據集。實驗設置的訓練配置為：批次大小為16，學習率為2e-5，訓練2個周期。為了減少對原始模型的影響，實驗中僅對鍵值層進行訓練。這樣的配置旨在保持實驗的公平性和一致性，使得結果具有較高的可信度。

微調性能比較

在訓練過程中，TransMLA模型的損失下降速度明顯快于原始GQA模型。這表明，TransMLA在數據擬合方面具有明顯優勢。在相同的訓練條件下，TransMLA能夠更快速地適應數據，并在較短的時間內達到更低的損失值。這一結果驗證了TransMLA在訓練效率上的改進，表明其優化方法在實際應用中具有較高的效能。

在具體的任務表現上，TransMLA模型也展現出了優越的性能。在數學任務（如GSM8K）和編程任務（如HumanEval）中，TransMLA模型的準確率顯著高于原始GQA模型。例如，在7B模型的實驗中，GSM8K的準確率從原始GQA模型的81%提升到了TransMLA模型的86%。這種顯著的性能提升，不僅表明TransMLA在特定任務中的適應性強，也展示了其在提升模型表達能力方面的潛力。

TransMLA的性能提升可以歸因于以下幾個方面。

增強的鍵值維度在模型表達能力上的提升起到了關鍵作用，通過調整權重矩陣，TransMLA擴大了鍵和值的維度，使得模型能夠捕捉到更多的特征和模式，從而提升了整體性能。

正交分解在優化模型結構方面也發揮了重要作用，實驗結果表明，僅增加參數量并不能帶來顯著的性能提升，正交分解方法通過優化模型的表示，使得模型在保持較低內存開銷的情況下，實現了性能的顯著提升。這一結果強調了結構優化在提升模型性能中的重要性。

綜上所述，TransMLA通過優化KV緩存的表示和引入先進的矩陣分解技術，實現了在不增加內存開銷的情況下，顯著提升模型性能的目標。這不僅為大規模語言模型的進一步發展提供了新的思路，也為未來的研究和應用指明了方向。

結論與展望

TransMLA成功地展示了多頭潛在注意力（MLA）的強大優勢，證明了其在表達能力上顯著優于組查詢注意力（GQA）。通過理論分析和實驗驗證，研究團隊展示了MLA在相同KV緩存開銷的情況下，能夠提供更強的模型性能。轉換后的MLA模型在多種下游任務中表現出色，為大型語言模型（LLMs）的注意力機制設計提供了全新的視角和思路。研究表明，TransMLA不僅有效地減少了KV緩存的內存需求，還大幅提升了模型的表達能力和推理速度。

TransMLA的研究不僅在技術層面上取得了重要突破，也為整個行業帶來了深遠的影響。

TransMLA顯著提升了資源效率，在不增加硬件負擔的情況下，通過優化注意力機制和鍵值緩存，提升了模型的性能。這為企業和研究機構在有限的計算資源下，進一步提高模型的表現提供了新的解決方案。

TransMLA對生態環境也具有友好性，通過減少模型對內存和計算資源的需求，降低了能源消耗和碳排放。這一特性使得TransMLA在綠色計算和可持續發展方面表現出色，符合當今社會對環保和節能的追求。

TransMLA提供了一種低成本的模型轉換方法，使現有的基于GQA的模型可以輕松遷移到MLA架構。這樣的遷移不僅能夠保留原有模型的優點，還能大幅提升其性能和效率。這為模型開發者和應用者提供了極大的便利，加速了新技術的推廣和應用。

研究團隊計劃在以下幾個方面繼續拓展和優化TransMLA。

團隊將擴展到更大規模的模型，包括LLaMA、Qwen、Mistral等大型語言模型，驗證TransMLA在更大規模下的有效性。這將進一步檢驗和提升MLA在處理超長序列和大規模數據上的表現。

研究團隊將使用DeepSeek R1蒸餾技術，進一步提升轉換后模型的性能。通過蒸餾技術，可以優化模型的參數和結構，使其在保持高性能的同時，具有更高的計算效率和穩定性。

針對MLA的特性，研究團隊計劃開發特定的推理加速策略。通過設計高效的推理算法，最大程度地發揮MLA的優勢，保持模型的低延遲和高響應速度。這將為實時應用和高頻交互場景中的LLMs提供更加優質的性能支持。

總的來說，TransMLA的成功不僅為現有的大模型優化提供了有效的解決方案，也為未來的研究和應用開辟了新的路徑。隨著技術的不斷進步和優化，相信MLA將在更多領域發揮重要作用，推動人工智能技術的發展和創新。

參考資料：???https://arxiv.org/pdf/2502.07864??

本文轉載自??獨角噬元獸??，作者： FlerkenS ????