MHA란 무엇인가? — 헤드를 8개로 나누는 이유와 PyTorch에서 조용히 망가지는 3가지 함정 - Insights

이 글을 다 읽고 나면, num_heads를 몇 개로 설정해야 할지, 그리고 nn.MultiheadAttention을 결과물을 망가뜨리지 않게 호출하기 위한 세 가지 설정을 직접 코드로 결정할 수 있게 됩니다.

PyTorch에서 nn.Linear나 nn.Module을 작성해 본 경험이 있는 사람
Self-Attention이라는 이름은 알고 있지만, 헤드를 나누는 의미를 설명할 수 없는 사람
num_heads를 선택하는 기준과, 파라미터 수가 헤드 수에 따라 변하지 않는 사실
에러가 발생하지 않도록 하는 세 가지 함정의 재현 코드와 구별 방법
nn.MultiheadAttention을 호출하기 전 체크리스트와, 에러 메시지별 대처표
Transformer 전체 구조(잔차 연결(Residual Connection), LayerNorm, FFN)
속도 비교 및 SDPA를 이용한 자체 구현이나 GQA로의 확장 (후속편에서 다룹니다)

지금까지 Self-Attention이나 Flash Attention에 대한 글을 써왔지만, 그 사이를 잇는 '헤드를 나누는' 단계를 제대로 다루지 못했습니다. PyTorch 2.14 공식 문서를 다시 읽어보고 놀란 점은, nn.MultiheadAttention이 자신을 '기초 이해를 위한 참조 구현(reference implementation)'이라고 설명하고 있다는 것이었습니다.

편리한 기성 모듈일수록 기본값이 자신의 전제와 달라도 에러가 발생하지 않는 한 알아차리기 어렵습니다. 이 글에서는 기성품의 내부 구조를 직접 손으로 재조립하여 확인해 볼 가치가 있다는 점을 함께 고민해 주시면 좋겠습니다. 검증은 GPU가 없는 Xeon 2코어, 메모리 8GB의 Linux 환경에서 Python 3.13과 PyTorch 2.14.0을 CPU로 실행했습니다.

먼저, 글 전체에서 사용할 비유를 정하겠습니다. 용어는 말미의 용어집에도 정리되어 있습니다.

이 섹션에서 알 수 있는 것: MHA(Multi-Head Attention, 멀티헤드 어텐션)가 무엇을 하는지, 그리고 왜 한 명이 아닌 여러 명으로 읽어야 하는지.

Attention은 문장의 각 단어가 '다른 어떤 단어를 얼마나 참고할지'를 결정하는 메커니즘입니다. 이를 원고지를 읽는 편집자 한 명의 작업에 비유해 보겠습니다.

편집자가 가진 '무엇을 찾고 있는지' 메모가 Query(질의)입니다.
각 단어에 붙은 '나는 무슨 이야기를 하는지' 제목표가 Key(키)입니다.
제목표 뒷면에 적힌 단어 내용이 Value(값)입니다.

편집자는 질의와 제목표를 대조하며, 유사해 보이는 단어일수록 빨간 펜으로 진하게 밑줄을 긋고, 그 농도에 따라 내용을 모읍니다.

MHA는 이 편집자를 8명으로 늘리는 메커니즘입니다. 한 명은 문법적 의존 관계를 담당하고, 다른 한 명은 대명사가 가리키는 대상을 담당하는 식으로, 각각 다른 색안경을 쓰고 같은 원고지를 읽습니다. 마지막에 편집장이 8명의 메모를 한 장에 정리합니다.

왜 한 명이 안 될까요? 원 논문 "Attention Is All You Need"에는 이렇게 쓰여 있습니다.

원문: "With a single attention head, averaging inhibits this."

해석: 어텐션 헤드가 하나만 있으면, 평균화가 이것(여러 표현 부분 공간의 정보에 동시에 주의를 기울이는 것)을 방해한다.

출처: Vaswani et al. 2017, §3.2.2

