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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)驗(yàn)證：GTA 的卓越性能與效率

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

模型有效性驗(yàn)證

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

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

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

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

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

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

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

盡管 GTA-1B 在設(shè)計(jì)上采用了更小的緩存尺寸，但其損失軌跡卻與 GQA-1B 高度匹配，這一事實(shí)有力地證明了 GTA 內(nèi)存高效架構(gòu)的有效性，即在減少資源消耗的同時(shí)不犧牲模型學(xué)習(xí)能力。在多項(xiàng)嚴(yán)苛的基準(zhǔn)測(cè)試中，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)大的競(jìng)爭(zhēng)力。這些詳盡的實(shí)驗(yàn)結(jié)果均已在下方的表格中呈現(xiàn)：

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

效率評(píng)估

理論效率分析

從理論層面分析，GTA 在計(jì)算復(fù)雜度和內(nèi)存使用方面均實(shí)現(xiàn)了顯著的效率提升。其 KV 緩存尺寸從 MHA 的 大幅減少至 GTA 的，其中 且 ，對(duì)于參數(shù)量龐大的大型模型而言，這意味著一個(gè)極其顯著的縮減因子，將有效緩解內(nèi)存壓力。同時(shí)，注意力計(jì)算量也從 MHA 的顯著降低到 GTA 的 ，這直接轉(zhuǎn)化為推理速度的顯著提升。這些嚴(yán)謹(jǐn)?shù)睦碚摲治龊蛯?duì)比數(shù)據(jù)均已在下方的表格中詳細(xì)列出：

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

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

實(shí)際部署表現(xiàn)

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

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

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

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

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

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

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