在剛剛舉辦的計(jì)算語言學(xué)和自然語言處理領(lǐng)域的頂級國際會議 ACL 上，由 DeepSeek 創(chuàng)始人梁文鋒親自署名的論文從 8000 多篇論文中脫穎而出，成為本屆大會的最佳論文之一。

（來源：資料圖）

該論文的其他合作者分別來自北京大學(xué)和美國華盛頓大學(xué)，擔(dān)任第一作者的是 Jingyang Yuan。

圖 | Jingyang Yuan（來源：https://ieeexplore.ieee.org/autho）

在這篇論文中，他們提出了一種名為 NSA 的本機(jī)可訓(xùn)練的稀疏注意力機(jī)制，該機(jī)制將算法創(chuàng)新與硬件優(yōu)化相結(jié)合，以實(shí)現(xiàn)高效的長上下文建模。

NSA 采用了一種動態(tài)分層稀疏策略，將粗粒度 token 壓縮與細(xì)粒度 token 選擇相結(jié)合，以同時保持全局上下文感知和局部精確性。

這一方法通過以下兩項(xiàng)關(guān)鍵創(chuàng)新推動了稀疏注意力設(shè)計(jì)的發(fā)展：

一方面，他們通過設(shè)計(jì)算術(shù)強(qiáng)度平衡的算法，并結(jié)合針對現(xiàn)代硬件的實(shí)現(xiàn)優(yōu)化，實(shí)現(xiàn)了顯著的速度提升。

另一方面，他們實(shí)現(xiàn)了端到端的訓(xùn)練，在不犧牲模型性能的前提下減少了預(yù)訓(xùn)練的計(jì)算量。

如下圖所示，實(shí)驗(yàn)表明，在通用基準(zhǔn)測試、長上下文任務(wù)和基于指令的推理任務(wù)中，使用 NSA 預(yù)訓(xùn)練的模型表現(xiàn)與全注意力模型相當(dāng)或更優(yōu)。

（來源：https://arxiv.org/pdf/2502.11089）

同時，在解碼、前向傳播和后向傳播方面，NSA 在 64k 長度序列上相較于全注意力機(jī)制實(shí)現(xiàn)了顯著加速，驗(yàn)證了其在整個模型生命周期中的效率。

圖 | 相關(guān)論文（來源：https://arxiv.org/pdf/2502.11089）

長上下文建模是下一代大模型的關(guān)鍵能力

研究界日益認(rèn)識到，長上下文建模是下一代大模型的關(guān)鍵能力，其推動因素是現(xiàn)實(shí)世界中的各種應(yīng)用，包括深度推理、庫級代碼生成和多輪自主代理系統(tǒng)。然而，標(biāo)準(zhǔn)注意力機(jī)制的高計(jì)算成本帶來了巨大的計(jì)算挑戰(zhàn)。稀疏注意力機(jī)制為在保持模型性能的同時提高效率提供了一個有前景的方向。

最近領(lǐng)域內(nèi)的一些突破，包括 OpenAI 的 o 系列模型、DeepSeek-R1 和 Gemini 1.5 Pro，使模型能夠處理整個代碼庫、長文檔，在數(shù)千個 token 上保持連貫的多輪對話，并在較長依賴關(guān)系中進(jìn)行復(fù)雜推理。

然而，隨著序列長度的增加，基礎(chǔ)注意力機(jī)制的高度復(fù)雜性成為關(guān)鍵的延遲瓶頸。

理論估計(jì)表明，在解碼 64k 長度上下文時，采用 softmax 架構(gòu)進(jìn)行注意力計(jì)算占總延遲的 70% 至 80%，這凸顯出人們迫切需要更高效的注意力機(jī)制。

實(shí)現(xiàn)高效長上下文建模的一種自然方法是利用 softmax 注意力機(jī)制的固有稀疏性，即選擇性地計(jì)算關(guān)鍵的查詢-鍵對，這可以在保持性能的同時顯著降低計(jì)算開銷。近期的研究進(jìn)展通過多種策略展示了這一潛力：鍵值（KV）緩存淘汰方法、分塊 KV 緩存選擇方法，以及基于采樣、聚類或哈希的選擇方法。

盡管這些策略頗具前景，但現(xiàn)有的稀疏注意力方法在實(shí)際部署中往往表現(xiàn)不佳。許多方法未能實(shí)現(xiàn)與其理論增益相當(dāng)?shù)乃俣忍嵘?；此外，大多?shù)方法缺乏有效的訓(xùn)練時間支持，無法充分利用注意力機(jī)制的稀疏模式。

為了克服這些局限性，部署有效的稀疏注意力機(jī)制必須應(yīng)對兩個關(guān)鍵挑戰(zhàn)：

