曾撼動(dòng)Transformer統(tǒng)治地位的Mamba作者之一Tri Dao，剛剛帶來(lái)新作——

提出兩種專(zhuān)為推理“量身定制”的注意力機(jī)制。

在保持模型性能不變的情況下，將解碼速度和吞吐量最高提升2倍，大大優(yōu)化了模型的長(zhǎng)上下文推理能力。

這項(xiàng)研究的三位作者均來(lái)自普林斯頓大學(xué)，論文主要有兩大貢獻(xiàn)：

其一，提出Grouped-Tied Attention（GTA），與已集成到LLaMA 3的注意力機(jī)制GQA質(zhì)量相當(dāng)，但KV緩存用量減少約50%。

其二，提出Grouped Latent Attention（GLA），與DeepSeek所使用的注意力機(jī)制MLA質(zhì)量匹配，但解碼速度更快，某些情況下比FlashMLA快2倍。

按照作者之一Ted Zadouri的總結(jié)：

GTA是GQA的有效替代品，而GLA是MLA的實(shí)用替代品。

一言以蔽之，通過(guò)優(yōu)化注意力機(jī)制的內(nèi)存使用和計(jì)算邏輯，在不犧牲模型生成質(zhì)量的前提下，可顯著提升大語(yǔ)言模型的推理效率和硬件資源利用率，尤其在長(zhǎng)上下文場(chǎng)景中優(yōu)勢(shì)更為突出。

相關(guān)論文公布后，一眾研究者也趕來(lái)祝賀~

那么，這項(xiàng)研究具體講了些啥？

引入推理感知注意力機(jī)制

概括而言，論文核心引入了推理感知注意力機(jī)制，即針對(duì)模型推理階段的內(nèi)存冗余、計(jì)算低效、長(zhǎng)上下文瓶頸等問(wèn)題，重新設(shè)計(jì)注意力機(jī)制。

據(jù)Tri Dao介紹，這項(xiàng)研究的起點(diǎn)始于一個(gè)想法：

在推理驅(qū)動(dòng)AI發(fā)展的時(shí)代，“理想”架構(gòu)應(yīng)該是什么樣子？

尤其在涉及長(zhǎng)上下文推理時(shí)，當(dāng)前的大語(yǔ)言模型（LLM）面臨內(nèi)存訪(fǎng)問(wèn)瓶頸和并行性限制兩大難題。

就是說(shuō)，模型生成文字時(shí)，每次都要從內(nèi)存里調(diào)取大量“歷史記錄”，不僅導(dǎo)致每個(gè)字生成變慢，而且只能按順序生成、沒(méi)法讓多個(gè)芯片同時(shí)干活。

對(duì)此，團(tuán)隊(duì)打算從兩個(gè)方向重新設(shè)計(jì)注意力機(jī)制：

• 更高的硬件效率：通過(guò)增加 “每字節(jié)內(nèi)存加載的計(jì)算量”（算術(shù)強(qiáng)度），減少對(duì)內(nèi)存帶寬的依賴(lài)；
• 保持并行可擴(kuò)展性：在不犧牲模型并行訓(xùn)練 / 推理能力的前提下優(yōu)化解碼速度。

而最終提出的GTA和GLA，在減少KV緩存用量的同時(shí)，模型質(zhì)量保持與現(xiàn)有方案相當(dāng)，且解碼速度顯著提升。

這里提到的“現(xiàn)有方案”，主要指早已聞名學(xué)術(shù)界的兩種方法：

一是分組查詢(xún)注意力（GQA）機(jī)制，它通過(guò)分組共享KV緩存減少內(nèi)存占用，在視覺(jué)Transformer（ViT）等任務(wù)中表現(xiàn)良好，適用于大規(guī)模數(shù)據(jù)處理，目前已應(yīng)用于Llama 3等開(kāi)源模型。

二是多頭潛在注意力（MLA）機(jī)制，最早可追溯到《Attention Is All You Need》這篇論文，后被DeepSeek再次帶火。它關(guān)注的是在不同層之間如何融合注意力信息，能減少每一層的冗余計(jì)算。

