Transformer Block 만들기

흩어져 있던 부품을 하나의 블록으로

지금까지 우리는 GPT를 이루는 부품들을 하나씩 직접 만들어 보았습니다. 토큰 임베딩과 위치 임베딩으로 입력을 만들었고, Self-Attention으로 문맥을 보게 했고, Multi-Head Attention으로 여러 관점을 동시에 사용하게 했습니다.

하지만 아직은 부품들이 작업대 위에 따로 놓여 있는 상태입니다. 이제 이 부품들을 하나의 단단한 모듈로 묶어야 합니다. 그 모듈이 바로 Transformer Block입니다.

Transformer Block은 이름만 들으면 거창해 보이지만, 구조는 생각보다 명확합니다. 한 블록 안에는 크게 두 개의 흐름이 들어 있습니다.

1. Attention으로 문맥 정보를 섞는다.
2. Feed Forward Network로 각 토큰 표현을 더 풍부하게 바꾼다.

그리고 이 두 흐름을 안정적으로 작동시키기 위해 LayerNorm과 Residual Connection을 함께 사용합니다.

우리가 만들 구조는 다음과 같습니다.

x = x + attention(layernorm(x))
x = x + ffn(layernorm(x))

이 두 줄이 이번 장의 핵심입니다.

첫 번째 줄은 정규화된 입력을 Multi-Head Attention에 넣고, 그 결과를 원래 입력 x에 더합니다. 두 번째 줄은 다시 정규화한 뒤 Feed Forward Network에 넣고, 그 결과를 또 원래 흐름에 더합니다.

이렇게 하면 모델은 매 블록마다 문맥을 한 번 섞고, 각 토큰 표현을 한 번 더 가공합니다. GPT는 이 블록을 한 번만 쓰지 않습니다. 같은 구조의 블록을 여러 층 쌓습니다. 그래서 GPT는 특별한 마법 모듈 하나가 아니라, 잘 설계된 Transformer Block을 반복해서 쌓은 구조라고 볼 수 있습니다.

Transformer Block의 입력과 출력

먼저 shape부터 다시 붙잡고 가겠습니다. 이전 장의 MultiHeadAttention은 입력과 출력의 모양이 같았습니다.

입력 x:  (B, T, n_embd)
출력 out: (B, T, n_embd)

여기서 각 기호의 의미는 이제 익숙할 것입니다.

B      : batch size, 한 번에 처리하는 샘플 수
T      : sequence length, 한 샘플 안의 토큰 수
n_embd : 각 토큰을 표현하는 벡터 차원

Transformer Block도 입력과 출력의 shape를 그대로 유지해야 합니다.

Transformer Block 입력:  (B, T, n_embd)
Transformer Block 출력:  (B, T, n_embd)

왜 그럴까요? 블록을 여러 개 쌓아야 하기 때문입니다. 첫 번째 블록의 출력이 (B, T, n_embd)라면, 두 번째 블록도 같은 형식의 입력을 받을 수 있습니다. 두 번째 블록의 출력도 같은 모양이면, 세 번째 블록을 또 붙일 수 있습니다.

레고 블록을 떠올리면 쉽습니다. 위아래 결합부의 모양이 같아야 계속 쌓을 수 있습니다. Transformer Block도 마찬가지입니다. 내부에서 Attention을 하든, Feed Forward Network를 거치든, 바깥으로 나올 때는 다시 같은 shape로 돌아와야 합니다.

이 성질은 Residual Connection에서도 중요합니다. 다음 코드를 보겠습니다.

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

여기서 x와 self.sa(self.ln1(x))는 더해져야 합니다. 두 텐서를 더하려면 shape가 같아야 합니다.

x                    : (B, T, n_embd)
self.sa(self.ln1(x)) : (B, T, n_embd)
결과                 : (B, T, n_embd)

따라서 Transformer Block을 구현할 때는 각 단계의 shape가 계속 유지되는지 확인하는 습관이 중요합니다. 모델이 커질수록 오류 메시지는 길어지고, 원인은 아주 작은 차원 불일치인 경우가 많습니다.

