원시 희소 어텐션: 하드웨어에 맞게 정렬되고 원시적으로 학습 가능한 희소 어텐션

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를 제안합니다. NSA는 전역 문맥 인식과 지역 정밀도를 보존하기 위해 거친 토큰 압축과 세밀한 토큰 선택을 결합하는 동적 계층적 희소 전략을 사용합니다. 우리의 접근 방식은 두 가지 주요 혁신을 통해 희소 어텐션 설계를 발전시킵니다: (1) 최신 하드웨어에 대한 구현 최적화를 통해 산술 강도 균형 조절 알고리즘 설계를 통해 상당한 가속을 달성합니다. (2) 모델 성능을 희생하지 않고 사전 훈련 계산을 줄이는 엔드-투-엔드 훈련을 가능하게 합니다. 그림 1에 나타난 바와 같이 실험 결과, NSA로 사전 훈련된 모델은 일반 벤치마크, 장거리 작업 및 명령 기반 추론에서 Full Attention 모델을 유지하거나 능가합니다. 한편, NSA는 디코딩, 순방향 전파 및 역방향 전파에서 64k 길이 시퀀스에서 Full Attention에 비해 상당한 가속을 달성하여 모델 수명 주기 전반에 걸쳐 효율성을 입증합니다.

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.