기술 요약: 확률적 어텐션

문제 제기

효율적인 어텐션 메커니즘은 트랜스포머 모델의 확장성에 필수적이지만, 기존 접근법들은 계산 비용과 수용 영역 (receptive field) 커버리지 간의 트레이드오프에 직면해 있습니다. **슬라이딩 윈도우 어텐션 (SWA)**은 토큰을 크기 $w$의 로컬 윈도우로 제한하여 레이어당 복잡도를 $O(nw)$로 줄입니다. 그러나 이러한 결정론적 국소성은 수용 영역을 선형 성장 ($\ell$레이어 후$\ell w$) 으로 제한하여,$w \ll n$일 때 시퀀스의 상당 부분을 도달 불가능하게 만듭니다. 글로벌 토큰, 수동으로 설계된 희소 패턴, 또는 블록 단위 라우팅과 같은 해결책들이 존재하지만, 이들은 종종 아키텍처의 복잡성을 증가시키거나 증명 가능한 커버리지 보장을 제공하지 못합니다. 반면, 전체 어텐션은 글로벌 연결성을 제공하지만 prohibitive 한 $O(n^2)$ 비용을 초래합니다.

본 논문은 약 130,000 개의 뉴런으로 구성된 연결 확률이 0.02% 에 불과한 **초파리 (Drosophila melanogaster) 의 연결체 (connectome)**에서 영감을 얻었습니다. 이러한 희소성에도 불구하고, 이 네트워크는 약 4.4 홉의 짧은 평균 경로 길이를 달성합니다. 이러한 "소셜 월드" 토폴로지는 밀집된 로컬 클러스터링과 광범위하게 분포된 확률적 쇼트컷으로 기능하는 장거리 연결 간의 상호작용에서 비롯됩니다. 저자들은 효율적인 글로벌 정보 흐름이 밀집된 모든-대-모든 연결성을 요구하는 것이 아니라, 분산된 확률적 장거리 연결로 보완된 구조화된 로컬 계산에서 발생할 수 있다고 주장합니다.

방법론: 확률적 어텐션 (SA)

저자들은 SWA 의 연결체 라우팅 전략을 모방하는 파라미터 없는 향상책인 **확률적 어텐션 (SA)**을 제안합니다.

핵심 메커니즘

SA 는 표준 SWA 계산 전후에 적용되는 3 단계 과정을 통해 고정된 로컬 윈도우를 확률적 글로벌 윈도우로 변환합니다:

1. 순열 (Permute): 각 레이어마다 독립적으로 무작위 순열 $\sigma \sim \text{Uniform}(S_n)$이 추출됩니다. 토큰 시퀀스 표현 ($Q, K, V$) 이 $\sigma$에 따라 재배열됩니다.
2. 윈도우 어텐션: 순열된 공간에서 표준 SWA 가 적용됩니다. 순열이 무작위이므로, 고정된 크기 $w$의 윈도우는 원래 시퀀스의 무작위 부분집합을 포괄하게 됩니다.
3. 순열 취소 (Undo Permute): 출력에 $\sigma^{-1}$를 적용하여 원래 순서를 복원합니다.

원래 토큰 공간에서 위치 $i$는 원래 거리 $|i-j|$와 관계없이 전체 시퀀스에 균일하게 분포된 기대 크기 $w$의 무작위 이웃 $\tilde{N}^\sigma_w(i)$에 어텐션합니다. 이는 연결체의 분산된 장거리 쇼트컷과 유사하게 임의의 토큰 쌍에 대해 약 $w/n$의 쌍별 연결 확률을 생성합니다.

이론적 특성

• 수용 영역 성장: SWA 의 수용 영역이 선형적으로 성장하여 (전체 커버리지를 위해 $O(n/w)$ 레이어 필요) SA 는 레이어 간 독립적으로 샘플링된 순열로 인해 지수적으로 성장하는 수용 영역을 제공합니다. 논문은 전체 시퀀스 커버리지가 $O(\log_w n)$ 레이어에서 달성됨을 증명합니다.
• 복잡도: 순열 및 역순열 연산은 $O(n)$ 인덱스 재배열입니다. 어텐션 계산은 $O(nw)$로 유지됩니다. 따라서 SA 는 표준 SWA 와 동일한 레이어당 $O(nw)$ 예산을 유지합니다.
• 인과적 제약: 자기회귀 설정에서 인과적 마스크는 순열된 공간에 적용되어, 토큰이 인과적으로 접근 가능한 선행 토큰에만 어텐션하도록 보장하면서도 확률적 장거리 연결을 가능하게 합니다.

게이트드 SA 와 SWA 결합

