Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

다음은 논문 "Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection"에 대한 상세한 기술 요약입니다.

1. 문제 제기

대형 언어 모델 (LLM) 은 프리필 (prefill) 단계에서 어텐션 메커니즘의 **2 차 복잡도 ($O(L^2)$)**로 인해 긴 컨텍스트 추론에서 근본적인 병목 현상에 직면해 있습니다. FlashAttention 과 같은 하드웨어 최적화는 메모리 I/O 오버헤드를 줄이지만, 알고리즘적 복잡도를 변경하지는 않습니다.

기존 가속화 방법들은 두 가지 주요 한계를 겪고 있습니다:

• 구조적 희소성 (Structured Sparsity): Minference 나 FlexPrefill 과 같은 방법들은 블록 단위 희소성을 적용합니다. 이는 중요한 토큰과 같은 블록에 묶여 있는 경우 불필요한 토큰을 유지하게 되어, 최적에 미치지 못하는 희소성을 초래합니다.
• 토큰 제거 (Token Eviction): FastKV 나 GemFilter 와 같은 방법들은 초기 레이어에서 토큰을 영구적으로 폐기합니다. 이는 토큰의 중요도가 정적이라고 가정하며, 서로 다른 레이어와 어텐션 헤드 간에 토큰 관련성의 동적 성질을 무시합니다. 일단 제거된 토큰은 다시 고려될 수 없으므로 정확도 저하로 이어집니다.

2. 방법론: Token Sparse Attention

저자들은 경량, 동적, 그리고 가역적인 토큰 단위 희소화 메커니즘인 Token Sparse Attention을 제안합니다. 제거 기반 방법들과 달리, 이 방법은 토큰을 영구적으로 제거하지 않고 계산 공간을 압축한 후 시퀀스를 복원합니다.

핵심 아키텍처

이 방법은 어텐션 헤드당 두 단계로 작동합니다:

1. 압축 (QKV 선택):
   • 각 어텐션 헤드 $h$에 대해, 중요도 점수를 기반으로 토큰 인덱스의 부분집합 $S_h$가 선택됩니다.
   • Query($Q$), Key($K$), Value($V$) 텐서가 수집되어 축소된 시퀀스 길이 $L' \ll L$를 가진 압축 텐서 $\hat{Q}, \hat{K}, \hat{V}$를 형성합니다.
   • 주요 특징: 선택은 헤드별로 독립적으로 수행되어, 서로 다른 헤드가 토큰의 서로 다른 부분집합에 집중할 수 있게 하며, 헤드별 이질성을 해결합니다.
   • 압축된 텐서는 밀집되고 연속적으로 유지되므로 FlashAttention 과 같은 최적화된 커널을 사용할 수 있습니다.