Feed Forward Network 구현하기

먼저 블록 안에 들어갈 Feed Forward Network를 만들겠습니다. Attention이 토큰들 사이의 관계를 섞는 역할이라면, Feed Forward Network는 각 토큰의 표현을 더 풍부하게 변환하는 역할을 합니다.

조금 더 구체적으로 말하면, Attention은 같은 문장 안의 다른 위치를 참고합니다. 반면 Feed Forward Network는 각 위치의 벡터에 독립적으로 적용됩니다.

Attention:
각 토큰이 다른 토큰들을 본다.

Feed Forward Network:
각 토큰의 벡터를 더 깊게 가공한다.

코드는 간단합니다.

class FeedForward(nn.Module):
    # 각 토큰 위치에 독립적으로 적용되는 작은 신경망

    def __init__(self, n_embd, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

입력 x의 shape가 (B, T, n_embd)라고 해 봅시다. 첫 번째 선형층은 마지막 차원만 바꿉니다.

(B, T, n_embd) -> (B, T, 4 * n_embd)

그다음 ReLU는 비선형성을 추가합니다. 비선형성은 단순한 직선 변환만으로는 표현하기 어려운 패턴을 배우게 해 주는 장치입니다. 마지막 선형층은 다시 원래 차원으로 되돌립니다.

(B, T, 4 * n_embd) -> (B, T, n_embd)

그래서 Feed Forward Network의 전체 입출력은 다음처럼 유지됩니다.

입력: (B, T, n_embd)
출력: (B, T, n_embd)

중간에 차원을 4 * n_embd로 키우는 이유는 표현 공간을 잠시 넓혀 주기 위해서입니다. 좁은 책상 위에서만 조립하던 부품을 잠깐 넓은 작업대로 옮겨 놓는다고 생각하면 됩니다. 넓은 공간에서 더 다양한 조합을 만든 뒤, 다시 원래 크기의 상자에 담아 다음 블록으로 넘깁니다.

실제 대형 모델에서는 ReLU 대신 GELU나 다른 활성화 함수를 사용하는 경우도 많습니다. 하지만 지금은 원리를 이해하는 것이 목표이므로 단순한 ReLU로 시작하겠습니다. 나중에 nn.ReLU()를 nn.GELU()로 바꾸어 학습 결과를 비교해 보는 것도 좋은 실험입니다.

Residual Connection은 왜 필요한가

Transformer Block의 핵심 코드는 다음 두 줄입니다.

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

여기서 x + ... 부분이 Residual Connection입니다. 한국어로는 잔차 연결이라고 부르기도 합니다.

Residual Connection의 생각은 단순합니다.

새로 계산한 결과로 원래 정보를 완전히 덮어쓰지 말고,
원래 정보에 새 정보를 더하자.

예를 들어 어떤 토큰 표현 x가 이미 유용한 정보를 담고 있다고 해 봅시다. Attention을 통과한 결과가 항상 더 좋다고 보장할 수는 없습니다. 특히 학습 초반에는 가중치가 거의 무작위이기 때문에 Attention 결과가 엉뚱할 수 있습니다.

이때 출력이 오직 Attention 결과만으로 정해진다면, 기존 정보가 쉽게 망가질 수 있습니다.

x = self.sa(x)

반면 Residual Connection을 사용하면 원래 흐름이 살아 있습니다.

x = x + self.sa(x)

이 구조에서는 Attention이 해야 할 일이 조금 달라집니다. Attention은 전체 표현을 처음부터 다시 만들 필요가 없습니다. 원래 표현에 더할 수정값 또는 추가 정보를 만들면 됩니다.

이는 깊은 모델을 학습할 때 특히 중요합니다. Transformer Block을 여러 층 쌓으면 입력은 많은 계산을 거쳐야 합니다. Residual Connection은 정보와 그래디언트가 더 잘 흐르도록 도와줍니다. 그래디언트는 학습 중 손실을 줄이기 위해 각 파라미터를 어느 방향으로 바꿔야 하는지 알려 주는 신호입니다. 이 신호가 깊은 층까지 잘 전달되어야 모델이 안정적으로 학습됩니다.

LayerNorm은 어디에 들어갈까

LayerNorm은 각 토큰 벡터의 값 분포를 정리해 주는 역할을 합니다. 값이 너무 커지거나 작아지는 것을 완화하고, 각 층이 더 안정적으로 계산할 수 있게 도와줍니다.

이번 구현에서는 LayerNorm을 Attention과 Feed Forward Network 앞에 둡니다.

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

이런 방식을 흔히 Pre-LN 구조라고 부릅니다. LN은 LayerNorm을 뜻하고, Pre는 주요 계산 전에 정규화한다는 뜻입니다.

흐름을 풀어 쓰면 다음과 같습니다.

1. x를 LayerNorm으로 정리한다.
2. 정리된 x를 Multi-Head Attention에 넣는다.
3. Attention 결과를 원래 x에 더한다.
4. 다시 x를 LayerNorm으로 정리한다.
5. 정리된 x를 Feed Forward Network에 넣는다.
6. Feed Forward 결과를 원래 x에 더한다.

코드 한 줄 안에 많은 일이 들어 있지만, 하나씩 보면 어렵지 않습니다.

먼저 첫 번째 sublayer입니다.

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

self.ln1(x)는 입력을 정규화합니다. self.sa(...)는 Multi-Head Attention입니다. 이 결과는 다시 원래 x에 더해집니다. 따라서 이 단계에서는 문맥 정보가 섞입니다.

두 번째 sublayer입니다.

x = x + self.ffwd(self.ln2(x))

이번에는 self.ln2(x)로 한 번 더 정규화하고, Feed Forward Network로 각 토큰 표현을 가공합니다. 그리고 그 결과를 다시 원래 흐름에 더합니다.

이렇게 Transformer Block은 두 개의 sublayer로 구성됩니다.

이제 이 구조를 그대로 코드로 옮기겠습니다.

Transformer Block 코드 작성하기

이전 장에서 만든 MultiHeadAttention 클래스를 그대로 사용합니다. 여기에 방금 만든 FeedForward, 그리고 PyTorch의 nn.LayerNorm을 결합합니다.

class Block(nn.Module):
    # Transformer block: communication followed by computation

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

        self.sa = MultiHeadAttention(
            n_embd=n_embd,
            n_head=n_head,
            block_size=block_size,
            dropout=dropout,
        )
        self.ffwd = FeedForward(n_embd, dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

코드는 짧지만, 지금까지 배운 핵심이 모두 들어 있습니다.

self.sa는 Multi-Head Attention입니다. 토큰들이 서로를 참고하면서 문맥 정보를 섞습니다. 물론 causal mask가 적용되어 있으므로 현재 위치는 미래 토큰을 볼 수 없습니다. 이전 장에서 만든 Head 안의 tril 마스크가 그 역할을 합니다.

self.ffwd는 Feed Forward Network입니다. Attention이 섞어 놓은 각 토큰 표현을 더 깊게 변환합니다.

self.ln1과 self.ln2는 LayerNorm입니다. Attention 앞과 Feed Forward Network 앞에서 각각 입력을 정리합니다.

마지막으로 x + ...는 Residual Connection입니다. 원래 정보를 유지하면서 새로 계산한 정보를 더합니다.

이제 GPT의 핵심 부품 하나가 완성되었습니다.

Block = Multi-Head Attention + Feed Forward Network + LayerNorm + Residual Connection

복잡해 보이던 GPT 구조가 조금씩 조립식 장난감처럼 보이기 시작합니다.

블록 하나를 따로 실행해 보기

새 모듈을 만들었으니 바로 언어 모델에 넣기 전에 shape를 확인해 보겠습니다. 모델 디버깅에서 가장 좋은 습관은 작은 입력으로 먼저 실행해 보는 것입니다.

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)
block = Block(
    n_embd=n_embd,
    n_head=n_head,
    block_size=block_size,
    dropout=dropout,
)

out = block(x)

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

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

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

이 결과는 아주 중요합니다. 블록 안에서는 Attention, concat, projection, Feed Forward Network, LayerNorm, Dropout이 모두 작동했습니다. 그런데 바깥에서 보면 입력과 출력의 shape는 그대로입니다.

이 성질 덕분에 우리는 블록을 여러 개 쌓을 수 있습니다.

blocks = nn.Sequential(
    Block(n_embd, n_head, block_size, dropout),
    Block(n_embd, n_head, block_size, dropout),
    Block(n_embd, n_head, block_size, dropout),
)

out = blocks(x)
print(out.shape)

출력은 여전히 같습니다.

torch.Size([4, 8, 64])

만약 여기서 shape 오류가 난다면, 대부분 다음 중 하나입니다.

• n_embd가 n_head로 나누어떨어지지 않는다.

• MultiHeadAttention의 projection 결과가 n_embd 차원으로 돌아오지 않는다.

• Feed Forward Network의 마지막 선형층이 n_embd를 출력하지 않는다.

• block_size보다 긴 입력 T를 넣었다.

오류가 날 때는 전체 코드를 한 번에 의심하지 말고, 각 모듈의 입력과 출력 shape를 하나씩 출력해 보면 됩니다.

shape를 따라가며 블록 내부 이해하기

n_embd = 64, n_head = 4인 예를 다시 보겠습니다. 이때 각 head의 크기는 16입니다.

head_size = n_embd // n_head = 64 // 4 = 16

Transformer Block 하나에서 shape는 다음처럼 흐릅니다.

x
(B, T, 64)

ln1(x)
(B, T, 64)

Multi-Head Attention
각 head: (B, T, 16)
concat:  (B, T, 64)
proj:    (B, T, 64)

Residual Add
(B, T, 64) + (B, T, 64) = (B, T, 64)

ln2(x)
(B, T, 64)

Feed Forward Network
Linear 1: (B, T, 256)
ReLU:    (B, T, 256)
Linear 2: (B, T, 64)

Residual Add
(B, T, 64) + (B, T, 64) = (B, T, 64)

이 흐름에서 가장 눈여겨볼 부분은 두 번의 확장과 복귀입니다.

첫 번째는 Multi-Head Attention 안에서 일어납니다. 여러 head가 각각 작은 차원의 출력을 만들고, 이를 이어 붙여 다시 n_embd 차원으로 만듭니다.

두 번째는 Feed Forward Network 안에서 일어납니다. n_embd 차원을 잠시 4 * n_embd로 넓혔다가 다시 n_embd로 줄입니다.

겉으로는 shape가 계속 같아 보이지만, 내부에서는 정보가 계속 섞이고 변환됩니다. 이것이 Transformer Block의 묘미입니다. 모양은 유지하되, 내용은 점점 더 풍부해집니다.

언어 모델에 Transformer Block 넣기

이제 이전 장의 언어 모델에 Block을 넣겠습니다. 달라지는 핵심은 단일 MultiHeadAttention을 직접 호출하는 대신, 여러 개의 Block을 통과시킨다는 점입니다.

먼저 전체 모델 코드를 보겠습니다.

class TransformerLanguageModel(nn.Module):
    def __init__(
        self,
        vocab_size,
        block_size,
        n_embd,
        n_head,
        n_layer,
        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.blocks = nn.Sequential(*[
            Block(
                n_embd=n_embd,
                n_head=n_head,
                block_size=block_size,
                dropout=dropout,
            )
            for _ in range(n_layer)
        ])

        self.ln_f = nn.LayerNorm(n_embd)
        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, n_embd)
        pos = torch.arange(T, device=idx.device)
        pos_emb = self.position_embedding_table(pos)    # (T, n_embd)

        x = tok_emb + pos_emb                           # (B, T, n_embd)
        x = self.blocks(x)                              # (B, T, n_embd)
        x = self.ln_f(x)                                # (B, T, n_embd)
        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

가장 중요한 부분은 여기입니다.

self.blocks = nn.Sequential(*[
    Block(
        n_embd=n_embd,
        n_head=n_head,
        block_size=block_size,
        dropout=dropout,
    )
    for _ in range(n_layer)
])

n_layer는 Transformer Block을 몇 개 쌓을지 정하는 값입니다. n_layer = 4라면 같은 구조의 블록이 네 개 만들어집니다.

여기서 *는 리스트 안의 요소들을 nn.Sequential에 하나씩 펼쳐서 전달하는 역할을 합니다. 즉 다음 코드와 비슷합니다.

self.blocks = nn.Sequential(
    Block(...),
    Block(...),
    Block(...),
    Block(...),
)

다만 n_layer 값을 바꾸면 블록 개수가 자동으로 바뀌므로 리스트 컴프리헨션을 쓰는 편이 편리합니다.

마지막에 들어간 self.ln_f도 눈여겨보아야 합니다.

self.ln_f = nn.LayerNorm(n_embd)

각 블록 안에도 LayerNorm이 있지만, 여러 블록을 모두 통과한 뒤 마지막으로 한 번 더 정규화합니다. 그다음 lm_head가 각 위치에서 다음 토큰의 점수, 즉 logits를 만듭니다.

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

토큰 ID
-> 토큰 임베딩
-> 위치 임베딩 더하기
-> Transformer Block 1
-> Transformer Block 2
-> ...
-> Transformer Block n_layer
-> 마지막 LayerNorm
-> lm_head
-> logits

이 구조가 바로 우리가 앞으로 Mini GPT라고 부를 모델의 뼈대입니다.

n_layer가 바꾸는 것

n_layer는 모델의 깊이를 정합니다. 블록 하나는 문맥을 한 번 섞고, 각 토큰 표현을 한 번 가공합니다. 블록이 여러 개 쌓이면 이 과정이 반복됩니다.

예를 들어 n_layer = 1인 모델은 다음과 같습니다.

임베딩 -> Block -> logits

n_layer = 4인 모델은 다음과 같습니다.

임베딩 -> Block -> Block -> Block -> Block -> logits

층이 많아지면 모델은 더 복잡한 패턴을 표현할 수 있습니다. 첫 번째 블록에서는 가까운 문자 조합이나 단순한 문장 부호 패턴을 잡고, 더 깊은 블록에서는 조금 더 넓은 흐름을 다룰 가능성이 생깁니다. 물론 작은 문자 단위 모델에서 이런 역할이 깔끔하게 나뉜다고 단정할 수는 없습니다. 하지만 구조적으로는 더 여러 번 정보를 섞고 변환할 수 있습니다.

대신 비용도 증가합니다. n_layer를 늘리면 파라미터 수가 늘고, 학습 시간도 늘고, GPU 또는 CPU 메모리도 더 필요합니다.

처음 실습에서는 n_layer = 2 또는 n_layer = 4 정도를 추천합니다. 모델이 너무 작으면 생성 결과가 잘 개선되지 않을 수 있고, 너무 크면 학습이 오래 걸리거나 작은 데이터셋에 과하게 맞춰질 수 있습니다.

여기서 중요한 점은 층을 늘린다고 무조건 좋은 문장이 나오는 것은 아니라는 사실입니다. 모델 크기, 데이터 크기, 학습 시간, 학습률, dropout, block size가 함께 맞아야 합니다. 작은 데이터셋에 지나치게 큰 모델을 쓰면 훈련 데이터는 잘 외우지만 검증 데이터에서는 성능이 좋아지지 않는 과적합이 생길 수 있습니다.

모델 생성과 파라미터 수 확인하기

이제 실제 모델을 만들어 보겠습니다. 이전 장에서 사용한 데이터 준비 코드, CharTokenizer, get_batch, train_data, val_data, 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
n_layer = 4
dropout = 0.2

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

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

파라미터 수는 모델이 학습해야 하는 숫자의 개수입니다. n_embd, n_head, n_layer, vocab_size에 따라 달라집니다. 지금 만드는 모델은 아주 작습니다. 목적은 ChatGPT급 모델을 만드는 것이 아니라, ChatGPT의 핵심 원리를 이해하기 위한 작은 GPT를 직접 조립해 보는 것입니다.

n_layer를 바꾸어 파라미터 수가 어떻게 달라지는지 확인해 보아도 좋습니다.

for n_layer in [1, 2, 4, 6]:
    model = TransformerLanguageModel(
        vocab_size=vocab_size,
        block_size=block_size,
        n_embd=n_embd,
        n_head=n_head,
        n_layer=n_layer,
        dropout=dropout,
    )
    num_params = sum(p.numel() for p in model.parameters())
    print(n_layer, f'{num_params / 1e6:.3f}M parameters')

실행해 보면 n_layer가 커질수록 파라미터 수가 거의 일정한 폭으로 증가하는 것을 볼 수 있습니다. 블록을 하나 추가할 때마다 Multi-Head Attention, Feed Forward Network, LayerNorm의 파라미터가 함께 추가되기 때문입니다.

forward 흐름 점검하기

언어 모델 전체가 제대로 작동하는지 작은 입력으로 확인해 보겠습니다.

B = 4
T = 8
idx = torch.randint(0, vocab_size, (B, T), device=device)
targets = torch.randint(0, vocab_size, (B, T), device=device)

logits, loss = model(idx, targets)

print('idx:', idx.shape)
print('logits:', logits.shape)
print('loss:', loss.item())

출력은 다음과 비슷해야 합니다.

idx: torch.Size([4, 8])
logits: torch.Size([4, 8, vocab_size])
loss: 4.1234

여기서 vocab_size 부분에는 실제 어휘 사전 크기가 들어갑니다. 예를 들어 문자 종류가 65개라면 logits의 shape는 다음과 같습니다.

torch.Size([4, 8, 65])

이 shape는 다음 뜻을 가집니다.

4  : 샘플 4개
8  : 각 샘플의 토큰 위치 8개
65 : 각 위치에서 다음 문자 후보 65개에 대한 점수

즉 모델은 한 번의 forward에서 모든 위치에 대해 다음 토큰 후보 점수를 동시에 계산합니다. 첫 번째 위치도 다음 토큰을 예측하고, 두 번째 위치도 다음 토큰을 예측하고, 여덟 번째 위치도 다음 토큰을 예측합니다.

손실 계산에서는 logits와 targets를 펼쳐서 cross entropy에 넣습니다.

logits_flat = logits.reshape(B * T, V)
targets_flat = targets.reshape(B * T)
loss = F.cross_entropy(logits_flat, targets_flat)

이 구조는 이전 모델들과 같습니다. 달라진 것은 logits를 만들기 전에 더 강력한 Transformer Block들이 들어갔다는 점입니다.

중간 shape를 직접 출력해 보기

복잡한 모델을 만들 때는 중간 shape를 눈으로 확인하는 습관이 큰 도움이 됩니다. 다음처럼 간단한 디버깅 함수를 만들어 볼 수 있습니다.

@torch.no_grad()
def debug_shapes(model, idx):
    model.eval()

    B, T = idx.shape
    print('idx:', idx.shape)

    tok_emb = model.token_embedding_table(idx)
    print('tok_emb:', tok_emb.shape)

    pos = torch.arange(T, device=idx.device)
    pos_emb = model.position_embedding_table(pos)
    print('pos_emb:', pos_emb.shape)

    x = tok_emb + pos_emb
    print('after add:', x.shape)

    for i, block in enumerate(model.blocks):
        x = block(x)
        print(f'after block {i}:', x.shape)

    x = model.ln_f(x)
    print('after ln_f:', x.shape)

    logits = model.lm_head(x)
    print('logits:', logits.shape)

    model.train()

실행해 봅니다.

idx = torch.randint(0, vocab_size, (2, block_size), device=device)
debug_shapes(model, idx)

출력은 다음과 비슷합니다.

idx: torch.Size([2, 16])
tok_emb: torch.Size([2, 16, 64])
pos_emb: torch.Size([16, 64])
after add: torch.Size([2, 16, 64])
after block 0: torch.Size([2, 16, 64])
after block 1: torch.Size([2, 16, 64])
after block 2: torch.Size([2, 16, 64])
after block 3: torch.Size([2, 16, 64])
after ln_f: torch.Size([2, 16, 64])
logits: torch.Size([2, 16, 65])

이 출력에서 가장 중요한 패턴은 블록을 몇 개 지나도 (B, T, n_embd)가 유지된다는 점입니다.

after block 0: (B, T, n_embd)
after block 1: (B, T, n_embd)
after block 2: (B, T, n_embd)
after block 3: (B, T, n_embd)

마지막 lm_head만 n_embd를 vocab_size로 바꿉니다. 왜냐하면 최종 목표는 각 위치에서 다음 토큰 후보들의 점수를 내는 것이기 때문입니다.

이런 shape 출력은 초보자용 임시 코드처럼 보일 수 있지만, 실제 개발에서도 매우 유용합니다. 특히 모델 구조를 조금이라도 바꿀 때는 shape를 먼저 확인하는 편이 좋습니다.

학습 루프는 그대로 사용할 수 있다

모델 구조는 바뀌었지만 학습 루프는 거의 그대로 사용할 수 있습니다. get_batch, estimate_loss, optimizer를 쓰는 흐름은 이전 장과 같습니다.

@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.2103, val loss 4.2087
step 300: train loss 2.5901, val loss 2.6604
step 600: train loss 2.3829, val loss 2.5018
step 900: train loss 2.2490, val loss 2.4175

숫자가 이 예시와 다르더라도 괜찮습니다. 중요한 것은 손실이 전반적으로 내려가는지, 그리고 검증 손실이 훈련 손실과 너무 크게 벌어지지 않는지 확인하는 것입니다.

훈련 손실만 계속 내려가고 검증 손실이 올라간다면 모델이 훈련 데이터에 과하게 맞춰지고 있을 수 있습니다. 이럴 때는 dropout을 조금 키우거나, n_layer 또는 n_embd를 줄이거나, 데이터셋을 더 늘리는 방법을 생각해 볼 수 있습니다.

생성 결과 확인하기

학습이 어느 정도 진행되었다면 텍스트를 생성해 보겠습니다.

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))