"소셜 월드" 영역 (높은 로컬 클러스터링 + 짧은 경로) 을 재현하기 위해, 저자들은 SA 와 SWA 를 결합한 듀얼-패스 아키텍처를 제안합니다:
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
여기서 $Y_{sa}$와 $Y_{swa}$는 각각 확률적 경로와 슬라이딩 윈도우 경로의 출력입니다. 게이트 가중치 ($g_{sa}, g_{swa}$) 는 독립적인 시그모이드 프로젝션을 통해 계산되어, 모델이 로컬 일관성 (SWA 에서) 과 글로벌 커버리지 (SA 에서) 사이의 균형을 학습할 수 있도록 합니다. 이 결합은 무시할 수 있는 파라미터를 추가하고 $O(nw)$ 복잡도를 유지합니다.

주요 기여

1. 확률적 어텐션 (SA): 무작위 순열을 적용하여 동일한 $O(nw)$계산 예산 내에서 지수적 수용 영역 성장 ($O(\log_w n)$) 을 달성하는 SWA 를 위한 파라미터 없는 플러그인 향상책.
2. 게이트드 SA + SWA: 생물학적 연결체에서 관찰된 소셜 월드 영역을 이론적 및 실증적으로 재현하여 로컬 클러스터링과 확률적 쇼트컷 사이의 균형을 맞추는 듀얼-패스 아키텍처.
3. 실증적 검증:
   • 프리-트레이닝: 처음부터 훈련된 360M 파라미터 모델에서, 게이트드 SA + SWA 결합은 평균 제로샷 정확도에서 가장 우수한 성과를 거두어 순수 SWA 를 능가하고, 퍼플렉시티에서는 전체 어텐션과 일치하면서도 하류 작업에서는 이를 능가했습니다.
   • 프리-트레이닝 없는 추론: 사전 훈련된 Qwen3-8B 및 Qwen3-30B-A3B 모델에 적용된 SA 는 동등한 계산 예산에서 SWA 를 일관되게 능가하고 **블록 어텐션 믹서 (MoBA)**와 일치하거나 능가했습니다. 윈도우 크기가 증가함에 따라 전체 어텐션 품질을 훨씬 빠르게 회복했습니다.

결과

• 프리-트레이닝: SA + SWA 모델은 평균 제로샷 정확도 35.9를 달성했습니다 (SWA 는 35.1, 전체 어텐션은 34.9). 이는 광범위한 담화 문맥이 필요한 LAMBADA에서 성능을 크게 향상시키는 동시에 WikiText에서 낮은 퍼플렉시티를 유지했습니다. 순수 SA 는 더 높은 퍼플렉시티를 보였는데, 이는 언어 모델링을 위한 로컬 일관성의 필요성을 확인시켜 주었습니다.
• Qwen3 에 대한 추론:
  • 유효 윈도우 크기 128 에서, Qwen3-8B 의 SA 는 전체 어텐션 기준치와 1 점 이내인 **70.9%**의 평균 정확도에 도달한 반면, SWA 는 62.2% 로 뒤처졌습니다.
  • Qwen3-30B-A3B 에서 SA 는 윈도우 크기 64 에서 **73.2%**의 정확도에 도달하여 SWA (47.0%) 와 MoBA (66.3%) 를 크게 능가했습니다.
  • SA 는 SWA 성능이 붕괴된 매우 작은 윈도우 크기 (예: $w=32$) 에서도 지식 집약적 작업 (MMLU, BoolQ) 에 대해 견고성을 보여주었습니다.
• 효율성: SA 는 선형 스케일링 $O(nw)$을 유지합니다. A100 GPU 에서 SA 는 시퀀스 길이 2,048 에서 전체 어텐션 대비 1.5 배의 속도 향상을 보였으며, 32,768 에서는 28 배로 증가했습니다.

중요성 및 주장

본 논문은 연결체에서 영감을 받은 확률적 라우팅이 효율적 어텐션의 표현력을 향상시키는 실용적인 원시 (primitive) 라고 주장합니다. 저자들은 글로벌 정보 흐름이 밀집된 모든-대-모든 연결성을 요구하는 것이 아니라, 구조화된 로컬 계산과 깊이 축적된 희소하며 확률적인 장거리 쇼트컷 간의 상호작용에서 발생할 수 있다고 논합니다.

이 연구의 중요성은 다음과 같습니다:

1. 간결성: SA 는 어텐션 메커니즘 자체에서 학습 가능한 파라미터가 필요 없으며 $O(n)$ 인덱스 오버헤드만 요구하므로, Mistral, Gemma 2, gpt-oss 등의 기존 SWA 커널에 대한 "플러그인" 업그레이드가 가능합니다.
2. 이론적 보장: SWA 의 선형 커버리지에 비해 유의미한 개선인 증명 가능한 $O(\log_w n)$ 커버리지 깊이를 제공합니다.
3. 상호 보완성: 이는 기존 선형 및 희소 접근법과 상호 보완적으로 제시되며, 수동으로 설계된 희소 패턴이나 글로벌 토큰의 복잡성 없이 로컬 효율성과 글로벌 표현력 사이의 간극을 연결하는 새로운 설계 원칙을 제공합니다.