太平洋科技要聞

首頁 > 科技要聞 > 科技> 正文

DeepSeek革命性NSA注意力機(jī)制問世！梁文鋒上陣，長文本推理能力飆升

新智元整合編輯：太平洋科技發(fā)布于：2025-02-19 15:47

2月18日，DeepSeek官方再發(fā)布新論文！

在這篇論文中，來自DeepSeek、北大和華盛頓大學(xué)的研究人員提出了一種全新的注意力機(jī)制NSA。

具體來說，它是一種面向硬件且支持原生訓(xùn)練的稀疏注意力機(jī)制（Sparse Attention），專為超快長上下文訓(xùn)練與推理設(shè)計。其核心組成包括——

動態(tài)分層稀疏策略

粗粒度的token壓縮

細(xì)粒度的token選擇

通過針對現(xiàn)代硬件的優(yōu)化設(shè)計，NSA（原生可訓(xùn)練稀疏注意力機(jī)制）在保證性能的同時顯著提升了推理速度，并有效降低了預(yù)訓(xùn)練成本。

在通用基準(zhǔn)測試、長文本處理以及基于指令的推理任務(wù)中，它的表現(xiàn)均能達(dá)到甚至超越傳統(tǒng)全注意力（Full Attention）模型的水平。

值得一提的是，DeepSeek創(chuàng)始人梁文鋒這次不僅親自上陣，并且還親自提交了論文。

論文地址：https://arxiv.org/abs/2502.11089

目前，這成為了該論文的最大亮點之一。

網(wǎng)友們也不忘艾特整天放小道消息的奧特曼：看，DeepSeek又有新論文了。

部署高效稀疏注意力機(jī)制的關(guān)鍵挑戰(zhàn)

隨著序列長度的增加，注意力計算在整體計算成本中變得越來越占主導(dǎo)地位，這對長上下文處理提出了重大挑戰(zhàn)。

尤其在深入推理、代碼庫級別的代碼生成以及多輪自主智能體系統(tǒng)等多種現(xiàn)實應(yīng)用的推動下，這一趨勢尤為明顯。

近期OpenAI的o系列模型、DeepSeek-R1和Gemini 1.5 Pro，在數(shù)千個token的范圍內(nèi)保持連貫的多輪對話，同時執(zhí)行復(fù)雜的長距離依賴推理了。

然而，隨著序列長度的增加，傳統(tǒng)注意力機(jī)制的高計算復(fù)雜度，成為關(guān)鍵的時延瓶頸。

理論估算表明，在處理64k長度上下文時，使用softmax結(jié)構(gòu)的注意力計算占總時延遲的70–80%，更凸顯了對更高效注意力機(jī)制的迫切需求。

部署高效的稀疏注意力機(jī)制必須應(yīng)對兩個關(guān)鍵挑戰(zhàn)：硬件對齊的推理加速和訓(xùn)練感知的算法設(shè)計。

為此，研究團(tuán)隊提出了NSA，一種原生可訓(xùn)練的稀疏注意力架構(gòu)，并集成了分層token建模。同時，還設(shè)計了專用計算內(nèi)核，以最大化實際計算效率。

如圖2所示，NSA通過將鍵值對組織為時間塊（temporal blocks），并通過三條注意力路徑進(jìn)行計算，從而降低每個查詢的計算量。

NSA引入了兩個核心創(chuàng)新點：

硬件對齊系統(tǒng)：優(yōu)化塊級稀疏注意力，使其充分利用Tensor Core并優(yōu)化內(nèi)存訪問，從而實現(xiàn)平衡的算術(shù)強(qiáng)度。

訓(xùn)練感知設(shè)計：通過高效算法和反向傳播運算符實現(xiàn)穩(wěn)定的端到端訓(xùn)練，使NSA能夠同時支持高效推理與完整訓(xùn)練流程。

上圖左：該框架通過三個并行的注意力分支處理輸入序列。對于給定的查詢，前面的鍵和值被處理成壓縮注意力以用于粗粒度模式、選擇注意力以用于重要token塊，以及滑動注意力以用于局部上下文。

上圖右：可視化每個分支產(chǎn)生的不同注意力模式。綠色區(qū)域表示需要計算注意力分?jǐn)?shù)的區(qū)域，而白色區(qū)域表示可以跳過的區(qū)域。

革命性NSA，碾壓全注意力

接下來，研究人員從三個方面評估了NSA，將其與全注意力基線，以及SOTA稀疏注意力方法進(jìn)行比較：

通用基準(zhǔn)性能

長文本基準(zhǔn)性能

思維鏈推理性能

如下圖4所示，NSA和全注意力基線的預(yù)訓(xùn)練損失曲線，呈現(xiàn)出穩(wěn)定和平滑的下降，NSA始終優(yōu)于全注意力模型。

通用基準(zhǔn)評估

研究人員在一套全面的基準(zhǔn)測試中評估了預(yù)訓(xùn)練的NSA和全注意力基線，覆蓋了知識、推理和編碼能力。

