Native Sparse Attention: Hardware-afgestemde en van nature trainbare sparse attention

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

Samenvatting

Lang-context modellering is cruciaal voor de volgende generatie taalmodelen, maar de hoge rekenkosten van standaard aandachtmechanismen vormen aanzienlijke computationele uitdagingen. Sparse aandacht biedt een veelbelovende richting voor het verbeteren van de efficiëntie terwijl de modelcapaciteiten behouden blijven. Wij presenteren NSA, een Natively trainable Sparse Attention mechanisme dat algoritmische innovaties combineert met hardware-afgestemde optimalisaties om efficiënte lang-context modellering te bereiken. NSA maakt gebruik van een dynamische hiërarchische sparse strategie, waarbij grofkorrelige tokencompressie wordt gecombineerd met fijnkorrelige tokenselectie om zowel globaal contextbewustzijn als lokale precisie te behouden. Onze aanpak brengt sparse aandachtontwerp verder met twee belangrijke innovaties: (1) We behalen aanzienlijke snelheidswinsten door een rekenintensiteit-gebalanceerd algoritmeontwerp, met implementatieoptimalisaties voor moderne hardware. (2) We maken end-to-end training mogelijk, waardoor de rekentijd vooraf wordt verminderd zonder in te leveren op modelprestaties. Zoals weergegeven in Figuur 1, tonen experimenten aan dat het model dat met NSA is voorgetraind, presteert op of boven het niveau van Full Attention-modellen op algemene benchmarks, lang-context taken en instructiegebaseerd redeneren. Tegelijkertijd behaalt NSA aanzienlijke snelheidswinsten ten opzichte van Full Attention op sequenties van 64k lengte tijdens decodering, voorwaartse propagatie en achterwaartse propagatie, wat de efficiëntie gedurende de gehele modellevenscyclus bevestigt.

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.