여전히 결과는 완벽하지 않을 것입니다. 문자 단위 모델은 단어의 의미를 직접 다루지 않습니다. 데이터도 작고, 모델도 아주 작습니다. 그래서 단어가 깨지거나, 문장이 갑자기 이상해지거나, 같은 패턴이 반복될 수 있습니다.

하지만 내부 구조는 이전보다 훨씬 GPT에 가까워졌습니다. 이제 모델은 단순히 Bigram 확률을 보는 것도 아니고, 단일 Attention head만 쓰는 것도 아닙니다. 여러 head로 문맥을 보고, Feed Forward Network로 표현을 바꾸고, LayerNorm과 Residual Connection으로 깊은 흐름을 안정화합니다. 그리고 그 블록을 여러 층 쌓습니다.

이 차이는 생성 결과가 아주 극적으로 좋아지지 않더라도 중요한 의미가 있습니다. 우리는 이제 GPT의 핵심 설계를 코드로 갖게 되었습니다.

block_size, n_embd, n_head, n_layer를 함께 보기

Mini GPT를 만들다 보면 여러 하이퍼파라미터가 등장합니다. 각각 따로 보면 이해되지만, 함께 작동할 때 헷갈릴 수 있습니다.

처음부터 큰 값을 넣고 시작하면 오류가 났을 때 원인을 찾기 어렵습니다. 작은 값으로 구조가 작동하는지 확인한 뒤, 하나씩 키우는 편이 좋습니다.

