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

2025-02-16

AI Summary

Standard attention incurs prohibitive O(n²) computational complexity for long-context modeling. To address this, we propose Native Trainable Sparse Attention (NSA), a novel sparse attention mechanism with dynamic hierarchical sparsity: coarse-grained global compression coupled with fine-grained preservation of critical local interactions. NSA introduces the first arithmetic-intensity-balanced, hardware-aligned kernel design, enabling end-to-end differentiability without dense pretraining. Its native gradient propagation ensures that sparsity patterns adaptively optimize during training. Experiments on 64K-length sequences demonstrate substantial speedups across forward pass, backward pass, and autoregressive decoding. Pretrained models using NSA match or exceed full-attention baselines on general language understanding, long-context question answering, and reasoning benchmarks—achieving both high efficiency and strong representational capacity.

Technology Category

• Natural Language Processing: (Large) Language Models
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Hardware-aware ML

Application Category

• Search and Retrieval-Augmented AI: Personalized, context-aware and across-device search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Efficient long-context modeling
• Hardware-aligned sparse attention
• End-to-end trainable mechanism

Innovation

• Hardware-aligned sparse attention
• Dynamic hierarchical sparse strategy
• End-to-end trainable design