（1）解決硬件適配的推理加速挑戰(zhàn)：將理論計(jì)算減少轉(zhuǎn)化為實(shí)際速度提升，需要在預(yù)填充和解碼階段都進(jìn)行硬件友好的算法設(shè)計(jì)，以便緩解內(nèi)存訪問和硬件調(diào)度瓶頸；

（2）解決訓(xùn)練感知算法設(shè)計(jì)的挑戰(zhàn)：通過可訓(xùn)練算子實(shí)現(xiàn)端到端的計(jì)算，以便降低訓(xùn)練成本，同時還得保持模型性能。這些要求對于實(shí)際應(yīng)用實(shí)現(xiàn)快速長上下文推理或訓(xùn)練至關(guān)重要。當(dāng)同時考慮這兩個方面時，現(xiàn)有方法仍然存在顯著的差距。

因此，為了實(shí)現(xiàn)更有效且高效的稀疏注意力機(jī)制，他們提出了 NSA 這種原生可訓(xùn)練的稀疏注意力架構(gòu)，并集成了分層 token 建模。如下圖所示，NSA 通過將鍵和值組織成時間塊，并通過三條注意力路徑來處理它們，從而減少每次查詢的計(jì)算量：壓縮的粗粒度 token、選擇性保留的細(xì)粒度 token 以及用于獲取局部上下文信息的滑動窗口。

（來源：https://arxiv.org/pdf/2502.11089）

然后可以實(shí)現(xiàn)專門的核函數(shù)，以最大限度地提高其實(shí)用效率。針對上述關(guān)鍵需求，NSA 推出了兩項(xiàng)核心創(chuàng)新：

首先是硬件對齊系統(tǒng)：針對 Tensor Core 的利用率和內(nèi)存訪問，優(yōu)化塊級稀疏注意力機(jī)制，確保算術(shù)強(qiáng)度平衡。

其次是訓(xùn)練感知設(shè)計(jì)：通過高效算法和反向操作符實(shí)現(xiàn)穩(wěn)定的端到端訓(xùn)練。這一優(yōu)化使 NSA 能夠同時支持高效部署和端到端訓(xùn)練。

研究中，該團(tuán)隊(duì)通過在真實(shí)世界語言語料庫上的綜合實(shí)驗(yàn)來評估 NSA。在擁有 2600 億個 token 的 270 億參數(shù)的 Transformer 主干網(wǎng)絡(luò)上進(jìn)行預(yù)訓(xùn)練后，他們評估了 NSA 在通用語言評估、長上下文評估和思維鏈推理評估中的表現(xiàn)。并進(jìn)一步比較了在英偉達(dá) A100 GPU 上采用優(yōu)化 Triton 實(shí)現(xiàn)的核速度。實(shí)驗(yàn)結(jié)果表明，NSA 的性能與全注意力基線相當(dāng)或更優(yōu)，同時優(yōu)于現(xiàn)有的稀疏注意力方法。此外，與全注意力機(jī)制相比，NSA 在解碼、前向和后向階段均實(shí)現(xiàn)了顯著加速，且隨著序列長度的增加，加速比也隨之提高。這些結(jié)果驗(yàn)證了他們的分層稀疏注意力設(shè)計(jì)有效地平衡了模型能力和計(jì)算效率。

三個角度評估 NSA

實(shí)驗(yàn)中，他們從以下三個角度來評估 NSA：通用基準(zhǔn)性能、長上下文基準(zhǔn)性能以及思維鏈推理性能，并與全注意力基線以及當(dāng)前最先進(jìn)的稀疏注意力方法進(jìn)行比較。

遵循最先進(jìn)的大模型的常見做法，他們的實(shí)驗(yàn)采用了一種結(jié)合分組查詢注意力（GQA，Grouped-Query Attention）和混合專家（MoE，Mixture-of-Experts）的骨干網(wǎng)絡(luò)，總參數(shù)為 270 億，其中活躍參數(shù)為 30 億。該模型由 30 層組成，隱藏層維度為 2560。對于 GQA，他們將組數(shù)設(shè)置為 4，總共有 64 個注意力頭。對于 MoE，他們采用了 DeepSeek MoE 結(jié)構(gòu)，該結(jié)構(gòu)包含 72 個路由專家和 2 個共享專家，并將 top-k 專家數(shù)設(shè)置為 6。為確保訓(xùn)練穩(wěn)定性，第一層中的 MoE 被替換為 SwiGLU 形式的 MLP。

