Multi-Head Attention 구현하기

하나의 눈에서 여러 개의 눈으로

이전 장에서 우리는 단일 causal Self-Attention head를 직접 구현했습니다. 각 토큰은 Query, Key, Value를 만들고, Query와 Key를 비교해 앞쪽 토큰들에 대한 Attention weight를 계산했습니다. 그리고 그 가중치로 Value를 섞어 현재 위치의 새로운 표현을 만들었습니다.

이제 한 단계 더 나아가겠습니다. 이번 장의 목표는 여러 개의 Attention head를 묶어 Multi-Head Attention을 구현하는 것입니다.

단일 head는 문맥을 하나의 관점으로 바라봅니다. 예를 들어 어떤 위치의 토큰이 바로 앞 글자를 강하게 볼 수도 있고, 문장 부호를 볼 수도 있고, 공백이나 줄바꿈 같은 구조적 신호를 볼 수도 있습니다. 그런데 head가 하나뿐이면 이 모든 관계를 하나의 Attention 패턴 안에서 처리해야 합니다.

Multi-Head Attention은 이 문제를 단순한 방식으로 해결합니다.

하나의 큰 Attention을 억지로 쓰지 말고,
작은 Attention head 여러 개가 각자 다른 관점으로 문맥을 보게 하자.

비유하자면, 한 명의 독자가 글을 읽는 대신 여러 명의 독자가 동시에 같은 문장을 읽는 것입니다. 어떤 독자는 바로 앞 글자와의 연결을 보고, 어떤 독자는 띄어쓰기나 문장 부호를 보고, 어떤 독자는 조금 먼 위치의 반복 패턴을 볼 수 있습니다. 물론 지금 만드는 작은 모델에서 실제로 그렇게 뚜렷한 역할 분담이 생긴다고 단정할 수는 없습니다. 하지만 구조적으로는 여러 관점을 병렬로 학습할 수 있는 길이 열립니다.

중요한 점은 구현이 생각보다 단순하다는 것입니다. 이미 Head 클래스를 만들었기 때문입니다. 이제 해야 할 일은 그 Head를 여러 개 만들고, 각각의 출력을 이어 붙인 다음, 다시 원래 임베딩 차원으로 섞어 주는 것입니다.

작은 모듈을 하나씩 조립해 나가면 GPT 구조는 더 이상 거대한 블랙박스가 아닙니다. Head를 만들고, MultiHeadAttention을 만들고, 다음 장에서 TransformerBlock을 만들면 우리가 직접 만든 Mini GPT의 뼈대가 거의 완성됩니다.

이번 장에서 만들 것

이번 장에서는 이전 장의 단일 head를 재사용하되, 몇 가지를 보강합니다. 먼저 Head 클래스에 dropout을 추가합니다. 그런 다음 여러 개의 Head를 nn.ModuleList로 묶어 MultiHeadAttention 클래스를 만듭니다. 각 head의 출력은 마지막 차원 방향으로 이어 붙이고, projection layer를 통과시켜 다시 n_embd 차원으로 합칩니다.

전체 흐름은 다음과 같습니다.

입력 x: (B, T, n_embd)

head 1 출력: (B, T, head_size)
head 2 출력: (B, T, head_size)
head 3 출력: (B, T, head_size)
head 4 출력: (B, T, head_size)

concat 결과: (B, T, n_head * head_size)
projection 결과: (B, T, n_embd)

보통은 다음 조건을 맞춥니다.

head_size = n_embd // n_head
n_head * head_size = n_embd

예를 들어 n_embd = 64, n_head = 4라면 각 head의 크기는 16입니다. 네 head의 출력을 이어 붙이면 다시 64차원이 됩니다.

여기서 오해하기 쉬운 부분이 있습니다. head를 4개로 늘린다고 해서 최종 출력 차원이 4배로 커지는 것은 아닙니다. 전체 임베딩 차원 n_embd를 여러 head가 나누어 쓰는 방식에 가깝습니다. 즉 하나의 64차원 관점 대신, 16차원짜리 관점 네 개를 병렬로 사용하는 셈입니다.

Head 클래스에 dropout 추가하기

먼저 이전 장의 Head 클래스를 조금 수정하겠습니다. 핵심 계산은 그대로입니다. 달라지는 점은 Attention weight에 dropout을 적용한다는 것입니다.

