Multi-Head Attention

하나의 시선으로는 부족하다

이전 장에서 우리는 Query, Key, Value를 나누어 생각했다. Query는 "나는 무엇을 찾고 있는가", Key는 "나는 어떤 단서를 가지고 있는가", Value는 "내가 실제로 전달할 정보는 무엇인가"에 해당했다. 이 세 가지를 이용하면 한 토큰이 문장 안의 다른 토큰들을 얼마나 참고해야 하는지 계산할 수 있다.

그런데 여기서 자연스러운 질문이 생긴다. Attention을 한 번만 하면 충분할까?

예를 들어 다음 문장을 보자.

민지는 책을 다 읽고 그것을 가방에 넣었다.

이 문장에서 그것을 이해하려면 책을 봐야 한다. 하지만 문장의 전체 의미를 제대로 잡으려면 그보다 더 많은 관계가 필요하다. 누가 행동했는지 보려면 민지를 봐야 하고, 무엇을 했는지 보려면 읽고, 넣었다 같은 동사를 봐야 한다. 또 가방에는 물건이 어디로 이동했는지 알려준다.

하나의 Attention head는 이 관계들 중 일부를 잡아낼 수 있다. 하지만 하나의 시선이 문법 관계, 의미 관계, 위치 관계, 참조 관계를 모두 동시에 완벽하게 보기는 어렵다. 마치 한 사람이 회의록을 읽으면서 "누가 말했는지", "무슨 결론이 났는지", "어떤 숫자가 중요한지", "앞뒤 문맥이 어떻게 이어지는지"를 한 번에 모두 표시하려고 하는 것과 비슷하다. 가능은 하지만 부담스럽다.

Multi-Head Attention은 이 문제를 단순하면서도 강력한 방식으로 해결한다. 하나의 큰 시선으로 문장을 보는 대신, 여러 개의 작은 시선을 동시에 사용한다.

여러 head가 서로 다른 관계를 본다

Multi-Head Attention에서 head는 하나의 Attention 계산 단위다. 각 head는 자기만의 Query, Key, Value 변환을 가지고 같은 입력을 다른 관점으로 바라본다.

여기서 중요한 점은 사람이 미리 "1번 head는 주어를 봐라", "2번 head는 목적어를 봐라", "3번 head는 가까운 단어를 봐라" 하고 지정하지 않는다는 것이다. 모델은 학습 과정에서 자연스럽게 각 head의 역할을 조정한다. 어떤 head는 가까운 토큰에 민감해질 수 있고, 어떤 head는 멀리 떨어진 참조 관계를 잘 볼 수 있다. 또 어떤 head는 특정 문장 구조나 반복 패턴을 포착하는 데 유리해질 수 있다.

다시 민지는 책을 다 읽고 그것을 가방에 넣었다라는 문장을 떠올려 보자. 여러 head는 대략 다음과 같은 관점을 나누어 가질 수 있다.

• 한 head는 그것과 책 사이의 참조 관계를 강하게 볼 수 있다.

• 한 head는 민지와 읽고, 넣었다 사이의 행동 주체 관계를 볼 수 있다.

• 한 head는 현재 토큰 주변의 가까운 단어들을 주로 보며 자연스러운 문장 흐름을 잡을 수 있다.

• 한 head는 가방에처럼 장소나 방향을 나타내는 정보를 더 중요하게 다룰 수 있다.

이렇게 말하면 head마다 뚜렷한 직업이 정해져 있는 것처럼 들릴 수 있지만, 실제로는 훨씬 더 흐릿하고 연속적이다. 어떤 head가 항상 하나의 역할만 하는 것은 아니다. 다만 여러 head가 있으면 모델은 한 번의 Attention으로는 섞여 버릴 수 있는 다양한 관계를 나누어 표현할 여지를 얻는다.