한 편집자가 의존 관계와 대명사 지시를 모두 신경 쓰면, 빨간 펜의 진하기가 두 가지 목적의 평균이 되어 어느 쪽도 어중간해집니다. 그래서 담당을 나누는 것입니다.

앞 섹션에서 MHA란 8명의 편집자가 각각 다른 색안경으로 같은 원고지를 읽는 것이라고 정의했습니다. 비유만으로는 코드로 옮길 수 없으므로, 같은 내용을 공식으로 대체합니다.

이 섹션에서 알 수 있는 것: MHA의 세 가지 식이 편집자의 비유 중 어디에 해당하는지.

기호는 사용하기 직전에 하나씩 도입하겠습니다.

$E$(이, 임베딩 차원(embed_dim), 정수): 단어를 나타내는 벡터의 길이. 비유에서는 원고 1단어당 정보의 폭입니다. 인자명은 embed_dim
$h$(에이치, 헤드 수(num_heads), 정수): 편집자의 인원수입니다. 인자명은 num_heads
$d$(디, 헤드 차원(head_dim), 정수): 한 명의 편집자가 담당하는 폭으로, $d = E / h$

먼저, 한 사람 분량의 어텐션 공식입니다. $Q$, $K$, $V$(행렬, 행은 단어 수, 열은 $d$)는 질의 메모, 제목표, 내용을 단어별로 세로로 나열한 것입니다.

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right) V

이 수식은 질문과 헤드(head)의 내적을 통해 적합성을 측정하고, softmax를 거쳐 빨간 펜의 농도(총합 1)로 만든 다음, 그 농도를 이용해 내용을 혼합한다는 의미입니다. 원 논문 §3.2.1에 따르면 $\sqrt{d}$로 나누는 이유는 폭 $d$가 클 경우 내적이 너무 커져서 softmax의 기울기(gradient)가 극단적으로 작아지는 것을 방지하기 위함이라고 설명합니다.

다음은 $i$번째 편집자입니다. $W_i^Q$, $W_i^K$, $W_i^V$(행렬, $E \times d$)는 편집자 $i$의 색안경(perspective)입니다.

\mathrm{head}_i = \mathrm{Attention}(Q W_i^Q,\; K W_i^K,\; V W_i^V)

이 수식이 의미하는 바는, 동일한 원고를 편집자별로 다른 색안경을 통해 폭 $E$에서 폭 $d$로 좁혀 읽는다는 것입니다.

마지막으로, 편집장의 요약입니다. $W^O$(행렬, $E \times E$)는 편집장이 내용을 정리하는 방식입니다.

\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h)\ W^O

이 수식이 의미하는 바는, $h$명의 폭 $d$ 메모를 가로로 나열하여 다시 폭 $E$로 되돌린 후, 편집장이 하나의 요약본으로 재작성한다는 것입니다.

기호	PyTorch에서의 이름	비유적 역할
$E$	embed_dim	원고 1단어당 정보의 폭
...
이전 절에서 MHA는 '투영(Projection), 분할(Split), 주의(Attention), 연결(Concatenate), 출력 투영(Output Projection)'의 5단계로 수식을 통해 확인했습니다. 수식이 정확하더라도 PyTorch 구현이 같은 과정을 거치는지 여부는 별개의 문제입니다. 일일이 대조해 보면, 후반부 함정에서 어느 부분이 수작업과 다른지 특정할 수 있습니다.

이 섹션에서 알게 될 것: nn.MultiheadAttention가 수식대로 작동하는지, 그리고 헤드(head) 수를 변경해도 파라미터 수가 전혀 변하지 않는다는 것입니다.

먼저, 이후 실험에서도 사용할 도구들을 정리합니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
...

다음으로, 수식대로 직접 구현한 MHA와 nn.MultiheadAttention의 출력을 비교해 봅니다.

import torch
import torch.nn as nn
from mha_utils import project_qkv, split_heads, merge_heads, max_diff
...

