기술 요약: SSA (Sparse Sparse Attention)

1. 문제 정의

본 논문은 거대 언어 모델(LLM)에서 컨텍스트 길이에 따라 제곱 비례($O(N^2)$)로 증가하는 연산 병목 현상인 풀 셀프 어텐션(full self-attention) 문제를 다룹니다. 희소 어텐션(sparse attention) 메커니즘은 쿼리를 키(key)의 부분 집합으로 제한하여 복잡도를 하위 이차(sub-quadratic)로 줄이지만, 기존 방식들은 두 가지 근본적인 격차를 가집니다:

1. 어텐션 격차 (Attention Gap, 학습-추론 불일치): 풀 어텐션으로 학습된 모델에 학습 없이 희소 어텐션을 적용하는 방식(Full-Sparse)입니다. 이는 모델이 학습한 분포(풀 컨텍스트로 학습됨)와 추론 시의 분포(제한된 컨텍스트)가 일치하지 않기 때문에 성능 저하를 유발합니다. 본 논문은 이러한 설정에서의 근사 오차(approximation error)가 "드롭된 어텐션 질량(dropped attention mass, 희소 메커니즘에 의해 선택되지 않은 토큰에 할당된 어텐션 가중치)"에 선형적으로 비례함을 증명합니다.
2. 역량 격차 (Capability Gap, 그래디언트 결핍): 순수하게 희소 어텐션만으로 학습된 방식(Sparse-Sparse)은 분포 불일치는 피할 수 있으나 불완전한 학습 문제에 직면합니다. 희소 어텐션은 포워드 패스(forward pass) 중에 키를 누락시키기 때문에, 누락된 토큰들은 그래디언트를 전혀 받지 못합니다. 또한, 정규화(normalization)가 전체 집합이 아닌 선택된 부분 집합에 대해서만 수행되므로, 무관한 토큰을 억제하려는 경쟁적 압력이 약화됩니다. 결과적으로, 이러한 모델들은 종종 풀 어텐션 모델의 성능을 따라가지 못합니다.

2. 방법론: Sparse Sparse Attention (SSA)

저자들은 역량 격차와 어텐션 격차를 동시에 해결하기 위해 풀 어텐션 스트림과 희소 어텐션 스트림을 통합하는 단일 학습 프레임워크인 SSA를 제안합니다.

핵심 메커니즘

각 학습 단계에서 모델은 두 가지 모드 중 하나를 동일한 확률로 무작위 선택합니다:

• 희소 어텐션 스트림 (SA): 모델은 희소 추론(sparse inference)에 직접 최적화되도록 희소 어텐션(구체적으로 블록 희소 어텐션)을 계산하여, 학습-추론 간의 분포 불일치를 제거합니다.
• 풀 어텐션 스트림 (FA): 모델은 모든 토큰에 완전한 그래디언트 흐름을 보장하기 위해 풀 어텐션을 계산하며, 이를 통해 순수 희소 학습의 학습 결핍 문제를 해결합니다.

양방향 정렬 (Bidirectional Alignment)

두 스트림 사이의 격차를 메우고 본질적으로 더 희소한 분포를 강제하기 위해, SSA는 각 레이어에서 상대적 어텐션 정렬(counterpart attention alignment) 목적 함수를 도입합니다. 보조 출력은 정렬만을 목적으로 반대되는 어텐션 모드(예: 희소 스트림 중 풀 어텐션 출력 계산)를 사용하여 계산되며, 이후 레이어로 전파되지는 않습니다.

정렬 손실($L_{align}$)은 두 가지 구성 요소로 이루어집니다:

1. 희소성 손실 ($L_{sparsity}$): 풀 어텐션 출력이 희소 어텐션 출력을 모방하도록($h_{full} \approx h_{sparse}$) 유도합니다. 이는 풀 어텐션 메커니즘이 선택된 블록에 집중하도록 강제하여, 드롭된 어텐션 질량($\delta(t)$)을 줄이고 이론적 오차 범위를 좁힙니다.
2. 커밋 손실 ($L_{commit}$): 희소 어텐션 출력이 풀 어텐션 출력에 가깝게 유지되도록($h_{sparse} \approx h_{full}$) 규제합니다. 이는 희소 표현을 풀 어텐션 특징 공간에 고정시켜, 드리프트(drift)를 방지하고 희소 모델이 풀 모델의 역량을 유지하도록 보장합니다.

전체 목적 함수는 표준 교차 엔트로피 손실과 정렬 손실을 결합합니다:
$$\mathcal{L} = \mathbb{E}_{mode \sim \{full, sparse\}}[\mathcal{L}_{ce}^{mode}] + \alpha \mathcal{L}_{align}$$

