LLM Research2026.10.058 min read

Multi-Head Attention 직접 구현: head를 나누면 파라미터는 그대로인데 왜 더 잘 볼까

하나의 attention을 여러 개의 작은 attention(head)으로 쪼개 병렬로 돌리는 Multi-Head Attention을 PyTorch로 직접 만들어 본다. head 분할·병합이 사실은 reshape와 transpose일 뿐이라는 점, head 수를 바꿔도 파라미터 수가 같다는 점을 코드로 확인한다. 마지막에 head 수 1·4·8을 비교하는

By 송민성

오늘의 목표

Attention을 여러 개의 head로 나누는 이유를 설명할 수 있게 된다. 그리고 head 분할·병합, 파라미터 수 계산, "한 번의 matmul로 모든 head를 처리하는 방법"을 PyTorch 코드로 직접 구현해 head 수 1·4·8을 비교한다.

개념과 직관

전제 복습. 앞선 Day에서 만든 self-attention은 각 토큰이 Query(Q), Key(K), Value(V) 벡터를 만들고, Q와 K의 내적으로 "누구를 얼마나 볼지" 점수를 낸 뒤 softmax로 가중치를 만들어 V를 섞는 구조였다.

text

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

문제. attention 한 번은 토큰마다 "하나의 가중치 분포"만 만든다. 한 문장 안에는 보고 싶은 관계가 여러 종류일 수 있다. 예를 들어 "어떤 동사의 주어가 누구인가", "바로 앞 단어는 무엇인가", "같은 대상을 가리키는 단어는 어디인가"가 동시에 필요하다. 분포 하나로 이것들을 전부 섞어 처리하기는 어렵다.

해법: Multi-Head Attention(MHA, 여러 개의 작은 attention을 병렬로 돌린 뒤 결과를 이어 붙이는 구조).

비유: 회의록을 읽는 사람이 한 명이면 한 가지 관점으로만 읽는다. 여러 명이 각자 다른 색 형광펜으로 읽고 마지막에 메모를 합치면 더 다양한 정보를 얻는다. 이때 각 사람이 head이고, 형광펜 색이 head마다 다른 부분공간(subspace, 전체 벡터 공간의 일부 방향들로 이루어진 더 작은 공간)이다.

head 분할·병합. 모델 차원을 d_model, head 수를 h라고 하면 head 하나의 차원은 d_head = d_model / h이다. 예: d_model=64, h=8이면 d_head=8.

입력 x (길이 T, 차원 d_model)를 Q, K, V로 선형 변환한다.
각각을 h조각으로 나눈다. (T, d_model) → (h, T, d_head)
head마다 독립적으로 attention을 계산한다.
결과 h개를 다시 이어 붙인다(concat). (h, T, d_head) → (T, d_model)
마지막에 출력 선형 변환 W_o를 한 번 더 적용해 head들의 정보를 섞는다.

"나눈다"는 새 계산이 아니라 같은 숫자를 다른 모양으로 보는 것(reshape)이다.

파라미터 수는 그대로다. 선형 변환 W_q, W_k, W_v, W_o는 각각 d_model x d_model 크기이다. head 수 h는 이 행렬의 크기에 들어가지 않는다.

text

파라미터 수(bias 제외) = 4 * d_model * d_model

d_model=64라면 4 * 64 * 64 = 16,384개이고, h가 1이든 8이든 같다. head 수는 파라미터를 늘리는 장치가 아니라 같은 파라미터를 어떻게 나눠 쓰느냐를 정하는 설정이다.

왜 sqrt(d_k)가 d_head 기준일까. head마다 내적하는 벡터의 길이가 d_head이므로, 점수 크기를 맞추는 나눗셈도 sqrt(d_head)를 쓴다.

