[24’ ICML] ACT: Unveiling and Harnessing Hidden Attention Sinks
카테고리: Language
태그: Attention
🔍 Abstract
- Problem
- LLM이 발전함에 따라 성능은 향상되고 있으나, 아직 Attention Distribution에 대해 잘 이해하지 못하고 있다.
- 본 논문에서는 최근 발견된 Attention Sink 현상을 더 깊게 이해하고, 이를 활용하여 LLM의 성능을 향상시키는 방법을 제안한다.
- Insight
- Attention Sink는 First Token 뿐만 아니라 Less Semantic Token에도 발생한다.
- 기존 StreamingLLM에서는 Attention Sink가 LLM 성능에 중요하다고 주장하였지만, 저자들은 이러한 Attention Sink가 항상 좋은 것은 아니라고 지적한다.
- Method
- 따라서, 저자들은 ACT(Attention Calibration Technique)를 제안한다.
- 이는 Attention Sink를 줄였을 때 성능이 향상되는 Head에 대하여 Attention Distribution을 조작하는 방법이다.
- Result
- 많은 데이터셋에서 일관적으로 성능 향상이 관찰되었다.
1. Location of Attention Sink
먼저, 저자들은 다음과 같이 Attention Head의 Notation을 정하였다.
이때 저자들의 Notation이 조금 헷갈리는 관계로 Query, Key 순서대로 정리하여 $\mathbf{A}_ h ^ l \in \mathbb{R}^ {N \times N}$이라고 하면, Attention Score를 다음과 같이 정의할 수 있다.
\[a_ h ^ l = \left[ \sum_ {i=j} ^ N \mathbf{A}_ h ^ l [i, j] / (N - j), ~ j \in \lbrace 1, 2, \cdots, N \rbrace \right]\]따라서 $a_ h ^ l \in \mathbb{R}^ N$이며, 이는 Attention Head $h$, Layer $l$에서 각 Token에 대한 Attention Score를 의미한다. 참고로 Autoregressive한 상황에서는, $\mathbf{A}_ h ^ l \in \mathbb{R}^ {1 \times N}$이고,
\[a_ h ^ l = \left[ \mathbf{A}_ h ^ l [0, j], ~ j \in \lbrace 1, 2, \cdots, N \rbrace \right]\]이다. 이를 Attention Matrix로 보면,
Attention Matrix는 위와 같고 이때 각 Key를 기준으로 Column Mean을 한 것이 각 토큰의 Attention Score이다. 즉, $j$번째 토큰의 Attention Score는 $a_ h ^ l [j]$이다. 저자들은 여기서 Attention Score가 일정 수준 이상, 즉 $a_ h ^ l [j] \gt \alpha / N$인 경우를 Attention Sink라고 정의한다. 이들의 통계는 다음과 같다.
위의 그림에서 Attention Sink는 First Token에만 발생하는 것이 아니라 Less Semantic Token에도 발생한다는 것을 확인할 수 있고, 특히 Layer 0, 1 이후 Layer에서 발생하여 유지되는 것을 확인할 수 있다.
2. Effect of Attention Sink
저자들은 각 Head에서 Attention Sink가 발생하는 경우의 성능이 어떻게 되고, 이를 줄였을 때의 성능이 어떻게 되는지 확인하였다. 이를 Calibration이라고 하며, 과정은 다음과 같다.
먼저 Attention Sink Token Set $\mathcal{S}_ h ^ l$을 다음과 같이 정의한다. 구현상에서는 <s> Token이 있는 경우, 즉 $j=1$의 경우는 포함하지 않았다.
그리고 Attention Sink Token에 대해서는 $\beta = 0.4$만큼 Attention Score를 줄이는 Calibration을 진행한다. 즉, $j_ \mathcal{S} \in \mathcal{S}_ h ^ l$인 경우,
\[\hat{A_ h ^ l} [i, j_ \mathcal{S}] = A_ h ^ l [i, j_ \mathcal{S}] \times \beta\]가 되고, $j_ {\mathcal{S}^ c} \notin \mathcal{S}_ h ^ l$인 경우에는,
\[\hat{A_ h ^ l} [i, j_ {\mathcal{S}^ c}] = A_ h ^ l [i, j_ {\mathcal{S}^ c}] + \left( \frac {A_ h ^ l [i, j_ {\mathcal{S}^ c}]}{\sum_ {j _ {\mathcal{S}^ c}} A_ h ^ l [i, j_ {\mathcal{S}^ c}]} \right) \left(\sum _ {j_ \mathcal{S}} A_ h ^ l [i, j_ {\mathcal{S}}] \times (1 - \beta) \right)\]로 조정되어 합은 여전히 1이 된다. 이러한 Calibration을 통해 Attention Sink를 줄이고, 성능을 향상시킬 수 있다. 그러나, 이것이 모든 Head에 대해 성능을 향상시키는 것은 아니다. 저자들은 모든 Layer, 모든 Head에 대하여 조작 이후 성능을 비교했다.
그 결과는 위와 같다. 전체의 76.8%의 Head가 조작 후 성능이 향상되었다. 이는 Attention Sink가 항상 좋은 것은 아니라는 것을 의미한다.
3. ACT(Attention Calibration Technique)
따라서 저자들은 이들 중 효과적인 Head만을 골라 Calibration을 진행하는 ACT를 제안한다. 이때 Task 별로 Head를 최적화할 것이다. 예를 들어, Classification Task라면 이에 대한 Dataset만을 모아 최적화를 진행하고, Multiple Question Task라면 이에 대한 Dataset만을 모아 최적화를 진행할 것이다. 즉, task $\mathcal{T} = \lbrace \mathcal{D}_ 1, \mathcal{D}_ 2, \cdots, \mathcal{D}_ Q \rbrace$라고 하고, 이때 각 데이터셋의 일부 $M$개만큼의 sample만을 이용해 최적화한다. 수학적으로는 $\Vert \mathcal{C} \cap \mathcal{D}_ q \Vert = M, \forall \mathcal{D}_ q \in \mathcal{T}$이다. 이후 성능의 향상에 도움을 주는 Head set $\mathcal{H}$를 찾아 Calibration을 진행한다.
4. Experiments
4.1. Main Results
전체 결과는 다음과 같다. 전반적으로 모델이 커지면 성능이 향상되는 것으로 보이며, Open-ended QA에서의 성능 향상이 가장 크다.
4.2. Ablation Study
Attention Calibration. Temp는 Initial Token 외에 모두 $\theta = 1.1$을 취하여 Attention 값을 높인 것이고, Inv-temp는 반대로 $\theta = 1 / 1.1$을 취하여 Attention 값을 낮춘 것이다. Inv-ours는 기존 방법과 반대로 Attention Sink에 Attention 값을 재분배해준 것이다. 이때 성능 차이가 꽤나 나지만 Inv-ours의 경우 성능에 큰 차이가 없다. 이는 Attention Sink에 대한 해당 논문의 방법론의 효과에 의문을 가지게 한다.
Attention Distribution. 저자들은 Attention Sink Token에서 가져온 남은 Attention 값을 (1) 전체에 상수 값으로 더해준 경우, (2) Question에만 적용한 경우, (3) Choice에만 적용한 경우와 (4) ACT를 비교하였다. 이때 ACT가 가장 좋은 성능을 보였다.
Hyperparameter. $\alpha$는 Attention Sink Token을 정의하는 기준이고, $\beta$는 Calibration 시 Attention Score를 조정하는 정도이다. $\alpha$는 충분한 범위에서 실험되지 못해 Robust한 성능을 보이는 것처럼 보이고, $\beta$는 어느 정도의 Trade-off를 보인다. 한편 데이터셋의 크기 $M$은 클수록 좋았지만, $M=1000$으로도 충분한 성능을 보여주었다.
Subset of $\mathcal{H}$. ACT를 적용한 경우, 모든 Head에 대해 Calibration을 진행하는 것이 아니라 일부 Head에 대해서만 Calibration을 진행한 경우를 비교하였다. 이때 모든 Head에 대해 Calibration을 진행한 경우가 가장 좋은 성능을 보였다.
💡 Summary
해당 논문의 내용을 정리하면 다음과 같다.
- Attention Sink는 First Token 뿐만 아니라 Less Semantic Token에도 발생함을 밝힘
- 일부 Head에서는 Attention Sink를 줄이면 성능이 향상되는 것을 확인함
댓글 남기기