출력의 최대 오차는 1.19e-07이었고, 가중치(weight)의 최대 오차는 5.96e-08이었습니다. 이는 float32의 반올림 오차 범위 내에 있기 때문에, nn.MultiheadAttention은 이전 절의 수식 그 자체입니다. 파라미터 수는 다음과 같습니다.

num_heads	head_dim	파라미터 수
1	512	1,050,624
...
헤드 수를 1부터 16까지 변경해도 파라미터 수는 전혀 변하지 않습니다. 색안경 $W^Q, W^K, W^V$와 편집장의 $W^O$ 모두 $E \times E$이며, 바이어스(bias)가 $4E$이므로 총합은 $4E^2 + 4E$입니다. 헤드 수는 이 총액을 몇 명에게 나눌지를 결정하는 것일 뿐입니다.

즉, num_heads는 모델의 크기가 아니라 사람당 폭을 결정하는 조절 장치(knob)입니다. 원 논문은 $E = 512$, $h = 8$, $d = 64$를 사용했습니다. 선택의 출발점은 다음 두 가지입니다.

$d$가 64 근처가 되도록 $h$를 선택한다 (원 논문의 설정에 맞추기 위함)
$h$는 짝수(even)로 만든다. 설치된 소스 코드(torch/nn/modules/activation.py)를 읽어보면, 추론용 고속 경로인 fastpath 판별 조건 목록에는 문서의 조건 목록에는 없는 'num_heads가 짝수여야 한다'는 내용이 포함되어 있기 때문입니다.

이전 절에서 내용물이 수식대로 작동하는지 확인했습니다. 다음으로 문제가 되는 것은, 그 내용물에 어떤 형태로 데이터를 주입하느냐입니다. 이후 다룰 세 가지 함정은 모두 계산 과정이 아니라 입력 형태(entry format)에서 발생하기 때문에, 먼저 인수를 정리하겠습니다.

이 섹션에서 알게 될 것: 자주 사용되는 인수와 입출력 텐서의 형태.

引数	既定値	意味
embed_dim	なし	모델 전체의 차원 $E$
...
`forward`의 주요 인자입니다. $N$은 배치 크기(batch size), $L$은 쿼리(query) 측의 시퀀스 길이, $S$는 참조(reference) 측의 시퀀스 길이입니다.

引数	形	意味
query	(N, L, E)	쿼리의 메모(memo)
...
반환 값은 `(attn_output, attn_output_weights)`의 튜플입니다. 가중치(weight)는 기본적으로 헤드 평균화된 $(N, L, S)$ 형태입니다. 각 헤드별로 '빨간펜' 표시를 보고 싶다면 `average_attn_weights=False`로 설정하여 $(N, h, L, S)$ 형태로 만들 수 있습니다.

위치적 맥락에 대해서는 문서에 다음과 같은 문장이 있습니다.

원문: "The intent of this layer is as a reference implementation for foundational understanding"

번역: 이 계층의 의도는 기초 이해를 위한 참조 구현입니다.

출처: PyTorch 2.14 torch.nn.MultiheadAttention

새로운 구조에서는 구성 요소(component)부터 직접 조합하는 것을 권장하지만, 학습이나 실험 과정에서 가독성을 위해서는 기성품(pre-built module)을 사용하는 것이 편리합니다. 본 글에서는 이 기성품을 안전하게 사용하는 측면에 초점을 맞춥니다.

MHA는 같은 원고를 $h$명의 편집자가 폭 $d = E/h$의 색안경으로 읽고, 편집장이 정리하는 메커니즘입니다.
nn.MultiheadAttention은 공식대로 작동하며, 수동 구현과의 차이는 약 1e-7 정도에 불과합니다.
헤드 수를 변경해도 파라미터 수는 $4E^2 + 4E$로 유지됩니다. num_heads는 $d$를 결정하는 조절 장치(knob) 역할을 하며, 짝수로 설정해야 합니다.
기본값은 batch_first=False이며, key_padding_mask가 True인 경우 무시됩니다.