추천하는 첫 설정은 다음과 같습니다.

batch_size = 64
block_size = 16
n_embd = 64
n_head = 4
n_layer = 4
dropout = 0.2

CPU에서 너무 느리다면 다음처럼 줄일 수 있습니다.

batch_size = 32
block_size = 16
n_embd = 32
n_head = 4
n_layer = 2
dropout = 0.1

여기서 주의할 점은 n_embd % n_head == 0이어야 한다는 것입니다. 예를 들어 n_embd = 32, n_head = 4라면 각 head는 8차원을 사용하므로 괜찮습니다. 하지만 n_embd = 32, n_head = 5는 나누어떨어지지 않으므로 지금 구현에서는 오류가 납니다.

자주 만나는 오류들

Transformer Block을 처음 구현할 때는 대부분 shape 오류를 만납니다. 오류가 난다고 해서 구조 전체를 이해하지 못한 것은 아닙니다. 오히려 shape 오류는 모델을 직접 만들고 있다는 증거에 가깝습니다.

가장 흔한 오류는 다음과 같습니다.

RuntimeError: The size of tensor a (...) must match the size of tensor b (...)

이 오류는 보통 Residual Connection에서 발생합니다.

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

왼쪽 x와 오른쪽 Attention 출력의 shape가 다르면 더할 수 없습니다. 이때는 MultiHeadAttention의 마지막 projection이 n_embd를 출력하는지 확인해야 합니다.