2. 압축 해제 (출력 복원):
   • 어텐션 출력 $\hat{O}$ (형태 $L' \times d$) 는 헤드가 특정하는 인덱스를 사용하여 원래 시퀀스 길이 $L \times d$의 영 (zero) 으로 초기화된 텐서로 다시 분산 (scatter) 됩니다.
   • 주요 특징: 선택되지 않은 위치는 0 으로 설정되지만, 이전 레이어의 **잔차 연결 (residual connection)**이 건너뛰어진 토큰의 정보를 보존합니다. 이를 통해 모델은 후속 레이어에서 토큰 중요도를 "다시 방문"하고 재평가할 수 있습니다.

동적 토큰 커버리지

어떤 토큰을 유지할지 그리고 얼마나 유지할지 결정하기 위해, 논문은 **동적 토큰 커버리지 (Dynamic Token Coverage)**를 도입합니다:

• 점수 매기기: 최근 쿼리를 모든 키에 대해 사용하여 경량 프록시 어텐션 맵을 계산합니다. 토큰 중요도 점수는 어텐션 가중치를 세로로 합산하여 도출됩니다.
• 예산 할당: 상위 $k$개 토큰을 유지하는 대신, 이 방법은 토큰을 오름차순 중요도로 정렬한 후 "꼬리" (가장 중요도가 낮은 토큰) 를 잘라내어 잘라낸 토큰의 누적 질량이 커버리지 임계값 $\tau$를 초과할 때까지 제거합니다. 이는 어텐션 노이즈에 대한 구조적 정규화로 제거를 처리합니다.
• 헤드별 선택: 글로벌 예산을 결정한 후, 각 헤드는 자신의 특정 점수에 기반하여 상위 $k$개 토큰을 독립적으로 선택합니다.

희소 레이어 선택

저자들은 모든 레이어에 희소성을 적용하면 성능이 저하된다는 것을 관찰합니다. 그들은 레이어 간 토큰 표현의 안정성을 측정하기 위해 **레이어 간 표현 드리프트 (Inter-Layer Representation Drift, $R_\ell$)**를 도입합니다.

• 낮은 드리프트 (안정적인 표현) 를 가진 레이어는 희소화를 위해 선택됩니다.
• 높은 드리프트 (표현이 크게 변화하는) 를 가진 레이어는 모델 동작을 보존하기 위해 밀집 상태로 유지됩니다.
• 이 선택은 모델당 한 번의 전처리 단계입니다.

3. 주요 기여

1. 가역적 토큰 단위 희소화: 토큰을 영구적으로 제거하지 않고 계산 비용을 줄이며 레이어 간 토큰 중요도의 동적 재평가를 가능하게 하는 새로운 "압축 및 압축 해제" 설계.
2. 헤드별 이질성: 서로 다른 헤드가 서로 다른 의미적 특징을 우선시한다는 사실을 포착하기 위해 각 어텐션 헤드에 대한 독립적인 토큰 선택을 지원합니다.
3. 원활한 호환성: 이 접근 방식은 기존 밀집 (FlashAttention) 및 희소 (Minference, FlexPrefill) 커널과 완전히 호환되며, 대체제가 아닌 보완적 가속기로 작용합니다.
4. 동적 예산 책정: 고정된 비율이 아닌 토큰 중요도 분포와 레이어별 안정성에 기반하여 희소성 수준을 결정하는 데이터 기반 접근법.

4. 실험 결과

실험은 RULER, InfiniteBench, Needle-in-a-Haystack과 같은 벤치마크를 사용하여 LLaMA-3.1-8B-Instruct 및 Mistral-Nemo-12B-Instruct에서 수행되었습니다.

• 정확도 - 지연 시간 트레이드오프:
  • Token Sparse Attention 은 일관되게 속도 - 정확도 트레이드오프를 개선합니다.
  • 128K 컨텍스트에서 FlashAttention 대비 최대 $\times 3.23$의 어텐션 속도 향상을 달성하며 1% 미만의 정확도 저하만 보입니다.
  • FlexPrefill과 결합할 때, 평균 정확도 (87.27%) 를 동일하게 유지하면서 속도 향상을 $\times 2.44$에서 $\times 2.76$으로 개선합니다.
• 제거 방법과의 비교:
  • 토큰 제거 방법 (PyramidInfer, FastKV, GemFilter) 보다 현저히 우수합니다. 예를 들어, 128K 에서 Token Sparse Attention 은 유사한 속도 향상에서 86.84% 정확도를 달성하는 반면 PyramidInfer 는 **78.49%**를 기록합니다.
• 확장성:
  • 속도 향상 이득은 컨텍스트 길이가 증가함에 따라 커집니다 (예: 4K 에서는 미미한 이득, 128K 에서는 상당한 이득). 이는 시퀀스 길이가 증가함에 따라 어텐션 희소성이 증가하기 때문입니다.
• 오버헤드:
  • 점수 매기기 및 인덱스 관리에 대한 오버헤드는 최소이며, 높은 희소성 수준에서도 전체 어텐션 지연 시간의 11% 미만을 차지합니다.

5. 의의

이 논문은 영구적인 토큰 제거에서 동적, 인터리브드 희소화로 전환함으로써 긴 컨텍스트 LLM 추론의 중요한 격차를 해소합니다.

• 이론적 통찰: 토큰 중요도는 레이어와 헤드를 거쳐 매우 동적임을 보여주며, 정적 가지치기보다는 유연한 선택 메커니즘이 필요함을 입증합니다.
• 실용적 영향: 모델을 재학습시키거나 상당한 정확도를 희생하지 않고 긴 컨텍스트 시나리오에서 막대한 속도 향상을 달성하기 위해 기존 최첨단 어텐션 커널 (FlashAttention, Minference 등) 위에 레이어링할 수 있는 "플러그 앤 플레이" 최적화를 제공합니다.
• 미래 방향: 영구적인 토큰 제거에 비해 가역적 압축이 $O(L^2)$ 병목 현상을 처리하는 데 더 우월한 전략임을 시사하며, 확장 가능한 추론을 위한 새로운 설계 지점을 확립합니다.