階躍公開了自家新型注意力機制：KV快取消耗直降93.7%，效能不減反增
机器之心發表於2025-01-17
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
隨著當前大語言模型的廣泛應用和推理時擴充套件的新正規化的崛起，如何實現高效的大規模推理成為了一個巨大挑戰。特別是在語言模型的推理階段，傳統注意力機制中的鍵值快取（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。
研究團隊透過研究 Attention 機制的一般性設計和容量分析，明確了 Attention 機制的容量最相關的兩個維度，並且提出了一系列的分析方法和設計原則。
為了更好地理解注意力機制的本質，研究團隊提出了廣義多頭注意力（GMHA）的概念框架，為理解不同的 MHA 變種注意力機制提供了一個統一的視角。進一步地，研究團隊分別從推理角度研究鍵值的計算和儲存方式，從分解角度探討模型的容量特徵。這種創新的分析方法為理解不同策略如何在模型效能和計算效率之間取得平衡提供了全新視角。在此基礎上，研究者們確立了完全引數化雙線性注意力（FPBA）作為理論上的效能上限標準。他們發現，目前廣泛使用的 MHA 及其各種變體實際上都可以被視為 FPBA 的低秩分解版本。研究團隊在分析現有解決方案時，重點考察了兩個代表性的改進方案：多查詢注意力（MQA）和多頭潛在注意力（MLA）。對於 MQA，研究人員發現它採用了一種更激進的引數共享策略。不同於 MHA 在頭部之間保持獨立的引數，MQA 讓所有注意力頭共享同一組鍵值引數。這種設計雖然將記憶體使用降到了極低的水平，但可能會影響模型的表達能力。研究團隊指出，這種權衡實際上反映了一個更普遍的設計困境：如何在保持模型效能的同時減少資源消耗。而對於更新的 MLA 方案，研究人員發現它採用了更為複雜的架構設計。MLA 引入了一個共享的潛在空間，透過在這個空間中進行引數壓縮來節省記憶體。雖然這種設計看似提供了更大的靈活性（因為中間維度可以設定得更大），但研究團隊的理論分析揭示了一個關鍵問題：模型的實際表達能力仍然受限於最小維度的約束，這意味著 MLA 的設計中增加中間維度並不能真正提升模型的效能上限。在深入分析現有方案的基礎上，研究團隊提出了一個雄心勃勃的目標：開發一種新的注意力機制，既能最大限度地節省資源，又能儘可能接近理論效能上限。這個目標促使他們開發出了多矩陣分解注意力（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 生態中。這種簡單而強大的創新，必將加速大語言模型在更多場景中的落地應用。那麼你覺得這項關於注意力機制的研究如何呢？歡迎在評論區留言討論。