GTA 工作由中國科學(xué)院自動化研究所、倫敦大學(xué)學(xué)院及香港科技大學(xué)（廣州）聯(lián)合研發(fā)，提出了一種高效的大模型框架，顯著提升模型性能與計算效率。一作為自動化所的孫羅洋博士生，研究方向為：大模型高效計算與優(yōu)化，通訊作者為香港科技大學(xué)（廣州）的鄧程博士、自動化所張海峰教授及倫敦大學(xué)學(xué)院汪軍教授。該成果為大模型的優(yōu)化部署提供了創(chuàng)新解決方案。

Grouped-head latent Attention (GTA) 震撼登場！這項創(chuàng)新機(jī)制通過共享注意力矩陣和壓縮潛在值表示，將計算量削減 62.5%，KV 緩存縮減 70%，prefill 和 decode 速度提升 2 倍。無論是處理海量數(shù)據(jù)構(gòu)成的長序列任務(wù)，還是在計算資源極為有限的邊緣設(shè)備上運(yùn)行，GTA 都展現(xiàn)出無與倫比的效率和卓越的性能，無疑將成為大型語言模型優(yōu)化領(lǐng)域的新標(biāo)桿。

大型語言模型面臨的效率困局

近年來，Transformer 架構(gòu)的橫空出世極大地推動了自然語言處理領(lǐng)域的飛速發(fā)展，使得大型語言模型在對話生成、文本摘要、機(jī)器翻譯以及復(fù)雜推理等多個前沿領(lǐng)域?qū)覄?chuàng)佳績，展現(xiàn)出令人驚嘆的能力。然而，隨著模型參數(shù)量從數(shù)十億激增至上千億，傳統(tǒng)多頭注意力機(jī)制 (Multi-Head Attention, MHA) 所固有的弊端也日益凸顯，成為制約其廣泛應(yīng)用和進(jìn)一步發(fā)展的瓶頸。

首當(dāng)其沖的是計算冗余問題。在 MHA (多頭注意力) 架構(gòu)中，每個注意力頭都像一個獨(dú)立的 “工作單元”，各自獨(dú)立地計算查詢 (Query)、鍵 (Key) 和值 (Value) 向量，這導(dǎo)致了大量的重復(fù)計算。特別是在處理長序列任務(wù)時，浮點運(yùn)算次數(shù) (FLOPs) 會呈平方級增長，嚴(yán)重拖慢了模型的處理效率，使得原本復(fù)雜的任務(wù)變得更加耗時。

其次是內(nèi)存瓶頸。每個注意力頭都需要完整存儲其對應(yīng)的鍵值對 (KV) 緩存，這使得內(nèi)存需求隨序列長度和注意力頭數(shù)量的增加而快速膨脹。例如，在處理長序列時，KV 緩存的規(guī)模可以輕松突破數(shù) GB，如此龐大的內(nèi)存占用極大地限制了大型模型在智能手機(jī)、物聯(lián)網(wǎng)設(shè)備等邊緣設(shè)備上的部署能力，使其難以真正走進(jìn)千家萬戶。

最后是推理延遲問題。高昂的計算和內(nèi)存需求直接導(dǎo)致了推理速度的顯著下降，使得像語音助手實時響應(yīng)、在線翻譯無縫切換等對延遲敏感的實時應(yīng)用難以提供流暢的用戶體驗。盡管業(yè)界的研究者們曾嘗試通過 Multi-Query Attention (MQA) 和 Grouped-Query Attention (GQA) 等方法來優(yōu)化效率，但這些方案往往需要在性能和資源消耗之間做出艱難的權(quán)衡，難以找到理想的平衡點。面對這一系列嚴(yán)峻的挑戰(zhàn)，研究團(tuán)隊經(jīng)過不懈努力，最終推出了 Grouped-head latent Attention (GTA)，以其顛覆性的設(shè)計，重新定義了注意力機(jī)制的效率極限，為大型語言模型的未來發(fā)展開辟了全新的道路。

• 論文標(biāo)題: GTA: Grouped-head latenT Attention
• 論文鏈接: https://arxiv.org/abs/2506.17286
• 項目鏈接: https://github.com/plm-team/GTA

GTA 的核心創(chuàng)新機(jī)制

GTA 的卓越成功源于其兩大核心技術(shù)突破，它們精妙地協(xié)同作用，使得大型語言模型即使在嚴(yán)苛的資源受限場景下，也能展現(xiàn)出前所未有的高效運(yùn)行能力。

分組共享注意力矩陣機(jī)制

