Native Sparse Attention : Hardware-Aligned and NativelyTrainable Sparse Attention (deepseek)

 

 

Native Sparse Attention 논문 리뷰

본 포스트에서는 Native Sparse Attention (NSA) 논문을 심도 있게 분석합니다. NSA는 기존의 Full Attention 방식의 높은 계산 비용 문제를 해결하고, 하드웨어에 최적화된 알고리즘과 end-to-end 학습이 가능한 스파스 어텐션(sparse attention) 메커니즘을 제안합니다. 논문에서는 동적 계층적 스파스 전략을 통해 compressed (압축), selected (선택), sliding window (슬라이딩 윈도우) 세 가지 경로를 활용하여 글로벌 컨텍스트와 로컬 정밀도를 동시에 보존하는 방법을 소개합니다.

1. 핵심 요약

NSA는 기존 sparse attention 기법들이 inference에 한정되거나 학습 시 계산 그래프의 불연속성 문제로 인해 발생하는 한계를 극복하기 위해 제안되었습니다. 본 논문의 주요 기여는 다음과 같습니다:

• 하드웨어 정렬 (Hardware-aligned) 최적화: Tensor Core 활용과 메모리 접근 효율성을 극대화하는 알고리즘 디자인을 통해 실제 속도 향상을 달성.
• 학습-인식 (Training-aware) 설계: end-to-end 학습이 가능한 sparse operator를 도입하여 pretraining 단계의 계산 비용을 줄임.
• 동적 계층적 스파스 전략: 입력 시퀀스를 압축(Compression), 선택(Selection), 슬라이딩 윈도우(Sliding Window) 세 경로로 나누어 각 경로별로 다른 정보를 효과적으로 학습.

실험 결과 NSA는 27B-파라미터 Transformer 백본(pretrained on 260B tokens)을 대상으로, 일반 벤치마크, long-context 작업 및 chain-of-thought reasoning 평가에서 Full Attention 모델과 동등하거나 우수한 성능을 보이며, 64k-length 시퀀스 처리 시 디코딩, Forward, Backward 단계에서 각각 최대 11.6×, 9.0×, 6.0×의 속도 향상을 달성하였습니다.

Figure 1: Full Attention 모델 대비 NSA의 성능 및 속도 향상 결과 (Decoding: 11.6×, Forward: 9.0×, Backward: 6.0×)

2. 각 소주제별 요약 및 핵심 내용

2.1 연구 배경 및 필요성 (Background & Motivation)

현대 대규모 언어 모델(LLM)은 수천 토큰 이상의 긴 문맥(long-context) 이해가 필수적입니다. 그러나 기존의 vanilla Attention (Vaswani et al., 2017) 메커니즘은 시퀀스 길이에 따라 계산 복잡도가 기하급수적으로 증가하여 성능의 병목 현상을 초래합니다. 이에 따라 여러 sparse attention 방법들이 제안되었지만, 대부분은 inference 단계에만 적용되거나 학습 시 불연속적 연산으로 인해 최적의 sparse 패턴 학습에 한계를 보였습니다.

• 문제점 1: 기존 sparse attention 방법은 prefilling과 decoding 중 하나의 단계에서만 효율 개선을 보이는 경우가 많아 전체적인 추론(latency)에서 기대한 속도 향상을 달성하지 못함.
• 문제점 2: 학습 과정에서 discrete 연산(예: k-means clustering, SimHash 등)으로 인한 gradient 흐름 차단 문제로 인해 end-to-end 학습이 어려움.

이러한 문제점을 해결하기 위해 NSA는 하드웨어와 학습 양 측면 모두를 고려한 native sparse attention 설계를 제안합니다.

2.2 NSA의 주요 혁신 (Key Innovations)

NSA가 제시하는 두 가지 주요 혁신은 다음과 같습니다:

• Hardware-Aligned System: 블록 단위의 스파스 어텐션(blockwise sparse attention)을 통해 메모리 접근 패턴을 최적화하고, Tensor Core와 같은 최신 하드웨어의 계산 자원을 극대화함으로써 실제 연산 속도를 크게 향상.
• Training-Aware Design: 연속적인 학습이 가능한 스파스 연산자를 도입하여 pretraining 단계에서도 효율적인 계산을 지원하고, gradient 흐름의 안정성을 보장.

논문에서는 이러한 혁신을 바탕으로 기존 sparse 방법들이 가진 한계를 극복하고, 실제 어플리케이션에서 적용 가능한 실질적인 성능 향상을 달성함을 입증합니다.

2.3 NSA 아키텍처 및 알고리즘 (Architecture & Algorithm)

NSA는 입력 시퀀스에 대해 다음과 같은 세 가지 병렬 경로로 어텐션을 계산합니다:

