在剛剛舉辦的計算語言學(xué)和自然語言處理領(lǐng)域的頂級國際會議ACL上，由DeepSeek創(chuàng)始人梁文鋒親自署名的論文從8000多篇論文中脫穎而出，成為本屆大會的最佳論文之一。
（來源：資料圖）
該論文的其他合作者分別來自北京大學(xué)和美國華盛頓大學(xué)，擔(dān)任第一作者的是JingyangYuan。
圖|JingyangYuan（來源：https://ieeexplore.ieee.org/autho）
在這篇論文中，他們提出了一種名為NSA的本機可訓(xùn)練的稀疏注意力機制，該機制將算法創(chuàng)新與硬件優(yōu)化相結(jié)合，以實現(xiàn)高效的長上下文建模。
NSA采用了一種動態(tài)分層稀疏策略，將粗粒度token壓縮與細(xì)粒度token選擇相結(jié)合，以同時保持全局上下文感知和局部精確性。
這一方法通過以下兩項關(guān)鍵創(chuàng)新推動了稀疏注意力設(shè)計的發(fā)展：
一方面，他們通過設(shè)計算術(shù)強度平衡的算法，并結(jié)合針對現(xiàn)代硬件的實現(xiàn)優(yōu)化，實現(xiàn)了顯著的速度提升。
另一方面，他們實現(xiàn)了端到端的訓(xùn)練，在不犧牲模型性能的前提下減少了預(yù)訓(xùn)練的計算量。
如下圖所示，實驗表明，在通用基準(zhǔn)測試、長上下文任務(wù)和基于指令的推理任務(wù)中，使用NSA預(yù)訓(xùn)練的模型表現(xiàn)與全注意力模型相當(dāng)或更優(yōu)。
（來源：https://arxiv.org/pdf/2502.11089）
同時，在解碼、前向傳播和后向傳播方面，NSA在64k長度序列上相較于全注意力機制實現(xiàn)了顯著加速，驗證了其在整個模型生命周期中的效率。
圖|相關(guān)論文（來源：https://arxiv.org/pdf/2502.11089）
長上下文建模是下一代大模型的關(guān)鍵能力
研究界日益認(rèn)識到，長上下文建模是下一代大模型的關(guān)鍵能力，其推動因素是現(xiàn)實世界中的各種應(yīng)用，包括深度推理、庫級代碼生成和多輪自主代理系統(tǒng)。然而，標(biāo)準(zhǔn)注意力機制的高計算成本帶來了巨大的計算挑戰(zhàn)。稀疏注意力機制為在保持模型性能的同時提高效率提供了一個有前景的方向。
最近領(lǐng)域內(nèi)的一些突破，包括OpenAI的o系列模型、DeepSeek-R1和Gemini1.5Pro，使模型能夠處理整個代碼庫、長文檔，在數(shù)千個token上保持連貫的多輪對話，并在較長依賴關(guān)系中進行復(fù)雜推理。
然而，隨著序列長度的增加，基礎(chǔ)注意力機制的高度復(fù)雜性成為關(guān)鍵的延遲瓶頸。
理論估計表明，在解碼64k長度上下文時，采用softmax架構(gòu)進行注意力計算占總延遲的70%至80%，這凸顯出人們迫切需要更高效的注意力機制。
實現(xiàn)高效長上下文建模的一種自然方法是利用softmax注意力機制的固有稀疏性，即選擇性地計算關(guān)鍵的查詢-鍵對，這可以在保持性能的同時顯著降低計算開銷。近期的研究進展通過多種策略展示了這一潛力：鍵值（KV）緩存淘汰方法、分塊KV緩存選擇方法，以及基于采樣、聚類或哈希的選擇方法。
盡管這些策略頗具前景，但現(xiàn)有的稀疏注意力方法在實際部署中往往表現(xiàn)不佳。許多方法未能實現(xiàn)與其理論增益相當(dāng)?shù)乃俣忍嵘淮送?，大多?shù)方法缺乏有效的訓(xùn)練時間支持，無法充分利用注意力機制的稀疏模式。
為了克服這些局限性，部署有效的稀疏注意力機制必須應(yīng)對兩個關(guān)鍵挑戰(zhàn)：
（1）解決硬件適配的推理加速挑戰(zhàn)：將理論計算減少轉(zhuǎn)化為實際速度提升，需要在預(yù)填充和解碼階段都進行硬件友好的算法設(shè)計，以便緩解內(nèi)存訪問和硬件調(diào)度瓶頸；
（2）解決訓(xùn)練感知算法設(shè)計的挑戰(zhàn)：通過可訓(xùn)練算子實現(xiàn)端到端的計算，以便降低訓(xùn)練成本，同時還得保持模型性能。這些要求對于實際應(yīng)用實現(xiàn)快速長上下文推理或訓(xùn)練至關(guān)重要。當(dāng)同時考慮這兩個方面時，現(xiàn)有方法仍然存在顯著的差距。
因此，為了實現(xiàn)更有效且高效的稀疏注意力機制，他們提出了NSA這種原生可訓(xùn)練的稀疏注意力架構(gòu)，并集成了分層token建模。如下圖所示，NSA通過將鍵和值組織成時間塊，并通過三條注意力路徑來處理它們，從而減少每次查詢的計算量：壓縮的粗粒度token、選擇性保留的細(xì)粒度token以及用于獲取局部上下文信息的滑動窗口。
（來源：https://arxiv.org/pdf/2502.11089）
然后可以實現(xiàn)專門的核函數(shù)，以最大限度地提高其實用效率。針對上述關(guān)鍵需求，NSA推出了兩項核心創(chuàng)新：
首先是硬件對齊系統(tǒng)：針對TensorCore的利用率和內(nèi)存訪問，優(yōu)化塊級稀疏注意力機制，確保算術(shù)強度平衡。
其次是訓(xùn)練感知設(shè)計：通過高效算法和反向操作符實現(xiàn)穩(wěn)定的端到端訓(xùn)練。這一優(yōu)化使NSA能夠同時支持高效部署和端到端訓(xùn)練。
研究中，該團隊通過在真實世界語言語料庫上的綜合實驗來評估NSA。在擁有2600億個token的270億參數(shù)的Transformer主干網(wǎng)絡(luò)上進行預(yù)訓(xùn)練后，他們評估了NSA在通用語言評估、長上下文評估和思維鏈推理評估中的表現(xiàn)。并進一步比較了在英偉達(dá)A100GPU上采用優(yōu)化Triton實現(xiàn)的核速度。實驗結(jié)果表明，NSA的性能與全注意力基線相當(dāng)或更優(yōu)，同時優(yōu)于現(xiàn)有的稀疏注意力方法。此外，與全注意力機制相比，NSA在解碼、前向和后向階段均實現(xiàn)了顯著加速，且隨著序列長度的增加，加速比也隨之提高。這些結(jié)果驗證了他們的分層稀疏注意力設(shè)計有效地平衡了模型能力和計算效率。
三個角度評估NSA
實驗中，他們從以下三個角度來評估NSA：通用基準(zhǔn)性能、長上下文基準(zhǔn)性能以及思維鏈推理性能，并與全注意力基線以及當(dāng)前最先進的稀疏注意力方法進行比較。
遵循最先進的大模型的常見做法，他們的實驗采用了一種結(jié)合分組查詢注意力（GQA，Grouped-QueryAttention）和混合專家（MoE，Mixture-of-Experts）的骨干網(wǎng)絡(luò)，總參數(shù)為270億，其中活躍參數(shù)為30億。該模型由30層組成，隱藏層維度為2560。對于GQA，他們將組數(shù)設(shè)置為4，總共有64個注意力頭。對于MoE，他們采用了DeepSeekMoE結(jié)構(gòu)，該結(jié)構(gòu)包含72個路由專家和2個共享專家，并將top-k專家數(shù)設(shè)置為6。為確保訓(xùn)練穩(wěn)定性，第一層中的MoE被替換為SwiGLU形式的MLP。
該團隊在論文中指出，其所提出的架構(gòu)在計算成本和模型性能之間實現(xiàn)了有效的權(quán)衡。其將全注意力模型和稀疏注意力模型均在2700億個8k長度文本的token上進行預(yù)訓(xùn)練，隨后使用YaRN在32k長度文本上進行持續(xù)訓(xùn)練和監(jiān)督微調(diào)，以實現(xiàn)長上下文適應(yīng)。兩種模型都經(jīng)過充分訓(xùn)練以實現(xiàn)完全收斂，以確保公平比較。如下圖所示，NSA和全注意力（FullAttention）基線的預(yù)訓(xùn)練損失曲線呈現(xiàn)出穩(wěn)定且平滑的下降趨勢，且NSA模型的表現(xiàn)始終優(yōu)于全注意力模型。
（來源：https://arxiv.org/pdf/2502.11089）
除了與全注意力（FullAttention）方法進行比較外，他們還評估了幾種最先進的推理階段稀疏注意力方法：H2O、infLLM、Quest和Exact-Top，這些方法首先計算全注意力得分，并選擇與每個查詢相對應(yīng)的前幾個得分最高的鍵，然后計算這些位置上的注意力。與此同時，這些方法涵蓋了多種稀疏注意力范式。
在一般評估中，當(dāng)大多數(shù)樣本的長度位于稀疏注意力基線的局部上下文窗口內(nèi)時，這些方法實際上等同于全注意力方法。因此，在這種設(shè)置之下，該團隊僅展示了NSA與全注意力基線之間的比較結(jié)果。在長文本評估中，他們對所有基線方法進行了比較，并將所有稀疏注意力方法的稀疏度設(shè)置為相同，以便確保比較的公平性。對于需要進行長文本監(jiān)督微調(diào)的思維鏈推理評估，他們僅將比較范圍限定在全注意力模型上，因為稀疏注意力基線模型不支持訓(xùn)練。
在一系列涵蓋知識、推理和編碼能力的綜合基準(zhǔn)測試上，該團隊對預(yù)訓(xùn)練的NSA和全注意力基線進行了評估，這些基準(zhǔn)測試包括MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP和HumanEval。
圖|對比結(jié)果（來源：https://arxiv.org/pdf/2502.11089）
盡管NSA較為簡潔，但其整體表現(xiàn)卻十分出色，在9項指標(biāo)中有7項優(yōu)于包括全注意力機制在內(nèi)的所有基線模型。這表明，盡管在較短序列上，NSA可能無法充分發(fā)揮其效率優(yōu)勢，但其表現(xiàn)仍然強勁。值得注意的是，NSA在推理相關(guān)基準(zhǔn)測試中表現(xiàn)出了顯著的提升（DROP：+0.042，GSM8K：+0.034），這表明該團隊所使用的預(yù)訓(xùn)練有助于模型發(fā)展出專門的注意力機制。這種稀疏注意力預(yù)訓(xùn)練機制迫使模型聚焦于最重要的信息，通過過濾掉無關(guān)注意力路徑中的噪聲，從而能夠潛在地提升性能。在不同評估中的一致表現(xiàn)也驗證了NSA作為一種通用架構(gòu)的穩(wěn)健性。
同時，他們還進行了長上下文評估。如下圖所示，在64k上下文的“大海撈針”測試中，NSA在所有位置均實現(xiàn)了完美的檢索準(zhǔn)確率。
（來源：https://arxiv.org/pdf/2502.11089）
這種表現(xiàn)源于該團隊的分層稀疏注意力設(shè)計方案，該設(shè)計方案結(jié)合了壓縮token以便實現(xiàn)高效的全局上下文掃描，以及能夠通過選擇token實現(xiàn)精確的局部信息檢索。粗粒度壓縮則能以較低的計算成本識別出相關(guān)的上下文塊，而對選定token的token級注意力能夠確保關(guān)鍵細(xì)粒度信息的保留。這種設(shè)計使得NSA既能保持全局感知，又能確保局部精確性。
該團隊還在LongBench上對NSA與最先進的稀疏注意力方法和全注意力基線進行了評估。為確保稀疏性一致，他們將所有稀疏注意力基線中每個查詢激活的token設(shè)置為2560個，這對應(yīng)于在處理32k序列長度時，NSA中激活的token的平均數(shù)量。繼StreamLLM之后，該token預(yù)算包括前128個token和512個本地token。他們排除了LongBench中的某些子集，因為它們在所有模型中的得分都很低，可能無法提供有意義的比較。如下表所示，NSA取得了最高的平均分0.469，優(yōu)于所有基線方法（比全注意力機制高出0.032，比精確頂部方法高出0.046）。
（來源：https://arxiv.org/pdf/2502.11089）
這一優(yōu)異表現(xiàn)源于以下兩項關(guān)鍵創(chuàng)新：第一，他們獨創(chuàng)的稀疏注意力設(shè)計，能夠在預(yù)訓(xùn)練期間對稀疏模式進行端到端的優(yōu)化，促進稀疏注意力模塊與其他模型組件之間的同步適應(yīng)；第二，分層稀疏注意力機制實現(xiàn)了局部和全局信息處理之間的平衡。
值得注意的是，NSA在需要長上下文復(fù)雜推理的任務(wù)上表現(xiàn)出色，在多跳問答任務(wù)（HPQ和2Wiki）上相較于全注意力機制分別提升了+0.087和+0.051，在代碼理解任務(wù)上超越了基線模型（LCC：+0.069），并且在段落檢索任務(wù)上優(yōu)于其他方法（PassR-en：+0.075）。這些結(jié)果驗證了NSA在處理多樣化長上下文挑戰(zhàn)的能力，其原生預(yù)訓(xùn)練的稀疏注意力在學(xué)習(xí)任務(wù)最優(yōu)模式方面提供了額外優(yōu)勢。
在思維鏈推理評估中，為了評估NSA與先進下游訓(xùn)練范式的兼容性，該團隊研究了其通過后訓(xùn)練獲取思維鏈數(shù)學(xué)推理能力的能力。鑒于強化學(xué)習(xí)在較小規(guī)模模型上的效果有限，他們采用了DeepSeek-R1的知識蒸餾方法，利用100億個32k長度數(shù)學(xué)推理軌跡的token進行監(jiān)督微調(diào)（SFT，supervisedfine-tuning）。這產(chǎn)生了兩個可比較的模型：FullAttention-R（全注意力基線模型）和NSA-R（該團隊的稀疏變體模型）。他們在AIME24基準(zhǔn)測試中評估了這兩個模型。為了驗證推理深度的影響，他們在兩種上下文限制（8k和16k個token）下進行了實驗，以便衡量擴展推理鏈?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é)果驗證了原生稀疏注意力機制的兩大關(guān)鍵優(yōu)勢：（1）預(yù)訓(xùn)練的稀疏注意力模式能夠高效捕捉對復(fù)雜數(shù)學(xué)推導(dǎo)至關(guān)重要的長距離邏輯依賴關(guān)系；（2）本次架構(gòu)采用與硬件相匹配的設(shè)計，保持了足夠的上下文密度，以支持不斷增加的推理深度，同時避免災(zāi)難性遺忘。
這種在不同上下文長度上的一致性表現(xiàn)證明，當(dāng)稀疏注意力機制被原生整合到訓(xùn)練流程中時，它能夠勝任高級推理任務(wù)。
（來源：https://openreview.net/profile?id=~Jingyang_Yuan1）
另據(jù)悉，擔(dān)任本次論文第一作者兼通訊作者的JingyangYuan于2022年獲得北京大學(xué)計算機科學(xué)學(xué)士學(xué)位，預(yù)計將于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
運營/排版：何晨龍
