DeepSeek 新論文來了！相關消息剛剛發(fā)布到 ?? 就吸引了大量用戶點贊、轉發(fā)、評論三連。

據介紹，DeepSeek 的這篇新論文提出了一種新的注意力機制 ——NSA。這是一個用于超快長上下文訓練和推斷的本地可訓練的稀疏注意力機制，并且還具有與硬件對齊的特點。

新研究發(fā)布兩個小時，就有近三十萬的瀏覽量。現在看來，DeepSeek 發(fā)布成果，比 OpenAI 關注度都高。

論文標題：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

論文鏈接：https://arxiv.org/abs/2502.11089

值得一提的是，幻方科技、DeepSeek 創(chuàng)始人梁文鋒也是論文的作者之一。這成了眾多網友討論的話題。

接下來，讓我們看下梁文鋒親自參與的研究，講了什么內容。

論文概覽

長上下文建模是下一代大型語言模型（LLM）的關鍵能力，這一需求源于多樣化的實際應用，包括深度推理、倉庫級代碼生成以及多輪自動智能體系統(tǒng)等。

最近大模型的突破 —— 如 OpenAI 的 o 系列模型、DeepSeek-R1 和 Gemini 1.5 Pro—— 已經能使得模型能夠處理整個代碼庫、長文檔、在數千個 token 上保持連貫的多輪對話，并在長距離依賴關系中進行復雜推理。然而，隨著序列長度的增加，普通注意力機制的高復雜性成為關鍵的延遲瓶頸。理論估計表明，在使用 softmax 架構進行 64k 長度上下文的解碼時，注意力計算占總延遲的 70-80%，這凸顯了對更高效注意力機制的迫切需求。

實現高效長上下文建模的自然方法是利用 softmax 注意力的固有稀疏性，通過選擇性計算關鍵 query-key 對，可以顯著減少計算開銷，同時保持性能。最近這一路線的進展包括多種策略：KV 緩存淘汰方法、塊狀 KV 緩存選擇方法以及基于采樣、聚類或哈希的選擇方法。盡管這些策略前景廣闊，現有的稀疏注意力方法在實際部署中往往表現不佳。許多方法未能實現與其理論增益相媲美的加速；此外，大多數方法主要關注推理階段，缺乏有效的訓練時支持以充分利用注意力的稀疏模式。

為了克服這些限制，部署有效的稀疏注意力必須應對兩個關鍵挑戰(zhàn)：

1、硬件對齊的推理加速：將理論計算減少轉化為實際速度提升，需要在預填充和解碼階段設計硬件友好的算法，以緩解內存訪問和硬件調度瓶頸；

2、訓練感知的算法設計：通過可訓練的操作符實現端到端計算，以降低訓練成本，同時保持模型性能。

這些要求對于實際應用實現快速長上下文推理或訓練至關重要。在考慮這兩方面時，現有方法仍顯不足。

為了實現更有效和高效的稀疏注意力，DeepSeek 研究人員提出了一種原生可訓練的稀疏注意力架構 NSA，它集成了分層 token 建模。

如圖 2 所示，NSA 通過將鍵和值組織成時間塊（temporal blocks）并通過三條注意力路徑處理它們來減少每查詢計算量：壓縮的粗粒度 token、選擇性保留的細粒度 token 以及用于局部上下文信息的滑動窗口。隨后，作者實現了專門的核以最大化其實際效率。

NSA 引入了兩個核心創(chuàng)新以對應于上述關鍵需求：

1、硬件對齊的系統(tǒng)：優(yōu)化塊狀稀疏注意力以利用 Tensor Core 和內存訪問，確保算術強度平衡；

2、訓練感知的設計：通過高效算法和反向操作符實現穩(wěn)定的端到端訓練。這一優(yōu)化使 NSA 能夠支持高效部署和端到端訓練。

研究通過對現實世界語言語料庫的綜合實驗來評估 NSA。在具有 260B token 的 27B 參數 Transformer 骨干上進行預訓練，作者評估了 NSA 在通用語言評估、長上下文評估和鏈式推理評估中的表現。作者還進一步比較了在 A100 GPU 上內核速度與優(yōu)化 Triton 實現的比較。實驗結果表明，NSA 實現了與 Full Attention 基線相當或更優(yōu)的性能，同時優(yōu)于現有的稀疏注意力方法。

此外，與 Full Attention 相比，NSA 在解碼、前向和后向階段提供了明顯的加速，且加速比隨著序列長度的增加而增加。這些結果驗證了分層稀疏注意力設計有效地平衡了模型能力和計算效率。

方法概覽

本文的技術方法涵蓋算法設計和內核優(yōu)化。作者首先介紹了方法背景，然后介紹了 NSA 的總體框架以及關鍵算法組件，最后詳細介紹了針對硬件優(yōu)化的內核設計，以最大限度地提高實際效率。

背景

注意力機制在語言建模中被廣泛使用，其中每個查詢 token q_??計算與所有前面鍵 k_:??的相關性分數，以生成值 v_:??的加權和。從形式上來說，對于長度為??的輸入序列，注意力操作定義如下：

