曾撼動Transformer主導地位的Mamba作者之一Tri Dao，剛帶來新發表——
提出兩種專為推論「量身打造」的注意力機制。
在保持模型效能不變的情況下，將解碼速度和吞吐量最高提升2倍，大幅優化了模型的長上下文推論能力。
這項研究的三位作者均來自普林斯頓大學，論文主要有兩大貢獻：
其一，提出Grouped-Tied Attention（GTA），與已整合到LLaMA 3的注意力機制GQA品質相當，但KV快取用量減少約50%。
其二，提出Grouped Latent Attention（GLA），與DeepSeek所使用的注意力機制MLA品質匹配，但解碼速度更快，某些情況下比FlashMLA快2倍。
按照作者之一Ted Zadouri的總結：
GTA是GQA的有效替代品，而GLA是MLA的實用替代品。
簡而言之，透過優化注意力機制的記憶體使用和計算邏輯，在不犧牲模型生成品質的前提下，可顯著提升大型語言模型的推論效率和硬體資源利用率，尤其在長上下文情境中，優勢更為突出。
相關論文公布後，許多研究人員也紛紛前來祝賀~
那麼，這項研究具體講了些什麼？
引入推論感知注意力機制
概括來說，論文核心引入了推論感知注意力機制，即針對模型推論階段的記憶體冗餘、計算低效、長上下文瓶頸等問題，重新設計注意力機制。
據Tri Dao介紹，這項研究的起點始於一個想法：
在推論驅動AI發展的時代，「理想」架構應該是什麼樣子？
尤其在涉及長上下文推論時，當前的大型語言模型（LLM）面臨記憶體存取瓶頸和平行性限制兩大難題。
也就是說，模型生成文字時，每次都要從記憶體裡提取大量「歷史紀錄」，不僅導致每個字生成變慢，而且只能按順序生成，沒法讓多個晶片同時運作。
對此，團隊打算從兩個方向重新設計注意力機制：
更高的硬體效率：透過增加「每位元組記憶體載入的計算量」（算術強度），減少對記憶體頻寬的依賴；
保持平行可擴展性：在不犧牲模型平行訓練/推論能力的前提下優化解碼速度。
而最終提出的GTA和GLA，在減少KV快取用量的同時，模型品質保持與現有方案相當，且解碼速度顯著提升。
這裡提到的「現有方案」，主要指早已聞名學術界的兩種方法：
一是分組查詢注意力（GQA）機制，它透過分組共享KV快取減少記憶體佔用，在視覺Transformer（ViT）等任務中表現良好，適用於大規模資料處理，目前已應用於Llama 3等開源模型。
二是多頭潛在注意力（MLA）機制，最早可追溯到《Attention Is All You Need》這篇論文，後被DeepSeek再次引發熱潮。它關注的是如何在不同層之間融合注意力資訊，能減少每一層的冗餘計算。
不過，由於GQA仍需為每組查詢頭儲存獨立KV、MLA平行優化不足，故仍需進一步改進。
下面分別展開團隊提出的新方法GTA和GLA。
分組綁定注意力機制GTA
GTA的核心設計思路是：將不同查詢頭的鍵（Key）與值（Value）狀態進行組合與重複使用，減少記憶體傳輸次數。
具體而言（右圖），它將多頭注意力的頭分為若干組（Group），每組內的頭共享相同的Key和Value參數。計算時，同一組內的頭使用相同的KV快取，僅查詢（Query）參數獨立。
相比之下，中間傳統的多頭注意力機制（MHA）每個查詢頭都有獨立的鍵和值，由於沒有共享，導致它需要更多的記憶體來儲存所有的鍵和值。
再對比GQA來看（左圖），GQA分組共享KV但每組仍獨立儲存，而GTA透過參數綁定實現了更徹底的KV重複利用。
分組潛在注意力機制GLA
而GLA的設計則採用了雙層結構：
潛在層（Latent Layer）：引入固定數量的潛在Tokens，作為全域上下文的壓縮表示，取代部分原始Token的KV快取；
分組頭機制：將查詢頭分組，每組頭共享潛在Token的KV，同時保留與原始Token的互動。
在解碼過程中，對比MLA（左圖），GLA透過共享聯合潛在表示減少了每個設備需要載入的KV快取量，從而減少了記憶體存取量。
並且由於每個設備上的KV快取量減少了，更多的請求也可以同時處理。
「GQA與MLA」的有效替代品
那麼，GTA和GLA的效果究竟如何呢？
團隊在四種規模的模型上進行了實驗，包括小型（183M）、中型（433M）、大型（876M）和XL（1471M）。這些模型基於FineWeb-Edu-100B資料集訓練，採用GPT-3架構和Llama 3分詞器。
測試的指標主要分為兩大類：
品質指標：困惑度（Perplexity）、下游任務準確率（Winogrande、SciQ等7個基準）；
效率指標：每個Token的解碼延遲、吞吐量、KV快取佔用量。
實驗對比了GQA、MLA、FlashMLA、傳統MHA等多種注意力機制。
困惑度實驗顯示，GTA在中大型模型上優於GQA，說明GTA可能更適合模型的進一步擴展；而GLA在多數情況下與MLA相當，說明GLA的設計是合理的，它能在平行計算和模型品質之間找到一個較好的平衡點。
幾種方案在下游任務中（涵蓋典型的常識推論、邏輯推論與知識問答等情境）的整體表現差異不大。
但從變化趨勢來看（下圖為從中型到大型），GTA和GLA可以保持或提高從中型到XL尺寸的下游任務效能。
KV快取方面，不犧牲模型品質的前提下，GTA相比GQA減少約50%的KV快取，驗證了「參數綁定+分組重複使用」的有效性。
同時，針對查詢長度為1的情況，MLA已接近計算瓶頸（達到610 TFLOPS/s），而GLA尚未使計算資源飽和（360 TFLOPS/s）。
且隨著序列長度從1K增加到64K，GLA的解碼速度比FlashMLA快2倍。
此外，在即時伺服器效能測試中，對於64個併發請求的輸出吞吐量（越高越好），相同平行方案下GLA的表現均優於MLA。
接下來，團隊還在DeepSeek Coder V2 Base（236B）模型上，當使用FP8精度時，對比了二者在不同預填充長度和解碼長度下的輸出吞吐量。
結果顯示，在預填充長度為32K和64K時，GLA-8的輸出吞吐量明顯高於MLA。這表明在處理長上下文時，GLA在吞吐量上優於MLA。
在處理不平衡負載時，GLA-8同樣展現出更高的輸出吞吐量。這表明GLA在處理不同長度的請求時，能夠更有效地利用資源，提升整體效能。
以上實驗均驗證了論文作者的說法，「GTA與GLA」是「GQA與MLA」的有效替代品。
論文作者均來自普林斯頓大學
論文作者包括Tri Dao在內一共三位，均來自普林斯頓大學。
Ted Zadouri，目前是普林斯頓大學博士生，研究方向為機器學習。
之前曾在英特爾有過兩段實習經歷（研究深度學習），還短暫在AI新創公司Cohere擔任研究員。
Hubert Strauss，普林斯頓大學研究工程師，研究方向為機器學習和模型深度學習。
大學畢業於法國知名工程學校Arts et Métiers，之後在佐治亞理工學院取得運籌學碩士學位。
畢業後曾有多段實習和工作經歷，成為普林斯頓大學工程師之前曾在一家公司擔任機器學習工程師，負責模型訓練與Transformer最佳化。
Tri Dao，目前是普林斯頓大學電腦科學助理教授，還是生成式AI新創公司Together AI的首席科學家。
他因提出一系列優化Transformer模型注意力機制的工作而聞名學界。
其中最具影響力的，是其作為作者之一提出了Mamba架構，這一架構在語言、音訊和基因組學等多種模態中都達到了SOTA效能。
尤其在語言建模方面，無論是預訓練還是下游評估，Mamba-3B模型都優於同等規模的Transformer模型，並能與兩倍於其規模的Transformer模型相媲美。
另外他還參與發表了FlashAttention 1-3 版，FlashAttention被廣泛用於加速Transformer，已經使注意力速度提高了4至8倍。
總之，回到這項研究，論文作者Ted Zadouri直言不諱：
這只是邁向測試時推論「理想」架構的第一步！
論文：https://arxiv.org/abs/2505.21487
程式碼：https://github.com/Dao-AILab/grouped-latent-attention
參考連結：
[1]https://x.com/tri_dao/status/1928170648863473892
[2]https://x.com/gm8xx8/status/1927572103806554262
[3]https://x.com/tedzadouri/status/1928167296821854363
— 完 —