Attention Sinks

어텐션 싱크는 내용이 질의와 관련이 없는데도 여러 헤드와 레이어에서 어텐션 가중치를 많이 받는 토큰으로, 대부분 시퀀스의 첫 토큰입니다. MIT, Meta AI, Carnegie Mellon University의 Guangxuan Xiao와 동료들은 ICLR 2024에서 발표한 어텐션 싱크를 활용한 효율적인 스트리밍 언어 모델에서 이 현상을 명명했습니다. 연구진은 Llama 2, MPT, Falcon, Pythia와 같은 대규모 언어 모델(LLM)이 슬라이딩 KV 캐시에서 첫 토큰을 제거하면 성능이 무너지고, 해당 토큰 네 개를 유지하면 생성 안정성이 회복된다는 사실을 보였습니다.

싱크가 형성되는 이유#

셀프 어텐션은 각 질의의 점수를 소프트맥스를 사용해 가중치로 변환하며, 가중치의 합은 항상 1입니다. 문맥에 헤드와 관련된 토큰이 없더라도 가중치는 어딘가에 배분되어야 합니다. 인과적 언어 모델에서는 첫 토큰이 이후 모든 토큰에 보이므로, 학습 과정에서 헤드가 필요로 하지 않는 가중치를 첫 토큰이 기본적으로 받게 됩니다. Llama 2에 관한 StreamingLLM 분석에서는 하위 두 레이어를 제외한 모든 레이어에서 모든 헤드의 어텐션이 초기 토큰에 집중되었습니다. 소프트맥스에 싱크 로짓을 추가하면 이러한 효과를 확인할 수 있습니다.

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

싱크가 없으면 헤드는 관련 없는 값 여섯 개에 전체 가중치를 분산하고 출력 노름은 0.83이 됩니다. 싱크 로짓이 3이면 싱크가 가중치의 77%를 차지하고 출력은 0.19로 감소하므로, 이 질의에 대해 헤드는 거의 아무 역할도 하지 않습니다. OpenAI의 오픈 웨이트 gpt-oss 모델은 모든 어텐션 헤드에 학습 가능한 싱크 로짓을 넣어 이를 구현합니다. 따라서 각 헤드는 출력에 영향을 주지 않는 남는 가중치를 배분할 위치를 갖습니다.

StreamingLLM을 사용한 스트리밍 생성#

무한히 생성하는 모델은 어떤 컨텍스트 윈도우도 초과하게 되므로 모든 토큰을 KV 캐시에 유지할 수 없습니다. 가장 최근 토큰만 유지하는 일반 윈도우 어텐션은 첫 토큰이 빠지는 순간 실패합니다. StreamingLLM은 초기 토큰 네 개의 키와 값을 최근 토큰의 롤링 윈도우 옆에 영구적으로 유지하고, 원래 텍스트 내 위치가 아닌 캐시 내 위치를 기준으로 위치를 지정합니다. 이때 매 단계마다 회전 위치 임베딩(RoPE)을 다시 적용합니다.

이 방식은 400만 토큰 이상에서도 언어 모델링을 안정적으로 유지했으며, 새 토큰이 추가될 때마다 캐시를 다시 계산하는 슬라이딩 윈도우보다 최대 22.2배 빠르게 실행되었습니다. 하지만 윈도우는 모델이 기억할 수 있는 범위를 여전히 제한합니다. 중간에서 제거된 토큰은 영구적으로 사라지므로 스트리밍 챗봇이나 AI 에이전트는 가장 최근 대화만 기억합니다. 저자들은 모든 샘플의 시작 부분에 학습 가능한 전용 플레이스홀더 토큰 하나를 넣어 모델을 사전 학습하기도 했으며, 이 토큰은 자체적으로 싱크 역할을 맡았습니다.

어텐션 싱크와 관련 용어 비교#

• 슬라이딩 윈도우 어텐션: 각 토큰을 최근 이웃으로 제한합니다. 전체 어텐션으로 학습된 모델에 추론 시 적용하면 싱크가 윈도우 밖으로 빠지는 순간 실패합니다.
• 컨텍스트 로트: 프롬프트가 컨텍스트 윈도우 안에 들어가는 상태로 길어질수록 정확도가 저하되는 현상입니다.

비전 모델의 싱크#

비전 트랜스포머(ViT)에서도 관련 현상이 나타납니다. 비전 트랜스포머에는 레지스터가 필요합니다에서 Meta와 Inria의 연구진은 지도 학습 및 자기지도 학습 ViT의 정보량이 적은 배경 패치에 비정상적으로 큰 노름을 가진 토큰이 나타나 내부 계산에 활용된다는 사실을 발견했습니다. 입력에 레지스터 토큰을 몇 개 추가하면 해당 역할을 대신하게 되어 특징 맵과 어텐션 맵의 아티팩트가 제거되고 조밀 예측 성능이 향상됩니다.

Ultralytics 패키지에서 실행되는 Meta의 SAM 2는 비슷한 방식으로 기준 프레임과 윈도우를 유지해 비디오 메모리를 관리합니다. 새 프레임마다 메모리 어텐션은 전체 비디오 동안 유지되는 사용자 프롬프트 프레임의 메모리 뱅크 항목과 가장 최근 프레임 6개의 항목을 읽으므로, 비디오 길이와 관계없이 프레임당 어텐션 비용이 일정하게 유지됩니다.