隨著當(dāng)前大語(yǔ)言模型的廣泛應(yīng)用和推理時(shí)擴(kuò)展的新范式的崛起，如何實(shí)現(xiàn)高效的大規(guī)模推理成為了一個(gè)巨大挑戰(zhàn)。特別是在語(yǔ)言模型的推理階段，傳統(tǒng)注意力機(jī)制中的鍵值緩存（KV Cache）會(huì)隨著批處理大小和序列長(zhǎng)度線性增長(zhǎng)，儼然成為制約大語(yǔ)言模型規(guī)?；瘧?yīng)用和推理時(shí)擴(kuò)展的「內(nèi)存殺手」。

雖然業(yè)界已有 MQA、GQA、MLA 等多種注意力機(jī)制變體試圖解決這一問題，但這些方案要么難以在嚴(yán)格的顯存限制下保持理想性能，要么在模型結(jié)構(gòu)上引入額外復(fù)雜度，帶來(lái)巨大的工程挑戰(zhàn)和生態(tài)兼容性問題。

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

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

實(shí)驗(yàn)中，MFA 和 MFA-KR 不僅超越了 MLA 的性能，還在減少了高達(dá) 93.7% 的 KV Cache 使用量的情況下，與傳統(tǒng)的 MHA 性能相當(dāng)。于此同時(shí)，MFA 實(shí)現(xiàn)簡(jiǎn)單，容易復(fù)現(xiàn)，對(duì)超參敏感度低，且兼容各種 Pos-embedding。

MFA 方法與分析

研究團(tuán)隊(duì)通過研究 Attention 機(jī)制的一般性設(shè)計(jì)和容量分析，明確了 Attention 機(jī)制的容量最相關(guān)的兩個(gè)維度，并且提出了一系列的分析方法和設(shè)計(jì)原則。

為了更好地理解注意力機(jī)制的本質(zhì)，研究團(tuán)隊(duì)提出了廣義多頭注意力（GMHA）的概念框架，為理解不同的 MHA 變種注意力機(jī)制提供了一個(gè)統(tǒng)一的視角。進(jìn)一步地，研究團(tuán)隊(duì)分別從推理角度研究鍵值的計(jì)算和存儲(chǔ)方式，從分解角度探討模型的容量特征。這種創(chuàng)新的分析方法為理解不同策略如何在模型性能和計(jì)算效率之間取得平衡提供了全新視角。

在此基礎(chǔ)上，研究者們確立了完全參數(shù)化雙線性注意力（FPBA）作為理論上的性能上限標(biāo)準(zhǔn)。他們發(fā)現(xiàn)，目前廣泛使用的 MHA 及其各種變體實(shí)際上都可以被視為 FPBA 的低秩分解版本。研究團(tuán)隊(duì)在分析現(xiàn)有解決方案時(shí)，重點(diǎn)考察了兩個(gè)代表性的改進(jìn)方案：多查詢注意力（MQA）和多頭潛在注意力（MLA）。

對(duì)于 MQA，研究人員發(fā)現(xiàn)它采用了一種更激進(jìn)的參數(shù)共享策略。不同于 MHA 在頭部之間保持獨(dú)立的參數(shù)，MQA 讓所有注意力頭共享同一組鍵值參數(shù)。這種設(shè)計(jì)雖然將內(nèi)存使用降到了極低的水平，但可能會(huì)影響模型的表達(dá)能力。研究團(tuán)隊(duì)指出，這種權(quán)衡實(shí)際上反映了一個(gè)更普遍的設(shè)計(jì)困境：如何在保持模型性能的同時(shí)減少資源消耗。

而對(duì)于更新的 MLA 方案，研究人員發(fā)現(xiàn)它采用了更為復(fù)雜的架構(gòu)設(shè)計(jì)。MLA 引入了一個(gè)共享的潛在空間，通過在這個(gè)空間中進(jìn)行參數(shù)壓縮來(lái)節(jié)省內(nèi)存。雖然這種設(shè)計(jì)看似提供了更大的靈活性（因?yàn)橹虚g維度可以設(shè)置得更大），但研究團(tuán)隊(duì)的理論分析揭示了一個(gè)關(guān)鍵問題：模型的實(shí)際表達(dá)能力仍然受限于最小維度的約束，這意味著 MLA 的設(shè)計(jì)中增加中間維度并不能真正提升模型的性能上限。

在深入分析現(xiàn)有方案的基礎(chǔ)上，研究團(tuán)隊(duì)提出了一個(gè)雄心勃勃的目標(biāo)：開發(fā)一種新的注意力機(jī)制，既能最大限度地節(jié)省資源，又能盡可能接近理論性能上限。這個(gè)目標(biāo)促使他們開發(fā)出了多矩陣分解注意力（MFA）。