dropout은 학습 중 일부 값을 무작위로 0으로 만드는 기법입니다. 처음 보면 일부러 정보를 버리는 이상한 행동처럼 보입니다. 하지만 이 덕분에 모델은 특정 경로나 특정 패턴에만 지나치게 의존하지 않게 됩니다. 작은 데이터셋으로 모델을 학습할 때는 특히 이런 정규화가 도움이 될 수 있습니다.

먼저 필요한 모듈을 가져옵니다.

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

이제 Head 클래스를 다시 작성합니다.

class Head(nn.Module):
    # 단일 head causal self-attention

    def __init__(self, n_embd, head_size, block_size, dropout):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.dropout = nn.Dropout(dropout)

        self.register_buffer(
            'tril',
            torch.tril(torch.ones(block_size, block_size))
        )

    def forward(self, x, return_weights=False):
        B, T, C = x.shape

        k = self.key(x)      # (B, T, head_size)
        q = self.query(x)    # (B, T, head_size)
        v = self.value(x)    # (B, T, head_size)

        head_size = k.shape[-1]

        wei = q @ k.transpose(-2, -1) * head_size**-0.5  # (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)                     # (B, T, T)

        dropped_wei = self.dropout(wei)
        out = dropped_wei @ v                            # (B, T, head_size)

        if return_weights:
            return out, wei

        return out

이전 장과 비교해 가장 중요한 변화는 이 두 줄입니다.

self.dropout = nn.Dropout(dropout)
dropped_wei = self.dropout(wei)

wei는 softmax를 거친 Attention weight입니다. 각 위치가 앞쪽의 어떤 토큰을 얼마나 볼지 정한 가중치입니다. 여기에 dropout을 적용하면 학습할 때 일부 Attention 연결이 무작위로 약해지거나 사라집니다.

예를 들어 어떤 위치가 원래 다음과 같은 가중치로 과거 토큰을 보고 있었다고 해 봅시다.

[0.10, 0.20, 0.70, 0.00]

dropout이 적용되면 학습 중에는 이 중 일부 연결이 꺼질 수 있습니다. 그러면 모델은 항상 세 번째 토큰만 믿는 대신, 다른 경로로도 정보를 가져오는 법을 배워야 합니다.

여기서 return_weights로 반환하는 값은 dropout이 적용되기 전의 wei입니다. 나중에 Attention을 시각화할 때는 원래의 Attention 분포를 보는 편이 이해하기 쉽기 때문입니다. 학습 중 dropout이 적용된 뒤의 가중치는 행의 합이 1처럼 보이지 않을 수 있습니다. 시각화할 때는 model.eval()을 호출해 dropout을 끄고 확인하는 것이 좋습니다.

MultiHeadAttention 클래스 만들기

이제 여러 개의 Head를 묶어 보겠습니다. PyTorch에서 여러 하위 모듈을 리스트처럼 보관할 때는 일반 파이썬 리스트가 아니라 nn.ModuleList를 사용해야 합니다.

왜 그럴까요? PyTorch는 모델 안에 들어 있는 학습 파라미터를 추적해야 합니다. nn.ModuleList에 담긴 모듈들은 PyTorch가 정식 하위 모듈로 인식합니다. 그래서 model.parameters()에 포함되고, model.to(device)를 호출했을 때 함께 이동하며, optimizer가 제대로 업데이트할 수 있습니다.

구현은 다음과 같습니다.