한 문장은 하나의 길만 가진 도로가 아니다. 문법의 길, 의미의 길, 위치의 길, 참조의 길이 겹쳐 있는 복잡한 지도다. Multi-Head Attention은 이 지도를 여러 렌즈로 동시에 읽는 장치다.

계산은 어떻게 이루어질까

이제 조금 더 구체적으로 들어가 보자. 입력 벡터들의 묶음을 X라고 하자. X는 문장에 들어 있는 토큰 임베딩들의 시퀀스다. 보통 구현에서는 다음과 같은 모양을 가진 텐서로 생각한다.

X: [B, T, C]

여기서 B는 배치 크기, T는 토큰 개수, C는 임베딩 차원이다. 예를 들어 한 번에 32개의 문장을 처리하고, 각 문장이 128개의 토큰으로 잘려 있으며, 각 토큰이 256차원 벡터라면 X의 모양은 [32, 128, 256]이 된다.

단일 Attention에서는 이 X에 대해 Query, Key, Value를 만들었다.

Q = XW_Q
K = XW_K
V = XW_V

Multi-Head Attention에서는 이 작업을 head마다 따로 한다. head가 4개라면 다음과 같이 네 벌의 Query, Key, Value가 생긴다.

Q1, K1, V1
Q2, K2, V2
Q3, K3, V3
Q4, K4, V4

각 head는 자신만의 가중치 행렬을 사용한다. 즉, 1번 head의 W_Q, W_K, W_V와 2번 head의 W_Q, W_K, W_V는 서로 다르다. 같은 입력을 받더라도 서로 다른 변환을 거치기 때문에 각 head는 다른 특징을 뽑아낼 수 있다.

각 head의 Attention 계산은 익숙한 형태다.

Attention(Q, K, V) = softmax(QKᵀ / √d_head) V

여기서 d_head는 head 하나가 사용하는 벡터 차원이다. QKᵀ는 토큰끼리의 관련도 점수이고, softmax는 그 점수를 확률처럼 바꾸며, 마지막으로 V의 가중합을 구해 새 표현을 만든다.

Causal Language Model, 즉 GPT처럼 왼쪽에서 오른쪽으로 다음 토큰을 예측하는 모델에서는 각 head에도 causal mask가 적용된다. 현재 토큰이 미래 토큰을 몰래 보면 안 되기 때문이다. 5번째 토큰을 예측하는 표현을 만들 때 6번째, 7번째 토큰을 참고한다면 시험지를 미리 본 것과 같다. 따라서 Multi-Head Attention이라고 해서 이 규칙이 사라지는 것은 아니다. 여러 head가 있더라도 모든 head는 같은 시간의 벽을 지켜야 한다.

concat과 output projection

각 head는 자기만의 Attention 결과를 만든다. head가 4개라면 결과도 4개다. 이제 이들을 하나로 합쳐야 한다. 이때 사용하는 방식이 concat, 즉 이어 붙이기다.

예를 들어 임베딩 차원 C가 256이고 head가 4개라면, 보통 각 head의 차원은 64가 된다.

C = 256
num_heads = 4
d_head = 64

각 head의 출력은 [B, T, 64] 모양이다. head 4개의 출력을 마지막 차원 방향으로 이어 붙이면 다음과 같이 된다.

head1: [B, T, 64]
head2: [B, T, 64]
head3: [B, T, 64]
head4: [B, T, 64]

concat 결과: [B, T, 256]

이제 다시 원래 임베딩 차원과 같은 크기가 되었다. 하지만 여기서 끝나지 않는다. 이어 붙인 결과에는 각 head가 만든 정보가 나란히 놓여 있을 뿐이다. 서로 섞이고 정리되는 과정이 필요하다. 이 역할을 하는 것이 output projection이다.

output = concat(head1, head2, ..., head_h) W_O

