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

提出兩種專為推理“量身定制”的注意力機制。

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

這項研究的三位作者均來自普林斯頓大學，論文主要有兩大貢獻：

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

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

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

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

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

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

那么，這項研究具體講了些啥？

引入推理感知注意力機制

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

據(jù)Tri Dao介紹，這項研究的起點始于一個想法：

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

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

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

對此，團隊打算從兩個方向重新設計注意力機制：

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

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

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

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

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

不過，由于GQA仍需為每組查詢頭存儲獨立KV、MLA并行優(yōu)化不足，故仍需進一步改進。

下面分別展開團隊提出的新方法GTA和GLA。

分組綁定注意力機制GTA

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

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

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

再對比GQA來看（左圖），GQA分組共享KV但每組仍獨立存儲，而GTA通過參數(shù)綁定實現(xiàn)了更徹底的KV重復利用。

分組潛在注意力機制GLA

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

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

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

并且由于每個設備上的KV緩存量減少了，更多的請求也可以同時處理。

「GQA和MLA」的有效替代品

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

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

測試的指標主要分為兩大類：

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

實驗對比了GQA、MLA、FlashMLA、傳統(tǒng)MHA等多種注意力機制。

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

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

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

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

同時，針對查詢長度為1的情況，MLA已接近計算瓶頸（達到610 TFLOPS/s ），而GLA尚未使計算資源飽和（360 TFLOPS/s ）。

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

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

接下來，團隊還在DeepSeek Coder V2 Base （236B）模型上，當使用FP8精度時，對比了二者在不同預填充長度和解碼長度下的輸出吞吐量。

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

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

以上實驗均驗證了論文作者的說法，「GTA和GLA」是「GQA和MLA」的有效替代品。

論文作者均來自普林斯頓大學

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

Ted Zadouri，目前是普林斯頓大學博士生，研究方向為機器學習。

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

Hubert Strauss，普林斯頓大學研究工程師，研究方向為機器學習和模型深度學習。

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

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

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

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

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

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

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

Anyway，回到這項研究，論文作者Ted Zadouri直言：

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

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