class MultiHeadAttention(nn.Module):
    # 여러 개의 self-attention head를 병렬로 실행한다

    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()

        if n_embd % n_head != 0:
            raise ValueError('n_embd must be divisible by n_head')

        head_size = n_embd // n_head

        self.heads = nn.ModuleList([
            Head(n_embd, head_size, block_size, dropout)
            for _ in range(n_head)
        ])

        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, return_weights=False):
        if return_weights:
            head_outputs = []
            head_weights = []

            for head in self.heads:
                out, wei = head(x, return_weights=True)
                head_outputs.append(out)
                head_weights.append(wei)

            out = torch.cat(head_outputs, dim=-1)
            out = self.dropout(self.proj(out))
            weights = torch.stack(head_weights, dim=1)

            return out, weights

        out = torch.cat([head(x) for head in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

짧은 코드지만 중요한 아이디어가 모두 들어 있습니다. 먼저 여러 head를 만듭니다.

self.heads = nn.ModuleList([
    Head(n_embd, head_size, block_size, dropout)
    for _ in range(n_head)
])

n_head = 4라면 Head 객체가 네 개 만들어집니다. 각각의 head는 자기만의 key, query, value 선형층을 가집니다. 즉 네 head가 같은 계산 구조를 쓰지만, 학습되는 가중치는 서로 다릅니다. 그래서 학습이 진행되면 서로 다른 Attention 패턴을 가질 가능성이 생깁니다.

그다음 forward에서는 모든 head를 같은 입력 x에 적용합니다.

[head(x) for head in self.heads]

각 head의 출력 모양은 (B, T, head_size)입니다. 이것들을 마지막 차원 방향으로 이어 붙입니다.

out = torch.cat([head(x) for head in self.heads], dim=-1)

여기서 dim=-1은 마지막 차원, 즉 channel 차원을 뜻합니다. 예를 들어 네 head가 각각 (B, T, 16)을 출력하면 concat 결과는 (B, T, 64)가 됩니다.

마지막으로 projection layer를 통과시킵니다.

out = self.dropout(self.proj(out))

self.proj는 여러 head에서 나온 정보를 다시 섞어 주는 선형층입니다. 각 head가 자기 관점으로 만든 정보를 그냥 이어 붙이기만 하면 head 사이의 정보가 충분히 섞이지 않습니다. projection layer는 이 정보들을 다시 조합해 다음 계층이 사용하기 좋은 표현으로 바꿉니다.

concat과 projection을 차원으로 따라가기

Multi-Head Attention은 코드는 짧지만, 차원 감각을 놓치면 헷갈리기 쉽습니다. 작은 예로 따라가 보겠습니다.

B = 4
T = 8
n_embd = 64
n_head = 4
head_size = n_embd // n_head  # 16

입력 x의 모양은 다음과 같습니다.

x: (4, 8, 64)

각 head는 입력을 받아 16차원 출력을 만듭니다.

head 0: (4, 8, 16)
head 1: (4, 8, 16)
head 2: (4, 8, 16)
head 3: (4, 8, 16)

이 네 개를 마지막 차원으로 이어 붙이면 다음과 같습니다.

concat: (4, 8, 64)

그다음 projection layer를 통과합니다.

projection: (4, 8, 64)

입력도 (B, T, n_embd)이고 출력도 (B, T, n_embd)입니다. 이 성질은 앞으로 Transformer Block을 만들 때 매우 중요합니다. 입력과 출력의 모양이 같아야 Residual Connection으로 원래 입력을 더할 수 있기 때문입니다.

다음 장에서 우리는 이런 계산을 하게 됩니다.

x = x + self.sa(self.ln1(x))

이때 x와 self.sa(...)의 모양이 같아야 더할 수 있습니다. 그래서 Multi-Head Attention은 최종적으로 다시 n_embd 차원을 출력하도록 설계합니다.

n_embd, n_head, head_size의 조건

Multi-Head Attention에서 가장 자주 만나는 오류는 차원 불일치입니다. 특히 n_embd와 n_head가 나누어떨어지지 않으면 곤란합니다.

우리는 다음 방식으로 head_size를 정했습니다.

head_size = n_embd // n_head

따라서 n_embd는 n_head로 나누어떨어지는 값이어야 합니다.

n_embd = 64, n_head = 4  -> head_size = 16  가능
n_embd = 96, n_head = 6  -> head_size = 16  가능
n_embd = 128, n_head = 8 -> head_size = 16  가능
n_embd = 64, n_head = 3  -> 나누어떨어지지 않음  곤란

물론 이론적으로는 n_head * head_size가 꼭 n_embd와 같지 않아도 projection layer를 맞추면 구현할 수 있습니다. 하지만 지금 만드는 Mini GPT에서는 구조를 단순하게 유지하겠습니다. n_embd를 head 수로 정확히 나누고, concat 결과가 다시 n_embd가 되게 만드는 방식이 가장 이해하기 좋습니다.

그래서 코드에 다음 검사를 넣었습니다.

if n_embd % n_head != 0:
    raise ValueError('n_embd must be divisible by n_head')

실습 중 이 오류가 나오면 n_embd나 n_head 값을 바꾸면 됩니다. 처음에는 n_embd = 64, n_head = 4 조합을 추천합니다. 계산량도 크지 않고, 차원도 깔끔합니다.

왜 dropout을 두 번 적용할까

이번 구현에서는 dropout이 두 곳에 들어갑니다. 하나는 Attention weight 뒤이고, 다른 하나는 projection 뒤입니다.

# Head 안
wei = F.softmax(wei, dim=-1)
dropped_wei = self.dropout(wei)
out = dropped_wei @ v

# MultiHeadAttention 안
out = self.dropout(self.proj(out))

두 dropout은 역할이 조금 다릅니다.

Attention weight에 적용하는 dropout은 토큰 사이의 연결을 정규화합니다. 어떤 위치가 항상 특정 과거 위치에만 지나치게 의존하지 않도록 만듭니다. 학습 중 일부 연결이 꺼지면 모델은 다른 토큰의 정보도 활용하는 법을 배워야 합니다.

Projection 뒤에 적용하는 dropout은 여러 head의 정보를 합친 결과를 정규화합니다. head들의 출력이 projection layer를 거쳐 섞인 뒤, 그 표현의 일부를 학습 중 무작위로 줄입니다. 덕분에 다음 계층은 특정 차원의 신호에만 과하게 기대지 않게 됩니다.

dropout은 학습할 때만 작동합니다. model.train() 상태에서는 무작위로 값을 끄지만, model.eval() 상태에서는 꺼지지 않습니다. 그래서 평가나 생성, Attention 시각화를 할 때는 보통 다음처럼 합니다.

model.eval()

그리고 다시 학습을 이어 갈 때는 다음처럼 돌려놓습니다.

model.train()

작은 모델에서는 dropout 값을 너무 크게 잡으면 오히려 학습이 어려워질 수 있습니다. 처음에는 dropout = 0.2 정도로 시작해 보겠습니다. 데이터가 아주 작고 모델도 작다면 0.1이나 0.0으로 바꿔 비교해 보는 것도 좋은 실험입니다.

Multi-Head Attention만 따로 실행해 보기

언어 모델에 넣기 전에 MultiHeadAttention 모듈만 따로 실행해 보겠습니다. 새 모듈을 만들 때는 항상 작은 입력으로 모양부터 확인하는 습관을 들이는 것이 좋습니다.

torch.manual_seed(1337)

B = 4
T = 8
n_embd = 64
n_head = 4
block_size = 8
dropout = 0.2

x = torch.randn(B, T, n_embd)
sa = MultiHeadAttention(n_embd, n_head, block_size, dropout)

out = sa(x)
print('x:', x.shape)
print('out:', out.shape)

출력은 다음과 같아야 합니다.

x: torch.Size([4, 8, 64])
out: torch.Size([4, 8, 64])

입력과 출력의 모양이 같습니다. 하지만 내부에서는 네 개의 head가 각자 Attention을 계산했고, 그 결과가 concat과 projection을 거쳐 다시 64차원으로 합쳐졌습니다.

이번에는 head별 Attention weight의 모양을 확인해 보겠습니다.

out, weights = sa(x, return_weights=True)

print('out:', out.shape)
print('weights:', weights.shape)

출력은 다음과 같아야 합니다.

out: torch.Size([4, 8, 64])
weights: torch.Size([4, 4, 8, 8])

weights의 모양에서 두 번째 차원 4가 head 수입니다.

weights: (B, n_head, T, T)

즉 weights[0, 2]는 첫 번째 샘플에서 세 번째 head의 Attention 행렬입니다. 이제 우리는 하나의 Attention 그림이 아니라 head별 Attention 그림을 볼 수 있습니다.

언어 모델에 Multi-Head Attention 넣기

이제 MultiHeadAttention을 언어 모델 안에 넣겠습니다. 이전 장의 SingleHeadAttentionLanguageModel과 구조는 거의 같습니다. 달라지는 부분은 self.sa_head 대신 self.sa_heads를 쓴다는 점입니다.

class MultiHeadAttentionLanguageModel(nn.Module):
    def __init__(self, vocab_size, block_size, n_embd, n_head, dropout):
        super().__init__()
        self.block_size = block_size

        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.sa_heads = MultiHeadAttention(
            n_embd=n_embd,
            n_head=n_head,
            block_size=block_size,
            dropout=dropout
        )
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding_table(idx)       # (B, T, C)
        pos = torch.arange(T, device=idx.device)
        pos_emb = self.position_embedding_table(pos)    # (T, C)

        x = tok_emb + pos_emb                           # (B, T, C)
        x = self.sa_heads(x)                            # (B, T, C)
        logits = self.lm_head(x)                        # (B, T, vocab_size)

        loss = None
        if targets is not None:
            B, T, V = logits.shape
            logits_flat = logits.reshape(B * T, V)
            targets_flat = targets.reshape(B * T)
            loss = F.cross_entropy(logits_flat, targets_flat)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

forward의 큰 흐름은 이제 익숙할 것입니다.

토큰 ID
-> 토큰 임베딩
-> 위치 임베딩 더하기
-> Multi-Head Attention
-> vocab_size 차원의 logits
-> cross entropy loss

아직 Feed Forward Network도 없고, LayerNorm도 없고, Residual Connection도 없습니다. 따라서 실제 GPT의 Transformer Block과 비교하면 여전히 단순합니다. 하지만 Attention 부분만 놓고 보면 이제 단일 head가 아니라 여러 head가 병렬로 문맥을 보는 구조가 되었습니다.

학습 설정하기

이전 장에서 만든 get_batch, train_data, val_data, tokenizer, vocab_size는 그대로 사용합니다. 여기서는 모델과 하이퍼파라미터만 바꾸겠습니다.

torch.manual_seed(1337)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

batch_size = 64
block_size = 16
max_iters = 3000
eval_interval = 300
eval_iters = 200
learning_rate = 1e-3
n_embd = 64
n_head = 4
dropout = 0.2

model = MultiHeadAttentionLanguageModel(
    vocab_size=vocab_size,
    block_size=block_size,
    n_embd=n_embd,
    n_head=n_head,
    dropout=dropout
)
model = model.to(device)

num_params = sum(p.numel() for p in model.parameters())
print(f'{num_params / 1e6:.3f}M parameters')

평가 함수와 학습 루프는 이전 장과 거의 같습니다.

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()

    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)

        for k in range(eval_iters):
            X, Y = get_batch(split)
            X = X.to(device)
            Y = Y.to(device)

            logits, loss = model(X, Y)
            losses[k] = loss.item()

        out[split] = losses.mean().item()

    model.train()
    return out

학습 루프입니다.

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
history = []

for step in range(max_iters):
    if step % eval_interval == 0 or step == max_iters - 1:
        losses = estimate_loss()
        train_loss = losses['train']
        val_loss = losses['val']
        history.append((step, train_loss, val_loss))
        print(f'step {step}: train loss {train_loss:.4f}, val loss {val_loss:.4f}')

    X, Y = get_batch('train')
    X = X.to(device)
    Y = Y.to(device)

    logits, loss = model(X, Y)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

실행하면 손실이 점차 내려가는 것을 볼 수 있습니다. 정확한 값은 데이터셋, 실행 환경, 초기값에 따라 달라집니다.

step 0: train loss 4.1902, val loss 4.1876
step 300: train loss 2.7118, val loss 2.7564
step 600: train loss 2.4879, val loss 2.5681
step 900: train loss 2.3812, val loss 2.4770

숫자 자체보다 중요한 것은 흐름입니다. 모델이 여러 head를 통해 문맥 정보를 섞고, 각 위치에서 다음 토큰을 예측하며, 손실을 줄이는 방향으로 파라미터를 조정하고 있습니다.

텍스트 생성해 보기

학습이 끝났다면 텍스트를 생성해 보겠습니다.

model.eval()

context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=500)[0].tolist()

