DeepSeek全新注意力機制NSA發布,超快速長文訓練與推理 精華
DeepSeek官推發布了最新技術成果NSA:一種面向硬件且支持原生訓練的稀疏注意力機制,專為超快長上下文訓練與推理設計。
NSA的核心組成:
- 動態分層稀疏策略
- 粗粒度的token壓縮
- 細粒度的token選擇
?? NSA針對現代硬件進行了優化設計,顯著提升了推理速度,并有效降低了預訓練成本——同時不損失性能。在通用基準測試、長文本任務和基于指令的推理任務中,NSA的表現均能達到甚至超越傳統全注意力模型的水平。
1. 動態分層稀疏策略
NSA的核心之一是動態分層稀疏策略,它結合了粗粒度的Token壓縮和細粒度的Token選擇。這種策略通過分層處理,既保證了模型對全局上下文的感知能力,又兼顧了局部信息的精確性。具體來說,NSA通過三個并行的注意力分支處理輸入序列:壓縮注意力、選擇注意力和滑動窗口注意力。這種設計使得模型能夠在不同粒度上捕捉信息,同時顯著降低計算量。
2. 粗粒度Token壓縮
粗粒度Token壓縮是NSA優化計算效率的關鍵技術之一。它通過將多個相鄰的Token合并為一個“超級Token”,減少了處理單元的數量,從而降低了計算量。NSA采用基于信息熵的方法,優先合并信息量較低的Token,最大限度地減少信息損失。例如,在處理新聞文章時,模型可以將常見的詞匯組合(如“的”“是”等)合并為一個超級Token,而保留關鍵的名詞和動詞。這種壓縮方式不僅提高了推理速度,還減少了存儲需求。
3. 細粒度Token選擇
在粗粒度壓縮的基礎上,NSA進一步引入了細粒度Token選擇機制。這一機制允許模型在壓縮后的“超級Token”中,根據任務需求動態選擇關鍵的子單元進行進一步處理。這種動態選擇機制類似于“二次篩選”,確保了模型在壓縮過程中不會丟失關鍵信息。例如,在處理問答任務時,模型可以優先選擇與問題相關的Token進行處理,而在生成文本時,則可以關注那些與上下文連貫性相關的Token。這種動態性不僅提高了模型的靈活性,還進一步優化了推理效率。
長文本基準測試
在長文本建模方面,NSA展現了其強大的能力。在64k上下文長度的“大海撈針”測試中,NSA實現了完美的檢索準確率。
在LongBench基準測試中,NSA獲得了最高平均分數0.469,優于所有基線,包括全注意力機制。
思維鏈推理性能
在思維鏈推理任務中,NSA同樣表現出色。研究人員通過從DeepSeek-R1進行知識蒸餾,使用100億個32k長度的數學推理軌跡進行監督微調,生成了兩個模型:全注意力模型和NSA稀疏變體。
在AIME 24基準測試中,NSA稀疏變體在8k上下文設置下比全注意力模型高出0.075的準確率,并在16k上下文中保持了這一優勢。這表明NSA能夠高效捕獲長距離邏輯依賴關系,并在推理深度增加時保持足夠的上下文密度。
效率性能分析
NSA在長文本建模中表現出顯著的加速效果,尤其是在64k上下文長度下,NSA實現了高達9.0倍的前向加速和6.0倍的反向加速。
解碼速度提升11.6倍:在處理超長文本時,NSA通過稀疏注意力機制顯著減少了內存訪問瓶頸,大幅降低了解碼延遲。
https://arxiv.org/abs/2502.11089
Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
本文轉載自??PaperAgent??
