Atención Dispersa Nativa: Atención Dispersa Alineada con el Hardware y Entrenable de Forma Nativa

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

February 16, 2025

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

Resumen

El modelado de contexto largo es crucial para la próxima generación de modelos de lenguaje, sin embargo, el alto costo computacional de los mecanismos de atención estándar plantea desafíos significativos. La atención dispersa ofrece una dirección prometedora para mejorar la eficiencia sin comprometer las capacidades del modelo. Presentamos NSA, un mecanismo de Atención Dispersa Nativamente Entrenable que integra innovaciones algorítmicas con optimizaciones alineadas al hardware para lograr un modelado eficiente de contexto largo. NSA emplea una estrategia jerárquica dinámica de dispersión, combinando compresión de tokens a nivel grueso con selección de tokens a nivel fino para preservar tanto la conciencia del contexto global como la precisión local. Nuestro enfoque avanza en el diseño de atención dispersa con dos innovaciones clave: (1) Logramos aceleraciones significativas mediante un diseño algorítmico equilibrado en intensidad aritmética, con optimizaciones de implementación para hardware moderno. (2) Habilitamos el entrenamiento de extremo a extremo, reduciendo el cómputo de preentrenamiento sin sacrificar el rendimiento del modelo. Como se muestra en la Figura 1, los experimentos demuestran que el modelo preentrenado con NSA mantiene o supera a los modelos de Atención Completa en benchmarks generales, tareas de contexto largo y razonamiento basado en instrucciones. Además, NSA logra aceleraciones sustanciales sobre la Atención Completa en secuencias de longitud 64k durante la decodificación, propagación hacia adelante y propagación hacia atrás, validando su eficiencia a lo largo del ciclo de vida del modelo.

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.