print(tokenizer.decode(generated))

또는 학습 데이터의 앞부분을 시작 문맥으로 사용할 수도 있습니다.

model.eval()

context = train_data[:block_size].unsqueeze(0).to(device)
generated = model.generate(context, max_new_tokens=500)[0].tolist()

print(tokenizer.decode(generated))

출력은 아직 완벽하지 않을 것입니다. 문자 단위 모델이기 때문에 단어가 깨질 수 있고, 같은 표현이 반복될 수 있으며, 문장이 갑자기 이상한 방향으로 흘러갈 수 있습니다. 괜찮습니다. 지금 목표는 ChatGPT급 문장을 만드는 것이 아닙니다. 이 책에서 만드는 것은 ChatGPT가 아니라, ChatGPT의 핵심 원리를 이해하기 위한 작은 GPT입니다.

그래도 이전보다 조금 더 그럴듯한 패턴이 보일 수 있습니다. 공백의 위치가 더 자연스러워지거나, 자주 등장하는 글자 조합이 더 잘 이어지거나, 문장 부호 주변의 흐름이 조금 나아질 수 있습니다. 물론 작은 데이터셋과 작은 모델에서는 차이가 크지 않을 수도 있습니다. 실험 결과가 기대보다 미묘하더라도 실망할 필요는 없습니다.