또 다른 흔한 오류는 다음입니다.

RuntimeError: mat1 and mat2 shapes cannot be multiplied

이 오류는 nn.Linear에 들어간 마지막 차원이 예상과 다를 때 자주 발생합니다. Feed Forward Network의 첫 번째 nn.Linear(n_embd, 4 * n_embd)에 들어가는 텐서의 마지막 차원이 정말 n_embd인지 확인하면 됩니다.

block_size 관련 오류도 자주 만납니다. 위치 임베딩 테이블은 block_size 길이까지만 준비되어 있습니다.

self.position_embedding_table = nn.Embedding(block_size, n_embd)

그런데 forward에 들어온 T가 block_size보다 크면 위치 ID가 범위를 벗어날 수 있습니다. 생성 함수에서 다음 코드가 필요한 이유가 바로 이것입니다.

idx_cond = idx[:, -self.block_size:]

생성된 전체 문장은 점점 길어지지만, 모델에 넣을 때는 마지막 block_size개 토큰만 잘라서 넣습니다. 이것이 지금 모델의 컨텍스트 윈도우입니다.

Transformer Block을 만들었다는 것의 의미

이번 장에서 만든 코드는 길지 않습니다. FeedForward 클래스는 몇 줄이고, Block 클래스도 매우 짧습니다. 그런데 이 짧은 코드 안에 GPT의 중요한 설계가 들어 있습니다.

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

