Natürliche Sparse Attention: Hardware-ausgerichtet und nativ trainierbare Sparse Attention

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

February 16, 2025

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

Zusammenfassung

Die Modellierung von langen Kontexten ist entscheidend für Sprachmodelle der nächsten Generation, jedoch stellt der hohe Rechenaufwand herkömmlicher Aufmerksamkeitsmechanismen bedeutende Rechenherausforderungen dar. Sparse Attention bietet eine vielversprechende Möglichkeit, die Effizienz zu verbessern, während die Modellfähigkeiten erhalten bleiben. Wir präsentieren NSA, einen nativ trainierbaren Sparse Attention Mechanismus, der algorithmische Innovationen mit hardwareorientierten Optimierungen integriert, um eine effiziente Modellierung langer Kontexte zu erreichen. NSA verwendet eine dynamische hierarchische Sparse-Strategie, die eine grobkörnige Token-Kompression mit einer feinkörnigen Token-Auswahl kombiniert, um sowohl das Bewusstsein für den globalen Kontext als auch die lokale Präzision zu erhalten. Unser Ansatz verbessert das Design von Sparse Attention durch zwei Schlüsselinnovationen: (1) Wir erzielen erhebliche Beschleunigungen durch eine algorithmische Gestaltung mit ausgewogener arithmetischer Intensität und Implementierungsoptimierungen für moderne Hardware. (2) Wir ermöglichen ein End-to-End-Training, das die Vorab-Berechnung reduziert, ohne die Modellleistung zu beeinträchtigen. Wie in Abbildung 1 gezeigt, zeigen Experimente, dass das mit NSA vorab trainierte Modell bei allgemeinen Benchmarks, langen Kontextaufgaben und instruktionsbasiertem Denken die Full Attention-Modelle beibehält oder übertrifft. Gleichzeitig erzielt NSA erhebliche Beschleunigungen gegenüber Full Attention bei Sequenzen mit einer Länge von 64k über Dekodierung, Vorwärts- und Rückwärtspropagation und bestätigt so seine Effizienz im gesamten Modelllebenszyklus.

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.