MFA 的設(shè)計(jì)體現(xiàn)了研究團(tuán)隊(duì)的三個(gè)關(guān)鍵創(chuàng)新：

• 首先，他們突破傳統(tǒng)設(shè)計(jì)的局限，顯著增加了注意力頭的數(shù)量和維度，極大增加了注意力模塊的模型容量。
• 其次，研究團(tuán)隊(duì)在矩陣分解方面實(shí)現(xiàn)了創(chuàng)新性突破，采用激進(jìn)的低秩分解策略，成功地在擴(kuò)展模型注意力頭的數(shù)量和維度時(shí)保持了極高的參數(shù)效率。
• 最后，研究團(tuán)隊(duì)采用單鍵值頭設(shè)計(jì)，這個(gè)設(shè)計(jì)確保了即使在增加模型復(fù)雜度的情況下，內(nèi)存使用仍然保持在最低水平。

為了進(jìn)一步分析 MFA 和其他注意力機(jī)制的不同，研究團(tuán)隊(duì)提供了一個(gè)清晰的對(duì)比表格。研究人員首先引入兩個(gè)關(guān)鍵指標(biāo)來(lái)度量 GMHA 系列模型的容量：模型總有效秩 TER （ Total Effective Rank ） 和 共享隱空間維度 SLSD（Shared latent subspace dim）?？傆行е?TER 定義為注意力頭數(shù)量與每個(gè)頭部分解秩（Factorization rank per head，簡(jiǎn)稱 FRH）的乘積，而共享隱空間維度 SLSD 則代表了所有注意力頭共同使用的隱空間維度，注意到 TER 和  SLSD 越大，模型具有更高的容量。同時(shí)注意到每個(gè)頭的分解秩（FRH）不超過 SLSD，而 KV Cache 的占用受制于 FRH 和 SLSD 之間的較大值，這構(gòu)成了容量和效率之間的關(guān)鍵權(quán)衡。

通過這個(gè)框架進(jìn)行分析，可以發(fā)現(xiàn)與 MQA 相比，MFA 同時(shí)實(shí)現(xiàn)了更高的 SLSD 和更高的 TER；與 MLA 相比，在相似的參數(shù)預(yù)算下，MFA 不僅達(dá)到了更小的 KV cache 尺寸，還實(shí)現(xiàn)了更高的 TER，同時(shí)保持了相當(dāng)?shù)?SLSD；與傳統(tǒng)的 MHA 相比，雖然 MFA 的 SLSD 較小，但它的 TER 更高，這解釋了為什么在實(shí)驗(yàn)中 MFA 能夠獲得更好的性能。

實(shí)驗(yàn)結(jié)果

為了研究新架構(gòu)在可能的更大規(guī)模上的表現(xiàn)，研究團(tuán)隊(duì)開展了一系列深入的擴(kuò)展性實(shí)驗(yàn)，系統(tǒng)地測(cè)試了從 1B 到 7B 參數(shù)的不同規(guī)模的模型，訓(xùn)練數(shù)據(jù)量從 10B 擴(kuò)展到 1T。在性能擴(kuò)展方面，研究團(tuán)隊(duì)的 MFA 方案展現(xiàn)出與傳統(tǒng) MHA 完全相當(dāng)?shù)臄U(kuò)展能力。這意味著，即使在更大規(guī)模下，MFA 依然能保持其優(yōu)異的性能。而 MFA-KR 雖然在性能上略有降低，但其擴(kuò)展趨勢(shì)與 MHA 保持一致。而隨著模型規(guī)模的增加，MFA 和 MFA-KR 在內(nèi)存節(jié)省方面的優(yōu)勢(shì)不僅得到保持，還呈現(xiàn)出進(jìn)一步擴(kuò)大的趨勢(shì)。在最大規(guī)模模型上，MFA 實(shí)現(xiàn)了 87.5% 的內(nèi)存節(jié)省，而 MFA-KR 更是將內(nèi)存使用降低到原來(lái)的 6.25%。

研究人員還進(jìn)行了一系列的消融實(shí)驗(yàn)，證明 MFA 和 MFA-KR 的設(shè)計(jì)的有效性，并在其他的主流的位置編碼上也驗(yàn)證了 MFA 和 MFA-KR 的性能優(yōu)勢(shì)。

展望

MFA 用最簡(jiǎn)潔的設(shè)計(jì)帶來(lái)了最顯著的提升，在不增加額外工程復(fù)雜度的前提下，優(yōu)雅地解決了 LLM 的高效推理顯存瓶頸問題，并能無(wú)縫集成到現(xiàn)有的 Transformer 生態(tài)中。這種簡(jiǎn)單而強(qiáng)大的創(chuàng)新，必將加速大語(yǔ)言模型在更多場(chǎng)景中的落地應(yīng)用。