하나의 matmul로 처리. head마다 별도의 W_q를 두고 for문으로 h번 계산할 필요가 없다. d_model x d_model짜리 큰 행렬 하나로 한 번에 Q를 만든 뒤, 결과를 reshape하면 h개 head의 Q가 동시에 나온다. 큰 행렬의 열을 d_head개씩 끊어 보면 head별 W_q가 이어 붙어 있는 것과 수학적으로 같다. GPU는 작은 연산 여러 개보다 큰 행렬 곱 하나를 훨씬 잘 처리한다.

코드로 직접 해보기

1단계: 분할·병합이 reshape임을 눈으로 확인

python

import torch

B, T, d_model, h = 2, 5, 8, 4
d_head = d_model // h

x = torch.arange(B * T * d_model, dtype=torch.float32).reshape(B, T, d_model)

# 분할: (B, T, d_model) -> (B, h, T, d_head)
xh = x.reshape(B, T, h, d_head).transpose(1, 2)
print(xh.shape)  # torch.Size([2, 4, 5, 2])

# 병합: 역순
back = xh.transpose(1, 2).reshape(B, T, d_model)
print(torch.equal(x, back))  # True

값은 하나도 바뀌지 않고 모양만 바뀐다.

2단계: MHA 모듈

python

import math
import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, causal=False):
        super().__init__()
        assert d_model % n_heads == 0, "d_model은 n_heads로 나누어떨어져야 한다"
        self.h = n_heads
        self.d_head = d_model // n_heads
        self.causal = causal
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def split(self, x):
        B, T, _ = x.shape
        return x.reshape(B, T, self.h, self.d_head).transpose(1, 2)  # (B,h,T,d_head)

    def forward(self, x):
        B, T, _ = x.shape
        q, k, v = self.split(self.W_q(x)), self.split(self.W_k(x)), self.split(self.W_v(x))

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)  # (B,h,T,T)
        if self.causal:
            mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=1)
            scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = attn @ v                                  # (B,h,T,d_head)

        out = out.transpose(1, 2).reshape(B, T, -1)     # 병합 (B,T,d_model)
        return self.W_o(out)

3단계: 모양과 파라미터 수 확인

python

x = torch.randn(2, 10, 64)
for h in [1, 4, 8]:
    m = MultiHeadAttention(64, h)
    n_params = sum(p.numel() for p in m.parameters())
    print(h, m(x).shape, n_params)
# 출력 shape은 모두 (2, 10, 64), 파라미터 수는 모두 16640 (bias 포함: 4*(64*64+64))

bias까지 포함하면 4 * (64*64 + 64) = 16,640이다. 세 값이 모두 같게 나오는지 직접 확인한다.

4단계: PyTorch 공식 구현과 비교(선택)

원리를 직접 만들어 봤으니, 가중치를 맞춰 nn.MultiheadAttention과 결과를 비교해 보면 검증이 된다. 다만 내부 가중치가 q, k, v를 하나로 합친 in_proj_weight로 저장되므로 복사 코드가 필요하다. 오늘은 shape과 파라미터 수가 같은지만 비교해도 충분하다.

오늘의 실험

목표: head 수 1·4·8에 따라 작은 시퀀스 태스크 성능이 어떻게 달라지는지 본다.

태스크 예시(직접 만들 수 있는 합성 태스크): 길이 16의 정수 시퀀스(어휘 크기 10)가 주어졌을 때, 각 위치에서 "이 시퀀스에서 현재 토큰과 같은 값이 앞에 몇 번 나왔는지" 또는 "시퀀스를 뒤집은 결과"를 맞추는 문제를 쓴다. 한 위치가 여러 위치를 참고해야 해서 attention의 차이가 드러나기 쉽다.

모델: 임베딩(d_model=64) → MHA 1층 → 선형층 → 위치별 분류. 위치 정보가 필요하면 학습 가능한 position embedding을 더한다.

통제 조건(바꾸는 것은 head 수 하나뿐):

d_model=64, 학습 데이터, 학습 스텝 수, 학습률, 초기화 시드를 모두 동일하게 고정
같은 조건으로 시드 3개 이상 반복(한 번의 결과로 결론 내리지 않기)

