Нативное разреженное внимание: аппаратно-ориентированное и изначально обучаемое разреженное внимание

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

February 16, 2025

Авторы: 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

Аннотация

Моделирование длинных контекстов имеет ключевое значение для языковых моделей следующего поколения, однако высокая вычислительная стоимость стандартных механизмов внимания создает значительные вычислительные трудности. Разреженное внимание предлагает перспективное направление для повышения эффективности при сохранении возможностей модели. Мы представляем NSA (Natively trainable Sparse Attention) — механизм разреженного внимания, который интегрирует алгоритмические инновации с оптимизациями, ориентированными на аппаратное обеспечение, для достижения эффективного моделирования длинных контекстов. NSA использует динамическую иерархическую стратегию разрежения, сочетая грубое сжатие токенов с точным выбором токенов для сохранения как глобальной осведомленности о контексте, так и локальной точности. Наш подход продвигает дизайн разреженного внимания с двумя ключевыми инновациями: (1) Мы достигаем значительного ускорения за счет сбалансированного по арифметической интенсивности проектирования алгоритма с оптимизациями для современного оборудования. (2) Мы обеспечиваем сквозное обучение, сокращая вычисления на этапе предварительного обучения без ущерба для производительности модели. Как показано на Рисунке 1, эксперименты демонстрируют, что модель, предварительно обученная с использованием NSA, сохраняет или превосходит модели с полным вниманием на общих бенчмарках, задачах с длинным контекстом и рассуждениях на основе инструкций. При этом NSA достигает значительного ускорения по сравнению с полным вниманием на последовательностях длиной 64k токенов на этапах декодирования, прямого и обратного распространения, подтверждая свою эффективность на протяжении всего жизненного цикла модели.

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.