W_O는 여러 head의 결과를 다시 한 번 선형 변환하는 가중치 행렬이다. 쉽게 말해 각 head가 가져온 보고서를 한곳에 모은 뒤, 최종 보고서로 편집하는 단계다. 1번 head의 정보와 3번 head의 정보를 섞어 더 유용한 특징을 만들 수도 있고, 덜 중요한 정보를 줄일 수도 있다.

이 구조 때문에 Multi-Head Attention은 단순히 Attention을 여러 번 하고 끝내는 장치가 아니다. 여러 관점으로 문맥을 본 뒤, 그 관점들을 다시 결합해 다음 층으로 넘길 표현을 만드는 장치다.

head 수와 임베딩 차원의 관계

Multi-Head Attention을 처음 보면 이런 생각이 들 수 있다.

head가 많을수록 무조건 좋은 것 아닐까?

직관적으로는 그럴듯하다. 관점이 많아지면 더 많은 관계를 볼 수 있을 것 같다. 하지만 실제로는 그렇게 단순하지 않다. head 수는 임베딩 차원, 연산량, 메모리 사용량과 함께 생각해야 한다.

보통 전체 임베딩 차원 C를 head 수로 나누어 각 head의 차원 d_head를 정한다.

d_head = C / num_heads

예를 들어 C = 256일 때 head 수가 4개라면 head 하나의 차원은 64다. head 수를 8개로 늘리면 head 하나의 차원은 32가 된다. 전체 차원은 그대로인데 관점의 수를 늘렸으므로, 각 관점이 사용할 수 있는 공간은 줄어든다.

이것은 카메라 여러 대를 설치하는 상황과 비슷하다. 카메라가 한 대뿐이면 한 각도에서만 볼 수 있다. 카메라를 여러 대 두면 다양한 각도를 볼 수 있다. 하지만 예산이 고정되어 있다면 카메라 수가 늘어날수록 카메라 한 대당 성능은 낮아질 수 있다. Multi-Head Attention에서도 비슷한 균형이 필요하다.

head 수를 늘릴 때의 장점은 분명하다. 모델이 여러 관계를 병렬로 포착할 수 있고, 단일 head가 놓치기 쉬운 패턴을 다른 head가 보완할 수 있다. 특히 긴 문맥 안에서 서로 다른 종류의 연결을 잡아야 할 때 도움이 된다.

하지만 단점도 있다. head가 너무 많아지면 head 하나의 차원이 지나치게 작아져 각 head가 충분히 풍부한 정보를 담기 어려워질 수 있다. 또한 head 수가 늘어나면 Attention 점수 행렬을 head마다 계산해야 하므로 구현에 따라 메모리 사용량과 계산 부담이 커진다. 특히 시퀀스 길이 T가 길어질수록 Attention 계산은 무거워진다. 각 head는 대체로 [T, T] 크기의 관계 점수 행렬을 만들기 때문이다.

따라서 head 수는 "많을수록 좋다"가 아니라 "모델 크기와 데이터, 목적에 맞게 적당히 정해야 하는 값"이다.

작은 모델에서는 더 조심해야 한다

이 책에서 나중에 만들 Mini GPT는 원리 이해를 위한 작은 모델이다. ChatGPT 같은 상용 규모 모델을 만드는 것이 아니라, Transformer가 어떻게 작동하는지 손으로 만져 보듯 이해하는 것이 목적이다. 작은 모델에서는 head 수를 정할 때 더 조심해야 한다.

예를 들어 임베딩 차원이 128인 작은 모델을 생각해 보자. head를 4개로 두면 head 하나는 32차원을 쓴다. 꽤 무난하다. 그런데 head를 16개로 늘리면 head 하나는 8차원만 가진다. 관점은 많아졌지만, 각 관점이 표현할 수 있는 내용은 매우 얇아진다.

