剛剛，DeepSeek又發(fā)高質(zhì)量論文，梁文鋒合著！

這篇論文的第一作者在實(shí)習(xí)期間完成了研究。

原標(biāo)題：剛剛，DeepSeek又發(fā)高質(zhì)量論文，梁文鋒合著！
文章來源：智東西
內(nèi)容字?jǐn)?shù)：6730字

DeepSeek發(fā)布新型稀疏注意力機(jī)制NSA：長上下文訓(xùn)練推理提速利器

近日，DeepSeek團(tuán)隊發(fā)布論文，介紹了一種名為NSA（Novel Sparse Attention）的改進(jìn)型稀疏注意力機(jī)制，該機(jī)制在長上下文訓(xùn)練和推理中展現(xiàn)出顯著的速度提升，尤其在解碼階段最高可達(dá)11.6倍。論文第一作者Jingyang Yuan在實(shí)習(xí)期間完成了這項研究，DeepSeek創(chuàng)始人兼CEO梁文鋒也參與其中。

1. 現(xiàn)有稀疏注意力機(jī)制的缺陷與NSA的優(yōu)勢

傳統(tǒng)的注意力機(jī)制在處理長序列時計算復(fù)雜度極高，限制了長文本建模的發(fā)展。現(xiàn)有的稀疏注意力機(jī)制雖然在理論上降低了計算復(fù)雜度，但在實(shí)際應(yīng)用中往往未能有效降低延遲，部分原因在于它們只在推理階段或預(yù)填充階段應(yīng)用稀疏性，缺乏對訓(xùn)練階段的支持，且難以適應(yīng)現(xiàn)代高效的解碼架構(gòu)。NSA則旨在解決這些問題，通過在訓(xùn)練和推理階段都應(yīng)用稀疏性，并針對現(xiàn)代硬件進(jìn)行優(yōu)化，實(shí)現(xiàn)全階段加速。

2. NSA的三大核心組件及軟硬件協(xié)同優(yōu)化

NSA的核心思想是通過動態(tài)分層稀疏策略，結(jié)合粗粒度token壓縮和細(xì)粒度token選擇，在提升效率的同時保留模型對全局長上下文感知能力和局部精確性。其三大核心組件包括：

1. 動態(tài)分層稀疏策略：根據(jù)不同層級的需求動態(tài)調(diào)整稀疏程度。
2. 粗粒度token壓縮：將鍵值對聚合成塊級表示，捕捉高層語義信息。
3. 精粒度token選擇：根據(jù)重要性分?jǐn)?shù)選擇關(guān)鍵的細(xì)粒度信息。

此外，NSA在Triton上實(shí)現(xiàn)了硬件對齊的稀疏注意力內(nèi)核，并針對共享KV緩存的架構(gòu)（如GQA和MQA）進(jìn)行了優(yōu)化，通過以組為中心的數(shù)據(jù)加載、共享KV加載和網(wǎng)格循環(huán)調(diào)度等策略，實(shí)現(xiàn)了近乎最優(yōu)的計算強(qiáng)度平衡。

3. NSA的性能表現(xiàn)：訓(xùn)練提速6-9倍，推理提速高達(dá)11.6倍

DeepSeek使用一個27B參數(shù)的LLM模型（其中3B為活躍參數(shù)）對NSA進(jìn)行了評估。結(jié)果顯示，在8卡A100計算集群上，NSA的前向傳播和反向傳播速度分別比全注意力快9倍和6倍。在長序列解碼時，NSA相較于全注意力模型速度顯著提升，最高可達(dá)11.6倍。在通用基準(zhǔn)、長上下文任務(wù)和基于指令的推理上，采用NSA機(jī)制的模型與全注意力模型相當(dāng)或表現(xiàn)更優(yōu)。

4. 總結(jié)

NSA作為一種新型的稀疏注意力機(jī)制，通過軟硬件協(xié)同深度優(yōu)化，有效解決了現(xiàn)有稀疏注意力機(jī)制的缺陷，在長上下文訓(xùn)練和推理中實(shí)現(xiàn)了顯著的加速效果，為下一代大語言模型的發(fā)展提供了新的可能性。其在訓(xùn)練階段的應(yīng)用以及對現(xiàn)代硬件的優(yōu)化，使其成為一個極具性價比的解決方案。

聯(lián)系作者

文章來源：智東西
作者微信：
作者簡介：智能產(chǎn)業(yè)新媒體！智東西專注報道人工智能主導(dǎo)的前沿技術(shù)發(fā)展，和技術(shù)應(yīng)用帶來的千行百業(yè)產(chǎn)業(yè)升級。聚焦智能變革，服務(wù)產(chǎn)業(yè)升級。