DeepSeek新注意力機制引熱議!梁文鋒親自提交預印本,目標明確降低計算成本
DeepSeek新注意力機制論文一出,再次引爆討論熱度。
依然是熟悉的畫風,熟悉的味道——
那邊馬斯克瘋狂燒了20萬張卡訓出Grok 3,這廂DeepSeek重點關注的依然是壓縮計算和推理成本。
具體來說,新論文提出了一種可原生訓練的稀疏注意力機制,名為NSA(Native Sparse Attention)。
目的很明確:解決大模型上下文建模中,標準注意力機制造成的計算成本高昂的問題。
效果也很明顯:
實驗表明,在解碼64k上下文時,基于softmax架構的注意力計算占到總延遲的70%-80%。而在不影響性能的前提下,NSA在64k上下文的解碼、前向傳播和反向傳播中均實現了顯著加速。
其中前向傳播速度最高可提升9倍,反向傳播速度最高可提升6倍,解碼速度提升可達11.6倍。
正如不少網友提到的,NSA意味著DeepSeek找到了優化注意力的方法,可以用更少的算力更加高效地訓練大模型,并且,他們還把這些“秘籍”公開了出來。
剛剛加入OpenAI不久的ViT核心作者Lucas Beyer也不禁開麥:粉了粉了。
另一點受到關注的是,這篇論文的作者可以說是DeepSeek系列大模型原班人馬,梁文鋒本人亦在其列——
論文還是梁文鋒親自傳上arXiv的。
實現超快速長上下文訓練推理
來看論文細節。
NSA的核心方法包括:
- 動態分層稀疏策略
- 粗粒度Token壓縮
- 細粒度Token選擇
整體框架上,NSA是用更緊湊的鍵值對替代原始注意力中的鍵值對,通過壓縮、選擇和滑動窗口(sliding window)三種映射策略得到優化注意力輸出,保持高稀疏率。
采取分層設計的好處是:
- 減少計算量的同時,能有效避免信息丟失,是模型在處理長序列時既高效又精準。
- 能夠在保證全局上下文感知的同時,精確捕捉局部信息,提升模型的表現。
- 可根據任務需求和上下文動態調整稀疏程度
具體到算法設計上,粗粒度Token壓縮通過將連續的Token聚合成塊級表示,可以捕獲全局語義信息,同時減少注意力的計算負擔。
細粒度Token選擇從序列中選擇最重要的Token,保留關鍵的局部信息。
滑動窗口則避免了局部模式的過度優勢——在注意力機制中,局部模式往往會主導學習過程,阻止模型有效地從壓縮和選擇Token中學習。
為了實現以上創新,DeepSeek官方還強調了兩方面的關鍵工作:
- 硬件對齊系統:針對張量核心利用率和內存訪問優化塊級稀疏注意力,確保均衡的算術強度。
- 訓練感知設計:通過高效算法和向后運算符實現穩定的端到端訓練。
研究人員用27B參數規模的模型進行了實驗,結果顯示,全注意力機制和NSA在預訓練損失方面,都表現出了穩定的收斂性,并且NSA實現了更低的損失值。
在包含知識、推理和編碼能力的多個通用基準測試中,與全注意力模型相比,NSA模型性能不降反超,在推理任務DROP中提升尤為明顯。
長上下文方面,64k上下文的“大海撈針”測試里,NSA完美過關。
在需要復雜長下文推理的各項任務中,NSA的表現也基本超過了包括全注意力模型在內的基線模型。
而在思維鏈推理評估中,通過知識蒸餾和監督微調,在8k和16k上下文設置下,AIME任務中NSA-R的得分均超過了全注意力模型。
這表明,NSA預訓練的稀疏注意力模式能有效捕捉長距離邏輯依賴,且其硬件對齊設計可支持不斷增加的推理深度。
效率方面,在8-GPU A100系統上,NSA的訓練加速效果會隨上下文長度的增加而增強。在64k上下文長度時,前向傳播速度最高可提升9倍,反向傳播速度最高可提升6倍,解碼速度提升可達11.6倍。
有意思的是,在馬斯克推出Grok 3炸場之時,不少人感慨:“大力出奇跡”在大模型訓練里仍然奏效——
在Grok 3成為大模型競技場首個突破1400分模型的背后,是驚人的10萬卡(后來擴展到20萬)H100集群。
但現在,DeepSeek為代表的效率派們顯然在引發另一種方向上的思考:通往AGI的路,可以更高效,更本地化,更人人可用。