在傳統(tǒng)的 MHA 架構(gòu)中，每個注意力頭都被視為一個獨(dú)立的 “獨(dú)行俠”，各自計算并維護(hù)自己的注意力分?jǐn)?shù)。這種分散式的計算模式雖然賦予了模型捕捉多種復(fù)雜依賴關(guān)系的能力，但同時也帶來了顯著的計算冗余。以一個包含 16 個注意力頭的 MHA 為例，當(dāng)每個頭獨(dú)立處理輸入時，會生成 16 組獨(dú)立的注意力矩陣，這使得總體的計算開銷隨注意力頭數(shù)量的增加而呈線性增長，效率低下。

與此形成鮮明對比的是，GTA 采用了全新的 “團(tuán)隊協(xié)作” 策略。該機(jī)制將注意力頭巧妙地劃分為若干個邏輯組，例如，每 4 個注意力頭可以組成一個小組，而這個小組內(nèi)部的成員將共享一張統(tǒng)一的注意力矩陣。這種創(chuàng)新的共享設(shè)計意味著，我們僅需對注意力分?jǐn)?shù)進(jìn)行一次計算，然后便可將其高效地分配給組內(nèi)所有注意力頭使用，從而大幅度減少了浮點運(yùn)算次數(shù) (FLOPs)。

實驗數(shù)據(jù)有力地證明，這一精巧的設(shè)計能夠?qū)⒖傆嬎懔肯鳒p，為處理超長序列任務(wù)帶來了顯著的推理加速效果。這恰如一位經(jīng)驗豐富的主廚，統(tǒng)一備齊所有食材，再分發(fā)給不同的助手進(jìn)行精細(xì)加工，既極大地節(jié)省了寶貴的時間，又確保了最終產(chǎn)出的高質(zhì)量和一致性。

壓縮潛在值表示技術(shù)

MHA 架構(gòu)的另一個關(guān)鍵痛點在于其 KV 緩存對內(nèi)存的巨大占用。由于每個注意力頭的值 (Value) 向量都需要被完整地存儲下來，導(dǎo)致模型的內(nèi)存需求會隨著輸入序列長度和注意力頭數(shù)量的增加而快速膨脹，成為部署大型模型的嚴(yán)重障礙。GTA 通過其獨(dú)創(chuàng)的 “壓縮 + 解碼” 巧妙設(shè)計，徹底解決了這一難題。

這項技術(shù)首先將所有注意力頭的值向量高效地壓縮為一個低維度的潛在表示 (Latent Representation)，從而極大地減少了所需的存儲空間。隨后，通過一個輕量級且高效的 WaLU (Weighted additive Linear Unit) 非線性解碼器，模型能夠根據(jù)每一組注意力頭的具體需求，從這個緊湊的潛在表示中動態(tài)地、定制化地生成所需的完整值向量。

這種創(chuàng)新的方法不僅顯著節(jié)省了寶貴的內(nèi)存資源，同時還巧妙地保留了每個注意力頭所特有的獨(dú)特表達(dá)能力，避免了信息損失。實驗結(jié)果令人鼓舞，GTA 的 KV 緩存規(guī)模成功縮減了 70%，這一突破性進(jìn)展為大型語言模型在性能受限的邊緣設(shè)備上的廣泛部署鋪平了道路，使其能夠更普惠地服務(wù)于各類應(yīng)用場景。

實驗驗證：GTA 的卓越性能與效率

研究團(tuán)隊通過一系列嚴(yán)謹(jǐn)而全面的實驗，對 Grouped-head Latent Attention (GTA) 在不同模型規(guī)模、輸入序列長度以及多樣化硬件平臺上的性能和效率進(jìn)行了深入評估。實驗結(jié)果令人信服地表明，GTA 在大幅度提升計算效率和內(nèi)存利用率的同時，不僅成功保持了，甚至在某些關(guān)鍵指標(biāo)上超越了現(xiàn)有主流注意力機(jī)制的模型性能，展現(xiàn)出其強(qiáng)大的實用價值和廣闊的應(yīng)用前景。

模型有效性驗證

為了確保實驗結(jié)果的公平性和準(zhǔn)確性，研究團(tuán)隊在實驗設(shè)計中采取了嚴(yán)格的控制變量法：所有非注意力相關(guān)的模型參數(shù)（例如隱藏層維度、多層感知機(jī) MLP 的大小等）都被固定不變，從而確保模型參數(shù)量的任何變化都僅僅來源于注意力機(jī)制自身的創(chuàng)新設(shè)計。

160M 參數(shù)模型表現(xiàn)