중요한 변화는 모델 내부 구조에 있습니다. 이제 모델은 하나의 Attention 관점이 아니라 여러 Attention 관점을 동시에 사용합니다.

단일 head와 multi-head 비교 실험

이제 작은 비교 실험을 해 보겠습니다. 단일 head 모델과 multi-head 모델을 같은 조건에서 학습시키고, 검증 손실과 생성 결과를 비교합니다. 이 실험은 엄밀한 벤치마크가 아닙니다. 데이터도 작고 모델도 작기 때문에 결과가 실행할 때마다 달라질 수 있습니다. 목표는 Multi-Head Attention이 어떤 차이를 만드는지 감각적으로 확인하는 것입니다.

먼저 학습을 함수로 묶어 보겠습니다.

def train_steps(model, max_iters=1000):
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    for step in range(max_iters):
        X, Y = get_batch('train')
        X = X.to(device)
        Y = Y.to(device)

        logits, loss = model(X, Y)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    return model

검증 손실을 간단히 재는 함수도 만듭니다.

@torch.no_grad()
def eval_model(model, eval_iters=100):
    model.eval()
    losses = torch.zeros(eval_iters)

    for k in range(eval_iters):
        X, Y = get_batch('val')
        X = X.to(device)
        Y = Y.to(device)

        logits, loss = model(X, Y)
        losses[k] = loss.item()

    model.train()
    return losses.mean().item()