結(jié)果如表1所示，盡管NSA具有稀疏性，但它在總體性能上表現(xiàn)優(yōu)越，在9個指標(biāo)中有7個超過了包括全注意力在內(nèi)的所有基線。

這表明，盡管NSA在較短序列上可能未能充分利用其效率優(yōu)勢，但它仍然表現(xiàn)出強(qiáng)大的性能。

值得注意的是，NSA在推理相關(guān)的基準(zhǔn)測試中顯示出顯著的提升（DROP: +0.042，GSM8K: +0.034）。

這進(jìn)一步證明了，預(yù)訓(xùn)練有助于模型發(fā)展專門的注意力機(jī)制。

NSA稀疏注意力預(yù)訓(xùn)練機(jī)制迫使模型集中關(guān)注最重要的信息，可能通過過濾掉來自無關(guān)注意力路徑的噪聲來增強(qiáng)性能。

不同評估中的一致表現(xiàn)，也驗證了NSA作為通用架構(gòu)的穩(wěn)健性。

長文本評估

如圖5所示，NSA在64k上下文中的「大海撈針」測試中，在所有位置上都實現(xiàn)了完美的檢索準(zhǔn)確率。

這一性能得益于層次化的稀疏注意力設(shè)計，該設(shè)計結(jié)合了壓縮token以實現(xiàn)高效的全局上下文掃描，以及用于精確局部信息檢索的選擇token。

粗粒度的壓縮在低計算成本下識別相關(guān)的上下文塊，而對選擇token的token級注意力則確保了關(guān)鍵細(xì)粒度信息的保留。

這種設(shè)計，使NSA能夠同時保持全局感知能力和局部精確性。

此外，研究人員還在LongBench基準(zhǔn)上，對NSA進(jìn)行了評估。

如表2所示，NSA達(dá)下了最高的平均分0.469，優(yōu)于所有基準(zhǔn)（比全注意力高0.032，比Exact-Top高0.046）。

思維鏈推理評估

考慮到RL在小規(guī)模模型上的效果有限，研究人員采用了從DeepSeek-R1進(jìn)行知識蒸餾的方法，使用100億個32k長度的數(shù)學(xué)推理軌跡進(jìn)行監(jiān)督微調(diào)（SFT）。

由此，產(chǎn)生了兩個模型：Full Attention-R（全注意力基準(zhǔn)）和NSA-R（稀疏變體）。

在AIME 24基準(zhǔn)測試上，研究人員評估了這兩個模型。

結(jié)果如下表3所示，在8k上下文設(shè)置下，NSA-R比Full Attention-R獲得了更高的準(zhǔn)確率（高出0.075），這種優(yōu)勢在16k上下文中依然保持（高出0.054）。

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

預(yù)訓(xùn)練的稀疏注意力模式，能夠有效捕獲復(fù)雜數(shù)學(xué)推導(dǎo)中至關(guān)重要的長程邏輯依賴關(guān)系；

架構(gòu)采用硬件對齊設(shè)計，在增加推理深度的同時保持足夠的上下文密度，避免了災(zāi)難性遺忘。

在不同上下文長度下的持續(xù)優(yōu)異表現(xiàn)，證實了當(dāng)稀疏注意力原生集成到訓(xùn)練流程中時，它在高級推理任務(wù)中的可行性。

64k上下文，前向傳播9倍速飆升

研究團(tuán)隊在配備8個A100 GPU的系統(tǒng)上評估了NSA相對于全注意力機(jī)制（Full Attention）的計算效率。

訓(xùn)練速度

為了評估訓(xùn)練速度，研究團(tuán)隊將基于Triton的NSA實現(xiàn)與基于Triton的FlashAttention-2進(jìn)行了比較，以確保在同一后端進(jìn)行公平的速度比較。

實驗結(jié)果表明，隨著上下文長度的增加，NSA實現(xiàn)了逐漸提高的加速，在64k上下文長度下實現(xiàn)了高達(dá)9.0倍的前向加速和6.0倍的反向加速。值得注意的是，隨著序列長度的增加，速度優(yōu)勢變得更加明顯。

這種加速源于NSA的硬件對齊算法設(shè)計，旨在最大限度地提高稀疏注意架構(gòu)的效率：分塊內(nèi)存訪問模式通過合并加載最大限度地利用了Tensor Core。內(nèi)核中精細(xì)的循環(huán)調(diào)度消除了冗余的KV傳輸。

下圖6比較了Triton實現(xiàn)的NSA和全注意力與Triton實現(xiàn)的FlashAttention-2，以確保在同一后端進(jìn)行公平的速度比較。

解碼速度

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

隨著解碼長度的增加，本文中的方法顯著降低了延遲，在64k上下文長度下實現(xiàn)了高達(dá)11.6倍的加速。

如下表4所示NSA在長序列解碼時具有顯著的效率優(yōu)勢。通過減少內(nèi)存訪問量，NSA能夠顯著提高解碼速度，尤其是在處理較長的上下文時。

這種內(nèi)存訪問效率的優(yōu)勢也會隨著序列的增長而放大。