不過(guò)，由于GQA仍需為每組查詢(xún)頭存儲(chǔ)獨(dú)立KV、MLA并行優(yōu)化不足，故仍需進(jìn)一步改進(jìn)。

下面分別展開(kāi)團(tuán)隊(duì)提出的新方法GTA和GLA。

分組綁定注意力機(jī)制GTA

GTA的核心設(shè)計(jì)思路是：將不同查詢(xún)頭的鍵（Key）和值（Value）狀態(tài)進(jìn)行組合與重用，減少內(nèi)存?zhèn)鬏敶螖?shù)。

具體而言（右圖），它將多頭注意力的頭分為若干組（Group），每組內(nèi)的頭共享相同的Key和Value參數(shù)。計(jì)算時(shí)，同一組內(nèi)的頭使用相同的KV緩存，僅查詢(xún)（Query）參數(shù)獨(dú)立。

相比之下，中間傳統(tǒng)的多頭注意力機(jī)制（MHA）每個(gè)查詢(xún)頭都有獨(dú)立的鍵和值，由于沒(méi)有共享，導(dǎo)致它需要更多的內(nèi)存來(lái)存儲(chǔ)所有的鍵和值。

再對(duì)比GQA來(lái)看（左圖），GQA分組共享KV但每組仍獨(dú)立存儲(chǔ)，而GTA通過(guò)參數(shù)綁定實(shí)現(xiàn)了更徹底的KV重復(fù)利用。

分組潛在注意力機(jī)制GLA

而GLA的設(shè)計(jì)則采用了雙層結(jié)構(gòu)：

• 潛在層（Latent Layer）：引入固定數(shù)量的潛在Tokens，作為全局上下文的壓縮表示，替代部分原始Token的KV緩存；
• 分組頭機(jī)制：將查詢(xún)頭分組，每組頭共享潛在Token的KV，同時(shí)保留與原始Token的交互。

在解碼過(guò)程中，對(duì)比MLA（左圖），GLA通過(guò)共享聯(lián)合潛在表示減少了每個(gè)設(shè)備需要加載的KV緩存量，從而減少了內(nèi)存訪(fǎng)問(wèn)量。

并且由于每個(gè)設(shè)備上的KV緩存量減少了，更多的請(qǐng)求也可以同時(shí)處理。

「GQA和MLA」的有效替代品

那么，GTA和GLA的效果究竟如何呢？

團(tuán)隊(duì)在四種規(guī)模的模型上進(jìn)行了實(shí)驗(yàn)，包括小型（183M）、中型（433M）、大型（876M）和XL（1471M）。這些模型基于FineWeb-Edu-100B數(shù)據(jù)集訓(xùn)練，采用GPT-3架構(gòu)和Llama 3分詞器。

測(cè)試的指標(biāo)主要分為兩大類(lèi)：

• 質(zhì)量指標(biāo)：困惑度（Perplexity）、下游任務(wù)準(zhǔn)確率（Winogrande、SciQ等7個(gè)基準(zhǔn)）；
• 效率指標(biāo)：每Token解碼延遲、吞吐量、KV緩存占用量。

實(shí)驗(yàn)對(duì)比了GQA、MLA、FlashMLA、傳統(tǒng)MHA等多種注意力機(jī)制。

困惑度實(shí)驗(yàn)顯示，GTA在中大型模型上優(yōu)于GQA，說(shuō)明GTA可能更適合模型的進(jìn)一步擴(kuò)展；而GLA在多數(shù)場(chǎng)景下與MLA相當(dāng)，說(shuō)明GLA的設(shè)計(jì)是合理的，它能在并行計(jì)算和模型質(zhì)量之間找到一個(gè)較好的平衡點(diǎn)。

幾種方案在下游任務(wù)中（涵蓋典型常識(shí)推理、邏輯推理和知識(shí)問(wèn)答等場(chǎng)景）的整體表現(xiàn)差距不大。

但從變化趨勢(shì)來(lái)看（下圖為從中型到大型），GTA和GLA可以保持或提高從中型到XL尺寸的下游任務(wù)性能。

KV緩存方面，不犧牲模型質(zhì)量的前提下，GTA相比GQA減少約50%的KV緩存，驗(yàn)證了 “參數(shù)綁定+分組重用” 的有效性。

