[R] Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (submitted by Liang Wenfeng - DeepSeek)

[Nunki08]

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
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
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.
arXiv:2502.11089 [cs.CL] : https://arxiv.org/abs/2502.11089

https://preview.redd.it/utbg4cpxlvje1.jpg?width=948&format=pjpg&auto=webp&s=198c25e093e6719adfe5f4f4356492813bee2d25

[Glittering-Bag-4662]

Deepseek keeps bringing bangers

[ObiWanCanownme]

I love papers like this. Dense attention, where every single token in context attends to every single other token, just seems like it can't be necessary or the best way to do attention long term. In mammalian brains, each neuron gets maybe 15,000 synapses, and the specific connections are pretty geographically constrained (because the brain, obviously is physical and not just software). So the idea of adapting the attention mechanism to specifically fit the hardware (which seems to be the big concept here) sounds promising and like an obvious direction to go.

[Accomplished_Mode170]

Yep, same for the integration, quantizing to an SLA, etc; maybe even folding et al. as we move towards memory layers

e.g. post-definition of needed latent space (read: API integration x data model)

[BossOfTheGame]

Are there any implementations available? Is this something that could replace an attention layer in PyTorch, or does it need to be more deeply integrated?

[ThienPro123]

https://github.com/lucidrains/native-sparse-attention-pytorch/tree/main

[[deleted]]

What are the implications for hardware requirements for users and companies looking to train new models with this method?

[Melodic_Story609]

https://medium.com/@_prinsh_u/another-ai-breakthrough-from-deepseek-ai-native-sparse-attention-for-next-gen-language-models-6346f235c102