該團(tuán)隊(duì)在論文中指出，其所提出的架構(gòu)在計(jì)算成本和模型性能之間實(shí)現(xiàn)了有效的權(quán)衡。其將全注意力模型和稀疏注意力模型均在 2700 億個 8k 長度文本的 token 上進(jìn)行預(yù)訓(xùn)練，隨后使用 YaRN 在 32k 長度文本上進(jìn)行持續(xù)訓(xùn)練和監(jiān)督微調(diào)，以實(shí)現(xiàn)長上下文適應(yīng)。兩種模型都經(jīng)過充分訓(xùn)練以實(shí)現(xiàn)完全收斂，以確保公平比較。如下圖所示，NSA 和全注意力（Full Attention）基線的預(yù)訓(xùn)練損失曲線呈現(xiàn)出穩(wěn)定且平滑的下降趨勢，且 NSA 模型的表現(xiàn)始終優(yōu)于全注意力模型。

（來源：https://arxiv.org/pdf/2502.11089）

除了與全注意力（Full Attention）方法進(jìn)行比較外，他們還評估了幾種最先進(jìn)的推理階段稀疏注意力方法：H2O、infLLM、Quest 和 Exact-Top，這些方法首先計(jì)算全注意力得分，并選擇與每個查詢相對應(yīng)的前幾個得分最高的鍵，然后計(jì)算這些位置上的注意力。與此同時，這些方法涵蓋了多種稀疏注意力范式。

在一般評估中，當(dāng)大多數(shù)樣本的長度位于稀疏注意力基線的局部上下文窗口內(nèi)時，這些方法實(shí)際上等同于全注意力方法。因此，在這種設(shè)置之下，該團(tuán)隊(duì)僅展示了 NSA 與全注意力基線之間的比較結(jié)果。在長文本評估中，他們對所有基線方法進(jìn)行了比較，并將所有稀疏注意力方法的稀疏度設(shè)置為相同，以便確保比較的公平性。對于需要進(jìn)行長文本監(jiān)督微調(diào)的思維鏈推理評估，他們僅將比較范圍限定在全注意力模型上，因?yàn)橄∈枳⒁饬€模型不支持訓(xùn)練。

在一系列涵蓋知識、推理和編碼能力的綜合基準(zhǔn)測試上，該團(tuán)隊(duì)對預(yù)訓(xùn)練的 NSA 和全注意力基線進(jìn)行了評估，這些基準(zhǔn)測試包括 MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP 和 HumanEval。

圖 | 對比結(jié)果（來源：https://arxiv.org/pdf/2502.11089）

盡管 NSA 較為簡潔，但其整體表現(xiàn)卻十分出色，在 9 項(xiàng)指標(biāo)中有 7 項(xiàng)優(yōu)于包括全注意力機(jī)制在內(nèi)的所有基線模型。這表明，盡管在較短序列上，NSA 可能無法充分發(fā)揮其效率優(yōu)勢，但其表現(xiàn)仍然強(qiáng)勁。值得注意的是，NSA 在推理相關(guān)基準(zhǔn)測試中表現(xiàn)出了顯著的提升（DROP：+0.042，GSM8K：+0.034），這表明該團(tuán)隊(duì)所使用的預(yù)訓(xùn)練有助于模型發(fā)展出專門的注意力機(jī)制。這種稀疏注意力預(yù)訓(xùn)練機(jī)制迫使模型聚焦于最重要的信息，通過過濾掉無關(guān)注意力路徑中的噪聲，從而能夠潛在地提升性能。在不同評估中的一致表現(xiàn)也驗(yàn)證了 NSA 作為一種通用架構(gòu)的穩(wěn)健性。

同時，他們還進(jìn)行了長上下文評估。如下圖所示，在 64k 上下文的“大海撈針”測試中，NSA 在所有位置均實(shí)現(xiàn)了完美的檢索準(zhǔn)確率。

（來源：https://arxiv.org/pdf/2502.11089）

這種表現(xiàn)源于該團(tuán)隊(duì)的分層稀疏注意力設(shè)計(jì)方案，該設(shè)計(jì)方案結(jié)合了壓縮 token 以便實(shí)現(xiàn)高效的全局上下文掃描，以及能夠通過選擇 token 實(shí)現(xiàn)精確的局部信息檢索。粗粒度壓縮則能以較低的計(jì)算成本識別出相關(guān)的上下文塊，而對選定 token 的 token 級注意力能夠確保關(guān)鍵細(xì)粒度信息的保留。這種設(shè)計(jì)使得 NSA 既能保持全局感知，又能確保局部精確性。

