原生稀疏注意力:硬件對齊且可原生訓練的稀疏注意力機制Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
長上下文建模對於下一代語言模型至關重要,然而標準注意力機制的高計算成本帶來了顯著的計算挑戰。稀疏注意力提供了一條提升效率同時保持模型能力的有前景的方向。我們提出了NSA,一種原生可訓練的稀疏注意力機制,它將算法創新與硬件對齊的優化相結合,以實現高效的長上下文建模。NSA採用了一種動態分層稀疏策略,結合了粗粒度的令牌壓縮與細粒度的令牌選擇,以保持全局上下文感知和局部精確性。我們的方法在稀疏注意力設計上取得了兩項關鍵創新:(1) 通過算術強度平衡的算法設計,並針對現代硬件進行實現優化,實現了顯著的加速。(2) 我們實現了端到端的訓練,在不犧牲模型性能的情況下減少了預訓練的計算量。如圖1所示,實驗表明,使用NSA預訓練的模型在通用基準測試、長上下文任務和基於指令的推理中保持或超越了全注意力模型。同時,NSA在64k長度序列的解碼、前向傳播和反向傳播上相較於全注意力實現了顯著的加速,驗證了其在模型整個生命週期中的高效性。