剛剛，DeepSeek發新成果，梁文鋒親自參與，實習生挑大樑，顯著加速AI訓練推理

智東西2月18日報導，今天下午，DeepSeek團隊發佈一篇新論文，介紹了一種改進的稀疏注意力機制NSA，可用於超快速的長上下文訓練與推理。NSA以性價比極高的方式，罕見地在訓練階段應用稀疏性，在訓推場景中均實現速度的明顯提升，特別是在解碼階段實現了高達11.6倍的提升。

讓人眼前一亮的是，DeepSeek創始人兼CEO梁文鋒這次出現在了合著名單之中，在作者排名中位列倒數第二。這意味著他作為項目管理者，參與了一線的研究工作。另外，這篇論文的第一作者Jingyang Yuan是在實習期間完成的這項研究。

據DeepSeek介紹，NSA具有三大核心組件：動態分層稀疏策略、粗粒度token壓縮、精粒度token選擇。通過三大組件的配合，既提升了效率，也保留了模型對全局長上下文的感知能力和局部精確性。

這一機制專門針現代硬件進行優化設計，原生支持模型訓練，在加速推理的同時降低預訓練成本，對性能也無明顯影響。採用NSA機制的模型在通用基準、長上下文任務和基於指令的推理上，與全注意力模型相當或表現更優。

在8卡A100計算集群上，NSA的前向傳播和反向傳播速度分別比全注意力快9倍和6倍，由於減少了內存訪問量，NSA在長序列解碼時相較於全注意力模型速度顯著提升。

01.現有稀疏注意力機制存在明顯缺陷，DeepSeek希望填補空白

長文本建模是下一代語言模型的關鍵能力，但傳統注意力機制的高複雜度限制了其在長序列上的應用。

例如，在解碼64k長度的上下文時，注意力計算佔據了總延遲的70%至80%。因此，稀疏注意力機制應運而生，通過選擇性計算關鍵的查詢鍵對來減少計算開銷。

然而，儘管許多稀疏注意力方法在理論上減少了計算複雜度，但這些方法在實際推理中未能顯著降低延遲。

一些方法僅在自回歸解碼階段應用稀疏性，而預填充階段仍需進行密集計算（如H2O）；另一些方法僅關注預填充階段的稀疏性（如MInference），導致在某些工作負載下無法實現全階段加速。

還有部分稀疏方法無法適應現代高效的解碼架構（如MQA和GQA），導致KV緩存訪問量仍然較高，無法充分發揮稀疏性優勢。

此外，現有的稀疏注意力方法大多僅在推理階段應用稀疏性，缺乏對訓練階段的支持。

NSA旨在通過針對硬件特點的推理加速和適用於訓練的算法設計，填補這一空白。DeepSeek推出NSA主要希望解決兩大問題：

一是事後稀疏化導致性能退化，如預訓練模型的檢索頭易被剪枝；

二是現有稀疏方法難以應對長序列訓練的效率需求。現有方法存在非可訓練組件和低效反向傳播等問題，阻礙了高效訓練和長上下文模型的發展。

02.軟硬件協同深度優化，無限逼近計算強度最優解

NSA的核心思想是通過動態分層稀疏策略，結合粗粒度的token壓縮和細粒度的token選擇，以保留全局上下文感知能力和局部精確性。

下方是NSA架構的概覽，左側，NSA將輸入序列通過三個並行的注意力分支處理：壓縮注意力（compressed attention）、選擇性注意力（selected attention）和滑動窗口注意力（sliding attention）。右側是對每個分支產生的不同注意力模式的可視化。綠色區域表示需要計算注意力分數的區域，而白色區域表示可以跳過的區域。

其中，壓縮注意力通過將鍵（key）和值（value）聚合成塊（block）級表示來捕捉粗粒度的語義信息。這種壓縮表示能夠捕捉更粗粒度的高層語義信息，並減輕注意力計算的負擔。

不過，僅使用壓縮後的鍵（key）和值（value）可能會丟失重要的細粒度信息，DeepSeek引入了選擇性注意力，通過塊選擇機制保留重要的細粒度信息。

他們給每個塊分配了重要性分數，根據塊的重要性分數選擇排名前n的塊，並將這些塊中的標記用於注意力計算。這種方法在保留關鍵信息的同時，顯著降低了計算負擔。

在注意力機制中，局部模式通常會快速適應並主導學習過程，可能會阻礙模型從壓縮和選擇token中有效學習。滑動窗口注意力可以回應這一問題，通過專注於局部上下文信息，防止模型過度依賴局部模式。

為了實現高效的稀疏注意力計算，NSA還針對現代硬件進行了優化。