在針對 160M 參數(shù)規(guī)模模型的測試中，無論輸入序列長度是 2048 還是 4096 個 token，GTA 都持續(xù)展現(xiàn)出卓越的性能優(yōu)勢。具體而言，采用 GTA2 配置的模型在 2048 token 序列長度下，成功實現(xiàn)了比傳統(tǒng) MHA、GQA 和 MLA 更低的評估損失，并獲得了更優(yōu)異的 Wikitext 困惑度（PPL）表現(xiàn)。此外，GTA1 配置的模型在多項下游任務(wù)中取得了更高的平均準(zhǔn)確率，彰顯了其在實際應(yīng)用中的有效性。尤為值得強(qiáng)調(diào)的是，GTA 在內(nèi)存效率方面表現(xiàn)出類拔萃，其每層所需的 KV 緩存大小僅為 MHA 的 12.5%（具體數(shù)據(jù)為 192 維度對比 MHA 的 1536 維度），這一顯著的縮減充分突顯了 GTA 在內(nèi)存優(yōu)化方面的強(qiáng)大能力。實驗結(jié)果詳細(xì)呈現(xiàn)在下方的表格中：

500M 參數(shù)模型表現(xiàn)

將模型規(guī)模擴(kuò)展至 500M 參數(shù)時，GTA 依然保持了其在性能上的領(lǐng)先地位。在 2048 token 序列長度的測試中，GTA 不僅實現(xiàn)了更低的評估損失，還在下游任務(wù)中取得了更高的平均準(zhǔn)確率，同時其 Wikitext 困惑度與 MHA 和 GQA 等主流模型保持在同等甚至更優(yōu)的水平。GTA 持續(xù)展現(xiàn)出其獨(dú)有的內(nèi)存優(yōu)勢，其 KV 緩存大小僅為 MHA 的 12.5%（具體為 320 維度對比 MHA 的 2560 維度），即使在采用更小緩存（例如 192 維度，僅為 MHA 的 7.5%）的情況下，GTA 也能獲得可比擬的性能表現(xiàn)，充分印證了其在內(nèi)存效率與性能之間取得的完美平衡。在處理 4096 token 長序列的任務(wù)中，GTA 不僅能夠與 MHA 的評估損失持平，更在 Wikitext 困惑度和下游任務(wù)的平均準(zhǔn)確率上提供了更優(yōu)異的表現(xiàn)。這些詳盡的實驗數(shù)據(jù)均已在下方的表格中列出：

1B 參數(shù)語言模型擴(kuò)展性

為了進(jìn)一步驗證 GTA 在大規(guī)模模型上的卓越擴(kuò)展能力和穩(wěn)定性，研究團(tuán)隊特意訓(xùn)練了 1B 參數(shù)級別的 GTA-1B 和 GQA-1B 模型。下圖清晰地展示了 GTA-1B 和 GQA-1B 在長達(dá) 50,000 訓(xùn)練步中的損失曲線和梯度范數(shù)曲線，從中可以觀察到兩者均展現(xiàn)出令人滿意的穩(wěn)定收斂趨勢。

盡管 GTA-1B 在設(shè)計上采用了更小的緩存尺寸，但其損失軌跡卻與 GQA-1B 高度匹配，這一事實有力地證明了 GTA 內(nèi)存高效架構(gòu)的有效性，即在減少資源消耗的同時不犧牲模型學(xué)習(xí)能力。在多項嚴(yán)苛的基準(zhǔn)測試中，GTA-1B（包括經(jīng)過 SFT 微調(diào)的版本）均展現(xiàn)出與 GQA-1B 相當(dāng)甚至更為優(yōu)異的性能，尤其在平均準(zhǔn)確率上取得了顯著提升。這充分表明，GTA 即使在資源受限的環(huán)境下，也能通過微調(diào)有效泛化到各種復(fù)雜任務(wù)，保持強(qiáng)大的競爭力。這些詳盡的實驗結(jié)果均已在下方的表格中呈現(xiàn)：

綜合來看，GTA-1B 無論是在基礎(chǔ)模型狀態(tài)還是經(jīng)過微調(diào)后，都成功實現(xiàn)了與 GQA-1B 相當(dāng)?shù)淖吭叫阅堋Ｅc此同時，其 KV 緩存尺寸僅為 GQA-1B 的 30%，而自注意力計算成本更是低至 37.5%。這些令人矚目的數(shù)據(jù)有力地強(qiáng)調(diào)了內(nèi)存和計算高效架構(gòu)在大型語言模型擴(kuò)展應(yīng)用方面的巨大潛力，預(yù)示著未來 AI 發(fā)展將更加注重可持續(xù)性和資源效率。

效率評估

理論效率分析

通過 LLM-Viewer 進(jìn)行經(jīng)驗基準(zhǔn)測試