3. 주요 기여

• 이론적 프레임워크: 본 논문은 순수 풀 학습이 어텐션 격차(분포 불일치)를 유발하고, 순수 희소 학습이 역량 격차(그래디언트 결핍)를 유발함을 확립합니다. 또한, 희소와 풀 어텐션 사이의 근사 오차가 드롭된 어텐션 질량에 선형적으로 비례함을 증명합니다.
• 통합 학습 프레임워크: SSA는 양방향 정렬을 통해 풀 어텐션과 희소 어텐션을 공동 학습하는 최초의 프레임워크로, 더 희소한 어텐션 분포를 명시적으로 유도하면서도 풀 어텐션의 역량을 유지합니다.
• 경험적 우월성: SSA는 MoBA 및 NSA와 같은 베이스라인 모델과 비교하여, 희소 및 풀 추론 모드 모두에서 최첨단(SOTA) 성능을 달하며 뛰어난 롱 컨텍스트 능력과 다양한 희소 예산에 대한 견고성을 입증했습니다.

4. 실험 결과

실험은 100B 토큰(8k 컨텍스트)으로 사전 학습된 1B 파라미터 모델(Llama-3.2-1B 아키텍처 기반)을 사용하여 32k 컨텍스트까지 지속 학습(continual training)하며 수행되었습니다.

• 상식 추론 및 퍼플렉서티 (Perplexity): SSA는 PIQA, HellaSwag, ARC 및 WikiText 퍼플렉서티에서 베이스라인(FullAttn, MoBA, NSA)을 일관되게 능가합니다. 특히, 희소 추론(Receptive Field = 256) 환경에서 SSA는 어텐션 격차로 인해 고통받는 FullAttn보다 우수한 성능을 보입니다.
• 롱 컨텍스트 성능:
  • Needle-in-a-Haystack (NIAH): SSA는 풀 어텐션 추론 시 8k 컨텍스트 내에서 100% 정확도를 달성하며, 다양한 수용 영역(receptive field)에서 MoBA 대비 우수한 성능을 유지합니다.
  • LongBench: SSA는 모든 추론 모드와 컨텍스트 확장 설정(외삽 및 지속 학습)에서 최고의 성능을 달성합니다.
  • 외삽 (Extrapolation): SSA는 수용 영역이 증가함에 따라 단조적인 성능 향상을 보이는 반면, MoBA는 퍼플렉서티와 LongBench에서 효과적인 외삽에 실패합니다.
• 효율성: 추론 시 SSA는 MoBA와 동일한 연산을 수행하며, 128k 컨텍스트에서 풀 어텐션 대비 2배의 속도 향상을 달성합니다. 보조 정렬 계산으로 인한 훈련 오버헤드는 미미합니다 (~17% 증가).

5. 의의 및 주장

본 논문은 희소 어텐션 방식이 효과적인 학습에 실패하거나(그래디언트 결핍 때문), 희소 추론에 대한 일반화에 실패하는(분포 불일치 때문) 역설을 SSA가 해결한다고 주장합니다.

• 내재적 희소성 (Inherent Sparsity): 풀 어텐션을 희소 출력에 정렬함으로써, SSA는 "본질적으로 더 희소한" 어텐션 분포를 학습합니다. 이는 드롭된 어텐션 질량을 줄이며, 이론적으로 근사 오차의 경계(bound)를 좁힙니다.
• 이중 모드 견고성 (Dual-Mode Robustness): 이 프레임워크는 모델이 효율성을 위한 희소 추론과 검색 집약적 작업을 위한 풀 추론 모두에서 잘 작동하게 하여, 성능 저하 없이 사용자에게 유연성을 제공합니다.
• 롱 컨텍스트 외삽: 저자들은 희소 어텐션으로 학습된 모델(특히 SSA)이 풀 어텐션 모델보다 RoPE 스케일링에 대한 강력한 견고성과 더 나은 롱 컨텍스트 외삽 능력을 보임을 관찰했습니다. 이는 더 적고 국소적인 토큰에 주의를 기울이는 것이 위치적 왜곡(positional distortions)에 대한 민감도를 낮추기 때문인 것으로 보입니다.

이 연구는 훈련 중에 높은 어텐션 희소성을 장려하는 것이 희소 추론뿐만 아니라 풀 추론에도 도움이 된다는 것을 보여주며, 확장 가능하고 효율적인 롱 컨텍스트 LLM을 설계하기 위한 새로운 방향을 제시합니다.