이제 두 모델을 준비합니다. 단일 head 모델은 이전 장에서 만든 SingleHeadAttentionLanguageModel을 사용합니다. Multi-head 모델은 방금 만든 클래스를 사용합니다.

torch.manual_seed(1337)

single_model = SingleHeadAttentionLanguageModel(
    vocab_size=vocab_size,
    block_size=block_size,
    n_embd=n_embd
)

multi_model = MultiHeadAttentionLanguageModel(
    vocab_size=vocab_size,
    block_size=block_size,
    n_embd=n_embd,
    n_head=n_head,
    dropout=dropout
)

single_model = train_steps(single_model, max_iters=1000)
multi_model = train_steps(multi_model, max_iters=1000)

single_val_loss = eval_model(single_model)
multi_val_loss = eval_model(multi_model)

print('single head val loss:', single_val_loss)
print('multi head val loss:', multi_val_loss)

생성 결과도 비교해 봅니다.

@torch.no_grad()
def generate_text(model, max_new_tokens=300):
    model.eval()
    context = train_data[:block_size].unsqueeze(0).to(device)
    generated = model.generate(context, max_new_tokens=max_new_tokens)[0].tolist()
    return tokenizer.decode(generated)

print('--- single head ---')
print(generate_text(single_model))

print('--- multi head ---')
print(generate_text(multi_model))

어떤 실행에서는 multi-head 모델의 검증 손실이 조금 더 낮을 수 있습니다. 어떤 실행에서는 차이가 거의 없거나, 오히려 단일 head가 비슷하게 나올 수도 있습니다. 이것은 이상한 일이 아닙니다. 지금 모델에는 아직 Transformer Block의 다른 중요한 부품들이 빠져 있습니다. 특히 Feed Forward Network, LayerNorm, Residual Connection이 없기 때문에 Multi-Head Attention의 장점이 충분히 드러나지 않을 수 있습니다.

하지만 생성 결과를 찬찬히 보면 몇 가지 차이를 관찰할 수 있습니다. multi-head 모델이 반복을 조금 덜 하거나, 가까운 글자 조합을 조금 더 안정적으로 이어 가거나, 특정 구두점 주변의 패턴을 더 잘 흉내 낼 수 있습니다. 반대로 데이터가 너무 작으면 두 모델 모두 비슷하게 어설플 수 있습니다.

