STUDYLLM
‹ Lecture 3 06 / 21

Multi-Head Attention

앞의 single head는 여러 Values를 하나의 weighted mixture로 만들었다. 한 문장 안에서도 위치 관계, 문법 관계 등 읽어야 할 정보는 여러 가지일 수 있다. Multi-Head Attention은 같은 입력에 서로 다른 learned projections를 적용하여, 여러 subspaces에서 Attention을 계산한 결과를 함께 사용한다. Head의 역할을 사람이 미리 하나씩 지정하는 방식은 아니다.

  1. XX
  2. QQ/KK/VV projection
  3. reshape into heads
  4. per-head Attention
  5. concat
  6. WOW_{O}

Projection & Reshape

GPT-2 Small은 dmodel=768d_{model}=768, h=12, dhead=64d_{head}=64다. QQ/KK/VV projection은 전체 768 input dimensions를 보고 [n768]\begin{bmatrix}n&768\end{bmatrix}을 만든 뒤 12×[n64]\begin{bmatrix}n&64\end{bmatrix}로 나눈다. Head outputs를 concat해 [n768]\begin{bmatrix}n&768\end{bmatrix}로 만들고 WO  [768768]W_{O}\;\begin{bmatrix}768&768\end{bmatrix}로 projection한다.

Raw X[:,0:64]X[:,0:64]만 잘라 첫 head에 주면 그 head는 나머지 704 dimensions를 직접 보지 못한다. 반면 XWQ(1)XW_Q^{(1)}의 WQ(1)W_Q^{(1)}는 전체 768 input dimensions를 64차원으로 섞을 수 있다. 전체 projection을 만든 뒤 reshape하는 구현은 이 개별 head projections를 함께 계산한 것과 같다.

각 head에서는 앞의 Scaled Dot-Product Attention을 수행한다. Outputs를 이어 붙이는 concat은 서로 다른 결과를 나란히 놓는 단계이고, WOW_O가 그 feature들을 다시 섞어 block의 output으로 만든다.

MHA⁡(X)=Concat⁡(O1,…,O12)WO\operatorname{MHA}(X) =\operatorname{Concat}(O_1,\ldots,O_{12})W_O

총 width가 같으면 12×2n2×64=2n2×76812\times2n^2\times64=2n^2\times768로 leading score matmul FLOPs는 한 개의 768-wide head와 비슷하다. Heads가 12개라는 이유만으로 비용이 12배가 되는 것은 아니다. 다만 실제 실행에는 reshape와 kernel 등의 overhead가 있어 latency까지 완전히 같다고 단정하지 않는다.

Head Patterns

일부 BERT heads에서 relative-position, dependency-aligned, broad [CLS] patterns가 관측됐다. Explicit syntax/coreference labels 없이 학습 중 나타난 것이다. Arrow는 Query→Key 선택 방향이지만 Value 정보는 Key 위치에서 Query output으로 흐른다.

Induction head [AA][B]…[AA]→[B]는 causal LM의 일부 ICL 패턴과 연결될 수 있다. 모든 head의 role을 미리 정하거나 관측 pattern을 인과 설명으로 간주하지 않는다.

상대 위치에 집중하는 head는 바로 앞이나 뒤의 token을 주로 읽고, dependency-aligned head는 목적어와 관련 동사 같은 관계를 따라 읽을 수 있다. Broad-attention pattern에서는 [CLS] Query가 여러 positions를 넓게 읽는다. 이런 패턴은 self-supervised pre-training 중 explicit syntax/coreference labels 없이 관측되었다는 점이 흥미롭다.

그러나 “head 1은 언제나 문법 담당”처럼 고정된 기능표를 붙일 수는 없다. Data와 checkpoint에 따라 패턴이 다르고, weight가 큰 위치를 그렸다고 그 위치가 최종 결정의 인과적 원인이라고 증명한 것도 아니다. Arrow가 Query → Key인 것은 선택 방향이며, 실제 Value 정보는 선택된 Key 위치에서 Query output으로 들어온다는 점도 구분한다.

Induction 예시에서는 앞서 본 Harry → Potter 관계 뒤에 다시 Harry가 오면 이전 continuation인 Potter를 읽어 이어 쓸 수 있다. [A][B] … [A] → [B] 패턴이 일부 In-context Learning을 뒷받침할 수 있지만, 모든 ICL을 단일 head pattern으로 설명한다는 뜻은 아니다.

Head가 12개라고 전체 fixed-width matmul 비용이 무조건 12배가 되는 것은 아니다.

전체 dims projection 후 head로 reshape한다.

Check

MHA는 XX의 각 64dimension을 미리 잘라서 head에 넣는가?

Answer

아니다. Learned projection이 전체 768 dims를 본 후 출력을 reshape한다. 각 head가 원래 모든 입력 features를 혼합할 수 있다.

Attention arrow와 실제 Value 정보가 흐르는 방향은 어떻게 다른가?

Answer

Arrow는 query가 어떤 key를 보는지 표시한다. 선택된 Value는 그 key 위치에서 query의 output으로 전달된다.

Contents