7 月 30 日,ACL(國際計(jì)算語言學(xué)年會)公布了 2025 年的獲獎?wù)撐摹A钊梭@喜的是,這些論文里的中國作者比例超過 51%,排在第二的美國僅為 14%。
其中,一篇由 DeepSeek 梁文鋒作為通訊作者、與北京大學(xué)等聯(lián)合發(fā)表的論文不僅拿下 Best Paper 獎,相關(guān)成果也引發(fā)熱議。
現(xiàn)場講座中,該論文的第一作者袁境陽透露,這項(xiàng)技術(shù)可以把上下文長度擴(kuò)展到 100 萬 tokens,并將應(yīng)用在他們的下一個(gè)前沿模型中。據(jù)了解,袁境陽當(dāng)時(shí)寫這篇論文時(shí)還只是 Deepseek 的實(shí)習(xí)生。
引入兩大核心技術(shù)創(chuàng)新
長上下文建模對于下一代語言模型至關(guān)重要,但標(biāo)準(zhǔn)注意力機(jī)制的高計(jì)算成本帶來了顯著的計(jì)算挑戰(zhàn)。隨著序列長度的增加,延遲瓶頸問題愈發(fā)凸顯。理論估算表明,在解碼 64k 長度的上下文時(shí),采用 softmax 架構(gòu)的注意力計(jì)算占總延遲的 70%–80%,這凸顯了對更高效注意力機(jī)制的迫切需求。
為解決這些局限性,有效的稀疏注意力機(jī)制在實(shí)際應(yīng)用中必須應(yīng)對兩項(xiàng)關(guān)鍵挑戰(zhàn):與硬件適配的推理加速,要將理論上的計(jì)算量減少轉(zhuǎn)化為實(shí)際的速度提升,就需要在預(yù)填充和解碼階段都采用硬件友好型的算法設(shè)計(jì),以緩解內(nèi)存訪問和硬件調(diào)度方面的瓶頸;兼顧訓(xùn)練的算法設(shè)計(jì),通過可訓(xùn)練算子實(shí)現(xiàn)端到端計(jì)算,在維持模型性能的同時(shí)降低訓(xùn)練成本。
綜合考慮這兩個(gè)方面,現(xiàn)有方法仍存在明顯差距。該團(tuán)隊(duì)認(rèn)為,稀疏注意力為在保持模型能力的同時(shí)提高效率提供了一個(gè)很有前景的方向。
在獲獎?wù)撐闹校麄兲岢隽?NSA,這是一種可原生訓(xùn)練的稀疏注意力(Natively trainable Sparse Attention)機(jī)制。它將算法創(chuàng)新與硬件對齊優(yōu)化相結(jié)合,以實(shí)現(xiàn)高效的長上下文建模。據(jù)介紹,NSA 采用動態(tài)分層稀疏策略,結(jié)合粗粒度的 token 壓縮和細(xì)粒度的 token 選擇,以同時(shí)保留全局上下文感知和局部精度。
具體來說,NSA 引入了兩項(xiàng)核心創(chuàng)新。
通過算術(shù)強(qiáng)度平衡的算法設(shè)計(jì)實(shí)現(xiàn)了顯著的加速,并針對現(xiàn)代硬件進(jìn)行了實(shí)現(xiàn)優(yōu)化:優(yōu)化塊式稀疏注意力,以提高張量核利用率和內(nèi)存訪問,確保均衡的算術(shù)強(qiáng)度。
通過高效算法和反向算子實(shí)現(xiàn)穩(wěn)定的端到端訓(xùn)練,在不犧牲模型性能的情況下減少了預(yù)訓(xùn)練計(jì)算量。
上下文處理速度狂飆,
準(zhǔn)確率堪稱“完美”
在真實(shí)世界語言語料庫上進(jìn)行綜合實(shí)驗(yàn)評估后,NSA 由于稀疏性過濾掉更多噪聲,在基準(zhǔn)測試中產(chǎn)生更好的準(zhǔn)確率。據(jù)悉,該團(tuán)隊(duì)在一個(gè)擁有 270 億參數(shù)的 Transformer 骨干網(wǎng)絡(luò)(其中激活參數(shù)為 30 億)上,使用 2600 億個(gè) token 進(jìn)行預(yù)訓(xùn)練,并從通用語言評估、長上下文評估和思維鏈推理評估三個(gè)方面評估了 NSA 的性能,還在 A100 GPU 上將其內(nèi)核速度與經(jīng)過優(yōu)化的 Triton 實(shí)現(xiàn)作了進(jìn)一步比較。
實(shí)驗(yàn)結(jié)果表明,NSA 的整體性能與全注意力模型相當(dāng)甚至更優(yōu)。在 9 項(xiàng)指標(biāo)中的 7 項(xiàng)上,NSA 均超過了包括全注意力模型在內(nèi)的所有基線。這表明,盡管 NSA 在較短序列上可能無法充分發(fā)揮其效率優(yōu)勢,但它展現(xiàn)出了強(qiáng)勁的性能。
值得注意的是,NSA 在推理相關(guān)的基準(zhǔn)測試中取得了顯著提升(DROP:+0.042,GSM8K:+0.034),這說明該團(tuán)隊(duì)的預(yù)訓(xùn)練有助于模型發(fā)展出專門的注意力機(jī)制。這種稀疏注意力預(yù)訓(xùn)練機(jī)制迫使模型聚焦于最重要的信息,通過過濾無關(guān)注意力路徑中的噪聲,可能會提升性能。在各類評估中表現(xiàn)出的一致性,也驗(yàn)證了 NSA 作為通用架構(gòu)的穩(wěn)健性。
在 64k 上下文的“大海撈針”測試中,NSA 在所有位置都實(shí)現(xiàn)了完美的檢索準(zhǔn)確率。此外,與全注意力相比,NSA 在解碼、前向傳播和反向傳播方面都實(shí)現(xiàn)了顯著的速度提升,且序列越長,提速比例越大。
據(jù)該團(tuán)隊(duì)稱,這一性能正是得益于其分層稀疏注意力設(shè)計(jì),該設(shè)計(jì)結(jié)合了用于高效全局上下文掃描的 token 壓縮和用于精確局部信息檢索的 token 選擇。粗粒度的 token 壓縮以較低的計(jì)算成本識別相關(guān)的上下文塊,而對 token 選擇的標(biāo)記級注意力則確保保留關(guān)鍵的細(xì)粒度信息。
同時(shí),NSA 優(yōu)于多種現(xiàn)有的稀疏注意力方法,包括 H2O、infLLM、Quest 以及 Exact-Top。
值得注意的是,NSA 在需要對長上下文進(jìn)行復(fù)雜推理的任務(wù)上表現(xiàn)出色,在多跳問答任務(wù)(HPQ 和 2Wiki)上比全注意力模型分別提升 0.087 和 0.051,在代碼理解任務(wù)(LCC)上超出基線模型 0.069,在段落檢索任務(wù)(PassR-en)上優(yōu)于其他方法 0.075。這些結(jié)果也驗(yàn)證了 NSA 處理各種長上下文挑戰(zhàn)的能力,其原生預(yù)訓(xùn)練的稀疏注意力在學(xué)習(xí)任務(wù)最優(yōu)模式方面帶來了額外優(yōu)勢。
為評估 NSA 與先進(jìn)下游訓(xùn)練范式的兼容性,該團(tuán)隊(duì)研究了其通過后期訓(xùn)練獲得思維鏈數(shù)學(xué)推理能力的潛力。鑒于強(qiáng)化學(xué)習(xí)在較小規(guī)模模型上的效果有限,其采用來自 DeepSeek-R1 的知識蒸餾,使用 100 億個(gè) 32k 長度的數(shù)學(xué)推理軌跡進(jìn)行有監(jiān)督微調(diào)(SFT)。這產(chǎn)生了兩個(gè)可比較的模型:全注意力 - R(全注意力基線模型)和 NSA-R(稀疏變體)。
接著,他們在具有挑戰(zhàn)性的美國數(shù)學(xué)邀請賽(AIME 24)基準(zhǔn)上對這兩個(gè)模型進(jìn)行了評估,使用 0.7 的采樣溫度和 0.95 的核采樣值,為每個(gè)問題生成 16 個(gè)回答并取平均分。并且,為驗(yàn)證推理深度的影響,他們在兩種生成上下文序列下進(jìn)行了實(shí)驗(yàn)。結(jié)果顯示,NSA-R 在 8k 和 16k 序列長度下的表現(xiàn)均優(yōu)于全注意力 - R。
這些結(jié)果驗(yàn)證了原生稀疏注意力的兩項(xiàng)關(guān)鍵優(yōu)勢:(1)預(yù)訓(xùn)練的稀疏注意力模式能夠高效捕捉復(fù)雜數(shù)學(xué)推導(dǎo)所必需的長程邏輯依賴關(guān)系;(2)我們架構(gòu)的硬件對齊設(shè)計(jì)保持了足夠的上下文密度,以支持推理深度的增加,同時(shí)避免災(zāi)難性遺忘。在不同上下文長度下的持續(xù)優(yōu)異表現(xiàn)證實(shí),當(dāng)稀疏注意力被原生整合到訓(xùn)練流程中時(shí),其在高級推理任務(wù)中具有切實(shí)可行性。
計(jì)算效率方面,該團(tuán)隊(duì)將基于 Triton 實(shí)現(xiàn)的 NSA 注意力機(jī)制和全注意力機(jī)制,與基于 Triton 的 FlashAttention-2 在 8-GPU A100 系統(tǒng)進(jìn)行了比較,以確保在相同后端下進(jìn)行公平的速度對比。
結(jié)果表明,隨著上下文長度的增加,NSA 實(shí)現(xiàn)了越來越顯著的速度提升。在 64k 上下文長度下,前向速度提升高達(dá) 9.0 倍,反向速度提升高達(dá) 6.0 倍。值得注意的是,序列越長,速度優(yōu)勢就越明顯。隨著解碼長度的增加,NSA 的方法延遲顯著降低,在 64k 上下文長度下提速高達(dá) 11.6 倍,且這種內(nèi)存訪問效率方面的優(yōu)勢也會隨著序列變長而進(jìn)一步擴(kuò)大。
值得一提的是,這篇論文早在今年 2 月就對外公布,而相關(guān)研究成果至今還沒有出現(xiàn)在任何 DeepSeek 模型中。不過,根據(jù)論文一作袁境陽的說法,DeepSeek 下一代模型就將應(yīng)用這項(xiàng)技術(shù),這也讓許多網(wǎng)友對 DeepSeek V4 的發(fā)布更加期待,畢竟其與 DeepSeek R2 的發(fā)布計(jì)劃似乎也有很大關(guān)聯(lián)。
早在今年 4 月,就有“DeepSeek R2 提前泄露”的傳言在 AI 圈刷屏。源頭是來自 Hugging Face CEO 發(fā)布的一條耐人尋味的帖子,配圖是 DeepSeek 在 Hugging Face 的倉庫鏈接,接著引發(fā)不少關(guān)于 R2 發(fā)布時(shí)間和技術(shù)細(xì)節(jié)的各類傳播。但對此,DeepSeek 官方一直未作出回應(yīng)。
前不久,有外媒報(bào)道稱,DeepSeek R2 可能繼續(xù)推遲。遲遲未發(fā)布的內(nèi)部原因是 DeepSeek 創(chuàng)始人梁文鋒對該模型當(dāng)前的性能不滿意,工程師團(tuán)隊(duì)仍在優(yōu)化和打磨。與此同時(shí),也有人這樣推測:R2 好歹要等 V4 出來再說,V3 可能已經(jīng)到達(dá)極限了。