기록 표:

head 수	d_head	파라미터 수	최종 train loss	검증 정확도(시드 평균)	편차
1	64
4	16
8	8

관찰 포인트:

파라미터 수가 같은데 성능 차이가 나는가?
차이가 시드 간 편차보다 큰가? 작다면 "차이가 있다"고 쓰지 않는다.
어느 head 수에서도 별 차이가 없다면 그것도 정상 결과이다. 태스크가 너무 쉬우면 head 수의 영향이 드러나지 않는다. 태스크를 어렵게 만들어 다시 돌려 본다.
선택: 학습이 끝난 뒤 attention 가중치(B,h,T,T)를 head별로 히트맵으로 그려, head마다 다른 패턴을 보이는지 눈으로 확인한다.

왜 이걸 배우나 (LLM / Post-training 연결)

GPT 계열 같은 디코더 기반 LLM은 이 MHA 블록을 층마다 쌓은 구조이므로, 모델의 파라미터·메모리·속도 논의는 거의 항상 head 수와 d_head로 돌아온다. LLM Research Engineer로서 모델 설정(config)을 읽고 "이 모델은 head가 몇 개이고 d_head가 얼마인지"를 바로 해석할 수 있어야 하며, 이후 Post-training(사전학습된 모델을 파인튜닝·정렬하는 단계)에서 어떤 가중치(W_q, W_k, W_v, W_o)를 어떻게 건드리는지 이해할 때 오늘의 구조가 기반이 된다.

초보자가 막히는 지점

`d_model % n_heads != 0` 에러: d_model=64에 head 5개처럼 나누어떨어지지 않으면 d_head를 정할 수 없다. assert로 먼저 막아 두자.
reshape 순서 실수: reshape(B, h, T, d_head)로 곧바로 바꾸면 모양은 맞아도 값의 배치가 뒤섞인다. 반드시 reshape(B, T, h, d_head) 후 transpose(1, 2)를 쓴다. 이 실수는 에러가 나지 않고 성능만 나빠져서 발견하기 어렵다. 1단계의 "분할 후 병합하면 원본과 같은가" 테스트가 이를 잡아 준다.
`transpose` 뒤 `view` 에러: transpose 후에는 메모리가 연속적이지 않아 view가 실패할 수 있다. reshape를 쓰거나 .contiguous()를 붙인다.
softmax 방향: dim=-1(Key 축)이어야 각 Query의 가중치 합이 1이 된다. attn.sum(-1)이 1에 가까운지 확인한다.
"head가 많을수록 좋다"는 오해: head를 늘리면 d_head가 작아져 head 하나가 표현할 수 있는 정보가 줄어든다. 오늘 실험이 이 트레이드오프를 직접 보는 연습이다.
head마다 다른 역할을 "학습한다"는 말: 초기화가 다르기 때문에 서로 다른 패턴으로 갈라질 여지가 생길 뿐, 역할이 보장되지는 않는다. 비슷한 head가 여러 개 생기기도 하므로 히트맵으로 직접 보기 전에는 단정하지 않는다.

알고리즘 (병행 1시간)

LeetCode Medium 2문제(그래프). 예: BFS/DFS로 풀 수 있는 연결 요소·섬 개수류, 위상 정렬류에서 고른다. 문제 선택은 직접 하되, 한 문제당 30분을 넘기면 힌트를 보고 이유를 기록한다.

풀이 후 기록할 것:

접근: 그래프를 어떻게 표현했는가(인접 리스트/격자), BFS와 DFS 중 무엇을 왜 골랐는가
시간복잡도·공간복잡도: 노드 수 V, 간선 수 E 기준으로 직접 유도해 적기
틀린 이유: 방문 체크 누락, 방문 표시 시점 실수, 경계 조건, 사이클 처리 등 막힌 원인을 한 줄로