Atención Dispersa Nativa: Atención Dispersa Alineada con el Hardware y Entrenable de Forma NativaNative Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
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.