Attenzione Sparse Nativa: Attenzione Sparse Allineata all'Hardware e Addestrabile in Modo Nativo

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

February 16, 2025

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

Abstract

La modellazione a lungo contesto è cruciale per le prossime generazioni di modelli linguistici, tuttavia l'elevato costo computazionale dei meccanismi di attenzione standard presenta sfide computazionali significative. L'attenzione sparsa offre una direzione promettente per migliorare l'efficienza mantenendo le capacità del modello. Presentiamo NSA, un meccanismo di Attenzione Sparsa Nativamente addestrabile che integra innovazioni algoritmiche con ottimizzazioni allineate all'hardware per ottenere una modellazione efficiente a lungo contesto. NSA utilizza una strategia sparsa gerarchica dinamica, combinando la compressione dei token a livello grossolano con la selezione dei token a livello fine per preservare sia la consapevolezza del contesto globale che la precisione locale. Il nostro approccio avanza il design dell'attenzione sparsa con due innovazioni chiave: (1) Otteniamo accelerazioni significative attraverso un design algoritmico bilanciato in termini di intensità aritmetica, con ottimizzazioni di implementazione per l'hardware moderno. (2) Abilitiamo l'addestramento end-to-end, riducendo il calcolo di pre-addestramento senza sacrificare le prestazioni del modello. Come mostrato nella Figura 1, gli esperimenti dimostrano che il modello pre-addestrato con NSA mantiene o supera i modelli con Attenzione Completa su benchmark generali, task a lungo contesto e ragionamento basato su istruzioni. Nel frattempo, NSA ottiene accelerazioni significative rispetto all'Attenzione Completa su sequenze di lunghezza 64k in decodifica, propagazione in avanti e propagazione all'indietro, validando la sua efficienza lungo l'intero ciclo di vita del modello.

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.