후반부에서는 이러한 기본값과 마스크 방향이 초래하는 함정들을 실제로 깨보면서 확인합니다.

앞 섹션 마지막에서 batch_first의 기본값이 False라는 점, 그리고 마스크는 True를 무시한다는 점을 확인했습니다. 둘 다 형태(shape)만 맞으면 오류가 발생하지 않습니다. 오류가 없는 설정 실수는 학습이 진행되지 않거나 정확도가 약간 낮은 등의 증상으로만 나타나기 때문에 인지하는 것이 가장 늦습니다.

이 섹션에서 알게 될 것: 결과물을 조용히 망가뜨리는 3가지 설정 실수와, 각각을 구별하는 방법입니다.

$(N, L, E)$ 형태의 데이터를 batch_first=False 상태로 전달하면, PyTorch는 이를 $(L, N, E)$로 해석합니다. 예를 들어, 서로 다른 원고에서 같은 위치에 있는 단어끼리 비교하게 됩니다.

import torch
import torch.nn as nn
from mha_utils import max_diff
...

batch_first	샘플 1만 변경했을 때, 샘플 0의 출력 변화	가중치 형태
False(기본값)	4.804e-01	(10, 4, 4)
True	0.000e+00	(4, 10, 10)

샘플 1만 변경했는데도 batch_first=False 상태에서는 샘플 0의 출력이 0.48만큼 변했습니다. 이는 배치 내 다른 원고의 정보가 새어 들어왔다는 증거이며, 오류는 한 번도 발생하지 않았습니다. 솔직히 이렇게 조용하게 작동할 줄은 몰랐습니다.

구별하는 방법은 가중치의 형태입니다. $N=4$, $L=10$이라면 올바른 형태는 $(4, 10, 10)$이어야 하며, $(10, 4, 4)$가 반환되었다면 배치와 시퀀스 차원이 뒤바뀐 것입니다.

어텐션 계산만 수행하는 최적화된 함수 F.scaled_dot_product_attention(이하 SDPA)는 bool 마스크의 True 의미가 nn.MultiheadAttention과 반대입니다. 문서에도 주의사항이 있습니다.

원문: "If migrating from MHA, ensure you invert your boolean mask"

번역: MHA에서 마이그레이션 하는 경우, bool 마스크를 반드시 반전시키십시오.

출처: PyTorch 2.14 torch.nn.functional.scaled_dot_product_attention

nn.MultiheadAttention의 key_padding_mask는 True를 무시하는 반면, SDPA의 attn_mask는 True가 참여합니다. 비유하자면, 한쪽은 "읽지 마" 스티커이고 다른 쪽은 "읽어라" 스티커입니다. 같은 마스크를 전달하여 비교해 보겠습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
...

조건	nn.MultiheadAttention과의 최대 오차
같은 마스크를 그대로 SDPA에 전달	1.35e+00
~mask로 반전하여 전달	1.19e-07

반전하는 것을 잊으면, 읽어야 할 단어를 읽지 않고 패딩만 읽는 계산이 됩니다. 형태가 맞기 때문에 여기에서도 오류는 발생하지 않습니다.

위 코드의 후반부는 특정 샘플의 모든 키(key)를 무시시켰을 때의 동작입니다.

모드	출력에 NaN을 포함하는가
train	True
eval	True

모든 제목표에 "읽지 마" 스티커가 붙으면, softmax에 들어가는 값이 모두 음의 무한대가 되고, 0 나누기 0 형태가 되어 빨간 펜의 농도를 정의할 수 없게 됩니다. 빈 샘플이 하나만 섞여도 해당 배치(batch)의 손실이 NaN이 될 수 있습니다. 데이터 측에서 빈 샘플을 제외하거나, 최소한 1개의 토큰은 남겨둡니다.

지금까지 num_heads 선택 방법과, 조용히 망가지는 3가지 함정을 실측으로 확인했습니다. 이 내용들을 호출하기 전에 위에서부터 순서대로 확인할 수 있도록 정리합니다.

