DeepSeek官推發(fā)布了最新技術(shù)成果NSA：一種面向硬件且支持原生訓(xùn)練的稀疏注意力機(jī)制，專為超快長(zhǎng)上下文訓(xùn)練與推理設(shè)計(jì)。

NSA的核心組成：

• 動(dòng)態(tài)分層稀疏策略
• 粗粒度的token壓縮
• 細(xì)粒度的token選擇

?? NSA針對(duì)現(xiàn)代硬件進(jìn)行了優(yōu)化設(shè)計(jì)，顯著提升了推理速度，并有效降低了預(yù)訓(xùn)練成本——同時(shí)不損失性能。在通用基準(zhǔn)測(cè)試、長(zhǎng)文本任務(wù)和基于指令的推理任務(wù)中，NSA的表現(xiàn)均能達(dá)到甚至超越傳統(tǒng)全注意力模型的水平。

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

NSA的核心之一是動(dòng)態(tài)分層稀疏策略，它結(jié)合了粗粒度的Token壓縮和細(xì)粒度的Token選擇。這種策略通過分層處理，既保證了模型對(duì)全局上下文的感知能力，又兼顧了局部信息的精確性。具體來說，NSA通過三個(gè)并行的注意力分支處理輸入序列：壓縮注意力、選擇注意力和滑動(dòng)窗口注意力。這種設(shè)計(jì)使得模型能夠在不同粒度上捕捉信息，同時(shí)顯著降低計(jì)算量。

2. 粗粒度Token壓縮

粗粒度Token壓縮是NSA優(yōu)化計(jì)算效率的關(guān)鍵技術(shù)之一。它通過將多個(gè)相鄰的Token合并為一個(gè)“超級(jí)Token”，減少了處理單元的數(shù)量，從而降低了計(jì)算量。NSA采用基于信息熵的方法，優(yōu)先合并信息量較低的Token，最大限度地減少信息損失。例如，在處理新聞文章時(shí)，模型可以將常見的詞匯組合（如“的”“是”等）合并為一個(gè)超級(jí)Token，而保留關(guān)鍵的名詞和動(dòng)詞。這種壓縮方式不僅提高了推理速度，還減少了存儲(chǔ)需求。

3. 細(xì)粒度Token選擇

在粗粒度壓縮的基礎(chǔ)上，NSA進(jìn)一步引入了細(xì)粒度Token選擇機(jī)制。這一機(jī)制允許模型在壓縮后的“超級(jí)Token”中，根據(jù)任務(wù)需求動(dòng)態(tài)選擇關(guān)鍵的子單元進(jìn)行進(jìn)一步處理。這種動(dòng)態(tài)選擇機(jī)制類似于“二次篩選”，確保了模型在壓縮過程中不會(huì)丟失關(guān)鍵信息。例如，在處理問答任務(wù)時(shí)，模型可以優(yōu)先選擇與問題相關(guān)的Token進(jìn)行處理，而在生成文本時(shí)，則可以關(guān)注那些與上下文連貫性相關(guān)的Token。這種動(dòng)態(tài)性不僅提高了模型的靈活性，還進(jìn)一步優(yōu)化了推理效率。

長(zhǎng)文本基準(zhǔn)測(cè)試

在長(zhǎng)文本建模方面，NSA展現(xiàn)了其強(qiáng)大的能力。在64k上下文長(zhǎng)度的“大海撈針”測(cè)試中，NSA實(shí)現(xiàn)了完美的檢索準(zhǔn)確率。

在LongBench基準(zhǔn)測(cè)試中，NSA獲得了最高平均分?jǐn)?shù)0.469，優(yōu)于所有基線，包括全注意力機(jī)制。

思維鏈推理性能

在思維鏈推理任務(wù)中，NSA同樣表現(xiàn)出色。研究人員通過從DeepSeek-R1進(jìn)行知識(shí)蒸餾，使用100億個(gè)32k長(zhǎng)度的數(shù)學(xué)推理軌跡進(jìn)行監(jiān)督微調(diào)，生成了兩個(gè)模型：全注意力模型和NSA稀疏變體。

在AIME 24基準(zhǔn)測(cè)試中，NSA稀疏變體在8k上下文設(shè)置下比全注意力模型高出0.075的準(zhǔn)確率，并在16k上下文中保持了這一優(yōu)勢(shì)。這表明NSA能夠高效捕獲長(zhǎng)距離邏輯依賴關(guān)系，并在推理深度增加時(shí)保持足夠的上下文密度。

效率性能分析

NSA在長(zhǎng)文本建模中表現(xiàn)出顯著的加速效果，尤其是在64k上下文長(zhǎng)度下，NSA實(shí)現(xiàn)了高達(dá)9.0倍的前向加速和6.0倍的反向加速。

解碼速度提升11.6倍：在處理超長(zhǎng)文本時(shí)，NSA通過稀疏注意力機(jī)制顯著減少了內(nèi)存訪問瓶頸，大幅降低了解碼延遲。

https://arxiv.org/abs/2502.11089
 Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

本文轉(zhuǎn)載自??PaperAgent??