具體來看，DeepSeek在Triton上實現了硬件對齊的稀疏注意力內核。鑒於多頭自注意力（MHA）內存密集且解碼效率低，他們專注於共享KV緩存的架構，如分組查詢注意力（GQA）和多查詢注意力（MQA），這些架構與當前最先進的LLMs一致。

DeepSeek的關鍵優化策略是採用不同的查詢分組策略，通過以下關鍵特性實現了近乎最優的計算強度平衡：

1、以組為中心的數據加載：在每個內循環中，加載組內所有頭的查詢及其共享的稀疏KV塊索引。

2、共享KV加載：在內循環中，連續加載KV塊以最小化內存加載。

3、網格循環調度：由於內循環長度在不同查詢塊中幾乎相同，將查詢/輸出循環放在Triton的網格調度器中，以簡化和優化內核。

03.超越多款基線模型，訓練提速6-9倍、推理最高提速11.6倍

為測試NSA機制在實際訓練、推理場景中的表現，DeepSeek用當前最先進的LLM的常見實踐，使用了一個結合分組查詢注意力（GQA）和混合專家（MoE）的骨幹架構作為樣本模型。這一模型總參數量為27B，其中3B為活躍參數。

在這一模型的基礎上，DeepSeek使用了NSA、全注意力以及其它注意力機制，並進行了評估。

在多個通用基準測試中，採用NSA的模型儘管具有稀疏性，但其總體性能優於所有基線模型，包括全注意力模型，在9項指標中有7項表現最佳。

NSA在較短序列上可能未能充分發揮其效率優勢，但其性能依然強勁。值得注意的是，NSA在推理相關基準測試中取得了顯著提升，這表明NSA的預訓練機制有助於模型開發專門的注意力機制，能迫使模型專注於最重要的信息，通過過濾掉無關的注意力路徑中的噪聲，潛在地提升了性能。

在長上下文任務中，NSA在64k上下文的「大海撈針（neddle in a haystack）」測試中實現了超強的的檢索精度。這得益於其分層稀疏注意力設計，通過粗粒度的壓縮token實現高效的全局上下文掃瞄，並通過細粒度的選擇標記保留關鍵信息，從而在全局感知和局部精確性之間取得平衡。

在LongBench上，NSA在多跳QA任務和代碼理解任務中表現優於所有基線，還顯示出在複雜長文本推理任務上的優勢。這些結果表明，NSA的原生稀疏注意力機制不僅提升了模型性能，還為長文本任務提供了更優的解決方案。

NSA機制還能與推理模型進行結合，適配前沿的後訓練方式。DeepSeek使用了從DeepSeek-R1蒸餾獲得的知識和監督微調（SFT）的方式，使採用NSA的模型在32k長度的數學推理任務上獲得鏈式數學推理能力。

實驗中，NSA-R（稀疏注意力變體）和全注意力-R（基線模型）在具有挑戰性的AIME 24基準測試上進行了對比。結果顯示，NSA-R在8k和16k上下文設置下均顯著優於全注意力-R（分別高出0.075和0.054），驗證了其在複雜推理任務中的優勢。

DeepSeek還在8-GPU A100系統上，對NSA的計算效率與全注意力機制進行了對比。

在訓練速度方面，隨著上下文長度的增加，NSA的加速效果愈發顯著。在64k上下文長度時，NSA的前向傳播速度提升了9倍，反向傳播速度提升了6倍。

這種加速主要得益於NSA的硬件對齊設計：塊狀的內存訪問模式通過合併加載最大化了Tensor Core的利用率，內核中精細的循環調度消除了冗餘的KV傳輸。

在解碼速度方面，注意力機制的解碼速度主要受限於KV緩存加載的內存瓶頸。隨著解碼長度的增加，NSA的延遲顯著降低，在64k上下文長度時實現了高達11.6倍的速度提升。這種內存訪問效率的優勢隨著序列長度的增加而更加明顯。

04.結語：DeepSeek持續給開源AI驚喜

儘管NSA取得了顯著的成果，但DeepSeek研究團隊也指出了一些可能的改進方向。例如，進一步優化稀疏注意力模式的學習過程，以及探索更高效的硬件實現方式。

正如DeepSeek之前發佈的所有技術報告那樣，這篇詳解NSA機制的論文內容詳實，對NSA機制中涉及的技術細節闡釋清晰，可操作性強，是DeepSeek給開源AI研究貢獻的最新成果。

本文來自微信公眾號「智東西」（ID：zhidxcom），作者：陳駿達，編輯：Panken，36氪經授權發佈。