Published on

階躍發布新型注意力機制：KV 快取消耗銳減 93.7%，效能不減反增

Authors
  • avatar
    Name
    Ajax
    Twitter

引言

隨著大型語言模型（LLM）的廣泛應用和推理時擴展的新範式的出現，高效的大規模推理成為一個關鍵挑戰。尤其是在語言模型的推理階段，傳統的注意力機制中的鍵值快取（KV Cache）會隨著批次大小和序列長度線性增長，這成為了制約大型語言模型規模化應用和推理時擴展的瓶頸，被稱為“記憶體殺手”。

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

MFA：多矩陣分解注意力

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

核心發現：

  • MFA 和 MFA-KR 在實驗中超越了 MLA 的效能。
  • 在減少高達 93.7% 的 KV 快取使用量的情況下，MFA 和 MFA-KR 與傳統的 MHA 效能相當。
  • MFA 實現簡單，容易復現，對超參數敏感度低，並且相容各種位置編碼 (Pos-embedding)。

MFA 方法與分析

研究團隊透過深入研究注意力機制的設計和容量分析，明確了注意力機制容量相關的兩個關鍵維度，並提出了一系列分析方法和設計原則。

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

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

  • MQA： 採用更激進的參數共享策略，讓所有注意力頭共享同一組鍵值參數。雖然降低了記憶體使用，但可能影響模型的表達能力。
  • MLA： 引入共享的潛在空間進行參數壓縮，節省記憶體。然而，模型的實際表達能力仍然受限於最小維度的約束，增加中間維度並不能真正提升模型的效能上限。

基於對現有方案的深入分析，研究團隊提出了一個目標：開發一種新的注意力機制，既能最大限度地節省資源，又能盡可能接近理論效能上限。由此，他們開發出了多矩陣分解注意力（MFA）。

MFA 的設計體現了三個關鍵創新：

  • 突破傳統設計： 顯著增加了注意力頭的數量和維度，極大提升了注意力模組的模型容量。
  • 矩陣分解創新： 採用激進的低秩分解策略，在擴展模型注意力頭的數量和維度的同時，保持了極高的參數效率。
  • 單鍵值頭設計： 確保即使在增加模型複雜度的情況下，記憶體使用仍然保持在最低水平。

MFA 與其他注意力機制的對比

為了進一步分析MFA和其他注意力機制的不同，研究團隊引入了兩個關鍵指標來度量GMHA系列模型的容量：

  • 模型總有效秩 (TER)： 注意力頭數量與每個頭部分解秩（FRH）的乘積。
  • 共享隱空間維度 (SLSD)： 代表所有注意力頭共同使用的隱空間維度。

TER 和 SLSD 越大，模型具有更高的容量。每個頭的分解秩（FRH）不超過 SLSD，而 KV 快取的佔用受制於 FRH 和 SLSD 之間的較大值，這構成了容量和效率之間的關鍵權衡。

透過這個框架分析：

  • 與 MQA 相比，MFA 同時實現了更高的 SLSD 和更高的 TER。
  • 與 MLA 相比，在相似的參數預算下，MFA 不僅達到了更小的 KV 快取尺寸，還實現了更高的 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 生態中。這種簡單而強大的創新，有望加速大型語言模型在更多場景中的落地應用。