ICML 2025 | 全局池化+局部保留，CCA-Attention為LLM長文本建模帶來突破性進(jìn)展

2025-06-09 09:23:05

作者提出了一種即插即用的高效長文本上下文建模方法——關(guān)鍵上下文感知注意力機(jī)制（CCA-Attention）。

琶洲實驗室、華南理工大學(xué)聯(lián)合推出關(guān)鍵上下文感知注意力機(jī)制（CCA-Attention），實現(xiàn)超長文本的高效上下文建模。在 128K 超長序列上下文建模任務(wù)中，CCA-Attention 的推理速度是標(biāo)準(zhǔn)自注意力機(jī)制的 7.9 倍，同時鍵值緩存（KV Cache）顯存占用減少 93%，性能全面優(yōu)于現(xiàn)有高效注意力方法。

論文標(biāo)題：Core Context Aware Transformers for Long Context Language Modeling
論文鏈接：https://arxiv.org/pdf/2412.12465
代碼鏈接：https://github.com/chenyaofo/CCA-Attention
發(fā)布時間：2024年12月17日

該成果已被 ICML 2025 接收，最早于 2024 年 12 月 17 日提交至 ArXiv，早于 DeepSeek NSA 和 Kimi MoBA 公開。CCA-Attention 不僅速度快、資源占用低，更在上下文建模的精準(zhǔn)度和效率上樹立了新標(biāo)桿，為長文本處理注入全新動力。

引言

近期研究 [1, 2, 3] 發(fā)現(xiàn)，LLMs 中的大多數(shù)層的注意力權(quán)重主要集中在少數(shù) token 上，表現(xiàn)出顯著的稀疏性（見圖 1）。這一發(fā)現(xiàn)啟示我們可以借助這種稀疏特性，降低注意力機(jī)制的計算復(fù)雜度。

圖 1： LLaMA2-7B 模型中注意力權(quán)重的可視化，陰影越深表示注意力權(quán)重越高。最后一個 token 僅對上下文少數(shù)幾個 token 有著較高的注意力權(quán)重，即注意力權(quán)重具有顯著的稀疏性。

現(xiàn)有稀疏注意力方法 [5, 6, 7] 通常通過預(yù)定義的稀疏模式來降低計算成本。然而，在問答任務(wù)中，關(guān)鍵信息可能分布在上下文的不同位置，模型需要能夠訪問任意位置的信息，作者稱這一特性為「可達(dá)性」。已有方法往往忽視了保持 token 之間可達(dá)性的重要性，可能導(dǎo)致信息傳遞受限，從而影響模型在長序列和復(fù)雜任務(wù)中的表現(xiàn)。

為解決這一問題，作者提出了一種即插即用的高效長文本上下文建模方法——關(guān)鍵上下文感知注意力機(jī)制（CCA-Attention），其特點(diǎn)如下：

高效長文本建模： 通過全局池化注意力與局部保留注意力的協(xié)同設(shè)計，在顯著降低計算量的同時保持對長距離依賴的建模能力。
線性計算復(fù)雜度：通過引入 core token 聚焦關(guān)鍵上下文，大幅提高計算效率。
可即插即用集成：無需修改模型結(jié)構(gòu)和從頭訓(xùn)練，可以輕松集成到預(yù)訓(xùn)練的 LLM 中，僅需少量微調(diào)即可實現(xiàn)性能優(yōu)化。

對比 DeepSeek 發(fā)布的 NSA [8] 需引入額外的壓縮模塊并從頭訓(xùn)練 LLMs，CCA-Attention 無需引入額外參數(shù)和修改模型結(jié)構(gòu)，可以無縫替換現(xiàn)有 LLMs 中的標(biāo)準(zhǔn)自注意力模塊。對比月之暗面發(fā)布的 MoBA [9] 通過門控機(jī)制丟棄不相關(guān)塊，CCA-Attention 通過動態(tài)聚合關(guān)鍵上下文為核心 token 的方式，在降低計算量的同時，確保所有 token 的信息交互，保留了完整的全局建模能力。

CCA-Attention：革新性的解決方案

圖 2： CCA-Attention 示意圖

全局感知池化：降低計算維度的智慧之舉

標(biāo)準(zhǔn)自注意力計算量隨序列長度呈平方級增長，長序列處理計算開銷極大。大量研究發(fā)現(xiàn)注意力權(quán)重的分布并不均勻，絕大部分注意力權(quán)重被分配給了少數(shù)重要 token，其余部分貢獻(xiàn)有限，屬于冗余上下文。

受此啟發(fā)，作者提出全局感知池化模塊。具體而言，將輸入序列，分成互不重疊的個組，g 為分組大小。對于第 i 組，使用該組最后一個 token 的 query 向量與組內(nèi)所有 token 的 key 向量計算重要性分?jǐn)?shù)，并獲得該組核心：

其中，是第 i 組的最后一個 token 對應(yīng)的 query 向量，是第 i 組的 key 矩陣，和是可學(xué)習(xí)的參數(shù)。將各組 core token 拼接起來得到 core token 序列。

為減少冗余，作者使用 core token 序列代替原始 token 進(jìn)行注意力計算，將維度從降至，從而降低了計算和存儲復(fù)雜度。通過 core token 序列計算得到的鍵值矩陣表示為：

其中和是可學(xué)習(xí)參數(shù)。

局部保留模塊：捕捉局部依賴的關(guān)鍵