이 두 줄은 다음 네 가지 아이디어를 동시에 담고 있습니다.

• Multi-Head Attention으로 토큰들이 서로의 문맥을 참고한다.

• Feed Forward Network로 각 토큰 표현을 비선형적으로 변환한다.

• LayerNorm으로 각 단계의 입력을 안정화한다.

• Residual Connection으로 원래 정보를 보존하며 새 정보를 더한다.

GPT는 이 블록을 한 번 통과하고 끝나는 모델이 아닙니다. 같은 구조를 여러 번 반복하면서 입력 토큰들의 표현을 조금씩 바꿉니다. 처음에는 단순한 문자 임베딩이었던 벡터가, 블록을 지날수록 문맥이 반영된 표현으로 변합니다. 마지막에는 그 표현을 사용해 다음 토큰의 확률 분포를 만듭니다.

이제 우리는 Mini GPT의 거의 모든 핵심 부품을 갖추었습니다. 다음 단계는 이 구조를 본격적으로 학습시키고, 학습 과정에서 손실이 어떻게 변하는지, 생성 결과가 어떻게 달라지는지 관찰하는 것입니다.

미니 GPT를 직접 만들어보면, LLM은 더 이상 마법이 아닙니다. 다음 장에서는 지금 만든 Transformer Block 기반 모델을 하나의 Mini GPT로 정리하고, 실제로 학습시키는 과정을 차근차근 살펴보겠습니다.