討論

研究團(tuán)隊反思了NSA的開發(fā)過程，并討論了在探索不同稀疏注意力策略過程中獲得的關(guān)鍵見解。

雖然本文中的方法展示了有希望的結(jié)果，但了解替代策略遇到的挑戰(zhàn)并分析注意力分布模式，為未來的研究方向提供了有價值的背景。

研究團(tuán)隊首先考察了促使他們做出設(shè)計選擇的替代token選擇策略所面臨的挑戰(zhàn)，然后通過可視化來深入了解注意力分布模式。

替代token選擇策略

研究人員考察了像ClusterKV這樣的基于聚類的策略。這些方法將來自同一集群的Key和Value存儲在連續(xù)的內(nèi)存區(qū)域中。雖然理論上對于訓(xùn)練和推理是可行的，但它們面臨著三個顯著的挑戰(zhàn)。

首先是動態(tài)聚類機(jī)制引入了不可忽略的計算開銷。其次，算子優(yōu)化困難，集群間的不平衡加劇了這種情況，特別是在混合專家（MoE）系統(tǒng)中，傾斜的專家并行（EP）組執(zhí)行時間導(dǎo)致了持續(xù)的負(fù)載不平衡。

最后是實現(xiàn)約束，需要強(qiáng)制性的周期性重新聚類和chunk-sequential訓(xùn)練協(xié)議。這些因素結(jié)合起來造成了巨大的瓶頸，嚴(yán)重限制了它們在實際部署中的有效性

研究人員還考慮了與NSA不同的分塊key、value選擇策略，例如Quest和InfLLM。

然而，現(xiàn)有方法面臨兩個關(guān)鍵問題。首先由于選擇操作是不可微的，因此基于神經(jīng)網(wǎng)絡(luò)的重要性得分計算依賴于輔助損失，這增加了算子開銷，并且通常會降低模型性能。其次啟發(fā)式無參數(shù)的重要性得分計算策略的召回率較低，導(dǎo)致性能欠佳。

研究人員在一個具有類似架構(gòu)的3B參數(shù)模型上評估了這兩種方法，并將它們的損失曲線與NSA和全注意力進(jìn)行了比較。

對于基于輔助損失的選擇方法，為每個塊引入了額外的查詢和代表性key，以估計塊重要性得分。這些得分由原始查詢和每個塊內(nèi)的key之間的平均注意力得分監(jiān)督。

對于啟發(fā)式無參數(shù)選擇方法，遵循Quest的策略，實現(xiàn)了使用查詢和key塊的坐標(biāo)式min-max乘積的直接選擇，而沒有引入額外的參數(shù)。

研究人員還探索了一種冷啟動訓(xùn)練方法，其中在過渡到啟發(fā)式分塊選擇之前，對初始的1000步應(yīng)用全注意力。這兩種方法的損失都較差。

下圖7展示了在3B參數(shù)模型上，使用全注意力和不同token選擇策略時的訓(xùn)練損失曲線比較。

可視化

研究人員對預(yù)訓(xùn)練的27B全注意力模型的注意力圖進(jìn)行了可視化（如下圖8所示）。

可視化結(jié)果揭示了一個有趣的模式：注意力分?jǐn)?shù)傾向于表現(xiàn)出分塊聚類的特性，即相鄰的key通常顯示出相似的注意力分?jǐn)?shù)。

這一觀察結(jié)果啟發(fā)了NSA的設(shè)計，表明基于空間連續(xù)性選擇key塊可能是一種很有前景的方法。分塊聚類現(xiàn)象表明，序列中相鄰的tokens可能與查詢tokens共享某些語義關(guān)系，盡管這些關(guān)系的確切性質(zhì)需要進(jìn)一步研究。

并促使研究人員探索一種在連續(xù)token塊上操作的稀疏注意力機(jī)制，而不是在單個token上操作，旨在提高計算效率并保留高注意力模式。

正如之前討論的，在設(shè)計NSA架構(gòu)時，研究人員面臨著基于Key-Clustering的策略和其他分塊選擇策略的挑戰(zhàn)，包括計算開銷、算子優(yōu)化困難和召回率低等問題。

結(jié)論

本文中，研究團(tuán)隊提出了NSA，一種面向硬件對齊的稀疏注意力架構(gòu)，用于高效的長上下文建模。

通過將分層token壓縮與塊級token選擇集成到一個可訓(xùn)練的架構(gòu)中，架構(gòu)在保持全注意力性能的同時，實現(xiàn)了加速的訓(xùn)練和推理。

該方法通過以下方式推進(jìn)了當(dāng)前最佳技術(shù)水平：在通用基準(zhǔn)測試中達(dá)到了全注意力的性能，長上下文評估中的建模能力更勝一籌，推理能力得到增強(qiáng)，同時計算延遲顯著降低，實現(xiàn)了可觀的加速。

參考資料：JHNYZ

https://arxiv.org/abs/2502.11089

本文來源：新智元

DeepSeek NSA 稀疏注意力機(jī)制

新智元

原創(chuàng)欄目