이 실험에서 얻어야 할 결론은 단순히 multi-head가 항상 이긴다는 것이 아닙니다. 더 중요한 결론은 이것입니다.

여러 head를 병렬로 붙이는 구조는 코드 몇 줄로 구현할 수 있고,
이 구조가 GPT의 핵심 부품 중 하나라는 사실을 직접 확인했다.

head별 Attention 그림 보기

Multi-Head Attention의 재미있는 점은 head별 Attention weight를 따로 볼 수 있다는 것입니다. 단일 head에서는 Attention 행렬이 하나뿐이었지만, 이제는 head 수만큼 Attention 행렬이 생깁니다.

먼저 검증 데이터에서 샘플 하나를 가져옵니다.

X, Y = get_batch('val')
idx = X[:1].to(device)

print(tokenizer.decode(idx[0].tolist()))

이제 모델 내부의 임베딩을 만든 뒤 sa_heads를 직접 호출합니다.

model.eval()

with torch.no_grad():
    B, T = idx.shape

    tok_emb = model.token_embedding_table(idx)
    pos = torch.arange(T, device=device)
    pos_emb = model.position_embedding_table(pos)
    x = tok_emb + pos_emb

    out, weights = model.sa_heads(x, return_weights=True)

print(weights.shape)

출력은 다음과 같아야 합니다.

torch.Size([1, 4, 16, 16])

block_size = 16, n_head = 4라면 (1, 4, 16, 16)입니다. 이제 각 head의 Attention 행렬을 그려 보겠습니다.

import matplotlib.pyplot as plt

att = weights[0].detach().cpu()  # (n_head, T, T)
labels = [tokenizer.decode([int(i)]) for i in idx[0].detach().cpu().tolist()]

fig, axes = plt.subplots(1, n_head, figsize=(4 * n_head, 4))

for h in range(n_head):
    ax = axes[h]
    ax.imshow(att[h], cmap='viridis')
    ax.set_title(f'head {h}')
    ax.set_xticks(range(T))
    ax.set_yticks(range(T))
    ax.set_xticklabels(labels, rotation=90)
    ax.set_yticklabels(labels)

plt.tight_layout()
plt.show()

그림을 보면 head마다 조금씩 다른 패턴이 보일 수 있습니다. 어떤 head는 자기 자신이나 가까운 과거 위치를 주로 볼 수 있고, 어떤 head는 더 넓게 퍼진 가중치를 가질 수 있습니다. 학습이 충분하지 않거나 데이터가 작으면 head들이 서로 비슷해 보일 수도 있습니다.

여기서도 해석에는 조심해야 합니다. Attention 그림이 곧 모델의 생각을 완벽히 설명해 주는 것은 아닙니다. 특히 지금처럼 작은 문자 단위 모델에서는 사람이 보기에 의미 있는 역할 분담이 뚜렷하게 드러나지 않을 수 있습니다. 그래도 head별 패턴을 관찰하는 것은 구현을 확인하고, Multi-Head Attention의 병렬 구조를 눈으로 이해하는 데 도움이 됩니다.

특히 오른쪽 위 삼각형 영역은 계속 어두워야 합니다. causal mask가 모든 head에 적용되기 때문입니다. 만약 어떤 head에서 미래 위치를 보는 값이 나타난다면 Head 클래스의 mask 적용 부분을 다시 확인해야 합니다.

Multi-Head Attention은 병렬 조립이다

여기까지 오면 Multi-Head Attention이 생각보다 단순한 조립이라는 사실이 보입니다. 수식으로 보면 복잡해 보이지만, 구현 관점에서는 다음 흐름입니다.

1. 같은 입력 x를 여러 Head에 넣는다.
2. 각 Head가 독립적으로 causal Self-Attention을 계산한다.
3. head별 출력을 마지막 차원으로 concat한다.
4. projection layer로 다시 n_embd 차원으로 섞는다.
5. dropout으로 학습을 정규화한다.

이 흐름은 Transformer를 이해할 때 매우 중요합니다. GPT는 단순히 하나의 거대한 연산으로 이루어진 모델이 아닙니다. 작은 모듈이 반복적으로 조립된 구조입니다. 토큰 임베딩, 위치 임베딩, Multi-Head Attention, Feed Forward Network, LayerNorm, Residual Connection이 일정한 패턴으로 쌓입니다.

