ネイティブ疎注意：ハードウェアに最適化された訓練可能な疎注意機構

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

February 16, 2025

著者: Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, Wangding Zeng

cs.AI

要旨

長文脈モデリングは次世代言語モデルにとって極めて重要であるが、標準的なアテンション機構の高い計算コストは重大な計算上の課題を引き起こしている。スパースアテンションは、モデルの能力を維持しつつ効率を向上させる有望な方向性を提供する。本論文では、アルゴリズムの革新とハードウェアに最適化された最適化を統合し、効率的な長文脈モデリングを実現するNatively trainable Sparse Attention（NSA）機構を提案する。NSAは、動的な階層的スパース戦略を採用し、粗粒度のトークン圧縮と細粒度のトークン選択を組み合わせることで、グローバルな文脈認識とローカルな精度の両方を保持する。我々のアプローチは、2つの重要な革新によりスパースアテンション設計を前進させる：(1) 算術強度バランスの取れたアルゴリズム設計により大幅な高速化を実現し、現代のハードウェア向けに実装を最適化する。(2) エンドツーエンドの学習を可能にし、モデルの性能を損なうことなく事前学習の計算量を削減する。図1に示すように、NSAで事前学習されたモデルは、一般的なベンチマーク、長文脈タスク、および指示に基づく推論において、Full Attentionモデルを維持または上回る性能を示す。一方、NSAは64k長のシーケンスにおいて、デコード、順伝播、逆伝播の全てにおいてFull Attentionに対して大幅な高速化を達成し、モデルライフサイクル全体での効率性を実証している。

English

Long-context modeling is crucial for next-generation language models, yet the high computational cost of standard attention mechanisms poses significant computational challenges. Sparse attention offers a promising direction for improving efficiency while maintaining model capabilities. We present NSA, a Natively trainable Sparse Attention mechanism that integrates algorithmic innovations with hardware-aligned optimizations to achieve efficient long-context modeling. NSA employs a dynamic hierarchical sparse strategy, combining coarse-grained token compression with fine-grained token selection to preserve both global context awareness and local precision. Our approach advances sparse attention design with two key innovations: (1) We achieve substantial speedups through arithmetic intensity-balanced algorithm design, with implementation optimizations for modern hardware. (2) We enable end-to-end training, reducing pretraining computation without sacrificing model performance. As shown in Figure 1, experiments show the model pretrained with NSA maintains or exceeds Full Attention models across general benchmarks, long-context tasks, and instruction-based reasoning. Meanwhile, NSA achieves substantial speedups over Full Attention on 64k-length sequences across decoding, forward propagation, and backward propagation, validating its efficiency throughout the model lifecycle.