一種 KV Cache 友好的大模型 Attention 機制。

原標題：階躍公開了自家新型注意力機制：KV緩存消耗直降93.7%，性能不減反增
文章來源：機器之心
內容字數：4981字

高效大語言模型推理：多矩陣分解注意力機制MFA的突破

機器之心AIxiv專欄報道了階躍星辰、清華大學等機構的研究成果：一篇關于新型注意力機制——多矩陣分解注意力（MFA）及其變體MFA-Key-Reuse的論文。該研究旨在解決大語言模型推理階段的內存瓶頸問題，即傳統注意力機制的鍵值緩存（KV Cache）隨批處理大小和序列長度線性增長的問題。

1. 問題與挑戰

現有解決KV Cache內存占用問題的注意力機制變體，例如MQA和MLA，存在性能與復雜度之間的權衡難題。MQA犧牲了模型表達能力以換取極低的內存使用；而MLA雖然結構靈活，但其性能上限受限于最小維度，增加中間維度并不能有效提升性能。

2. MFA機制的創新

MFA的提出旨在最大限度地節省資源并接近理論性能上限。研究團隊通過三個關鍵創新實現了這一目標：

1. 顯著增加注意力頭數量和維度：突破傳統設計局限，極大提升模型容量。

2. 創新的低秩分解策略：在擴展模型注意力頭數量和維度時保持高參數效率。

3. 單鍵值頭設計：即使增加模型復雜度，也能保持最低水平的內存使用。

3. 理論分析與比較

研究團隊提出了廣義多頭注意力（GMHA）的概念框架，并引入“總有效秩（TER）”和“共享隱空間維度（SLSD）”兩個指標來衡量GMHA系列模型的容量。通過此框架，MFA與MQA、MLA、MHA的比較分析表明，MFA在參數效率和性能之間取得了最佳平衡。MFA在更高的TER和SLSD下，實現了更低的KV Cache占用。

4. 實驗結果與結論

實驗結果顯示，MFA和MFA-KR在不同規模模型（1B到7B參數）和訓練數據量（10B到1T）下均表現出色。MFA與傳統MHA具有相當的性能擴展能力，并在內存節省方面展現出顯著優勢，最大規模模型上實現了高達87.5%的內存節省。消融實驗也驗證了MFA和MFA-KR設計的有效性。

5. 展望

MFA以簡潔的設計解決了LLM高效推理的顯存瓶頸問題，并能無縫集成到現有Transformer生態中。這項研究為大語言模型的廣泛應用提供了重要的技術支撐。

聯系作者

文章來源：機器之心
作者微信：
作者簡介：專業的人工智能媒體和產業服務平臺