Atenção Esparsa Nativa: Alinhada ao Hardware e Intrinsecamente Treinável com Esparsidade

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

Resumo

A modelagem de contexto longo é crucial para os modelos de linguagem de próxima geração, mas o alto custo computacional dos mecanismos de atenção padrão apresenta desafios computacionais significativos. A atenção esparsa oferece uma direção promissora para melhorar a eficiência enquanto mantém as capacidades do modelo. Apresentamos o NSA, um mecanismo de Atenção Esparsa Nativamente Treinável que integra inovações algorítmicas com otimizações alinhadas ao hardware para alcançar uma modelagem eficiente de contexto longo. O NSA emprega uma estratégia esparsa hierárquica dinâmica, combinando compressão de tokens em nível grosseiro com seleção de tokens em nível fino para preservar tanto a consciência do contexto global quanto a precisão local. Nossa abordagem avança o design da atenção esparsa com duas inovações principais: (1) Alcançamos acelerações substanciais por meio de um design algorítmico balanceado em intensidade aritmética, com otimizações de implementação para hardware moderno. (2) Habilitamos o treinamento de ponta a ponta, reduzindo a computação de pré-treinamento sem sacrificar o desempenho do modelo. Como mostrado na Figura 1, os experimentos demonstram que o modelo pré-treinado com NSA mantém ou supera os modelos de Atenção Completa em benchmarks gerais, tarefas de contexto longo e raciocínio baseado em instruções. Enquanto isso, o NSA alcança acelerações substanciais em relação à Atenção Completa em sequências de 64k de comprimento ao longo da decodificação, propagação direta e propagação reversa, validando sua eficiência em todo o ciclo de vida do 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.