논문 리뷰 / LLM / Attention Sink / Transformer / Language Model Pre-training
https://arxiv.org/abs/2410.10781
Xiangming Gu, Tianyu Pang, Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, Min Lin
When Attention Sink Emerges in Language Models: An Empirical View
Language Models (LMs) assign significant attention to the first token, even if it is not semantically important, which is known as attention sink. This phenomenon has been widely adopted in applications such as streaming/long context generation, KV cache o
arxiv.org
이 논문을 한 문장으로 요약하면, LLM이 의미적으로 중요하지 않은 첫 번째 token에 과도한 attention을 주는 attention sink 현상이 언제, 왜, 어떤 조건에서 등장하는지 실험적으로 분석하고, 그 원인이 softmax attention의 구조적 의존성과 key bias 역할에 있음을 보인 논문이다.
1. 왜 이 논문이 중요한가
대규모 언어모델을 분석하다 보면 매우 이상한 현상이 하나 관찰된다. 모델이 문장을 처리할 때, 의미적으로 특별히 중요하지 않은 첫 번째 token에 지나치게 큰 attention을 부여하는 경우가 많다는 것이다. 이를 attention sink라고 부른다. 예를 들어 문장의 첫 token이 단순한 BOS token이거나, 의미 없는 시작 token이거나, 심지어 random token이어도 여러 attention head가 이 위치에 강하게 attention을 몰아주는 현상이 나타난다.
attention sink는 단순히 “해석상 이상한 현상”에 그치지 않는다. 실제로 streaming generation, long-context inference, KV cache compression, inference acceleration, quantization 등 여러 응용에서 attention sink는 이미 실용적으로 활용되고 있다. 대표적으로 긴 문맥을 처리할 때 모든 과거 token의 KV cache를 유지하지 않고도, sink token 몇 개를 유지하면 성능 저하를 줄일 수 있다는 식의 아이디어가 여기에 연결된다.
그런데 흥미로운 점은, attention sink가 널리 활용되고 있음에도 불구하고 왜 생기는지, 언제 생기는지, 어떤 학습 조건에서 강해지거나 사라지는지에 대한 이해는 충분하지 않았다는 것이다. 이 논문은 바로 그 질문을 정면으로 다룬다. 저자들은 다양한 open-source LM과 직접 pre-training한 small LM을 이용해 attention sink의 발생 조건을 분석한다.
이 논문을 읽으면서 가장 먼저 드는 의문은 이것이다. “첫 번째 token이 중요해서 attention을 받는 것인가, 아니면 모델 구조상 attention을 저장할 장소가 필요해서 첫 번째 token을 이용하는 것인가?” 논문의 결론은 후자에 가깝다. 첫 번째 token은 의미 정보를 전달한다기보다, key bias처럼 작동하면서 extra attention score를 흡수하는 역할을 한다. 더 과감하게 말하면, attention sink token은 value computation에 실질적으로 기여하지 않는 attention 저장소처럼 행동할 수 있다.
Attention sink는 LLM의 long-context 처리와 KV cache 최적화에서 실용적으로 중요하지만,
그 발생 원인은 충분히 설명되지 않았다. 이 논문은 attention sink가 효과적인 LM pre-training, 충분한 데이터,
softmax attention의 정규화 구조, 그리고 key bias 역할과 밀접하게 연결되어 있음을 보인다.
2. Attention Sink란 무엇인가
attention sink는 Transformer language model에서 특정 token, 특히 첫 번째 token이 많은 attention head에서 과도하게 큰 attention score를 받는 현상이다. causal language model에서는 각 token이 자신보다 앞선 token들을 볼 수 있다. 이때 뒤쪽 token들이 실제 의미상 중요한 앞 token들보다 첫 번째 token에 반복적으로 큰 attention을 주는 경우가 관찰된다.
일반적인 직관으로 보면 이상하다. 첫 번째 token은 항상 문장의 핵심 정보를 담는 것이 아니다. 문서의 첫 단어가 의미상 중요할 수도 있지만, BOS token이나 EOS token처럼 단순한 boundary marker일 수도 있고, random sequence에서는 아무 의미도 없을 수도 있다. 그럼에도 모델은 첫 번째 token을 마치 특별한 anchor처럼 활용한다.
논문은 이 현상을 다음과 같은 관점에서 본다. 첫 번째 token이 semantic content를 많이 담아서 attention을 받는 것이 아니라, 모델이 softmax attention 구조 안에서 extra attention score를 저장하거나 흡수할 지점이 필요하고, 첫 번째 token이 그 역할을 하게 된다는 것이다.
여기서 중요한 점은 attention sink가 단순한 visualization artifact가 아니라는 것이다. 모델 내부의 attention distribution, hidden state, key/query geometry, pre-training dynamics와 연결되어 있다. 따라서 attention sink를 이해하는 것은 LLM의 내부 계산 방식, long-context behavior, KV cache 설계, quantization 안정성까지 이해하는 데 중요하다.
3. 논문의 질문
이 논문은 attention sink를 단순히 관찰하는 데서 그치지 않고, 다음 질문들을 체계적으로 던진다.
- Attention sink는 대형 모델에서만 나타나는가, 아니면 작은 LM에서도 나타나는가?
- Attention sink는 자연어 입력에서만 나타나는가, random token sequence에서도 나타나는가?
- Attention sink는 pre-training 중 언제 등장하는가?
- Learning rate, weight decay, batch size 같은 optimization 요소가 attention sink에 어떤 영향을 주는가?
- 데이터 양과 데이터 분포를 바꾸면 sink 위치가 바뀌는가?
- Loss function이나 attention mask를 바꾸면 sink가 이동하거나 사라지는가?
- Positional embedding, pre-norm/post-norm, attention design 중 무엇이 attention sink의 발생에 핵심적인가?
- Softmax attention이 아닌 sigmoid attention이나 다른 attention operation을 쓰면 어떻게 되는가?
이 질문들이 중요한 이유는, attention sink를 단순한 현상으로 보지 않고 언어모델 pre-training의 구조적 부산물로 분석하기 때문이다. 논문은 attention sink가 우연히 생기는 것이 아니라, 충분히 학습된 LM에서 자연스럽게 등장하는 현상이며, softmax attention의 정규화 구조와 강하게 연결되어 있다고 주장한다.
4. Transformer LM과 attention sink의 기본 구조
논문은 auto-regressive language model을 Transformer decoder block의 반복으로 정의한다. 입력 token sequence는 embedding layer를 거쳐 hidden states가 되고, 각 block은 multi-head self-attention과 feed-forward network로 구성된다. 대부분의 현대 LLM은 pre-norm 구조를 사용한다. 즉, attention이나 FFN에 들어가기 전에 layer normalization 또는 RMSNorm을 적용한다.
causal self-attention에서는 각 token이 자신보다 뒤에 있는 token을 볼 수 없다. 따라서 i번째 token은 1번째부터 i번째 token까지만 attend할 수 있다. 이 구조에서 첫 번째 token은 모든 뒤쪽 token들이 공통적으로 볼 수 있는 가장 오래된 token이다. attention sink는 바로 이 첫 번째 token에 attention이 몰리는 현상이다.
document packing도 중요하다. LM pre-training에서는 여러 문서를 이어 붙인 뒤 context length 단위로 자른다. 이때 chunk의 첫 token은 반드시 문서의 진짜 첫 token일 필요가 없다. 어떤 경우에는 문서 중간 token일 수도 있고, BOS 또는 EOS일 수도 있다. 그럼에도 attention sink가 첫 번째 위치에 나타난다면, 이는 그 token의 의미보다 위치와 학습 구조가 더 중요할 가능성을 시사한다.
5. 첫 번째 token은 key bias처럼 작동한다
논문에서 가장 중요한 분석 중 하나는 “첫 번째 token이 정확히 어떤 방식으로 sink가 되는가”이다. 기존 연구에서는 첫 번째 token의 hidden state norm이 비정상적으로 크다는 massive activation 관찰이 있었다. 이 논문도 LLaMA3-8B Base에서 첫 번째 token의 hidden state L2 norm이 다른 token보다 매우 커지는 현상을 확인한다.
그런데 저자들은 여기서 한 걸음 더 들어간다. hidden state norm이 크다고 해서 key norm이나 value norm도 반드시 큰 것은 아니다. 오히려 논문은 첫 번째 token의 key와 value norm이 다른 token보다 작을 수 있음을 보인다. 그렇다면 왜 attention이 첫 번째 token에 몰릴까?
핵심은 query-key 사이의 각도다. attention score는 단순히 key norm이 큰 token에만 커지는 것이 아니다. query와 key의 dot product는 norm product와 cosine similarity의 곱으로 볼 수 있다. 논문은 첫 번째 token의 key가 다른 token들의 query와 작은 각도를 이루면서 큰 cosine similarity를 만들고, 이로 인해 attention sink가 발생한다고 해석한다.
이 분석은 매우 중요하다. attention sink를 단순히 “첫 번째 token의 activation이 커져서 생긴다”고만 설명하면 부족하다. 실제로는 첫 번째 token의 key가 여러 query와 잘 맞는 방향으로 정렬되어, 마치 attention score를 흡수하는 bias vector처럼 작동한다. 따라서 논문은 attention sink가 value information을 전달하기보다 key bias 역할에 가깝다고 주장한다.
6. Attention Sink를 어떻게 측정하는가
attention sink는 시각적으로는 attention map의 vertical stripe로 쉽게 볼 수 있다. 하지만 논문은 다양한 모델과 학습 조건을 비교하기 위해 정량적 metric이 필요하다고 본다. 이를 위해 threshold-based metric을 정의한다.
기본 아이디어는 간단하다. 각 attention head에서 특정 token, 주로 첫 번째 token이 받는 attention 비율이 일정 threshold를 넘으면, 그 head는 attention sink를 가진 것으로 간주한다. 그런 다음 전체 block과 head 중 attention sink가 나타나는 비율을 모델 단위 metric으로 계산한다.
threshold가 낮으면 attention sink를 느슨하게 정의하고, threshold가 높으면 더 엄격하게 정의한다. 논문은 token length에 덜 민감하면서도 충분히 엄격한 threshold를 선택해 실험을 진행한다. 이 metric은 완전한 이론적 정의라기보다, 다양한 조건에서 attention sink가 얼마나 강하게 나타나는지를 비교하기 위한 실험적 도구에 가깝다.
7. Attention Sink는 입력 도메인과 무관하게 나타나는가
논문은 attention sink가 특정 자연어 도메인에서만 나타나는지 확인한다. 이를 위해 The Pile의 여러 도메인 데이터를 사용하고, 자연어가 아닌 random token sequence도 실험한다. 결과적으로 입력 도메인은 attention sink metric에 큰 영향을 주지 않는다. 즉, attention sink는 특정 문체나 주제 때문이라기보다 모델 구조와 학습 방식에서 비롯된 현상에 가깝다.
더 흥미로운 것은 random token sequence에서도 attention sink가 유지된다는 점이다. 자연어 문법이나 의미가 없는 random token 입력에서도 GPT2-XL, Mistral-7B, LLaMA2-7B, LLaMA3-8B 등은 여전히 첫 번째 token에 attention sink를 보인다. 이는 sink token이 의미 정보를 담아서 선택되는 것이 아니라는 강한 증거다.
반면 repeated token sequence에서는 Mistral과 LLaMA 계열에서 attention sink가 사라진다. 모든 token이 동일하게 반복되면 첫 token만 특별한 hidden state나 key geometry를 갖기 어려워지고, attention sink가 분산되거나 사라질 수 있다. 이 결과는 첫 번째 token의 uniqueness가 attention sink에 중요하다는 해석을 가능하게 한다.
| 모델 | Natural | Random | Repeat | 해석 |
|---|---|---|---|---|
| GPT2-XL | 77.00 | 70.29 | 62.28 | random sequence에서도 sink 유지 |
| Mistral-7B | 97.49 | 75.21 | 0.00 | 반복 token에서는 sink 소실 |
| LLaMA2-7B Base | 92.47 | 90.13 | 0.00 | 자연어 의미보다 위치/구조 효과가 큼 |
| LLaMA3-8B Base | 99.02 | 91.23 | 0.00 | 강한 attention sink 관찰 |
8. Chat 모델과 Base 모델의 차이는 큰가
저자들은 instruction tuning이 attention sink를 크게 바꾸는지도 확인한다. base model과 chat model을 비교한 결과, attention sink metric은 큰 차이를 보이지 않는다. 예를 들어 LLaMA3-8B는 Base 99.02, Chat 98.85로 거의 동일하다. Mistral-7B도 Base 97.49, Chat 88.34로 여전히 높은 sink metric을 유지한다.
이 결과는 attention sink가 instruction tuning 이후에 생기는 것이 아니라, pre-training 단계에서 이미 형성되는 현상임을 시사한다. 따라서 논문은 이후 분석의 초점을 LM pre-training에 맞춘다.
| 모델 | Base | Chat | 해석 |
|---|---|---|---|
| Mistral-7B | 97.49 | 88.34 | chat tuning 이후에도 sink 유지 |
| LLaMA2-7B | 92.47 | 92.88 | 거의 변화 없음 |
| LLaMA2-13B | 91.69 | 90.94 | instruction tuning 영향 작음 |
| LLaMA3-8B | 99.02 | 98.85 | pre-training에서 이미 형성된 현상 |
9. 작은 모델에서도 Attention Sink가 나타나는가
attention sink는 대규모 모델에서만 나타나는 현상처럼 보일 수 있다. 하지만 논문은 그렇지 않다고 말한다. LLaMA, GPT-2, OPT, Pythia 등 다양한 모델 family를 분석한 결과, attention sink는 작은 LM에서도 나타난다. 특히 Pythia-14M처럼 매우 작은 모델에서도 attention sink가 관찰된다.
다만 모델 family마다 scaling 양상은 다르다. Pythia family에서는 모델 크기가 커질수록 attention sink가 더 뚜렷해지는 경향이 있지만, 모든 모델 family에서 동일한 방식으로 scale-dependent하게 증가하는 것은 아니다. 즉, attention sink는 model size만으로 설명되지 않고, architecture, data, optimization, loss design이 함께 영향을 준다.
10. Attention Sink는 pre-training 중 언제 등장하는가
이 논문의 핵심 실험 중 하나는 직접 작은 LLaMA-style LM을 pre-training하면서 attention sink가 언제 등장하는지 추적한 것이다. 저자들은 약 60M parameter 규모의 LLaMA-style 모델을 구성하고, 5B token 규모의 The Pile 데이터를 사용해 학습한다. context length는 2048, batch size는 1M tokens, training step은 20k로 설정한다.
실험 결과 attention sink는 학습 초기부터 존재하는 것이 아니다. 일정한 optimization step이 진행되고 training loss와 validation loss가 충분히 낮아진 뒤에 등장한다. 기본 설정에서는 대략 1k에서 2k step 사이에 attention sink가 나타나기 시작하고, 학습이 진행될수록 더 뚜렷해진다.
이 결과는 attention sink가 랜덤 초기화나 우연한 구조 때문만은 아니라는 점을 보여준다. 모델이 실제로 language modeling objective를 효과적으로 최적화하면서, attention distribution을 효율적으로 구성하는 과정에서 sink가 형성된다고 볼 수 있다.
Attention sink는 학습 초기에 바로 나타나는 고정된 구조가 아니라,
LM이 충분히 최적화된 이후 등장하는 학습된 내부 행동에 가깝다.
11. Optimization: learning rate, weight decay, batch size의 영향
논문은 optimization 요소들이 attention sink에 어떤 영향을 미치는지도 분석한다. 먼저 learning rate를 줄이면 training loss가 낮아지는 속도가 느려지고, attention sink의 등장도 지연된다. 또한 작은 learning rate로 학습된 모델에서는 attention sink가 덜 뚜렷해지는 경향이 있다. 하지만 learning rate를 너무 작게 하면 최적화 자체가 어려워지고, 모델 성능도 떨어진다.
weight decay는 attention sink를 더 강하게 유도하는 경향을 보인다. weight decay ratio가 0.0에서 0.5로 증가할 때 attention sink metric은 15.20에서 41.08까지 증가한다. 그러나 weight decay가 지나치게 커지면 validation loss가 악화되고, attention sink도 오히려 사라진다. 즉, weight decay는 일정 범위 안에서는 sink를 촉진하지만, 과도하면 모델 최적화 자체를 망가뜨린다.
반면 batch size만 바꾸는 것은 attention sink에 큰 영향을 주지 않는다고 보고한다. 따라서 논문은 attention sink가 단순 batch noise의 산물이 아니라, effective optimization과 loss landscape의 구조적 결과에 더 가깝다고 해석한다.
| Weight Decay | Attention Sink Metric | Validation Loss | 해석 |
|---|---|---|---|
| 0.0 | 15.20 | 3.72 | sink는 존재하지만 약함 |
| 0.1 | 18.18 | 3.73 | 기본 설정 수준 |
| 0.5 | 41.08 | 3.80 | weight decay가 sink를 강하게 유도 |
| 1.0 | 37.71 | 3.90 | sink는 여전히 크지만 성능 악화 |
| 2.0 | 6.13 | 4.23 | 과도한 regularization으로 최적화 악화 |
| 5.0 | 0.01 | 5.24 | 학습 자체가 크게 망가지며 sink 소실 |
12. Data Distribution: 충분한 데이터가 있어야 sink가 생긴다
논문은 attention sink가 데이터 양과 어떤 관계를 가지는지도 분석한다. 기본 실험에서는 5B tokens를 사용하지만, 이를 2.5B, 500M, 200M, 100M, 50M으로 줄여가며 같은 compute budget에서 학습한다.
결과는 분명하다. 데이터가 충분하지 않으면 모델은 overfitting 경향을 보이고, attention sink도 사라진다. 이는 매우 흥미로운 결과다. attention sink가 단순한 overfitting artifact라면 데이터가 부족할수록 더 강해질 수도 있을 것이다. 하지만 실제로는 반대다. attention sink는 충분한 데이터 위에서 모델이 일반화 가능한 language modeling behavior를 학습할 때 등장한다.
즉, attention sink는 “나쁜 학습의 부산물”이라기보다, 충분한 데이터와 효과적인 최적화를 통해 모델이 softmax attention 구조 안에서 찾아낸 효율적인 계산 전략일 수 있다.
13. Sink 위치는 고정되어 있는가
attention sink는 보통 첫 번째 token에서 나타난다. 하지만 이 논문은 sink 위치가 절대적으로 고정된 것이 아니라, data distribution과 loss function에 의해 이동할 수 있음을 보인다.
예를 들어 document packing 후 chunk의 첫 token을 random하게 다시 샘플링하면 attention sink가 오히려 더 강해질 수 있다. 이는 sink token이 semantic information을 많이 담아서 attention을 받는 것이 아니라는 점을 다시 보여준다. 또한 첫 번째 token이 아닌 두 번째 token이 특정 조건에서 더 특별한 위치가 되도록 만들면 attention sink가 두 번째 token으로 이동할 수 있다.
더 직접적인 실험으로, 저자들은 특정 learnable token을 두 번째 또는 세 번째 위치에 고정한다. 그러면 attention sink는 첫 번째 token이 아니라 그 고정 token 위치로 이동한다. 즉, attention sink는 “첫 번째 위치 자체”의 필연적 성질이라기보다, 모델이 bias-like token을 학습할 수 있는 위치와 pre-training setup의 결과로 이해해야 한다.
Attention sink는 항상 첫 token에만 고정되는 성질이 아니다. 데이터 분포와 loss 설계가 바뀌면 sink 위치도 이동할 수 있다.
따라서 sink는 위치 자체보다 학습 과정에서 형성되는 key-bias 역할로 이해하는 편이 더 정확하다.
14. Loss Function: prefix LM과 shifted window attention
논문은 loss function과 attention mask가 attention sink 위치에 미치는 영향도 분석한다. 먼저 prefix language modeling을 고려한다. 일반적인 auto-regressive LM에서는 첫 번째 token이 prediction target이 아니다. 이를 일반화해 앞쪽 여러 token을 prefix로 두면, attention sink는 첫 번째 token 하나가 아니라 prefix token들 사이에서 나타난다.
이는 매우 중요한 관찰이다. sink token은 “예측되지 않는 token” 또는 “다른 token들이 공통적으로 참조할 수 있는 token”과 관련이 있다. prefix length를 늘리면 모델은 그 prefix token들 중 일부를 attention sink로 활용한다.
shifted window attention도 분석한다. window attention에서는 각 token이 전체 과거 문맥을 보는 것이 아니라 제한된 window 안의 token만 볼 수 있다. 이 경우 중요한 것은 “각 window 안의 상대적인 첫 token”이 아니라, 전체 sequence의 절대적인 첫 token이다. 논문은 attention sink가 relative first token이 아니라 absolute first token에 나타나는 경향이 있으며, window size가 작아지면 attention sink가 덜 나타난다고 보고한다.
15. Positional Embedding은 attention sink의 원인인가
첫 번째 위치에 attention sink가 나타난다면, 자연스럽게 positional embedding이 원인일 것이라고 의심할 수 있다. 저자들은 이를 확인하기 위해 Rotary, NoPE, absolute PE, learnable PE, relative PE, ALiBi 등을 비교한다.
결과적으로 positional embedding의 종류는 attention sink 발생 자체를 막지 못한다. Rotary에서도, ALiBi에서도, NoPE에서도 attention sink는 나타난다. relative PE는 학습이 어렵고 validation loss가 나빠지지만, 그 경우에도 attention sink metric 자체는 관찰된다.
| Positional Embedding | Attention Sink Metric | Validation Loss | 해석 |
|---|---|---|---|
| NoPE | 20.35 | 3.81 | 명시적 PE가 없어도 sink 발생 |
| Absolute PE | 32.73 | 3.74 | sink 유지 |
| Learnable PE | 33.13 | 3.79 | sink 유지 |
| Relative PE | 35.53 | 5.45 | 학습은 어려우나 sink는 관찰 |
| ALiBi | 20.78 | 3.71 | sink 유지 |
| Rotary | 18.18 | 3.73 | 기본 설정 |
이 결과는 attention sink를 단순히 positional embedding의 부산물로 설명할 수 없다는 점을 보여준다. 위치 정보는 sink 위치와 관계가 있을 수 있지만, attention sink 발생 자체의 핵심 원인은 positional embedding보다 attention operation과 optimization dynamics에 더 가깝다.
16. Pre-norm과 Post-norm 구조의 영향
현대 LLM은 주로 pre-norm 구조를 사용한다. pre-norm에서는 layer normalization을 attention이나 FFN에 들어가기 전에 적용하고, residual connection을 통해 hidden state가 다음 block으로 전달된다. 이 구조에서는 특정 block에서 massive activation이 생기면, residual path를 통해 이후 block까지 유지될 수 있다.
그렇다면 post-norm으로 바꾸면 attention sink가 사라질까? 논문은 그렇지 않다고 보고한다. post-norm 구조에서도 attention sink는 여전히 존재한다. 다만 massive activation이 최종 hidden state가 아니라 post-layer normalization 이전 feature에서 나타난다. 즉, normalization 위치를 바꿔도 모델은 attention sink를 만드는 다른 방식의 내부 표현을 형성한다.
17. Attention Sink는 key bias 역할을 한다
논문의 가장 중요한 결론 중 하나는 attention sink가 key bias처럼 작동한다는 것이다. 저자들은 learnable sink token, KV bias, K bias, V bias를 도입해 비교한다.
결과적으로 key bias가 있는 경우 attention sink는 첫 번째 token에서 bias 위치로 이동한다. KV bias나 K bias를 추가하면 모델 성능은 유지되면서도 attention sink가 첫 번째 token에서 사라지고, 새로 도입된 key bias 쪽으로 이동한다. 반면 value bias만 추가하는 경우에는 attention sink를 이동시키지 못한다.
이 결과는 매우 강력한 해석을 제공한다. Attention sink의 본질은 value information을 많이 전달하는 것이 아니라, attention score를 받아내는 key-side structure에 가깝다. 즉, sink token은 key를 통해 query들과 잘 정렬되고, extra attention score를 저장하는 bias vector처럼 작동한다.
| 설계 | Attention sink 변화 | 해석 |
|---|---|---|
| Learnable sink token | sink가 learnable token으로 이동 | 고정된 특수 token이 sink 역할을 대신함 |
| KV biases | 첫 token sink가 bias 쪽으로 이동 | key/value bias가 sink 역할 제공 |
| K biases | 첫 token sink가 거의 사라지고 key bias로 이동 | attention sink의 핵심은 key-side bias |
| V biases | first-token sink 유지 | value만으로는 sink 위치를 바꾸기 어려움 |
이 지점에서 Vision Transformer의 register token과도 비교할 수 있다. ViT register는 global information을 모으는 의미 있는 latent slot처럼 해석될 수 있다. 반면 이 논문에서의 attention sink는 value computation에 실질적으로 기여하지 않는 extra attention score 저장소에 더 가깝다. 따라서 LLM의 attention sink는 ViT register와 유사해 보이지만, 기능적으로는 다를 수 있다.
18. Softmax Attention이 핵심 원인인가
논문은 attention operation 자체를 바꾸는 실험을 통해 더 근본적인 질문을 던진다. attention sink가 softmax attention의 구조적 결과라면, softmax normalization을 제거하거나 다른 kernel을 사용했을 때 sink가 사라져야 한다.
softmax attention에서는 각 query가 과거 token들에 대해 attention score를 분배하고, 그 합은 1이 된다. 즉, attention score들은 서로 독립적이지 않다. 어떤 token에 attention을 덜 주려면 다른 token에 더 주어야 하고, 전체 probability mass는 항상 보존된다. 논문은 이 attention score 간 내부 의존성이 attention sink를 유발하는 핵심 요인 중 하나라고 본다.
이를 검증하기 위해 저자들은 sigmoid attention without normalization, ELU+1 kernel, identity kernel, MLP kernel 등 다양한 attention operation을 실험한다. 특히 normalization을 제거한 sigmoid attention에서는 validation loss가 comparable하면서도 attention sink가 거의 사라진다. 1B parameter 규모까지 scale-up해도 sigmoid attention without normalization에서는 massive activation과 attention sink가 나타나지 않는다.
| Attention Operation | Normalization | Attention Sink 경향 | 해석 |
|---|---|---|---|
| Softmax attention | 있음 | sink 발생 | attention score가 합 1로 묶임 |
| Sigmoid attention | 없음 | sink 거의 소실 | token 간 attention score 의존성 완화 |
| Sigmoid attention | 있음 | sink 다시 발생 | 정규화가 sink 발생에 중요함 |
| MLP kernel | 조건에 따라 다름 | sink 없음 | softmax식 score 경쟁 구조가 약화됨 |
이 결과는 논문의 가장 강한 메시지다. attention sink는 단순히 첫 token, BOS token, positional embedding, norm explosion만으로 설명되지 않는다. softmax attention이 모든 token의 attention weight를 하나의 확률분포로 정규화하고, token들이 제한된 attention mass를 경쟁적으로 나눠 갖는 구조 자체가 sink token의 필요성을 만든다는 해석이 가능하다.
19. 논문을 읽으면서 생겼던 의문과 해소 과정
이 논문을 읽으면서 가장 먼저 생기는 의문은 “attention sink가 정말 의미 없는 token에 대한 attention인가?”이다. 모델이 첫 token에 attention을 많이 준다면, 혹시 그 token이 실제로 중요한 global information을 담고 있는 것 아닐까? 논문은 이에 대해 key와 value를 분리해서 답한다. 첫 번째 token은 key bias로는 강하게 작동하지만, value computation에 의미 있게 기여한다고 보기 어렵다. 특히 K bias만으로 sink를 이동시킬 수 있다는 결과는 attention sink의 본질이 value 정보가 아니라 key-side score storage에 있음을 보여준다.
두 번째 의문은 “attention sink는 단순히 BOS token 때문에 생기는 것 아닌가?”이다. 논문은 random token sequence와 document packing, fixed token position 실험을 통해 이 의문을 해소한다. sink는 BOS token에만 의존하지 않는다. 첫 token이 random token이어도 sink가 나타나고, 고정 token을 다른 위치에 두면 sink가 그 위치로 이동한다. 따라서 attention sink는 특정 token semantic보다 학습 구조와 위치적 역할에 더 가깝다.
세 번째 의문은 “massive activation이 원인인가, 결과인가?”이다. 기존 연구는 massive activation에 주목했지만, 이 논문은 key geometry와 attention operation을 함께 본다. 첫 token hidden state norm은 크지만 key/value norm은 작을 수 있으며, query-key cosine similarity가 attention sink에 크게 기여한다. 또한 K bias를 도입하면 massive activation 없이도 sink 역할을 대체할 수 있다. 따라서 massive activation은 attention sink와 관련이 있지만, 단독 원인으로 보기에는 부족하다.
네 번째 의문은 “attention sink를 없애면 모델이 더 좋아지는가?”이다. 논문은 이 질문에 조심스럽게 접근한다. sigmoid attention without normalization은 sink를 제거하면서도 comparable한 validation loss를 보인다. 그러나 이것이 모든 실전 LLM에서 softmax를 대체해야 한다는 의미는 아니다. 오히려 논문의 핵심은 attention sink가 softmax attention의 구조적 성질과 연결되어 있다는 점을 실험적으로 보였다는 데 있다.
20. 이 논문의 강점
- 현상 분석이 매우 체계적이다. 단순히 attention map을 보여주는 수준이 아니라 입력 종류, 모델 크기, pre-training step, optimization, data distribution, loss function, architecture를 모두 분해해 분석한다.
- attention sink를 정량화한다. threshold-based metric을 통해 여러 모델과 실험 조건을 비교할 수 있게 했다.
- key bias 관점이 설득력 있다. attention sink가 value information을 전달한다기보다 key-side bias로 extra attention score를 저장한다는 해석을 실험적으로 뒷받침한다.
- softmax attention의 구조적 한계를 지적한다. attention score들이 합 1로 정규화되는 구조가 sink token의 필요성을 만든다는 관점을 제시한다.
- 실용적 응용과 연결된다. attention sink는 long-context streaming, KV cache compression, quantization 등 실제 LLM 시스템 최적화와 밀접하다.
21. 한계
첫 번째 한계는 연구 범위가 주로 첫 번째 token sink에 집중되어 있다는 점이다. 실제 LLM에서는 period, newline, delimiter 같은 특정 word token도 sink처럼 작동할 수 있다. 논문도 future work에서 이러한 non-first sink token을 추가로 분석할 필요가 있다고 언급한다.
두 번째로, attention sink를 줄이는 것이 항상 실용적으로 바람직한지는 별개의 문제다. attention sink는 해석상 이상해 보이지만, long-context inference나 KV cache compression에서는 오히려 유용하게 활용된다. 따라서 attention sink를 제거할지, 별도 bias로 분리할지, 유지하면서 최적화할지는 응용 목적에 따라 달라진다.
세 번째로, sigmoid attention without normalization이 1B 규모에서 sink를 제거했다는 결과는 매우 흥미롭지만, 대규모 frontier LLM 수준에서 같은 결론이 유지되는지는 추가 검증이 필요하다. validation loss가 comparable하더라도 downstream task, long-context reasoning, instruction following, safety alignment까지 모두 검증된 것은 아니다.
네 번째로, attention sink metric은 threshold 기반 경험적 지표다. 모델 간 비교를 가능하게 해주지만, attention sink의 모든 기능적 의미를 완전히 포착한다고 보기는 어렵다. 특히 attention score와 실제 causal contribution은 다를 수 있기 때문에, 후속 연구에서는 intervention이나 causal tracing 기반 분석이 필요하다.
22. 후속 연구 관점에서의 의미
이 논문은 LLM 내부 메커니즘 해석에서 중요한 방향을 제시한다. attention sink를 단순히 “첫 token에 attention이 많이 간다”는 표면적 현상이 아니라, softmax attention이 만들어내는 key-bias-like internal mechanism으로 분석했기 때문이다.
후속 연구에서는 몇 가지 방향이 자연스럽게 이어질 수 있다. 첫째, first-token sink 외에 newline, punctuation, delimiter, special token sink를 분석할 수 있다. 둘째, attention sink를 explicit key bias나 learnable sink token으로 분리해 long-context inference를 더 안정적으로 만들 수 있다. 셋째, softmax attention을 대체하는 attention operation이 attention sink뿐 아니라 LLM의 성능, 안정성, scaling law에 어떤 영향을 주는지 분석할 수 있다.
특히 quantization과 KV cache optimization 관점에서 이 논문은 중요하다. attention sink token은 attention 분포에서 매우 큰 비중을 차지하기 때문에, 해당 token의 key/value를 어떻게 보존하느냐가 성능에 영향을 줄 수 있다. 실제로 pivot token을 유지하거나 sink token을 보존하는 quantization/KV compression 기법들은 이 현상을 활용한다.
23. 정리
When Attention Sink Emerges in Language Models: An Empirical View는 attention sink 현상을 매우 폭넓게 분석한 논문이다. 이 논문은 attention sink가 대형 LLM에서만 나타나는 특이 현상이 아니라, 작은 LM에서도 나타나며, 자연어뿐 아니라 random token sequence에서도 관찰된다는 점을 보인다. 또한 instruction tuning보다 pre-training 단계에서 이미 형성되는 현상임을 보여준다.
논문의 핵심 결론은 세 가지로 정리할 수 있다. 첫째, attention sink는 충분한 데이터와 효과적인 최적화를 통해 등장한다. 둘째, sink 위치는 data distribution과 loss function에 의해 이동할 수 있으며, 반드시 첫 번째 token에 고정되는 것은 아니다. 셋째, attention sink는 value information을 많이 전달하는 token이라기보다 key bias처럼 작동하며, softmax attention의 정규화로 인한 attention score 간 의존성과 깊이 연결되어 있다.
개인적으로 이 논문에서 가장 흥미로운 지점은 attention sink를 “이상한 attention pattern”이 아니라 “모델이 softmax attention 안에서 만든 내부 계산 장치”처럼 해석한다는 점이다. 모델은 모든 attention score를 확률분포로 정규화해야 하고, 이 과정에서 실제 value computation에는 크게 기여하지 않으면서 attention mass를 흡수하는 key-side bias가 필요해질 수 있다. 첫 번째 token은 그 역할을 가장 자연스럽게 수행하는 위치가 된다.
이 논문은 LLM을 더 잘 이해하고 싶은 사람뿐 아니라, long-context inference, KV cache compression, quantization, efficient attention, mechanistic interpretability를 연구하는 사람에게도 중요한 참고점이 된다. attention sink는 단순한 시각화 현상이 아니라, 현대 Transformer LM의 계산 구조와 시스템 최적화를 연결하는 핵심 현상 중 하나라고 볼 수 있다.
이 논문은 LLM의 attention sink가 충분한 pre-training 이후 softmax attention 구조 속에서 자연스럽게 등장하는 key-bias-like 현상이며, 의미 정보 전달보다 extra attention score를 저장하는 내부 계산 장치에 가깝다는 점을 실험적으로 보여준다.
원문 링크: https://arxiv.org/abs/2410.10781