階躍公開了新型注意力機制：KV緩存消耗直降93.7%，性能不減反增

隨着當前大語言模型的廣泛應用和推理時擴展的新範式的崛起，如何實現高效的大規模推理成爲了一個巨大挑戰。特別是在語言模型的推理階段，傳統注意力機制中的鍵值緩存（KV Cache）會隨着批處理大小和序列長度線性增長，儼然成爲制約大語言模型規模化應用和推理時擴展的「內存殺手」。

雖然業界已有 MQA、GQA、MLA 等多種注意力機制變體試圖解決這一問題，但這些方案要麼難以在嚴格的顯存限制下保持理想性能，要麼在模型結構上引入額外複雜度，帶來巨大的工程挑戰和生態兼容性問題。

在近期由階躍星辰、清華大學等機構完成的論文《Multi-matrix Factorization Attention》中，研究人員提出新型注意力機制架構 —— 多矩陣分解注意力（MFA）及其變體 MFA-Key-Reuse，在大幅降低語言模型推理成本的同時，還實現了性能的顯著提升。

鏈接：https://arxiv.org/abs/2412.19255

實驗中，MFA 和 MFA-KR 不僅超越了 MLA 的性能，還在減少了高達 93.7% 的 KV Cache 使用量的情況下，與傳統的 MHA 性能相當。於此同時，MFA 實現簡單，容易復現，對超參敏感度低，且兼容各種 Pos-embedding。

MFA 方法與分析

研究團隊通過研究 Attention 機制的一般性設計和容量分析，明確了 Attention 機制的容量最相關的兩個維度，並且提出了一系列的分析方法和設計原則。

爲了更好地理解注意力機制的本質，研究團隊提出了廣義多頭注意力（GMHA）的概念框架，爲理解不同的 MHA 變種注意力機制提供了一個統一的視角。進一步地，研究團隊分別從推理角度研究鍵值的計算和存儲方式，從分解角度探討模型的容量特徵。這種創新的分析方法爲理解不同策略如何在模型性能和計算效率之間取得平衡提供了全新視角。

在此基礎上，研究者們確立了完全參數化雙線性注意力（FPBA）作爲理論上的性能上限標準。他們發現，目前廣泛使用的 MHA 及其各種變體實際上都可以被視爲 FPBA 的低秩分解版本。研究團隊在分析現有解決方案時，重點考察了兩個代表性的改進方案：多查詢注意力（MQA）和多頭潛在注意力（MLA）。

對於 MQA，研究人員發現它採用了一種更激進的參數共享策略。不同於 MHA 在頭部之間保持獨立的參數，MQA 讓所有注意力頭共享同一組鍵值參數。這種設計雖然將內存使用降到了極低的水平，但可能會影響模型的表達能力。研究團隊指出，這種權衡實際上反映了一個更普遍的設計困境：如何在保持模型性能的同時減少資源消耗。

而對於更新的 MLA 方案，研究人員發現它採用了更爲複雜的架構設計。MLA 引入了一個共享的潛在空間，通過在這個空間中進行參數壓縮來節省內存。雖然這種設計看似提供了更大的靈活性（因爲中間維度可以設置得更大），但研究團隊的理論分析揭示了一個關鍵問題：模型的實際表達能力仍然受限於最小維度的約束，這意味着 MLA 的設計中增加中間維度並不能真正提升模型的性能上限。

在深入分析現有方案的基礎上，研究團隊提出了一個雄心勃勃的目標：開發一種新的注意力機制，既能最大限度地節省資源，又能儘可能接近理論性能上限。這個目標促使他們開發出了多矩陣分解注意力（MFA）。

MFA 的設計體現了研究團隊的三個關鍵創新：

爲了進一步分析 MFA 和其他注意力機制的不同，研究團隊提供了一個清晰的對比表格。研究人員首先引入兩個關鍵指標來度量 GMHA 系列模型的容量：模型總有效秩 TER （ Total Effective Rank ） 和 共享隱空間維度 SLSD（Shared latent subspace dim）。總有效秩 TER 定義爲注意力頭數量與每個頭部分解秩（Factorization rank per head，簡稱 FRH）的乘積，而共享隱空間維度 SLSD 則代表了所有注意力頭共同使用的隱空間維度，注意到 TER 和 SLSD 越大，模型具有更高的容量。同時注意到每個頭的分解秩（FRH）不超過 SLSD，而 KV Cache 的佔用受制於 FRH 和 SLSD 之間的較大值，這構成了容量和效率之間的關鍵權衡。

通過這個框架進行分析，可以發現與 MQA 相比，MFA 同時實現了更高的 SLSD 和更高的 TER；與 MLA 相比，在相似的參數預算下，MFA 不僅達到了更小的 KV cache 尺寸，還實現了更高的 TER，同時保持了相當的 SLSD；與傳統的 MHA 相比，雖然 MFA 的 SLSD 較小，但它的 TER 更高，這解釋了爲什麼在實驗中 MFA 能夠獲得更好的性能。

實驗結果

爲了研究新架構在可能的更大規模上的表現，研究團隊開展了一系列深入的擴展性實驗，系統地測試了從 1B 到 7B 參數的不同規模的模型，訓練數據量從 10B 擴展到 1T。在性能擴展方面，研究團隊的 MFA 方案展現出與傳統 MHA 完全相當的擴展能力。這意味着，即使在更大規模下，MFA 依然能保持其優異的性能。而 MFA-KR 雖然在性能上略有降低，但其擴展趨勢與 MHA 保持一致。而隨着模型規模的增加，MFA 和 MFA-KR 在內存節省方面的優勢不僅得到保持，還呈現出進一步擴大的趨勢。在最大規模模型上，MFA 實現了 87.5% 的內存節省，而 MFA-KR 更是將內存使用降低到原來的 6.25%。

研究人員還進行了一系列的消融實驗，證明 MFA 和 MFA-KR 的設計的有效性，並在其他的主流的位置編碼上也驗證了 MFA 和 MFA-KR 的性能優勢。

展望

MFA 用最簡潔的設計帶來了最顯著的提升，在不增加額外工程複雜度的前提下，優雅地解決了 LLM 的高效推理顯存瓶頸問題，並能無縫集成到現有的 Transformer 生態中。這種簡單而強大的創新，必將加速大語言模型在更多場景中的落地應用。

那麼你覺得這項關於注意力機制的研究如何呢？歡迎在評論區留言討論。