Attention Sparse Native : Une Attention Sparse Alignée sur le Matériel et Entraînable de Manière Native

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

February 16, 2025

Auteurs: 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

Résumé

La modélisation de contexte long est cruciale pour les modèles de langage de nouvelle génération, mais le coût computationnel élevé des mécanismes d'attention standard pose des défis significatifs. L'attention parcimonieuse offre une direction prometteuse pour améliorer l'efficacité tout en conservant les capacités du modèle. Nous présentons NSA, un mécanisme d'Attention Parcimonieuse Natif (Natively trainable Sparse Attention) qui intègre des innovations algorithmiques avec des optimisations alignées sur le matériel pour atteindre une modélisation de contexte long efficace. NSA utilise une stratégie parcimonieuse hiérarchique dynamique, combinant une compression grossière des tokens avec une sélection fine des tokens pour préserver à la fois la conscience du contexte global et la précision locale. Notre approche fait progresser la conception de l'attention parcimonieuse avec deux innovations clés : (1) Nous obtenons des accélérations substantielles grâce à une conception algorithmique équilibrée en intensité arithmétique, avec des optimisations d'implémentation pour le matériel moderne. (2) Nous permettons un entraînement de bout en bout, réduisant le calcul de pré-entraînement sans sacrifier les performances du modèle. Comme le montre la Figure 1, les expériences montrent que le modèle pré-entraîné avec NSA maintient ou dépasse les modèles à Attention Complète sur des benchmarks généraux, des tâches à contexte long et des raisonnements basés sur des instructions. Par ailleurs, NSA obtient des accélérations substantielles par rapport à l'Attention Complète sur des séquences de longueur 64k, que ce soit en décodage, propagation avant ou propagation arrière, validant ainsi son efficacité tout au long du cycle de vie du modèle.

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.