DeepSeek 新論文來了！相關訊息剛剛釋出到 𝕏 就吸引了大量使用者點贊、轉發、評論三連。
據介紹，DeepSeek 的這篇新論文提出了一種新的注意力機制 ——NSA。這是一個用於超快長上下文訓練和推斷的本地可訓練的稀疏注意力機制，並且還具有與硬體對齊的特點。
新研究釋出兩個小時，就有近三十萬的瀏覽量。現在看來，DeepSeek 釋出成果，比 OpenAI 關注度都高。
論文標題：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
論文連結：https://arxiv.org/abs/2502.11089
值得一提的是，幻方科技、DeepSeek 創始人梁文鋒也是論文的作者之一。這成了眾多網友討論的話題。
接下來，讓我們看下樑文鋒親自參與的研究，講了什麼內容。
論文概覽
長上下文建模是下一代大型語言模型（LLM）的關鍵能力，這一需求源於多樣化的實際應用，包括深度推理、倉庫級程式碼生成以及多輪自動智慧體系統等。
最近大模型的突破 —— 如 OpenAI 的 o 系列模型、DeepSeek-R1 和 Gemini 1.5 Pro—— 已經能使得模型能夠處理整個程式碼庫、長文件、在數千個 token 上保持連貫的多輪對話，並在長距離依賴關係中進行復雜推理。然而，隨著序列長度的增加，普通注意力機制的高複雜性成為關鍵的延遲瓶頸。理論估計表明，在使用 softmax 架構進行 64k 長度上下文的解碼時，注意力計算佔總延遲的 70-80%，這凸顯了對更高效注意力機制的迫切需求。
實現高效長上下文建模的自然方法是利用 softmax 注意力的固有稀疏性，透過選擇性計算關鍵 query-key 對，可以顯著減少計算開銷，同時保持效能。最近這一路線的進展包括多種策略：KV 快取淘汰方法、塊狀 KV 快取選擇方法以及基於取樣、聚類或雜湊的選擇方法。儘管這些策略前景廣闊，現有的稀疏注意力方法在實際部署中往往表現不佳。許多方法未能實現與其理論增益相媲美的加速；此外，大多數方法主要關注推理階段，缺乏有效的訓練時支援以充分利用注意力的稀疏模式。
為了克服這些限制，部署有效的稀疏注意力必須應對兩個關鍵挑戰：
1、硬體對齊的推理加速：將理論計算減少轉化為實際速度提升，需要在預填充和解碼階段設計硬體友好的演算法，以緩解記憶體訪問和硬體排程瓶頸；
2、訓練感知的演算法設計：透過可訓練的運算子實現端到端計算，以降低訓練成本，同時保持模型效能。
這些要求對於實際應用實現快速長上下文推理或訓練至關重要。在考慮這兩方面時，現有方法仍顯不足。
為了實現更有效和高效的稀疏注意力，DeepSeek 研究人員提出了一種原生可訓練的稀疏注意力架構 NSA，它整合了分層 token 建模。
如圖 2 所示，NSA 透過將鍵和值組織成時間塊（temporal blocks）並透過三條注意力路徑處理它們來減少每查詢計算量：壓縮的粗粒度 token、選擇性保留的細粒度 token 以及用於區域性上下文資訊的滑動視窗。隨後，作者實現了專門的核以最大化其實際效率。
NSA 引入了兩個核心創新以對應於上述關鍵需求：
1、硬體對齊的系統：最佳化塊狀稀疏注意力以利用 Tensor Core 和記憶體訪問，確保算術強度平衡；
2、訓練感知的設計：透過高效演算法和反向運算子實現穩定的端到端訓練。這一最佳化使 NSA 能夠支援高效部署和端到端訓練。
研究透過對現實世界語言語料庫的綜合實驗來評估 NSA。在具有 260B token 的 27B 引數 Transformer 骨幹上進行預訓練，作者評估了 NSA 在通用語言評估、長上下文評估和鏈式推理評估中的表現。作者還進一步比較了在 A100 GPU 上核心速度與最佳化 Triton 實現的比較。實驗結果表明，NSA 實現了與 Full Attention 基線相當或更優的效能，同時優於現有的稀疏注意力方法。
此外，與 Full Attention 相比，NSA 在解碼、前向和後向階段提供了明顯的加速，且加速比隨著序列長度的增加而增加。這些結果驗證了分層稀疏注意力設計有效地平衡了模型能力和計算效率。
方法概覽
本文的技術方法涵蓋演算法設計和核心最佳化。作者首先介紹了方法背景，然後介紹了 NSA 的總體框架以及關鍵演算法元件，最後詳細介紹了針對硬體最佳化的核心設計，以最大限度地提高實際效率。
背景
注意力機制在語言建模中被廣泛使用，其中每個查詢 token q_𝑡計算與所有前面鍵 k_:𝑡的相關性分數，以生成值 v_:𝑡的加權和。從形式上來說，對於長度為𝑡的輸入序列，注意力操作定義如下：
其中 Attn 表示注意力函式。
這裡，𝛼_𝑡,𝑖 表示 q_𝑡 和 k_𝑖 之間的注意力權重，𝑑_𝑘是鍵的特徵維度。隨著序列長度的增加，注意力計算在總計算成本中變得越來越占主導地位，對長上下文處理提出了重大挑戰。
算術強度（Arithmetic Intensity）是計算操作與記憶體訪問的比率，本質上決定了硬體的演算法最佳化。每個 GPU 都有一個由其峰值計算能力和記憶體頻寬決定的臨界算術強度，以這兩個硬體限制的比率計算。對於計算任務，高於此臨界閾值的算術強度將構成計算限制（受限於 GPU FLOPS），而低於此閾值的算術強度將構成記憶體限制（受限於記憶體頻寬）。
具體來說，對於因果自注意力機制，在訓練和預填充階段，批次矩陣乘法和注意力計算表現出高算術強度，使得這些階段在現代加速器上計算受限。相反，自迴歸解碼會受到記憶體頻寬的限制，因為它每次前向傳遞都會生成一個 token，同時需要載入整個鍵值快取，從而導致算術強度較低。這樣就出現了不同的最佳化目標，即減少訓練和預填充期間的計算成本，同時減少解碼期間的記憶體訪問。
總體框架
為了充分利用具有自然稀疏模式的注意力機制的潛力，作者提出將 (1) 式中原始的鍵 - 值對替換成更加緊湊和資訊密集的表徵鍵 - 值對
。具體而言，可將最佳化的注意力輸出的形式定義成：
其中 是基於當前查詢 q_𝑡 和上下文記憶 k_:𝑡 , v_:𝑡 動態構建的。透過設計不同的對映策略，可以得到的不同類別，然後可將它們按以下方式組合起來
如圖 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 的做法，將時間連續的查詢塊載入到 SRAM 中，則會導致記憶體訪問效率低下，因為塊內的查詢可能需要不相交的 KV 塊。
為了解決這個問題，這裡的關鍵最佳化在於不同的查詢分組策略：對於查詢序列上的每個位置，將 GQA 組內的所有查詢頭（它們共享相同的稀疏 KV 塊）載入到 SRAM 中。圖 3 說明了其前向傳遞實現。
該設計能夠 (1) 透過分組共享消除冗餘的 KV 傳輸，以及 (2) 跨 GPU 流式多處理器平衡計算工作負載，由此實現了近乎最佳的算術強度。
NSA 的實驗表現
作者從三個角度對新提出的 NSA 進行了評估：一般基準效能、長上下文基準效能和思維鏈推理效能。
一般基準效能
該團隊在大量基準上對比了 NSA 與 Full Attention 的表現。結果見下表 1。
可以看到，儘管 NSA 比較稀疏，但它實現了卓越的整體效能，在 9 個指標中的 7 個上都優於包括 Full Attention 在內的所有基線。
這表明，儘管 NSA 可能無法充分利用其在較短序列上的效率優勢，但它依然表現出了強勁的效能。值得注意的是，NSA 在推理相關基準測試中表現出了顯著的提升（DROP：+0.042，GSM8K：+0.034），這表明 DeepSeek 的預訓練有助於模型發展出專門的注意力機制。透過過濾掉不相關的注意力路徑中的噪音，這種稀疏注意力預訓練機制可迫使模型專注於最重要的資訊，有可能提高效能。在不同評估中的一致表現也證明了 NSA 作為通用架構的穩健性。
長上下文基準效能
下圖 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，優於所有基線，其中比 Full Attention 高出 0.032，比 Exact-Top 高出 0.046。
這一改進源於兩個關鍵創新，分別是（1）原生的稀疏注意力設計，能夠在預訓練期間對稀疏模式進行端到端最佳化，促進稀疏注意力模組與其他模型元件之間的同步適應；（2）分層稀疏注意力機制實現了區域性和全域性資訊處理之間的平衡。
值得注意的是，NSA 在需要對長上下文進行復雜推理的任務上表現出色，在兩項多跳 QA 任務（HPQ 和 2Wiki）上比 Full Attention 實現了 + 0.087 和 +0.051 的改進，在程式碼理解方面超過了基線（LCC 上實現 + 0.069 的改進），並在段落檢索方面優於其他方法（PassR-en 上實現 + 0.075 的改進）。
這些結果驗證了 NSA 處理各種長上下文任務中的能力，其原生預訓練的稀疏注意力在學習任務最優模式方面提供了額外的助益。
思維鏈推理效能評估
為了評估 NSA 與前沿下游訓練正規化的相容性，作者研究了其透過後訓練獲得思維鏈數學推理能力的能力。
鑑於強化學習在較小模型上的有效性有限，作者從 DeepSeek-R1 進行知識蒸餾，用 100 億個 32k 長度的數學推理軌跡進行監督微調（SFT）。
這產生了兩個可比較的模型：Full Attention-R（全注意力基線）和 NSA-R（NSA 的稀疏變體）。
然後，作者在具有挑戰性的美國數學邀請賽（AIME 24）基準上評估這兩個模型。使用取樣溫度為 0.7 和 top-𝑝值為 0.95 的配置，為每個問題生成 16 個回答，並計算平均得分。為了驗證推理深度的影響，作者在兩種生成上下文限制下進行實驗：8k 和 16k token，測量擴充套件的推理鏈是否提高了準確性。模型預測的示例比較見附錄 A。
如表 3 所示，在 8k 上下文設定下，NSA-R 的準確性顯著高於 Full Attention-R（+0.075），這一優勢在 16k 上下文設定下仍然保持（+0.054）。
這些結果驗證了原生稀疏注意力的兩個關鍵優勢：
（1）預訓練的稀疏注意力模式能夠高效捕捉對複雜數學推導至關重要的長距離邏輯依賴關係；
（2）該架構的硬體對齊設計保持了足夠的上下文密度，以支援不斷增長的推理深度，而不會出現災難性遺忘。在不同上下文長度下的一致優勢證實了稀疏注意力在原生整合到訓練流程中時，對於高階推理任務的可行性。
效率分析
作者在一個 8-GPU A100 系統上評估了 NSA 相對於 Full Attention 的計算效率。
訓練速度
為了確保在相同後端下進行公平的速度比較，作者將基於 Triton 的 NSA 注意力和 Full Attention 實現與基於 Triton 的 FlashAttention-2 進行對比。如圖 6 所示，隨著上下文長度的增加，NSA 實現了越來越大的加速，在 64k 上下文長度下實現了 9.0 倍的前向加速和 6.0 倍的反向加速。
值得注意的是，序列越長，速度優勢就越明顯。這種加速源於 DeepSeek 的硬體對齊演算法設計，其能最大限度地提高稀疏注意力架構的效率：(1) 分塊式記憶體訪問模式透過合併載入最大限度地利用了 Tensor Core；(2) 核心中精細的迴圈排程消除了冗餘的 KV 傳輸。
解碼速度
注意力機制的解碼速度主要由記憶體訪問瓶頸決定，這與 KV 快取載入量密切相關。
在每一步解碼過程中，NSA 最多隻需要載入
個壓縮 token、𝑛𝑙′個選定的 token 以及𝑤個鄰近 token，其中𝑠是快取的序列長度。如表 4 所示，隨著解碼長度的增加，該方法在延遲方面表現出顯著的減少，在 64k 上下文長度下實現了高達 11.6 倍的加速。
這種記憶體訪問效率的優勢在序列長度增加時也會進一步放大。
關於此研究的更多內容，大家可以檢視原論文。