장문 컨텍스트 처리 기법: 어텐션은 O(n²)이므로 모델은 어떻게 책 전체를 읽을까? 윈도우, 싱크, 그리고 Lost-in-the-Middle
요약
본 글은 Self-attention의 계산 복잡도 O(n²) 문제를 다루며, 긴 컨텍스트 처리를 위한 세 가지 주요 기법을 소개합니다. 슬라이딩 윈도우 어텐션으로 비용을 선형적으로 줄이고, 희소/확장 어텐션을 통해 장거리 의존성을 유지하며, RoPE 스케일링과 YaRN 같은 방법을 사용해 모델의 학습 범위를 확장하는 방안을 설명합니다.
핵심 포인트
- Self-attention은 컨텍스트 길이(n)에 대해 O(n²) 비용 증가 문제가 있습니다.
- 슬라이딩 윈도우는 계산 복잡도를 O(n·w)로 낮춰 효율성을 높입니다.
- 희소/확장 어텐션은 로컬 정보와 글로벌 토큰을 결합하여 장거리 의존성을 유지합니다.
- RoPE 스케일링 및 YaRN 기법은 모델의 학습 범위를 효과적으로 확장하는 방법을 제시합니다.
Self-attention은 모든 토큰이 다른 모든 토큰을 바라보게 만듭니다. n개의 토큰에 대해 이는 n×n 점수 행렬을 의미하므로, 계산량과 메모리 모두 n²에 비례하여 증가합니다. 컨텍스트를 두 배로 늘리면 비용은 네 배가 됩니다. 이 단 하나의 사실 때문에 200k 토큰 창(window)은 어렵고, 전체 장문 컨텍스트 플레이북은 그 벽을 피하기 위한 일련의 트릭들입니다. 저는 인터랙티브 페이지를 만들었는데, 여기서 슬라이더를 움직이며 각 트릭이 어떻게 작동하는지 볼 수 있습니다. 모든 것은 어텐션 _패턴(patterns)_에 대해 클라이언트 측에서 시뮬레이션되며 모델 호출은 없습니다. 여기 투어 내용입니다.
이차적 벽 (The quadratic wall)
완전한 어텐션(Full attention)은 모든 쿼리 토큰을 모든 과거 키 토큰과 비교합니다. 컨텍스트 4k는 16M 점수이고, 128k는 16.8 _billion_입니다. 이는 단순히 확장되지 않으며, 아래의 모든 내용은 이를 우회하는 방법들입니다.
scores = Q @ K.T # shape (n, n) <- n×n 폭증
scores = scores.masked_fill(~causal, -inf) # 미래를 엿보지 않음
attn = softmax(scores, dim=-1) @ V # 비용 및 메모리 ~ O(n²)
슬라이딩-윈도우 (Sliding-window) 어텐션
가장 저렴한 해결책입니다: 각 토큰이 가장 최근의 w개의 키에만 주의를 기울이도록 합니다. 전체 삼각형 대신 대각선 _밴드(band)_를 사용하는 것입니다. 비용은 O(n²)에서 O(n·w)로 떨어지며, 이는 컨텍스트 대비 _선형적(linear)_입니다.
def streaming_mask(n, w, sink=4):
mask = sliding_window_mask(n, w)
mask[:, :sink] = True # 처음 sink개의 키는 모두에게 보이도록 유지
...
demo에서 이것을 지켜보는 것이 저를 사로잡은 부분입니다. 싱크(sinks)를 끄고, "스트리밍된 토큰" 슬라이더를 창 크기를 지나서 당기면, 싱크가 제거되는 순간 퍼플렉서티 곡선이 평평했던 ~8에서 260까지 급등합니다.
희소/확장 어텐션 (Sparse / dilated attention)
로컬 윈도우는 멀리 있지만 중요한 토큰을 인식하지 못합니다. 희소 어텐션은 저렴한 로컬 밴드(local band)를 유지하고 몇 개의 추가 링크만 더합니다: 모두가 볼 수 있는 글로벌 토큰, 그리고 시퀀스를 건너뛰는 스트라이드(dilated) 점프입니다. 이제 정보는 패턴을 ~O(n·w)로 유지하면서 전체 컨텍스트를 이동할 수 있습니다. 이것이 Longformer / BigBird의 방식입니다 — 로컬 + 글로벌 + 스트라이드.
def sparse_mask(n, w, stride=64, globals=(0,1,2,3)):
mask = sliding_window_mask(n, w)
for i in range(n):
...
RoPE 스케일링 — 학습된 길이를 넘어 확장하기
4k 토큰에서 학습된 모델에 30,000 위치를 입력하면 고장 납니다. 왜냐하면 그 회전 각도를 본 적이 없기 때문입니다. 해결책은 _외삽(extrapolate)하지 말고 보간(interpolate)하는 것_입니다. RoPE(Rotary embeddings)는 위치를 회전 각도로 인코딩합니다. 새로운 위치들을 모델이 이미 알고 있는 범위로 압축하면 각도가 친숙하게 유지됩니다. NTK-aware 스케일링과 YaRN이 이를 개선하며, 약간의 파인튜닝으로 품질을 복구할 수 있습니다. 이것이 4k 모델을 32k+ 모델로 만드는 방법입니다.
scale = train_len / target_len # 예: 4096 / 32768 = 1/8
pos = arange(target_len) * scale # 위치가 이제 [0, 4096)에 존재함
q, k = apply_rope(q, k, pos) # 스케일된 위치에 대한 회전 임베딩
Lost in the middle — 이 모든 것이 해결하지 못하는 문제
여기 가장 해결하기 어려운 문제가 있습니다. 심지어 어떤 사실이 컨텍스트에 포함될 수 있더라도, 모델은 그 정보가 시작 (최초 효과성, primacy) 또는 끝 (최신성, recency) 근처에 있을 때 훨씬 더 잘 검색합니다. 바늘을 중간에 묻으면 정확도가 U자 모양으로 떨어지며 — 컨텍스트가 커질수록 그 하락 폭은 더 깊어집니다. 이것은 패치로 고칠 수 있는 버그가 아니라, 근본적인 한계입니다. 긴 창(window)이 있다는 것은 토큰 하나하나가 읽힌다는 보장이 아니라, 단지 _용량_일 뿐입니다.
# 정확도는 위치에 따라 U자 모양을 띠므로 -> 최초 효과성 및 최신성을 활용:
ordered = [가장 관련성 높은 정보] + 중간 청크 + [두 번째로 관련성 높은 정보]
# 또한: 더 작은 top-k (잃어버릴 것이 적음) + 바늘을 가장자리로 재순위화.
데모에서 컨텍스트를 가로질러 바늘을 끌면 검색 곡선이 그 U자 모양으로 떨어지는 것을 볼 수 있습니다. 이는
AI 자동 생성 콘텐츠
본 콘텐츠는 Dev.to AI tag의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기