其中 Attn 表示注意力函數。

這里，??_??,?? 表示 q_?? 和 k_?? 之間的注意力權重，??_??是鍵的特征維度。隨著序列長度的增加，注意力計算在總計算成本中變得越來越占主導地位，對長上下文處理提出了重大挑戰(zhàn)。

算術強度（Arithmetic Intensity）是計算操作與內存訪問的比率，本質上決定了硬件的算法優(yōu)化。每個 GPU 都有一個由其峰值計算能力和內存帶寬決定的臨界算術強度，以這兩個硬件限制的比率計算。對于計算任務，高于此臨界閾值的算術強度將構成計算限制（受限于 GPU FLOPS），而低于此閾值的算術強度將構成內存限制（受限于內存帶寬）。

具體來說，對于因果自注意力機制，在訓練和預填充階段，批量矩陣乘法和注意力計算表現出高算術強度，使得這些階段在現代加速器上計算受限。相反，自回歸解碼會受到內存帶寬的限制，因為它每次前向傳遞都會生成一個 token，同時需要加載整個鍵值緩存，從而導致算術強度較低。這樣就出現了不同的優(yōu)化目標，即減少訓練和預填充期間的計算成本，同時減少解碼期間的內存訪問。

總體框架

為了充分利用具有自然稀疏模式的注意力機制的潛力，作者提出將 (1) 式中原始的鍵 - 值對替換成更加緊湊和信息密集的表征鍵 - 值對。具體而言，可將優(yōu)化的注意力輸出的形式定義成：

其中 是基于當前查詢 q_?? 和上下文記憶 k_:?? , v_:?? 動態(tài)構建的。通過設計不同的映射策略，可以得到的不同類別，然后可將它們按以下方式組合起來

如圖 2 所示，NSA 有三種映射策略 C = {cmp, slc, win}，分別表示鍵和值的壓縮、選取和滑動窗口。??^??_?? ∈ [0, 1] 是對應于策略 c 的門控分數，可通過 MLP 和 sigmoid 激活從輸入特征中得出。令 ??_?? 表示重新映射的鍵 / 值的總數

通過使 ??_?? ? ??，可保持較高的稀疏率。

接下來，DeepSeek 還介紹了重新映射策略的具體設計，涵蓋 token 壓縮、token 選取和滑動窗口。詳細的算法設計見原論文。下面來看看 NSA 為何具有 FlashAttention 相當的速度。

核設計

為了在訓練和預填充期間實現 FlashAttention 級別的加速，作者基于 Triton 實現了硬件對齊的稀疏注意力內核。

由于多頭注意力（MHA）會占用大量內存且解碼效率低下，因此該團隊選擇專注于遵循 SOTA LLM 的共享 KV 緩存架構，如 GQA 和 MQA。

雖然壓縮和滑動窗口注意計算與現有的 FlashAttention-2 內核很容易兼容，但他們卻引入了專門用于稀疏選擇注意的內核設計。如果這時候遵循 FlashAttention 的做法，將時間連續(xù)的查詢塊加載到 SRAM 中，則會導致內存訪問效率低下，因為塊內的查詢可能需要不相交的 KV 塊。

為了解決這個問題，這里的關鍵優(yōu)化在于不同的查詢分組策略：對于查詢序列上的每個位置，將 GQA 組內的所有查詢頭（它們共享相同的稀疏 KV 塊）加載到 SRAM 中。圖 3 說明了其前向傳遞實現。

該設計能夠 (1) 通過分組共享消除冗余的 KV 傳輸，以及 (2) 跨 GPU 流式多處理器平衡計算工作負載，由此實現了近乎最佳的算術強度。

NSA 的實驗表現

作者從三個角度對新提出的 NSA 進行了評估：一般基準性能、長上下文基準性能和思維鏈推理性能。

一般基準性能

該團隊在大量基準上對比了 NSA 與 Full Attention 的表現。結果見下表 1。

可以看到，盡管 NSA 比較稀疏，但它實現了卓越的整體性能，在 9 個指標中的 7 個上都優(yōu)于包括 Full Attention 在內的所有基線。

這表明，盡管 NSA 可能無法充分利用其在較短序列上的效率優(yōu)勢，但它依然表現出了強勁的性能。值得注意的是，NSA 在推理相關基準測試中表現出了顯著的提升（DROP：+0.042，GSM8K：+0.034），這表明 DeepSeek 的預訓練有助于模型發(fā)展出專門的注意力機制。通過過濾掉不相關的注意力路徑中的噪音，這種稀疏注意力預訓練機制可迫使模型專注于最重要的信息，有可能提高性能。在不同評估中的一致表現也證明了 NSA 作為通用架構的穩(wěn)健性。

長上下文基準性能

下圖 5 展示了 NSA 在 64k 上下文的大海撈針（needle-in-a-haystack） 測試中的結果，它在所有位置上都實現了完美的檢索準確率。