작은 모델에서 head 수를 늘리면 처음에는 결과가 좋아질 수 있다. 단일 head보다 여러 관계를 나누어 볼 수 있기 때문이다. 하지만 어느 지점을 넘으면 효과가 작아지거나 오히려 나빠질 수 있다. 모델의 전체 표현 능력은 제한되어 있는데, 그것을 너무 잘게 쪼개면 각 head가 의미 있는 특징을 충분히 학습하지 못할 수 있다.

또 하나의 현실적인 문제는 속도와 메모리다. 작은 실습 환경에서는 노트북 CPU나 보급형 GPU로 학습을 돌릴 가능성이 높다. head 수를 무작정 늘리면 학습 시간이 길어지고, 배치 크기나 컨텍스트 길이를 줄여야 할 수도 있다. 그러면 오히려 전체 실험 품질이 떨어진다.

Mini GPT를 만들 때는 처음부터 거창한 설정을 잡기보다 작고 안정적인 값으로 시작하는 편이 좋다. 예를 들어 n_embd = 128, n_head = 4처럼 전체 임베딩 차원이 head 수로 깔끔하게 나누어지는 설정이 이해하기 쉽다. 이후에 head 수를 2, 4, 8로 바꾸어 보며 생성 결과와 학습 속도가 어떻게 달라지는지 관찰하면 Multi-Head Attention의 감각이 훨씬 선명해진다.

왜 병렬이라는 점이 중요할까

Multi-Head Attention에서 "Multi"만큼 중요한 단어가 "병렬"이다. 여러 head가 순서대로 이어져 있는 것이 아니라, 같은 입력을 동시에 다른 방식으로 바라본다.

만약 1번 head가 먼저 문장을 보고, 그 결과를 2번 head가 보고, 다시 3번 head가 보는 식이라면 계산 흐름이 길어진다. 하지만 Multi-Head Attention은 같은 입력 X에서 여러 head의 Q, K, V를 한꺼번에 만든다. 각 head는 독립적으로 Attention을 계산하고, 마지막에 결과를 합친다.

이 구조는 Transformer가 RNN 계열 모델과 달리 병렬 계산에 강한 이유와도 연결된다. 물론 Attention은 시퀀스 길이가 길어질수록 계산량이 커지는 문제가 있지만, 한 층 안에서 여러 토큰과 여러 head의 계산을 행렬 연산으로 묶어 처리할 수 있다는 점은 큰 장점이다. GPU는 이런 큰 행렬 계산을 매우 잘 처리한다.

여기서 다시 중요한 그림을 떠올려 보자. Transformer는 문장을 한 토큰씩 천천히 읽으며 내부 상태 하나에 모든 것을 압축하는 방식이 아니다. 문장 안의 토큰들이 서로를 바라보며 표현을 갱신한다. Multi-Head Attention은 이 "서로 바라보기"를 여러 관점으로 확장한다.

Mini GPT에서는 어떻게 설계할까

나중에 Mini GPT를 구현할 때 Multi-Head Attention은 별도의 모듈로 분리하는 것이 좋다. 이유는 간단하다. Transformer Block 안에서 Attention은 핵심 부품이고, 여러 번 반복해서 사용되기 때문이다. 부품을 잘 분리해 두면 코드를 읽기도 쉽고 실험하기도 쉽다.

가장 이해하기 쉬운 설계는 Head 모듈과 MultiHeadAttention 모듈을 나누는 방식이다. Head는 Attention head 하나를 담당한다. 자기만의 Key, Query, Value 선형 변환을 가지고 Attention을 계산한다. MultiHeadAttention은 여러 Head를 모아 실행한 뒤, 결과를 concat하고 output projection을 적용한다.

대략적인 구조는 다음과 같다.

class Head:
    def __init__(self, n_embd, head_size):
        self.key = Linear(n_embd, head_size)
        self.query = Linear(n_embd, head_size)
        self.value = Linear(n_embd, head_size)

    def forward(self, x):
        k = self.key(x)
        q = self.query(x)
        v = self.value(x)

        scores = q @ k.transpose(-2, -1)
        scores = scores / sqrt(head_size)
        scores = apply_causal_mask(scores)
        weights = softmax(scores)

        out = weights @ v
        return out

