提出可直接训练的稀疏注意力机制,显著提升长文本建模效率
Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

- 采用分层动态稀疏策略,结合粗粒度压缩与细粒度选择
- 在64k序列上实现解码/前向/反向传播全链路提速
- 支持端到端训练,预训练成本降低且性能不降
长上下文建模对下一代语言模型至关重要,但标准注意力机制计算开销大。稀疏注意力为提升效率提供了可行路径。本文提出NSA(Natively Sparse Attention),一种算法创新与硬件对齐优化融合的原生可训练稀疏注意力机制,实现高效长上下文建模。NSA采用动态分层稀疏策略,结合粗粒度标记压缩与细粒度标记选择,兼顾全局上下文感知与局部精度。核心创新包括:(1) 通过算术强度平衡的设计实现显著加速,并针对现代硬件进行实现优化;(2) 支持端到端训练,减少预训练计算量而不牺牲模型性能。实验表明,使用NSA预训练的模型在通用基准、长上下文任务及指令推理任务中表现持平或优于全连接注意力模型。同时,在64k长度序列上,NSA在解码、前向传播和反向传播各阶段均实现显著提速,验证了其在整个模型生命周期中的高效性。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。