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

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

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

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

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

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

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

論文概覽

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

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

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

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

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

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

這些要求對于實際應(yīng)用實現(xiàn)快速長上下文推理或訓(xùn)練至關(guān)重要。在考慮這兩方面時，現(xiàn)有方法仍顯不足。

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

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

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

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

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

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

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

方法概覽

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

背景

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

其中 Attn 表示注意力函數(shù)。

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

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

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

總體框架

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

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

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

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

接下來，DeepSeek 還介紹了重新映射策略的具體設(shè)計，涵蓋 token 壓縮、token 選取和滑動窗口。詳細的算法設(shè)計見原論文。下面來看看 NSA 為何具有 FlashAttention 相當?shù)乃俣取?/span>

核設(shè)計

為了在訓(xùn)練和預(yù)填充期間實現(xiàn) FlashAttention 級別的加速，作者基于 Triton 實現(xiàn)了硬件對齊的稀疏注意力內(nèi)核。

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

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

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

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

NSA 的實驗表現(xiàn)

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

一般基準性能

該團隊在大量基準上對比了 NSA 與 Full Attention 的表現(xiàn)。結(jié)果見下表 1。

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

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

長上下文基準性能

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

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

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

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

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

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

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

思維鏈推理性能評估

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

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

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

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

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

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

（1）預(yù)訓(xùn)練的稀疏注意力模式能夠高效捕捉對復(fù)雜數(shù)學推導(dǎo)至關(guān)重要的長距離邏輯依賴關(guān)系；

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

效率分析

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

訓(xùn)練速度

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

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

解碼速度

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

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

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

關(guān)于此研究的更多內(nèi)容，大家可以查看原論文。