슬라이딩 윈도우 어텐션 (Sliding-window attention): 밴드 마스크(banded mask)를 통해 O(N²)을 O(N·W)로
요약
전체 셀프 어텐션의 O(N²) 복잡도를 O(N·W)로 줄이는 슬라이딩 윈도우 어텐션 기법을 설명합니다. 밴드 마스킹을 통해 계산 비용과 메모리를 획기적으로 절감하면서도, 모델의 깊이를 통해 장거리 문맥을 유지하는 원리를 다룹니다.
핵심 포인트
- 어텐션 복잡도를 O(N²)에서 O(N·W)로 최적화
- 밴드 마스킹을 통해 윈도우 밖의 점수를 -Infinity로 설정
- KV-캐시를 고정 크기 롤링 버퍼로 구현하여 메모리 효율 증대
- 모델의 깊이를 통해 장거리 도달 범위(receptive field) 확보
전체 셀프 어텐션 (Full self-attention)은 N개의 토큰 각각이 다른 모든 N개의 토큰을 참조하게 하므로, 비용과 메모리가 O(N²)처럼 폭발적으로 증가합니다. 즉, 컨텍스트(context)를 두 배로 늘리면 비용은 네 배가 됩니다. 이러한 쌍(pair)의 대부분은 낭비됩니다. 언어는 압도적으로 국소적(local)이며, 하나의 토큰이 수천 개의 위치만큼 떨어진 무언가를 직접 참조해야 하는 경우는 드뭅니다. 슬라이딩 윈도우 어텐션 (Sliding-window attention) (Longformer, Mistral, BigBird)은 각 토큰이 이전 W개의 토큰만을 참조하게 만듭니다. 이는 전체 사각형 대신 얇은 대각선 밴드(diagonal band)를 형성하여, 비용을 시퀀스 길이(sequence length)에 비례하는 O(N·W)로 낮춥니다. 그리고 장거리 도달 범위(long-range reach)를 잃는 것도 아닙니다. 깊이(depth)를 통해 이를 획득하게 됩니다. 저는 이 밴드와 확장되는 수용 영역 (receptive field)을 실시간으로 보여주는 데모를 제작했습니다. 작동 원리는 다음과 같습니다.
모든 셀을 결정하는 하나의 규칙
어텐션 행렬 (attention matrix)은 행(row)에 쿼리 (query)를, 열(column)에 키 (key)를 가집니다. 채워진 셀은 "이 쿼리가 저 키를 참조할 수 있음"을 의미합니다. 슬라이딩 윈도우는 키가 윈도우(window) 내에 있고 (디코더의 경우) 미래의 것이 아니라면 셀을 유지합니다.
function allowed(i, j){
if (causal && j > i) return false; // 미래를 엿보지 않음
const dist = i - j; // 키 j가 얼마나 뒤에 있는지
...
이것이 밴드의 전부입니다: 쿼리 i는 고정된 W 너비의 슬라이스인 키 [i−W+1 … i]를 봅니다.
마스킹(Mask), 그 후 윈도우에 대한 소프트맥스 (softmax)
불리언 마스크 (boolean mask)로 밴드를 구축하고, 소프트맥스 (softmax)를 적용하기 전에 윈도우를 벗어난 모든 점수(score)를 −∞로 설정합니다. 그러면 각 행은 오직 자신의 윈도우 내에서만 정규화됩니다. 즉, N개가 아닌 W개의 이웃들에 대한 가중치 혼합(weighted blend)이 됩니다.
if (!mask[i][j]) scores[i][j] = -Infinity; // 윈도우 밖의 쌍을 차단
const weights = scores.map(softmax); // 행의 합이 윈도우 내에서 1이 됨
절감 효과: O(N²) → O(N·W)
채워진 셀(cell)의 개수를 세어보면, 전체 어텐션 (full attention)은 N²이고 밴드 (band)는 대략 N·W이므로, W ≪ N일 때 절감 효과는 극적입니다. N=32k, W=4k인 경우 약 10억(1.0B) 개의 쌍이 약 1억 3,100만(131M) 개로 줄어들어 약 8배 가량 감소합니다. N=100k, W=4k일 때 윈도우는 전체 어텐션의 약 4%만을 계산합니다. 데모의 기본 설정인 N=24, W=4에서도 이미 6배의 차이를 보여줍니다. W를 N까지 올리면 밴드가 삼각형 전체를 채우게 되며, 이는 절감 효과 없이 다시 전체 어텐션으로 돌아감을 의미합니다. 이는 마스크 (mask)가 주장하는 대로 정확히 작동하고 있음을 보여주는 좋은 검증 (sanity check)입니다. 보너스로, 윈도우 기반 디코더 (windowed decoder)의 KV-캐시 (KV-cache)는 마지막 W개의 항목을 담는 고정 크기의 롤링 버퍼 (rolling buffer)가 됩니다. 즉, O(W) 메모리만 사용하며, 이미 생성한 양에 관계없이 무제한 길이에 대해 스트리밍 (streamable)이 가능합니다.
깊이가 거리를 확보한다 — ~L·W
당연한 반론이 있을 수 있습니다: 만약 토큰이 과거 W개만 본다면, 어떻게 장거리 구조 (long-range structure)를 학습할 수 있을까요? 이는 작은 컨볼루션 (convolution)들을 쌓아 큰 수용 영역 (receptive field)을 키우는 방식과 동일합니다. 레이어 (layer) ℓ의 윈도우가 레이어 ℓ+1의 윈도우로 정보를 전달하므로, 정보는 레이어당 하나의 윈도우만큼 더 멀리 도약합니다.
const field = Math.min(N, 1 + L * (W - 1)); // ≈ L·W <- 깊이를 통한 도달 거리
따라서 윈도우 크기가 W인 L개의 레이어는 컨텍스트 (context)가 약 L·W인 것처럼 동작하면서도, 모든 레이어는 O(N·W)의 저렴한 비용을 유지합니다. Mistral의 32개 레이어와 4k 윈도우는 4k 토큰을 훨씬 뛰어넘는 거리까지 도달합니다.
변형: 확장된 윈도우 (dilated windows)와 글로벌 토큰 (global tokens)
두 가지 확장 방식은 동일한 비용으로 도달 거리를 더 넓힙니다. 확장된 (dilated) 윈도우는 W개의 키 (key)를 유지하되 간격을 d만큼 띄워서, (W−1)·d + 1개의 토큰에 걸쳐 있게 합니다. 서로 다른 헤드 (head)들이 서로 다른 확장 간격 (dilation)을 사용합니다 (Longformer).
몇 개의 글로벌 토큰 (global tokens) (예: [CLS] 또는 태스크 토큰)은 모든 토큰을 어텐션하고 모든 토큰으로부터 어텐션을 받습니다 (전체 행과 열에 해당). 따라서 어떤 토큰이든 두 번의 도약 (two hops)만으로 다른 어떤 토큰에도 도달할 수 있습니다. 로컬 윈도우 (local window)와 글로벌 토큰의 조합은 정확히 Longformer/BigBird의 방식입니다.
슬라이딩 윈도우 어텐션 (Sliding-window attention)은 희소 근사 (sparse-approximate) 방식입니다. 즉, 계산되는 쌍(pairs)의 _대상_을 변경합니다. 이는 더 스마트한 IO 패턴을 통해 정확한 전체 어텐션 (full attention)을 계산하는 FlashAttention과는 상호 보완적입니다. FlashAttention은 계산 대상이 아닌 계산하는 _방법_을 변경하기 때문입니다. 따라서 실제 시스템에서는 이 둘을 함께 쌓아서 사용합니다. 예를 들어, Mistral은 하단에 Flash 커널을 두고 윈도우 패턴을 실행합니다.
제가 얻은 교훈은 다음과 같습니다: 이차 복잡도 (quadratic) 어텐션은 자연의 법칙이 아니라 기본 설정일 뿐입니다. 각 토큰을 로컬 밴드 (local band)로 제한하고, 정보가 이 밴드들을 가로질러 전달될 수 있도록 깊이 (depth)에 의존한다면, 긴 범위의 도달 거리 (long-range reach)를 가지면서도 선형 비용 (linear cost)을 얻을 수 있습니다.
W, N 및 레이어 수를 조절하며, 밴드는 얇아지는 반면 수용 영역 (receptive field)은 확장되는 모습을 확인해 보세요:
https://dev48v.infy.uk/ai/days/day48-sliding-window-attention.html
AI 자동 생성 콘텐츠
본 콘텐츠는 Dev.to AI tag의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기