盡管全局感知池化模塊能有效捕捉長距離依賴，但由于其壓縮特性，可能會忽略細(xì)粒度的局部上下文，而這些局部語義對于語言建模同樣至關(guān)重要。為此，作者進(jìn)一步提出局部保留模塊（Locality-preserving Module），為全局模塊提供有效互補(bǔ)信息。

具體來說，該模塊會確保每個 token 都能至少關(guān)注前面 w 個原始 token，以此來捕捉局部上下文信息，保留連續(xù)性語義信息：

為了應(yīng)對生成過程中標(biāo)記數(shù)量難以維持為組大小 g 的整數(shù)倍的問題，作者將局部窗口大小設(shè)置為，確保注意力窗口與組大小對齊，避免信息遺漏；是原始 token 序列經(jīng)過線性變換后的鍵值矩陣。

局部保留模塊與全局池化模塊共享線性變換參數(shù)，不會引入額外參數(shù)開銷。在實際推理中，局部模塊提供精細(xì)語義支持，彌補(bǔ)全局壓縮帶來的信息損失，共同構(gòu)成完整的上下文建模體系。

全局-局部模塊可微融合：打造全面可達(dá)性的橋梁

全局感知池化模塊和局部保留模塊在計算時都只涉及部分 token，導(dǎo)致注意力的可達(dá)性有限。為解決這個問題，作者采用全局-局部模塊可微融合策略。具體而言，該策略將兩種注意力模塊中的鍵值矩陣進(jìn)行組合，形成統(tǒng)一的鍵矩陣和值矩陣。由此，CCA-Attention 的最終輸出表示為：

其中，每個位置的輸出計算表達(dá)式如下：

基于 Triton 的底層加速：提升效率的強(qiáng)大動力

為了在訓(xùn)練、預(yù)填充、解碼期間實現(xiàn) FlashAttention 級別的加速，作者基于 Triton 實現(xiàn)了硬件對齊的 CCA-Attention 內(nèi)核。作者借鑒 FlashAttention 的設(shè)計思路，利用 Triton 進(jìn)行底層算子融合，將全局池化注意力和局部保留注意力整合為一個獨(dú)立且緩存友好的算子，有效消除冗余計算，并原生支持 KV 緩存技術(shù)，進(jìn)一步提升訓(xùn)練、預(yù)填充、解碼階段的計算效率。相比標(biāo)準(zhǔn)自注意力機(jī)制，CCA-Attention 在計算復(fù)雜度和 KV 緩存內(nèi)存占用方面具有顯著優(yōu)勢，從而在整體上實現(xiàn)了更快的運(yùn)行速度與更高的內(nèi)存利用效率。

實驗結(jié)果

實驗設(shè)置

作者將 CCA-Attention 應(yīng)用于 LLaMA2-7B-32K 和 LLaMA2-7B-80K 模型，并在 SlimPajama 數(shù)據(jù)集上微調(diào) 1,000 步。對比方法包括 StreamingLLM、LM-Infinite 和 MInference 等高效注意力方法。評估指標(biāo)涵蓋 LongBench 基準(zhǔn)測試和多文檔問答準(zhǔn)確匹配得分（EM Score）等，全面衡量模型在長文本任務(wù)中的性能表現(xiàn)。

長序列語言建模

在 LongBench-E 基準(zhǔn)測試中，CCA-LLM 取得了最高的平均得分。以 LLaMA2-7B-32K 模型為例，其得分顯著優(yōu)于 LM-Infinite 和 MInference；在 LLaMA2-7B-80K 模型上，CCA-Attention 依然表現(xiàn)出色，平均分?jǐn)?shù)與標(biāo)準(zhǔn)自注意力相當(dāng)，同時推理延遲和顯存占用大幅降低，展現(xiàn)出更強(qiáng)的長序列處理效率優(yōu)勢。

表 1： 長序列語言建模實驗

長文檔問答任務(wù)

在多文檔問答任務(wù)的 EM Score 評估中，CCA-LLM 在不同序列長度下均展現(xiàn)出優(yōu)異的表現(xiàn)，且其性能優(yōu)勢隨著上下文長度的增加而愈加明顯。在處理超長上下文（如 64K 和 128K）任務(wù)時，CCA-LLM 的 EM 得分超越了標(biāo)準(zhǔn)自注意力機(jī)制，同時推理速度也顯著提升——在 128K 上下文長度下，推理速度達(dá)到標(biāo)準(zhǔn)自注意力方法的 7.9 倍，展現(xiàn)出其在高效長文本建模方面的突出優(yōu)勢。

表 2： 長文檔問答實驗

計算和存儲效率對比

相比標(biāo)準(zhǔn)自注意力及其他高效注意力方法（如 MInference），CCA-Attention 在推理速度與內(nèi)存占用方面展現(xiàn)出顯著優(yōu)勢。不同于 MInference 等僅關(guān)注預(yù)填充（prefilling）階段加速的方法，CCA-Attention 能夠同時優(yōu)化預(yù)填充和解碼（decoding）兩個階段，實現(xiàn)端到端的全流程高效推理。

在 64K 上下文長度下，CCA-Attention 的推理速度達(dá)到標(biāo)準(zhǔn)自注意力的 5.7 倍，KV Cache 顯存占用也大幅降低；在 128K 上下文任務(wù)中，推理速度提升更是達(dá)到 7.9 倍，同時 KV Cache 顯存使用減少高達(dá) 93%，充分體現(xiàn)了其在長序列建模中的高效性與實用性。