Sliding Window Attention

슬라이딩 윈도우 어텐션(SWA)은 각 토큰이 전체 시퀀스에 대해 수행하는 일반 셀프 어텐션 대신, 고정된 수의 인접 토큰만 참고하도록 제한합니다. 인과적 언어 모델에서는 가장 최근의 w개 토큰이 이에 해당합니다. 그러면 어텐션 연산량은 시퀀스 길이에 따라 이차적으로 증가하는 대신 선형적으로 증가하며, 생성 모델은 KV 캐시에 최근 w개 토큰의 키와 값만 보관하면 됩니다. Longformer는 2020년에 슬라이딩 윈도우를 어텐션의 핵심으로 삼았고, Mistral AI의 Mistral 7B는 2023년에 윈도우 크기 4,096토큰을 사용해 이를 널리 사용되는 오픈 대규모 언어 모델(LLM)에 적용했습니다.

깊이에 따라 윈도우가 확장되는 방식#

SWA의 한 계층은 w개 토큰만 보지만 Transformer 계층은 여러 개가 쌓입니다. 2계층에서 토큰은 1계층의 w개 토큰을 이미 요약한 은닉 상태를 읽으므로, 합성곱 신경망(CNN)의 수용 영역이 합성곱을 쌓을수록 넓어지는 것처럼 각 계층마다 정보가 약 w − 1개 위치만큼 더 이동할 수 있습니다. PyTorch에서는 윈도우가 불리언 어텐션 마스크이며, 해당 마스크를 자기 자신과 곱하면 정보가 이동하는 범위를 추적할 수 있습니다.

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

16개 토큰에서 윈도우는 계층당 136개 대신 58개의 점수를 계산하며, 마지막 토큰의 도달 범위는 계층마다 3개 토큰씩 늘어나 4, 7, 10, 13이 됩니다. 동일한 마스크를 PyTorch의 scaled_dot_product_attention에 attn_mask으로 전달할 수 있습니다. Mistral 7B는 4,096토큰 윈도우를 사용하는 32개 계층으로 이론상 약 131K 토큰의 어텐션 범위를 갖지만, 멀리 떨어진 정보가 전달되려면 여러 단계를 거쳐야 합니다.

Mistral 7B는 또한 키와 값을 w개의 슬롯으로 이루어진 순환 버퍼에 저장합니다. 위치 i의 항목은 i mod w 슬롯에 기록되어 가장 오래된 항목을 덮어씁니다. 32K 토큰 시퀀스에서는 품질 저하 없이 캐시 메모리를 8분의 1로 줄였으며, 고정된 캐시 크기 덕분에 생성이 아무리 오래 지속되어도 메모리 사용량은 일정합니다.

로컬 레이어와 글로벌 레이어의 교차 배치#

순수 윈도우 방식은 멀리 떨어진 과거의 정보를 한 단계만으로 복사할 수 없으므로, 최신 모델은 SWA 레이어와 전체 글로벌 레이어를 혼합합니다. Google의 Gemma 2는 4,096토큰 윈도우로 두 레이어를 하나씩 번갈아 배치했고, Gemma 3는 1,024토큰에 불과한 윈도우에서 글로벌 레이어 하나마다 로컬 레이어 다섯 개를 사용하여 최대 128K 토큰의 컨텍스트 윈도우에서도 KV 캐시를 작게 유지합니다. 글로벌 레이어만 전체 컨텍스트의 키와 값을 저장하므로, 6개 레이어 중 5개는 각각 1,024토큰만 보유합니다.

비전 트랜스포머는 모든 패치가 토큰인 고해상도 이미지에서도 동일한 혼합 방식을 사용합니다. Ultralytics 패키지에서 이미지 세그멘테이션을 수행하는 Meta의 SAM 3는 이미지 인코더의 32개 블록 중 28개에서 패치로 이루어진 겹치지 않는 24 × 24 윈도우 내에 어텐션을 적용하고, 7, 15, 23, 31번 블록에서는 1,008픽셀 이미지의 전체 72 × 72 패치 그리드에 걸쳐 글로벌 어텐션을 적용합니다.

SWA와 관련 용어 비교#

• Swin Transformer: 이미지를 고정된 겹치지 않는 윈도우로 나누고, 각 윈도우 내의 모든 토큰이 해당 윈도우를 공유합니다. 레이어 사이에서 그리드를 이동시켜 인접한 윈도우끼리 정보를 교환합니다.
• 희소 어텐션: 고정 패턴을 사용하는 더 넓은 계열입니다. SWA는 이 계열에서 가장 단순한 방식이며, 다른 방식에는 확장 윈도우, 글로벌 토큰 또는 무작위 연결이 추가됩니다.
• 어텐션 싱크: 시퀀스의 첫 번째 토큰으로, 남는 어텐션 가중치를 흡수합니다. 전체 어텐션으로 학습된 모델에 추론 시 윈도우를 적용할 때 이 토큰을 제거하면 생성이 중단되며, StreamingLLM은 윈도우와 함께 이 토큰을 캐시에 유지합니다.
• 변형 어텐션: 각 쿼리가 오프셋을 학습하는 몇 개의 샘플링 지점에 어텐션을 적용하므로, 이웃 영역이 콘텐츠에 따라 달라집니다.
• Flash attention: 정확한 어텐션을 위한 GPU 커널입니다. 구현에서 윈도우 크기를 지정할 수 있으므로, SWA 모델은 윈도우 밖의 블록을 완전히 건너뜁니다.