** When Attention Sink Emerges in Language Models: An Empirical View (ICLR 2025)

https://www.dropbox.com/scl/fi/77jzbfvnx6c1algovtd8z/iclr25_Attention_Sink_Anatomy.pdf?rlkey=2br6s5zfe1x8mj62v3q4eyi45&dl=0

이 논문은 Attention Sink가 왜 발생하는가?, 언제 발생하는가?, **무엇이 Attention Sink를 결정하는가?**를 체계적으로 분석한 최초의 empirical study이다. 기존 연구들이 “Attention Sink를 활용”하는 데 초점을 맞췄다면, 이 논문은 pretraining 과정에서 Attention Sink가 생성되는 메커니즘을 규명하는 것이 목표이다.


1. 논문의 핵심 질문

저자들은 다음 네 가지 질문을 던진다.

  1. 언제 Attention Sink가 생기는가?
  2. 왜 항상 첫 번째 token에 나타나는가?
  3. 어떤 요소가 이를 결정하는가?
    • optimization
    • data distribution
    • loss
    • architecture
  4. Softmax와 어떤 관계가 있는가?

이를 위해 저자들은

  • GPT-2
  • OPT
  • Pythia
  • LLaMA2
  • LLaMA3
  • Mistral

등 다양한 autoregressive LM을 분석하였다.


2. 방법론

논문의 방법론은 새로운 모델을 제안하는 것이 아니라,

Attention Sink를 발생시키는 원인을 하나씩 제거(ablation)하면서 분석

하는 것이다.

전체 실험은 다음 4개의 축으로 구성된다.

Optimization
        ↓
Data Distribution
        ↓
Loss Function
        ↓
Architecture

3. Attention Sink 측정 방법

기존 연구들은 heatmap만 보여주었지만, 이 논문은 처음으로 정량 metric을 정의한다.

각 head에서

αkl,h=1Tk+1i=kTAi,kl,h\alpha_k^{l,h} = \frac1{T-k+1} \sum_{i=k}^{T} A_{i,k}^{l,h}

즉, token k가 평균적으로 얼마나 attention을 받는지를 계산한다.

특히 α1\alpha_1이 크면 첫 token이 sink이다.

그리고 threshold ε를 이용하여

Sinkkϵ=1Ll1HhI(αk>ϵ)\text{Sink}^{\epsilon}_k = \frac1L \sum_l \frac1H \sum_h I(\alpha_k>\epsilon)

을 정의한다.

“Attention Sink가 발생한 head의 비율”

이다.

논문에서는

  • ε=0.3
  • T=64

를 기본 설정으로 사용한다.


4. 왜 첫 번째 Token이 Sink인가?

이 부분이 논문의 가장 중요한 분석이다.

기존에는

첫 token의 hidden state norm이 크기 때문

이라고 알려져 있었다.

하지만 저자들은 이것이 충분한 설명이 아님을 보인다.

(1) Hidden state

첫 token은

  • hidden norm ↑↑

이다.

즉 massive activation.


(2) 그런데 Key norm은 작다

의외로 첫 token의

  • key norm ↓
  • value norm ↓

이다.

즉, hidden norm이 크다고 해서 attention score가 커지는 것이 아니다.


(3) 진짜 원인

Attention score는 qkT=||q||||k||cosθqk^T = ||q|| ||k|| \cos\theta 이다.

저자들은 첫 token은 norm이 아니라, cosine similarity가 매우 크다는 사실을 발견한다.

즉,

norm
↓↓

angle
↑↑↑↑

↓

attention sink

이다.

즉, 첫 token의 key가 다른 모든 query와 거의 같은 방향으로 정렬되어 있다.

이를

Key Bias

라고 해석한다.


5. Optimization이 Sink에 미치는 영향

저자들은 60M LLaMA를 처음부터 pretrain하였다.

(1) 언제 생기나?

Figure 4.

처음에는 Sink ≈ 0 이다.

그러다가 loss가 충분히 감소한 이후 갑자기 Sink ↑↑가 된다.

즉, Attention Sink는 초기 random initialization에서 생기는 것이 아니라

학습이 충분히 진행되어야 발생한다.


(2) Learning Rate

Learning rate를 줄이면

  • loss 감소 느림
  • sink도 늦게 생성

심지어 더 오래 학습해도 sink가 약하다.

결론

작은 learning rate는 sink 생성을 억제한다.


(3) Batch Size

Batch size는 거의 영향이 없다.


6. Data Distribution 실험

이 논문의 두 번째 핵심이다.

(1) 데이터가 적으면?

학습 token을

5B
↓
2.5B
↓
500M
↓
50M

으로 줄인다.

결과:

데이터가 적을수록 Attention Sink가 거의 생기지 않는다.

즉, 충분한 데이터가 필요하다.


(2) 첫 token을 Random으로 바꾸면?

매 sequence마다 첫 token을 random token으로 교체한다.

놀랍게도 Attention Sink는 여전히 발생한다.

심지어 더 강해진다.

즉, semantic information과 거의 관계없다.


(3) 첫 두 token을 Random으로 만들면?

x1 random, x2 random 으로 만들면 Sink가 첫 번째가 아니라 두 번째 token으로 이동한다.

즉, Attention Sink는 “첫 번째 위치”가 아니라 loss에서 특별한 역할을 하는 위치에 생긴다.


(4) 특정 위치 token 고정

항상 position 2를 같은 token으로 만들면, Sink도 position 2로 이동한다.