이 섹션에서 알게 될 것: nn.MultiheadAttention을 작성하기 전과, 결과가 이상할 때 확인할 순서.

마지막 상자가 가장 저렴하게 효과를 볼 수 있습니다. 단지 1배치만 흘려보내 attn_output_weights의 형태를 print 하는 것만으로도 함정 1은 확실히 잡아낼 수 있습니다.

이전 섹션의 체크리스트는 오류가 발생하지 않는 함정을 방지하기 위한 것이었습니다. 여기서는 반대로, 오류로 표출되는 것들을 모았습니다. 실제로 오류를 일으켜 메시지를 확인한 코드입니다.

import torch
import torch.nn as nn
E, H = 64, 8
...

증상	원인	대처
AssertionError: embed_dim must be divisible by num_heads	E가 h로 나누어 떨어지지 않음	E를 나누는 h를 선택함
...

"is_causal" 행은 인자 설명의 첫 문장만 읽으면 "지정하면 인과 마스크(causal mask)를 적용한다"고 읽을 수 있습니다. 실제로는, 같은 설명에 있는 경고와 같이 is_causal는 "전달한 attn_mask는 인과 마스크입니다"라는 힌트이며, attn_mask를 함께 전달하지 않으면 오류가 발생합니다.

MHA(Multihead Attention)는 동일한 원고를 h명의 편집자가 각기 다른 색안경으로 읽고, 편집장이 정리하는 구조였습니다. nn.MultiheadAttention은 공식대로 작동하며, 헤드 수를 바꿔도 파라미터 수는 변하지 않습니다. num_heads는 1인당 폭 d를 결정하는 조절 장치이며, d가 64 전후이면서 짝수일 때를 출발점으로 삼습니다.

반면, batch_first의 기본값, 마스크의 방향성, 모든 마스크 행의 NaN이라는 세 가지 함정은 모두 오류를 내지 않고 결과를 망칩니다. 편집부가 조용히 원고를 바꿔도, 편집장의 요약본은 그럴싸하게 완성됩니다. 그래서 1배치에서 가중치의 형태를 확인하는 습관이 가장 저렴한 보험입니다.

개인적으로는 문서에 없는 짝수 조건(even condition)을 소스 코드에서 발견한 것이 가장 큰 수확이었습니다. 기성품의 내부 구조는 한 번은 직접 조합해서 확인할 가치가 있습니다.

후속편: fastpath와 SDPA의 속도를 CPU에서 실측하고, SDPA로 자체 MHA와 GQA를 구성하는 글을 준비 중입니다(공개 후 링크 추가 예정).

다음으로 읽으신다면, 주의를 빠르게 계산하는 Flash Attention, 생성 시 메모리를 결정하는 KV 캐시, 주의를 사용하지 않는 Mamba-2 순서가 추천합니다.

MHA 정의 및 참고 자료

용어	의미	비유적 대응
MHA	주의(Attention)를 여러 헤드(head)로 나누어 병렬로 계산한 후, 마지막에 통합하는 메커니즘	8명의 편집자와 1명의 편집장

참고 자료 및 원문 출처:

PyTorch 2.14 문서: torch.nn.MultiheadAttention(인수, fastpath 조건, 참조 구현으로서의 위치가 기술된 일차 정보)
PyTorch 2.14 문서: torch.nn.functional.scaled_dot_product_attention(마스크 방향에 대한 주의사항)
Vaswani et al., "Attention Is All You Need", arXiv:1706.03762, 2017 (일본어 번역 제목: 필요한 것은 주의만으로 충분하다. MHA를 제안한 원 논문. §3.2.2에 $h=8$, $d=64$ 설정)

끝까지 읽어주셔서 감사합니다. X에서도 업데이트 정보를 전하고 있습니다.

MHA란 무엇인가? — 헤드를 8개로 나누는 이유와 PyTorch에서 조용히 망가지는 3가지 함정

요약

핵심 포인트

MHA 정의 및 참고 자료

댓글

MHA 정의 및 참고 자료