這一性能源于 DeepSeek 團隊的分層稀疏注意力設計，該設計結合壓縮 token 來實現高效的全局上下文掃描，以及結合選擇 token 來實現精確的局部信息檢索。粗粒度壓縮以較低的計算成本識別相關的上下文塊，而對選定 token 的 token 級注意力可確保關鍵細粒度信息的保留。這種設計使得 NSA 能夠同時保持全局意識和局部精度。

作者還在 LongBench 上對 NSA 進行了評估，并與 SOTA 稀疏注意力方法和 Full Attention 基線進行了比較。為了確保一致的稀疏性，他們將所有稀疏注意力基線中每個查詢激活的 token 設置為 2560 個 tokens，這對應于 NSA 在處理 32k 序列長度時激活的 token 的平均數量。按照 StreamLLM，此 token 預算包括前 128 個 tokens 和 512 個本地 tokens。

作者從 LongBench 中排除了某些子集，因為它們在所有模型中的得分都較低，可能無法提供有意義的比較。結果如下表 2 所示，NSA 獲得了最高平均分數 0.469，優(yōu)于所有基線，其中比 Full Attention 高出 0.032，比 Exact-Top 高出 0.046。

這一改進源于兩個關鍵創(chuàng)新，分別是（1）原生的稀疏注意力設計，能夠在預訓練期間對稀疏模式進行端到端優(yōu)化，促進稀疏注意力模塊與其他模型組件之間的同步適應；（2）分層稀疏注意力機制實現了局部和全局信息處理之間的平衡。

值得注意的是，NSA 在需要對長上下文進行復雜推理的任務上表現出色，在兩項多跳 QA 任務（HPQ 和 2Wiki）上比 Full Attention 實現了 + 0.087 和 +0.051 的改進，在代碼理解方面超過了基線（LCC 上實現 + 0.069 的改進），并在段落檢索方面優(yōu)于其他方法（PassR-en 上實現 + 0.075 的改進）。

這些結果驗證了 NSA 處理各種長上下文任務中的能力，其原生預訓練的稀疏注意力在學習任務最優(yōu)模式方面提供了額外的助益。

思維鏈推理性能評估

為了評估 NSA 與前沿下游訓練范式的兼容性，作者研究了其通過后訓練獲得思維鏈數學推理能力的能力。

鑒于強化學習在較小模型上的有效性有限，作者從 DeepSeek-R1 進行知識蒸餾，用 100 億個 32k 長度的數學推理軌跡進行監(jiān)督微調（SFT）。

這產生了兩個可比較的模型：Full Attention-R（全注意力基線）和 NSA-R（NSA 的稀疏變體）。

然后，作者在具有挑戰(zhàn)性的美國數學邀請賽（AIME 24）基準上評估這兩個模型。使用采樣溫度為 0.7 和 top-??值為 0.95 的配置，為每個問題生成 16 個回答，并計算平均得分。為了驗證推理深度的影響，作者在兩種生成上下文限制下進行實驗：8k 和 16k token，測量擴展的推理鏈是否提高了準確性。模型預測的示例比較見附錄 A。

如表 3 所示，在 8k 上下文設置下，NSA-R 的準確性顯著高于 Full Attention-R（+0.075），這一優(yōu)勢在 16k 上下文設置下仍然保持（+0.054）。

這些結果驗證了原生稀疏注意力的兩個關鍵優(yōu)勢：

（1）預訓練的稀疏注意力模式能夠高效捕捉對復雜數學推導至關重要的長距離邏輯依賴關系；

（2）該架構的硬件對齊設計保持了足夠的上下文密度，以支持不斷增長的推理深度，而不會出現災難性遺忘。在不同上下文長度下的一致優(yōu)勢證實了稀疏注意力在原生集成到訓練流程中時，對于高級推理任務的可行性。

效率分析

作者在一個 8-GPU A100 系統(tǒng)上評估了 NSA 相對于 Full Attention 的計算效率。

訓練速度

為了確保在相同后端下進行公平的速度比較，作者將基于 Triton 的 NSA 注意力和 Full Attention 實現與基于 Triton 的 FlashAttention-2 進行對比。如圖 6 所示，隨著上下文長度的增加，NSA 實現了越來越大的加速，在 64k 上下文長度下實現了 9.0 倍的前向加速和 6.0 倍的反向加速。

值得注意的是，序列越長，速度優(yōu)勢就越明顯。這種加速源于 DeepSeek 的硬件對齊算法設計，其能最大限度地提高稀疏注意力架構的效率：(1) 分塊式內存訪問模式通過合并加載最大限度地利用了 Tensor Core；(2) 內核中精細的循環(huán)調度消除了冗余的 KV 傳輸。

解碼速度

注意力機制的解碼速度主要由內存訪問瓶頸決定，這與 KV 緩存加載量密切相關。

在每一步解碼過程中，NSA 最多只需要加載個壓縮 token、????′個選定的 token 以及??個鄰近 token，其中??是緩存的序列長度。如表 4 所示，隨著解碼長度的增加，該方法在延遲方面表現出顯著的減少，在 64k 上下文長度下實現了高達 11.6 倍的加速。

這種內存訪問效率的優(yōu)勢在序列長度增加時也會進一步放大。

關于此研究的更多內容，大家可以查看原論文。