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

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

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

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

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

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

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

粗粒度的token壓縮

細(xì)粒度的token選擇

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

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

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

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

目前，這成為了該論文的最大亮點(diǎn)之一。

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

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

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

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

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

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

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

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

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

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

NSA引入了兩個(gè)核心創(chuàng)新點(diǎn)：

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

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

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

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

革命性NSA，碾壓全注意力

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

通用基準(zhǔn)性能

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

思維鏈推理性能

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

通用基準(zhǔn)評(píng)估

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

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

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

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

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

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

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

長(zhǎng)文本評(píng)估

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

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

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

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

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

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

思維鏈推理評(píng)估

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

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

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

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

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

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

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

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

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

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

訓(xùn)練速度

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

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

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

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

解碼速度

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

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

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

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

討論

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

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

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

替代token選擇策略

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

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

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

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

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

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

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

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

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

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

可視化

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

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

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

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

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

結(jié)論

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

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

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

參考資料：JHNYZ

https://arxiv.org/abs/2502.11089

本文來源：新智元

網(wǎng)友評(píng)論

聚超值•精選

推薦 手機(jī) 筆記本 影像 硬件 家居 商用 企業(yè) 出行 未來
  • 漢美馳李梟雄：擁抱AI浪潮，智能家居需求暴增勢(shì)不可擋
    漢美馳李梟雄：擁抱AI浪潮，智能家居需求暴增勢(shì)不可擋

    作為美國(guó)百年家電品牌，HamiltonBeach漢美馳在美國(guó)市場(chǎng)處于行業(yè)領(lǐng)導(dǎo)地位，在本屆AWE展會(huì)上，漢美馳與太平洋網(wǎng)絡(luò)合作，展出了他們的三個(gè)主打品類：破壁機(jī)、慢燉鍋和空氣炸鍋，其中，破壁機(jī)屬于攪拌類小家電，漢美馳是這個(gè)品類的開創(chuàng)者；慢燉鍋是美國(guó)家庭的必需品，就像我們中國(guó)家庭的電飯煲一樣，漢美馳慢燉鍋累計(jì)在美國(guó)的銷量超過3000萬臺(tái)，是這個(gè)品類的領(lǐng)導(dǎo)品牌；漢美馳的這款HALO空氣炸鍋剛剛獲得了全球頂級(jí)設(shè)計(jì)大獎(jiǎng)——iF設(shè)計(jì)獎(jiǎng)。 今年，漢美馳在國(guó)內(nèi)市場(chǎng)的動(dòng)作很多，包括：推出了家電行業(yè)第一款應(yīng)用chatGPT的AI牛排機(jī)，全球首發(fā)煙灶產(chǎn)品，全球首發(fā)中式廚電產(chǎn)品，自建抖音、淘寶直播間，與頭部主播烈兒寶貝等合作……這些經(jīng)營(yíng)動(dòng)作的背后，漢美馳有怎樣的戰(zhàn)略規(guī)劃？他們對(duì)中國(guó)市場(chǎng)有些什么樣的判斷？他們的優(yōu)勢(shì)有哪些？請(qǐng)看PConline獨(dú)家專訪漢美馳中國(guó)營(yíng)銷中心總經(jīng)理李梟雄先生。

    呼倫 呼倫 2023-05-05 00:03
  • 極氪“流血”上市，最急的是埃安、哪吒?jìng)?

    在經(jīng)歷了2018-2020年“蔚小理”的上市潮后，隨著國(guó)內(nèi)新能源汽車市場(chǎng)競(jìng)爭(zhēng)日益激烈，以及全球EV需求放緩，IPO的大門正在向造車新勢(shì)力們緩緩關(guān)閉。極氪的成功上市既是對(duì)新勢(shì)力的一次激勵(lì)，也是一次警示——后來者必須面對(duì)越來越嚴(yán)苛的上市條件。留給哪吒汽車、廣汽埃安們的機(jī)會(huì)可能不多了。

    劉凡 劉凡 2024-05-11 11:26
    極氪   吉利   哪吒   埃安
    • 二維碼 回到頂部