class MultiHeadAttention:
    def __init__(self, n_embd, num_heads):
        head_size = n_embd // num_heads
        self.heads = [Head(n_embd, head_size) for _ in range(num_heads)]
        self.proj = Linear(n_embd, n_embd)

    def forward(self, x):
        out = concat([head(x) for head in self.heads])
        out = self.proj(out)
        return out

실제 PyTorch 코드에서는 nn.Module, nn.Linear, torch.cat, masked_fill, F.softmax 같은 도구를 사용하게 된다. 또한 학습 안정화와 과적합 방지를 위해 dropout을 넣을 수 있다. 하지만 핵심 구조는 위와 같다.

실전 규모의 구현에서는 효율을 위해 head마다 따로 Linear를 만들기보다, 하나의 큰 선형 변환으로 Q, K, V를 한 번에 만든 뒤 텐서 모양을 바꾸어 여러 head로 나누는 경우가 많다. 이 방식이 더 빠르고 메모리 접근에도 유리하다. 하지만 처음 배우는 단계에서는 head를 명시적으로 나누어 구현하는 편이 훨씬 이해하기 쉽다. 먼저 구조를 눈으로 확인한 뒤, 나중에 최적화된 구현으로 넘어가면 된다.

Multi-Head Attention이 바꾸는 것

Self-Attention은 토큰들이 서로를 참고해 자기 표현을 갱신하는 방법이었다. Multi-Head Attention은 그 과정을 여러 관점으로 확장한다. 이 차이는 작아 보이지만 Transformer의 표현력을 크게 끌어올린다.

하나의 head만 있다면 모델은 문맥을 하나의 기준으로 압축해 보아야 한다. 반면 여러 head가 있으면 문맥 안의 다양한 관계를 나누어 볼 수 있다. 어떤 관계는 가까운 단어 사이에서 생기고, 어떤 관계는 멀리 떨어진 단어 사이에서 생긴다. 어떤 관계는 문법적이고, 어떤 관계는 의미적이다. Multi-Head Attention은 이 복잡한 관계들을 한 번에 하나의 렌즈로 뭉개지 않고, 여러 렌즈로 나누어 관찰한 뒤 다시 합친다.

정리하면 Multi-Head Attention의 핵심은 다음과 같다.

• 각 head는 자기만의 Q, K, V 투영을 통해 같은 입력을 다른 관점으로 본다.

• 각 head의 Attention 결과는 concat으로 이어 붙인 뒤 output projection으로 다시 섞인다.

• head 수가 늘어나면 다양한 관계를 볼 수 있지만, 고정된 임베딩 차원 안에서는 head 하나의 표현 공간이 줄어든다.

• 작은 모델에서는 head 수를 무작정 늘리기보다 임베딩 차원, 학습 속도, 메모리 사용량을 함께 고려해야 한다.

• Mini GPT 구현에서는 Multi-Head Attention을 독립 모듈로 분리하면 Transformer Block을 이해하고 확장하기 쉬워진다.

이제 우리는 Transformer의 핵심 렌즈를 하나 더 얻었다. 토큰은 임베딩 벡터가 되고, Self-Attention을 통해 서로를 바라보며, Multi-Head Attention을 통해 여러 관점에서 문맥을 읽는다.

하지만 아직 중요한 질문이 남아 있다. Attention은 토큰 사이의 관계를 계산하지만, 그 자체만으로는 "첫 번째 토큰", "두 번째 토큰", "마지막에 가까운 토큰" 같은 순서를 자연스럽게 알기 어렵다. 문장은 단어의 모음이 아니라 순서를 가진 흐름이다.

다음 장에서는 Transformer가 이 순서 정보를 어떻게 다루는지, 위치 임베딩이 왜 필요한지 살펴본다.