該團(tuán)隊(duì)還在 LongBench 上對 NSA 與最先進(jìn)的稀疏注意力方法和全注意力基線進(jìn)行了評估。為確保稀疏性一致，他們將所有稀疏注意力基線中每個查詢激活的 token 設(shè)置為 2560 個，這對應(yīng)于在處理 32k 序列長度時，NSA 中激活的 token 的平均數(shù)量。繼 StreamLLM 之后，該 token 預(yù)算包括前 128 個 token 和 512 個本地 token。他們排除了 LongBench 中的某些子集，因?yàn)樗鼈冊谒心Ｐ椭械牡梅侄己艿?，可能無法提供有意義的比較。如下表所示，NSA 取得了最高的平均分 0.469，優(yōu)于所有基線方法（比全注意力機(jī)制高出 0.032，比精確頂部方法高出 0.046）。

（來源：https://arxiv.org/pdf/2502.11089）

這一優(yōu)異表現(xiàn)源于以下兩項(xiàng)關(guān)鍵創(chuàng)新：第一，他們獨(dú)創(chuàng)的稀疏注意力設(shè)計(jì)，能夠在預(yù)訓(xùn)練期間對稀疏模式進(jìn)行端到端的優(yōu)化，促進(jìn)稀疏注意力模塊與其他模型組件之間的同步適應(yīng)；第二，分層稀疏注意力機(jī)制實(shí)現(xiàn)了局部和全局信息處理之間的平衡。

值得注意的是，NSA 在需要長上下文復(fù)雜推理的任務(wù)上表現(xiàn)出色，在多跳問答任務(wù)（HPQ 和 2Wiki）上相較于全注意力機(jī)制分別提升了+0.087 和+0.051，在代碼理解任務(wù)上超越了基線模型（LCC：+0.069），并且在段落檢索任務(wù)上優(yōu)于其他方法（PassR-en：+0.075）。這些結(jié)果驗(yàn)證了 NSA 在處理多樣化長上下文挑戰(zhàn)的能力，其原生預(yù)訓(xùn)練的稀疏注意力在學(xué)習(xí)任務(wù)最優(yōu)模式方面提供了額外優(yōu)勢。

在思維鏈推理評估中，為了評估 NSA 與先進(jìn)下游訓(xùn)練范式的兼容性，該團(tuán)隊(duì)研究了其通過后訓(xùn)練獲取思維鏈數(shù)學(xué)推理能力的能力。鑒于強(qiáng)化學(xué)習(xí)在較小規(guī)模模型上的效果有限，他們采用了 DeepSeek-R1 的知識蒸餾方法，利用 100 億個 32k 長度數(shù)學(xué)推理軌跡的 token 進(jìn)行監(jiān)督微調(diào)（SFT，supervised fine-tuning）。這產(chǎn)生了兩個可比較的模型：Full Attention-R（全注意力基線模型）和 NSA-R（該團(tuán)隊(duì)的稀疏變體模型）。他們在 AIME 24 基準(zhǔn)測試中評估了這兩個模型。為了驗(yàn)證推理深度的影響，他們在兩種上下文限制（8k 和 16k 個 token）下進(jìn)行了實(shí)驗(yàn)，以便衡量擴(kuò)展推理鏈?zhǔn)欠衲芴岣邷?zhǔn)確性。

如下表所示，在 8k 上下文設(shè)置下，NSA-R 的準(zhǔn)確率顯著高于全注意力-R（+0.075），且在 16k 上下文設(shè)置下，這一優(yōu)勢仍然存在（+0.054）。

（來源：https://arxiv.org/pdf/2502.11089）

這些結(jié)果驗(yàn)證了原生稀疏注意力機(jī)制的兩大關(guān)鍵優(yōu)勢：（1）預(yù)訓(xùn)練的稀疏注意力模式能夠高效捕捉對復(fù)雜數(shù)學(xué)推導(dǎo)至關(guān)重要的長距離邏輯依賴關(guān)系；（2）本次架構(gòu)采用與硬件相匹配的設(shè)計(jì)，保持了足夠的上下文密度，以支持不斷增加的推理深度，同時避免災(zāi)難性遺忘。

這種在不同上下文長度上的一致性表現(xiàn)證明，當(dāng)稀疏注意力機(jī)制被原生整合到訓(xùn)練流程中時，它能夠勝任高級推理任務(wù)。

（來源：https://openreview.net/profile?id=~Jingyang_Yuan1）

另據(jù)悉，擔(dān)任本次論文第一作者兼通訊作者的 Jingyang Yuan 于 2022 年獲得北京大學(xué)計(jì)算機(jī)科學(xué)學(xué)士學(xué)位，預(yù)計(jì)將于 2027 年從北京大學(xué)博士畢業(yè)，他的研究方向包括圖表示學(xué)習(xí)、神經(jīng)物理模擬、大模型等。

參考資料：

https://ieeexplore.ieee.org/author/37090050318

https://scholar.google.com/citations?user=mDwlqfkAAAAJ&hl=en

https://arxiv.org/pdf/2502.11089

運(yùn)營/排版：何晨龍