TR-02. 멀티-헤드 셀프 어텐션(Multi-head Self Attention)

멀티-헤드 어텐션은 트랜스포머의 핵심 부품으로, "여러 개의 다른 관점에서" 문장을 동시에 바라보는 메커니즘입니다.

'어텐션 헤드' 하나가 스케일드 닷-프로덕트 어텐션 계산기 한 세트라고 생각하시면 됩니다. 멀티-헤드는 이 계산기를 여러 개(예: 8개) 복제하여 병렬로 실행하는 방식입니다.

왜 '멀티-헤드'가 필요한가?

'싱글-헤드' 어텐션(어텐션을 한 번만 계산)은 문장 내 단어 간의 한 가지 종류의 관계만 학습하는 경향이 있습니다.

예를 들어, "그녀는 강아지에게 밥을 주었다"라는 문장에서 어텐션이 '주어-동사' 관계(그녀는-주었다)에만 집중하도록 학습될 수 있습니다.

하지만 문장에는 다양한 관계가 존재합니다.

• 주어-동사 관계 (그녀는 - 주었다)
• 동사-목적어 관계 (주었다 - 밥을)
• 수식 관계 (강아지 - 에게)
• 문맥적 유사 관계 등

멀티-헤드 어텐션은 모델이 이러한 다양한 문맥적 관계를 동시에 학습할 수 있게 해줍니다. 마치 8명의 전문가(Head)가 각자 다른 관점(예: 한 명은 문법 구조, 한 명은 의미적 연관성)으로 문장을 분석하고 그 결과를 합치는 것과 같습니다.

멀티-헤드 셀프 어텐션의 구조

멀티-헤드 어텐션(MHA)은 여러 단계를 거쳐 작동합니다. (여기서는 트랜스포머 인코더에서 사용되는 셀프 어텐션을 기준으로 설명하며, Q, K, V는 모두 같은 입력 문장으로부터 나옵니다.)

전체 입력 문장 행렬을 $X$라고 하고, 헤드의 개수($h$)를 8개라고 가정해 봅시다.

1단계: 입력 프로젝션 (Projection)

Q, K, V 얻기

1. 먼저, 입력 문장 행렬($X$)을 Q, K, V로 변환하기 위한 3개의 큰 가중치 행렬($W^Q, W^K, W^V$)을 준비합니다.
   
   • 이 큰 가중치 행렬을 8개의 작은 헤드용 가중치 행렬로 쪼갭니다.
     
     • $W^Q \rightarrow W^Q_1, W^Q_2, \dots, W^Q_8$
     • $W^K \rightarrow W^K_1, W^K_2, \dots, W^K_8$
     • $W^V \rightarrow W^V_1, W^V_2, \dots, W^V_8$
2. 이제 8개의 헤드가 각자 다른 가중치를 사용하여 자신만의 Q, K, V '조각'을 만듭니다.
   
   • Head 1: $Q_1 = XW^Q_1$, $K_1 = XW^K_1$, $V_1 = XW^V_1$
   • Head 2: $Q_2 = XW^Q_2$, $K_2 = XW^K_2$, $V_2 = XW^V_2$
   • ...
   • Head 8: $Q_8 = XW^Q_8$, $K_8 = XW^K_8$, $V_8 = XW^V_8$

이 과정을 통해 8개의 헤드는 원본 입력을 8개의 서로 다른 '관점' 또는 '부분 공간(subspace)'으로 투영(projection)시킵니다.

2단계: 병렬 어텐션 계산 (Parallel Attention)

어텐션 점수 행렬(Attention Score Matrix) 계산

어텐션 값(Attention Value) 계산

• 8개의 헤드가 동시에, 그리고 독립적으로 스케일드 닷-프로덕트 어텐션을 계산합니다.
• 각 헤드는 자신만의 어텐션 출력($Z$)을 만듭니다.

$$\text{Head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right)V_i$$

• Head 1 $\rightarrow Z_1$
• Head 2 $\rightarrow Z_2$
• ...
• Head 8 $\rightarrow Z_8$

여기서 $d_k$는 각 헤드의 차원입니다. (예: 전체 모델 차원이 512이고 헤드가 8개라면, 각 헤드의 $d_k$는 512 / 8 = 64가 됩니다.)

 

결합 (Concatenation)

• 8개의 헤드가 각각 계산한 출력 행렬($Z_1$ ~ $Z_8$)을 다시 하나의 큰 행렬로 이어 붙입니다.

$$\text{Concat}(Z_1, Z_2, \dots, Z_8)$$
최종 프로젝션 (Final Projection)

• 결합된 행렬은 아직 최종 출력이 아닙니다. 이 행렬을 마지막으로 또 다른 가중치 행렬 $W^O$ (Output Weight Matrix)와 곱해줍니다.
• 이 과정을 통해 8개 헤드의 서로 다른 관점에서 나온 정보들을 하나로 통합하고 정리하여, 모델의 다음 레이어가 처리하기 좋은 형태로 만들어 줍니다.

$$\text{MultiHead}(Q, K, V) = \text{Concat}(Z_1, \dots, Z_8) W^O$$

이 최종 출력이 바로 멀티-헤드 어텐션 레이어의 최종 결과물이며, 다음 레이어(예: Feed-Forward Network)의 입력으로 전달됩니다.

요약

멀티-헤드 셀프 어텐션은 다음과 같은 장점을 가집니다.

• 다양한 관점 학습: 여러 헤드가 문장 내의 다양한 종류의 관계(문법적, 의미적 등)를 동시에 학습할 수 있습니다.
• 풍부한 정보 추출: 단일 어텐션으로는 놓칠 수 있는 미묘한 문맥 정보를 여러 헤드가 포착하여 모델의 표현력을 높입니다.
• 안정적인 학습: 입력을 여러 조각(헤드)으로 나누어 처리함으로써, 각 헤드가 더 작고 집중된 부분 공간에서 학습하게 되어 전체적인 학습 과정을 안정화하는 데 도움을 줄 수 있습니다.