同時(shí)，針對(duì)查詢(xún)長(zhǎng)度為1的情況，MLA已接近計(jì)算瓶頸（達(dá)到610 TFLOPS/s ），而GLA尚未使計(jì)算資源飽和（360 TFLOPS/s ）。

且隨著序列長(zhǎng)度從1K增加到64K ，GLA的解碼速度比FlashMLA快2倍。

此外，在實(shí)時(shí)服務(wù)器性能測(cè)試中，對(duì)于64個(gè)并發(fā)請(qǐng)求的輸出吞吐量（越高越好），相同并行方案下GLA的表現(xiàn)均優(yōu)于MLA。

接下來(lái)，團(tuán)隊(duì)還在DeepSeek Coder V2 Base （236B）模型上，當(dāng)使用FP8精度時(shí)，對(duì)比了二者在不同預(yù)填充長(zhǎng)度和解碼長(zhǎng)度下的輸出吞吐量。

結(jié)果顯示，在預(yù)填充長(zhǎng)度為32K和64K時(shí)，GLA-8的輸出吞吐量明顯高于MLA。這表明在處理長(zhǎng)上下文時(shí)，GLA在吞吐量上優(yōu)于MLA。

在處理不均衡負(fù)載時(shí)，GLA-8同樣展現(xiàn)出更高的輸出吞吐量。這表明GLA在處理不同長(zhǎng)度的請(qǐng)求時(shí)，能夠更有效地利用資源，提高整體性能。

以上實(shí)驗(yàn)均驗(yàn)證了論文作者的說(shuō)法，「GTA和GLA」是「GQA和MLA」的有效替代品。

論文作者均來(lái)自普林斯頓大學(xué)

論文作者包括Tri Dao在內(nèi)一共三位，均來(lái)自普林斯頓大學(xué)。

Ted Zadouri，目前是普林斯頓大學(xué)博士生，研究方向?yàn)闄C(jī)器學(xué)習(xí)。

之前曾在英特爾有過(guò)兩段實(shí)習(xí)經(jīng)歷（研究深度學(xué)習(xí)），還短暫在AI創(chuàng)企Cohere擔(dān)任研究員。

Hubert Strauss，普林斯頓大學(xué)研究工程師，研究方向?yàn)闄C(jī)器學(xué)習(xí)和模型深度學(xué)習(xí)。

本科畢業(yè)于法國(guó)知名工程學(xué)校Arts et Métiers，之后在佐治亞理工學(xué)院取得運(yùn)籌學(xué)碩士學(xué)位。

畢業(yè)后曾有多段實(shí)習(xí)和工作經(jīng)歷，成為普林斯頓大學(xué)工程師之前曾在一家公司擔(dān)任機(jī)器學(xué)習(xí)工程師，負(fù)責(zé)模型訓(xùn)練和Transformer優(yōu)化。

Tri Dao，目前是普林斯頓大學(xué)計(jì)算機(jī)科學(xué)助理教授，還是生成式AI初創(chuàng)公司Together AI的首席科學(xué)家。

他因提出一系列優(yōu)化Transformer模型注意力機(jī)制的工作而聞名學(xué)界。

其中最有影響力的，是其作為作者之一提出了Mamba架構(gòu)，這一架構(gòu)在語(yǔ)言、音頻和基因組學(xué)等多種模態(tài)中都達(dá)到了SOTA性能。

尤其在語(yǔ)言建模方面，無(wú)論是預(yù)訓(xùn)練還是下游評(píng)估，Mamba-3B模型都優(yōu)于同等規(guī)模的Transformer模型，并能與兩倍于其規(guī)模的Transformer模型相媲美。

另外他還參與發(fā)表了FlashAttention1-3版本，F(xiàn)lashAttention被廣泛用于加速Transformers，已經(jīng)使注意力速度提高了4-8倍。

Anyway，回到這項(xiàng)研究，論文作者Ted Zadouri直言：

這只是邁向test-time推理“理想”架構(gòu)的第一步！

論文：
https://arxiv.org/abs/2505.21487
代碼：
https://github.com/Dao-AILab/grouped-latent-attention