為了將理論優(yōu)勢轉(zhuǎn)化為可量化的實際性能，研究團(tuán)隊利用 LLM-Viewer 框架，在配備 NVIDIA H100 80GB GPU 的高性能計算平臺上，對 GTA-1B 和 GQA-1B 進(jìn)行了全面的經(jīng)驗基準(zhǔn)測試。下圖清晰地展示了在不同配置下，兩種模型的預(yù)填充和解碼時間對比。從中可以明顯看出，GTA-1B 在計算密集型的預(yù)填充階段和 I/O 密集型的解碼階段都持續(xù)地優(yōu)于 GQA-1B，充分展現(xiàn)了其卓越的延遲特性和更高的運(yùn)行效率。

實際部署表現(xiàn)

為了更真實地評估 GTA-1B 在實際應(yīng)用場景中的性能，研究團(tuán)隊利用 transformers 庫，在多種異構(gòu)硬件平臺（包括服務(wù)器級的 NVIDIA H100、NVIDIA A800，消費(fèi)級的 RTX 3060，以及邊緣設(shè)備如 Apple M2 和 BCM2712）上進(jìn)行了深入的推理實驗。

上圖直觀地展示了在不同配置下，GTA-1B 與 GQA-1B 的預(yù)填充和解碼時間對比。GTA-1B（藍(lán)色實線）在所有測試平臺上都持續(xù)展現(xiàn)出優(yōu)于 GQA-1B（橙色虛線）的預(yù)填充時間，尤其是在處理 2k token 等更長輸入序列時，性能差距更為顯著，體現(xiàn)了其在處理長文本時的強(qiáng)大優(yōu)勢。在解碼階段，GTA-1B 同樣表現(xiàn)出卓越的性能，特別是在擴(kuò)展生成長度時，這種優(yōu)勢在所有硬件類型上都保持一致，充分突顯了其設(shè)計的魯棒性。

上圖進(jìn)一步展示了在啟用緩存卸載功能時的性能表現(xiàn)。在 NVIDIA H100 平臺上，GTA-1B 在處理更長輸入序列時依然保持了其預(yù)填充優(yōu)勢，并且在解碼時間上實現(xiàn)了比 GQA-1B 更大的改進(jìn)。這種在所有平臺上的持續(xù)趨勢，有力地突顯了 GTA-1B 在 I/O 密集型場景中的高效性，這類場景中緩存卸載需要 GPU 和 CPU 內(nèi)存之間頻繁的數(shù)據(jù)傳輸，而 GTA-1B 在這種復(fù)雜環(huán)境下依然表現(xiàn)出色。

綜上所述，GTA-1B 在各種硬件平臺下，無論是在預(yù)填充還是解碼時間上，均全面超越了 GQA-1B，并在處理更長輸入序列時展現(xiàn)出顯著的性能優(yōu)勢。它不僅在標(biāo)準(zhǔn)推理設(shè)置中表現(xiàn)出色，在啟用緩存卸載的 I/O 密集型條件下也同樣杰出，充分展現(xiàn)了其在不同硬件能力和批處理大小下的強(qiáng)大多功能性。這種卓越的適應(yīng)性使得 GTA-1B 成為服務(wù)器級和消費(fèi)級部署的理想解決方案，通過顯著降低計算復(fù)雜度和內(nèi)存需求，極大地提升了大型語言模型中注意力機(jī)制的整體效率。

技術(shù)局限與未來方向

盡管 Grouped-head latent Attention (GTA) 在效率和性能方面取得了令人矚目的突破，但作為一項新興技術(shù)，仍有一些關(guān)鍵的技術(shù)挑戰(zhàn)需要我們持續(xù)關(guān)注和深入探索。首先，非線性解碼器在進(jìn)行模型壓縮的過程中，可能會引入微小的近似誤差，這需要未來的研究在架構(gòu)設(shè)計和訓(xùn)練策略上進(jìn)一步優(yōu)化，確保模型輸出的準(zhǔn)確性。其次，當(dāng)前 GTA 的研究和驗證主要集中在自然語言處理任務(wù)上，其在計算機(jī)視覺或多模態(tài)任務(wù)中的適用性和性能表現(xiàn)，還需要進(jìn)行更廣泛和深入的探索與驗證。

盡管存在這些局限，研究團(tuán)隊已經(jīng)為 GTA 的未來發(fā)展制定了清晰且富有前景的后續(xù)研究方向。他們計劃持續(xù)改進(jìn)非線性解碼器的架構(gòu)設(shè)計，以期在保證高效解碼的同時，進(jìn)一步減少信息損失，提升模型性能上限。此外，研究團(tuán)隊還雄心勃勃地計劃將 GTA 應(yīng)用于更大規(guī)模的模型，以驗證其在超大規(guī)模場景下的可擴(kuò)展性和效率優(yōu)勢，推動大型語言模型向更廣闊的應(yīng)用領(lǐng)域邁進(jìn)。