즉, Sink 위치는 데이터 분포에 의해 결정된다.


7. Loss Function 실험

(1) Weight Decay

Weight decay를 증가시키면

γ ↑
↓
Sink ↑

가 된다.

그러나, 너무 크면 optimization이 망가지면서 Sink도 사라진다.


(2) Prefix LM

원래는 첫 token은 prediction 대상이 아니다.

Prefix LM에서는 앞쪽 p개의 token이 모두 prediction에서 제외된다.

결과:

Sink도

token1
↓
prefix 영역 전체

로 이동한다.

즉, “prediction되지 않는 token”들이 Sink가 된다.


(3) Sliding Window

Mistral처럼 window attention을 사용하면,

Sink는 window 시작점이 아니라 여전히 sequence의 절대 첫 token에 존재한다.

또한 window가 작을수록 Sink가 약해진다.


8. Architecture 분석

(1) Positional Embedding

다음을 비교한다.

  • NoPE
  • Absolute
  • Learnable
  • Relative
  • ALiBi
  • RoPE

결과:

모든 경우 Attention Sink가 존재한다.

즉, PE가 원인이 아니다.


(2) Pre-Norm vs Post-Norm

PostNorm에서도 Sink는 존재한다.

차이는 massive activation이 LN 이전에 숨어 있을 뿐이다.

즉, LayerNorm도 근본 원인이 아니다.


9. Key Bias 가설 검증

가장 흥미로운 실험이다.

저자들은 Transformer에 직접 learnable key bias를 넣는다.

query
↓
key bias
↓
attention

그러면 첫 token Sink가 사라지고 bias가 Sink가 된다.

반대로 value bias만 넣으면 아무 변화가 없다.

즉, Attention Sink는 value가 아니라 key 때문이다.


또한 K bias의 value 벡터 크기를 점점 키우면, 처음에는 bias가 attention을 흡수하지만 value의 영향이 커질수록 다시 첫 번째 token으로 Sink가 이동한다. 이는 모델이 “추가 attention 저장소”를 실제 출력(value)에 영향을 거의 주지 않는 형태로 유지하려 한다는 해석을 뒷받침한다.


10. Softmax가 원인인가?

논문의 가장 중요한 결론이다.

저자들은 Softmax를 다음으로 교체한다.

  • sigmoid
  • ELU+1
  • identity kernel
  • MLP kernel

그리고 Normalization을 제거한다.

결과:

Softmax를 제거하면 Attention Sink가 거의 사라진다.

Validation loss는 거의 유지된다.

즉, Attention Sink는 Softmax normalization 때문에 attention score들이 합 = 1 이어야 하기 때문에 생긴다.

즉, 모든 token이 “attention을 저장할 공간” 하나를 만들어야 한다.

그 역할을 첫 token이 수행한다.

논문은 이를

Attention Sink는 Key Bias처럼 extra attention score를 저장하는 공간이며, value computation에는 거의 기여하지 않는다.

라고 결론짓는다. 1B 규모 모델에서도 normalization 없는 sigmoid attention은 성능 저하가 매우 작으면서 Sink와 massive activation이 거의 나타나지 않았다.


11. 실험 결과 요약

실험결과의미
작은 모델(Pythia-14M 포함)Sink 존재대규모 모델 특유 현상이 아님
Chat vs Base거의 동일Instruction tuning 영향 미미
Random token 입력Sink 유지의미 정보와 무관
반복 토큰 입력LLaMA/Mistral에서 Sink 소실동일 hidden state가 Sink를 분산
학습 진행일정 단계 이후 Sink 생성충분한 최적화 필요
데이터 감소Sink 감소충분한 학습 데이터 필요
Weight decay 증가Sink 증가(과도하면 감소)정규화가 Sink 형성 촉진
Prefix LMPrefix 전체로 Sink 이동Loss 구조가 위치 결정
Sliding window작은 window에서 Sink 감소장거리 attention과 연관
모든 PESink 존재PE는 근본 원인 아님
K bias 추가Sink가 bias로 이동Sink는 Key Bias 역할
V bias만 추가거의 변화 없음Value는 핵심 원인 아님
Softmax 제거Sink 거의 소멸Softmax normalization이 핵심 요인

12. 논문의 의의와 한계

의의

  • Attention Sink를 **정량적으로 측정하는 지표(Sinkε)**를 제안했다.
  • Optimization → Data → Loss → Architecture의 네 측면에서 최초로 체계적인 원인 분석을 수행했다.
  • Attention Sink = Key Bias + Softmax Normalization의 부산물이라는 설득력 있는 메커니즘을 제시했다.
  • Streaming inference, KV cache compression, quantization 등 Sink를 활용하는 기존 연구들의 현상을 이론적으로 설명하는 기반을 제공했다.

한계

  • 대부분의 분석이 empirical observation에 기반하며, Sink 생성 메커니즘에 대한 엄밀한 수학적 증명은 제공하지 않는다.
  • Softmax가 중요한 요인임을 보였지만, 왜 최적화 과정에서 특정 key 방향이 bias처럼 정렬되는지에 대한 표현학습 관점의 설명은 충분하지 않다.
  • 분석 단위가 attention 수준에 머물러 있어, 어떤 attention head나 MLP neuron이 Sink를 생성하는지와 같은 mechanistic circuit 수준의 분석은 포함되어 있지 않다. 이는 이후의 EAP, activation patching, SAE 기반 회로 분석과 자연스럽게 연결될 수 있는 후속 연구 주제이다.

댓글

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다