Native Sparse Attention

Native Sparse Attention(NSA)은 장문맥 언어 모델의 효율성을 높이기 위해 설계된 희소 어텐션 메커니즘이다. 표준 어텐션은 시퀀스 길이에 제곱 비례하는 계산 비용을 가지므로 긴 문맥 처리에 한계가 있다. NSA는 이러한 문제를 해결하기 위해 동적 계층적 희소 전략을 도입하여, 전역 문맥 정보를 유지하면서도 계산량을 크게 줄인다. 이 메커니즘은 하드웨어 구조에 정렬된 최적화를 포함하여 실제 시스템 수준에서 성능 향상을 제공한다.

대규모 언어 모델(LLM)에서 장문맥 모델링은 심층 추론, 코드 생성, 다중 턴 대화 등 실세계 응용에서 필수적이다. 그러나 표준 어텐션 메커니즘은 모든 토큰 쌍 간의 관계를 계산하므로 $O(n^2)$의 시간 및 메모리 복잡도를 가진다. 시퀀스 길이가 64k(6만 4천 토큰) 수준에 도달하면 어텐션 연산이 전체 지연 시간의 약 70~80%를 차지하는 병목 현상이 발생한다. 기존의 희소 어텐션 기법들은 추론 단계에만 집중하거나 하드웨어 효율성이 떨어지는 한계가 있었으며, NSA는 이를 해결하기 위해 제안되었다.

NSA는 동적 계층적 희소 전략을 사용하며, 이는 크게 두 가지 구성 요소로 나뉜다.

• Coarse-grained 토큰 압축: 전체 시퀀스를 블록 단위로 압축하여 전역 문맥 정보를 효율적으로 포착한다.
• Fine-grained 토큰 선택: 중요한 개별 토큰을 선택하여 지역적 정밀도를 유지한다.

이러한 이중 접근 방식은 전역 인식과 지역 정밀도 사이의 균형을 제공하며, 연산 복잡도를 획기적으로 낮춘다.

NSA의 주요 혁신은 하드웨어 정렬 최적화에 있다. 산술 강도(arithmetic intensity)를 최적화하여 GPU와 같은 현대 가속기에서 실질적인 속도 향상을 달성한다. 이론적인 연산량 감소에 그치지 않고, 메모리 접근 패턴을 하드웨어 구조에 맞게 설계함으로써 실제 시스템에서의 실행 효율을 극대화하였다.

NSA는 학습 가능한 희소 패턴을 통해 사전 학습 단계에서부터 희소성을 내재화한다. 이를 통해 사전 학습 비용을 절감하면서도 Full Attention 대비 동등하거나 우수한 성능을 유지한다. 64k 시퀀스 환경에서 디코딩, 순전파, 역전파 모든 과정에서 유의미한 속도 향상을 증명하였으며, 긴 문맥에 대한 모델의 이해 능력을 보존한다.

NSA의 연구 결과는 ACL 2025에 게재되었으며, 여러 오픈소스 구현체가 존재한다.

특히 홍콩과학기술대학교와 카네기멜론대학교 연구진이 발표한 FSA(An Alternative Efficient Implementation of Native Sparse Attention Kernel)는 NSA 커널의 효율적인 대체 구현으로 평가받는다.