1. Token Compression (압축): 연속된 토큰 블록을 하나의 압축 토큰으로 집약하기 위해 learnable MLP와 intra-block position encoding을 사용합니다. (논문 Eq. (7))
2. Token Selection (선택): 블록 단위의 토큰 중요도를 평가하여 가장 중요한 블록(Top-n)을 선택합니다. 이는 attention score의 spatial continuity (공간 연속성)를 활용하며, 논문에서는 Eq. (8)~(11)로 수식화되어 있습니다.
3. Sliding Window (슬라이딩 윈도우): 최근 토큰을 별도의 창(window)으로 관리하여 로컬 컨텍스트 정보를 별도로 취급합니다. 이를 통해 압축 및 선택 경로가 전역 패턴에 집중할 수 있도록 합니다.

세 경로의 결과는 gating mechanism을 통해 통합되어 최종 어텐션 결과를 산출합니다. 또한, 논문에서는 기존 Full Attention과 비교하여 계산 복잡도를 대폭 줄이면서도, 정보 손실을 최소화하는 설계 전략을 상세히 설명합니다.

Figure 2: NSA의 전체 아키텍처. 왼쪽은 세 가지 병렬 어텐션 경로(Compression, Selection, Sliding Window)를, 오른쪽은 각 경로에서 계산되는 어텐션 패턴을 시각화.

2.4 하드웨어 최적화 및 커널 디자인 (Hardware-Aligned Kernel Design)

NSA는 학습과 추론 양 단계에서 FlashAttention-2와 견줄 만한 효율을 달성하기 위해, Triton 기반의 커스텀 커널을 구현합니다. 특히, GQA(Grouped-Query Attention) 및 MQA(Multiple-Query Attention)와 같은 최신 디코딩 아키텍처에 맞춰 KV-cache의 공유와 블록 단위 메모리 접근을 최적화하였습니다.

주요 최적화 전략은 쿼리 그룹화(query grouping)로, 이는 연속적이지 않은 메모리 접근 문제를 해결하여 전체 연산 속도를 높이는 데 크게 기여합니다.

2.5 실험 결과 및 평가 (Experimental Results & Evaluation)

NSA의 성능은 27B-파라미터 Transformer 모델을 대상으로, 260B tokens의 대규모 pretraining 후 평가되었습니다. 실험 결과는 다음과 같은 주요 성과를 보여줍니다:

• 전반적인 성능: 일반 벤치마크, long-context 작업, chain-of-thought reasoning 평가에서 Full Attention 모델과 동등하거나 그 이상의 성능을 달성.
• 계산 효율성: 64k-length 시퀀스 처리 시, 디코딩 단계에서 최대 11.6×, Forward 단계에서 9.0×, Backward 단계에서 6.0×의 속도 향상을 기록.

테이블: NSA와 Full Attention 비교 (속도 향상 지표)

Stage	Speedup Ratio
Decoding	11.6×
Forward Propagation	9.0×
Backward Propagation	6.0×

위 테이블은 NSA가 각 단계별로 Full Attention에 비해 얼마나 빠른 성능을 보이는지 명확하게 보여줍니다.

3. 결론 및 향후 연구 (Conclusion & Future Directions)

본 논문은 장문맥(long-context) 모델링의 병목 문제를 해결하기 위한 새로운 방향성을 제시합니다. NSA는 하드웨어와 학습 양 측면에서 기존 sparse attention 방식의 한계를 극복하며, 실제 응용 환경에서의 속도 향상과 성능 유지를 동시에 달성하였습니다.

향후 연구에서는 NSA의 동적 계층적 스파스 전략을 기반으로 보다 다양한 모델 아키텍처에의 확장 가능성, 그리고 다른 하드웨어 플랫폼에서의 최적화 방안을 추가적으로 탐구할 필요가 있을 것입니다.

부록: 논문 수식에 대한 간단한 설명

논문에서는 어텐션 메커니즘의 수학적 정의를 통해 기본 개념을 명확히 하고 있습니다. 예를 들어, 일반적인 어텐션 연산은 아래와 같이 정의됩니다:

oₜ = Attn(qₜ, k₁:ₜ, v₁:ₜ)
Attn(qₜ, k₁:ₜ, v₁:ₜ) = Σ (αₜ,ᵢ · vᵢ),  where  αₜ,ᵢ = exp(qₜᵀkᵢ/√dₖ)
  

NSA에서는 이 기본 어텐션 수식을 변형하여, 압축 및 선택 등의 전략을 도입한 새로운 mapping 함수 fₖ(·)와 fᵥ(·)를 정의합니다. 이를 통해 전체 토큰 집합 대신 정보 밀집도가 높은 토큰 집합을 구성하여 계산량을 줄이고, 동시에 모델 성능을 유지할 수 있습니다.