이 관점으로 보면 거대한 LLM도 조금 덜 낯설어집니다. 물론 실제 상용 LLM은 훨씬 더 큰 데이터와 더 많은 파라미터, 더 정교한 학습 기법을 사용합니다. 하지만 핵심 구조는 우리가 지금 손으로 구현하는 작은 코드와 연결되어 있습니다.

자주 만나는 오류

Multi-Head Attention을 구현할 때는 대부분 차원 문제에서 오류가 납니다. 오류 메시지가 길어도 겁먹을 필요는 없습니다. 거의 항상 어느 텐서의 마지막 차원이 예상과 다르다는 뜻입니다.

먼저 n_embd가 n_head로 나누어떨어지는지 확인하세요.

n_embd = 64
n_head = 4
assert n_embd % n_head == 0

n_head = 3처럼 설정하면 64를 정확히 나눌 수 없습니다. 이 경우 n_head를 2, 4, 8 등으로 바꾸거나 n_embd를 조정해야 합니다.

다음으로 concat 차원을 확인하세요.

out = torch.cat([head(x) for head in self.heads], dim=-1)

여기서 dim=-1을 써야 channel 차원으로 붙습니다. 만약 dim=1로 붙이면 시간 길이 T 방향으로 이어 붙게 되어 모양이 달라지고, 뒤의 projection layer에서 오류가 날 수 있습니다.

또 하나는 projection layer의 입력 차원입니다. 우리는 head_size = n_embd // n_head로 정했으므로 concat 결과가 n_embd가 됩니다. 따라서 projection layer는 다음처럼 만들 수 있습니다.

self.proj = nn.Linear(n_embd, n_embd)

만약 head size를 직접 지정하는 방식으로 바꾸었다면 n_head * head_size가 projection의 입력 차원이 되어야 합니다. 지금은 구조를 단순하게 유지하기 위해 n_head * head_size = n_embd가 되도록 만들었습니다.

마지막으로 device 문제를 확인하세요. tril은 register_buffer로 등록되어 있기 때문에 model.to(device)를 호출하면 함께 이동합니다. 하지만 입력 X, Y도 같은 device로 보내야 합니다.

X = X.to(device)
Y = Y.to(device)

모델은 GPU에 있고 입력은 CPU에 있으면 행렬 곱에서 오류가 납니다. 실습 중 device 관련 오류가 나오면 모델과 데이터가 같은 장치에 있는지 먼저 확인하세요.

정리

이번 장에서는 단일 Head 클래스를 재사용해 MultiHeadAttention 모듈을 만들었습니다. 여러 head를 nn.ModuleList로 구성했고, 각 head가 같은 입력을 독립적으로 바라보게 했습니다. 그런 다음 head별 출력을 torch.cat(..., dim=-1)로 이어 붙이고, projection layer로 다시 n_embd 차원으로 섞었습니다.

또 Attention weight와 projection 뒤에 dropout을 적용했습니다. Attention weight dropout은 토큰 사이의 특정 연결에 과하게 의존하지 않도록 돕고, projection 뒤 dropout은 여러 head가 합쳐진 표현을 정규화합니다. 학습 중에는 dropout이 작동하지만, 평가와 생성에서는 model.eval()을 통해 꺼진다는 점도 확인했습니다.

차원 조건도 중요했습니다. n_embd는 n_head로 나누어떨어져야 하며, 보통 head_size = n_embd // n_head로 둡니다. 예를 들어 n_embd = 64, n_head = 4라면 head 하나의 크기는 16이고, 네 head를 concat하면 다시 64차원이 됩니다.

마지막으로 단일 head와 multi-head 모델을 작은 실험으로 비교했습니다. 결과가 항상 극적으로 달라지지는 않지만, 이제 모델은 여러 관점으로 문맥을 볼 수 있는 구조를 갖게 되었습니다.

다음 장에서는 여기에 Feed Forward Network, LayerNorm, Residual Connection을 더해 하나의 Transformer Block을 만들겠습니다. 그 순간부터 우리가 만든 코드는 단순한 Attention 모델을 넘어, Mini GPT의 진짜 뼈대에 가까워집니다.