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

一種 KV Cache 友好的大模型 Attention 機(jī)制。

原標(biāo)題：階躍公開了自家新型注意力機(jī)制：KV緩存消耗直降93.7%，性能不減反增
文章來源：機(jī)器之心
內(nèi)容字?jǐn)?shù)：4981字

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

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

1. 問題與挑戰(zhàn)

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

2. MFA機(jī)制的創(chuàng)新

MFA的提出旨在最大限度地節(jié)省資源并接近理論性能上限。研究團(tuán)隊(duì)通過三個(gè)關(guān)鍵創(chuàng)新實(shí)現(xiàn)了這一目標(biāo)：

1. 顯著增加注意力頭數(shù)量和維度：突破傳統(tǒng)設(shè)計(jì)局限，極大提升模型容量。

2. 創(chuàng)新的低秩分解策略：在擴(kuò)展模型注意力頭數(shù)量和維度時(shí)保持高參數(shù)效率。

3. 單鍵值頭設(shè)計(jì)：即使增加模型復(fù)雜度，也能保持最低水平的內(nèi)存使用。

3. 理論分析與比較

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

4. 實(shí)驗(yàn)結(jié)果與結(jié)論

實(shí)驗(yàn)結(jié)果顯示，MFA和MFA-KR在不同規(guī)模模型（1B到7B參數(shù)）和訓(xùn)練數(shù)據(jù)量（10B到1T）下均表現(xiàn)出色。MFA與傳統(tǒng)MHA具有相當(dāng)?shù)男阅軘U(kuò)展能力，并在內(nèi)存節(jié)省方面展現(xiàn)出顯著優(yōu)勢(shì)，最大規(guī)模模型上實(shí)現(xiàn)了高達(dá)87.5%的內(nèi)存節(jié)省。消融實(shí)驗(yàn)也驗(yàn)證了MFA和MFA-KR設(shè)計(jì)的有效性。

5. 展望

MFA以簡(jiǎn)潔的設(shè)計(jì)解決了LLM高效推理的顯存瓶頸問題，并能無縫集成到現(xiàn)有Transformer生態(tài)中。這項(xiàng)研究為大語言模型的廣泛應(yīng)用提供了重要的技術(shù)支撐。

聯(lián)系作者

文章來源：機(jī)器之心
作者微信：
作者簡(jiǎn)介：專業(yè)的人工智能媒體和產(chǎn)業(yè)服務(wù)平臺(tái)