밑바닥부터 만드는 mini GPT 공부 시리즈 4편
self-attention 하나로도 각 토큰이 문맥을 참고할 수 있다. 그런데 GPT는 보통 attention을 하나만 쓰지 않고 여러 head로 나눈다.
핵심은 여러 관점이다. 같은 문장을 보더라도 한 head는 가까운 단어 관계를, 다른 head는 문장 전체 분위기를, 또 다른 head는 특정 패턴을 더 잘 보도록 학습될 수 있다.
이번 글에서 다루는 것
- head를 나눈다는 말의 shape 의미
- `d_model`과 `n_heads`, `head_dim`의 관계
- 각 head 결과를 다시 합치는 이유
- Multi-Head Attention을 shape 흐름으로 읽는 법
1. head는 embedding 차원을 나눈 작은 attention 공간이다
입력 벡터의 마지막 차원이 `d_model`이라고 하자. head를 4개로 나누면 각 head는 `d_model / 4` 크기의 작은 공간에서 attention을 계산한다.
32차원
32차원
32차원
32차원
그래서 `d_model`은 `n_heads`로 나누어떨어져야 한다. 나누어떨어지지 않으면 각 head가 같은 크기의 `head_dim`을 가질 수 없다.
2. shape는 이렇게 움직인다
Multi-Head Attention을 이해할 때는 수식보다 shape를 먼저 보는 편이 좋다. 입력은 `(B, T, C)`이고, head를 나누면 `(B, H, T, D)`가 된다.
x: (B, T, C)
q, k, v: (B, T, C)
split head: (B, H, T, D)
score: (B, H, T, T)
context: (B, H, T, D)
merge: (B, T, C)여기서 `T x T`는 각 위치가 다른 위치를 얼마나 볼지 나타내는 attention weight 행렬이다. causal mask도 바로 이 `T x T` 위에서 적용된다.
B = 2, T = 5, C = 128, H = 4
D = C / H = 32
x: (2, 5, 128)
split: (2, 4, 5, 32)
score: (2, 4, 5, 5)
merge: (2, 5, 128)head를 나눠도 최종 출력 차원은 원래 C로 돌아온다. 다음 Transformer block이 같은 형식의 입력을 받아야 하기 때문이다.
3. 여러 head 결과는 다시 하나의 벡터로 합쳐진다
head를 나누는 것은 모델 출력을 여러 갈래로 분리해서 끝내겠다는 뜻이 아니다. 각 head가 계산한 문맥 벡터를 다시 이어 붙여 원래 `d_model` 크기로 되돌린다.
head를 나누는 이유를 한 문장으로
Multi-Head Attention은 하나의 큰 attention으로 모든 관계를 한 번에 보려 하지 않고, embedding 차원을 여러 작은 관점으로 나누어 각 관점에서 문맥을 본 뒤 다시 합치는 방식이다.
구현할 때 자주 나는 실수
- `view` 뒤에 head 차원을 앞으로 옮기지 않으면 score shape가 맞지 않는다.
- `transpose` 이후에는 메모리가 연속적이지 않을 수 있어, 병합 전 `contiguous()`가 필요할 수 있다.
- head를 합친 뒤에는 output projection을 거쳐 head별 결과를 다시 섞는다.
스스로 점검
- `d_model`이 `n_heads`로 나누어떨어져야 하는 이유는 무엇인가?
- attention weight의 shape가 `(B, H, T, T)`가 되는 이유는 무엇인가?
- head를 나눈 뒤 왜 다시 합쳐야 하는가?
다음 글 예고
다음 글에서는 LayerNorm, GELU, FeedForward, residual connection을 묶어 GPT Block을 만드는 흐름을 본다.
한 줄 정리: Multi-Head Attention은 여러 관점으로 문맥을 본 뒤, 그 결과를 다시 하나의 표현으로 합치는 구조다.