Native Sparse Attention

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

대규모 언어 모델(LLM)에서 장문맥 모델링은 중요한 과제이다. 표준 어텐션 메커니즘은 모든 토큰 쌍 간의 관계를 계산하므로 O(n²)의 시간 및 메모리 복잡도를 가진다. 이는 문맥 길이가 증가함에 따라 실용적인 한계에 부딪힌다. 희소 어텐션은 이러한 비용을 줄이기 위한 유망한 방향으로 제안되었으나, 기존 방법들은 하드웨어 효율성이나 학습 가능성에서 한계를 보였다. NSA는 이러한 문제를 해결하기 위해 하드웨어 정렬 및 학습 가능한 희소 어텐션을 설계하였다.

NSA는 동적 계층적 희소 전략을 사용한다. 이 전략은 두 가지 주요 구성 요소로 이루어진다:

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

이러한 이중 접근 방식은 전역 인식과 지역 정밀도 사이의 균형을 제공한다. 또한 NSA는 하드웨어 정렬 최적화를 통해 GPU와 같은 현대 가속기에서 효율적으로 실행되도록 설계되었다.

NSA의 주요 혁신은 두 가지이다. 첫째, 산술 강도(arithmetic intensity)를 최적화하여 실질적인 속도 향상을 달성한다. 둘째, 학습 가능한 희소 패턴을 통해 사전 학습 단계에서부터 희소성을 내재화하여 추론 시에도 일관된 효율성을 제공한다. 기존의 희소 어텐션 방법들은 종종 사전 학습 후에 희소성을 적용하거나 하드웨어 비효율성을 겪었으나, NSA는 이러한 문제를 해결한다.

NSA의 공식 구현은 공개되지 않았으나, 여러 오픈소스 구현체가 존재한다. lucidrains가 PyTorch 기반 구현을 MIT 라이선스로 공개하였고, XunhaoLai가 Triton 기반의 효율적인 구현을 Apache 2.0 라이선스로 공개하였다. 또한 홍콩과학기술대학교와 카네기멜론대학교 연구진이 FSA(An Alternative Efficient Implementation of Native Sparse Attention Kernel)라는 대체 구현을 발표하였다. 이러한 구현체들은 NSA의 희소 어텐션 패턴을 다양한 프레임워크에서 사용할 수 있게 한다.

NSA는 장문맥 언어 모델의 효율성을 크게 향상시킬 잠재력을 가진다. 하드웨어 정렬 및 학습 가능한 설계는 실제 배포 환경에서 실용적인 이점을 제공한다. 이 연구는 희소 어텐션 분야에서 중요한 진전으로 평가되며, 향후 대규모 언어 모델의 장문맥 처리 능력 발전에 기여할 것으로 기대된다.