GPT Study

강의 목표

Andrej Karpathy의 강의는 Tiny Shakespeare 데이터와 character-level tokenizer에서 출발해, next-token prediction을 수행하는 작은 GPT를 PyTorch로 직접 구현한다. Bigram language model을 기준선으로 만든 뒤 causal self-attention과 Transformer block을 차례로 추가하면서 각 구성 요소가 필요한 이유를 확인한다.

GPT와 확률적 생성

GPT는 Generative Pre-trained Transformer의 약자이며, Transformer를 기반으로 다음 token의 확률분포를 예측한다. 이 강의에서 구현하는 작은 언어 모델도 같은 next-token prediction 원리를 따른다.

1. 텍스트를 학습 데이터로 바꾸기

강의에서는 데이터에 등장하는 고유 문자를 vocabulary로 만들고, 각 문자를 정수 ID와 양방향으로 매핑한다. 전체 텍스트를 token ID tensor로 변환한 뒤 train/validation split을 구성한다.

Character-level tokenization

Context와 batch

block_size는 모델이 한 번에 볼 수 있는 최대 context 길이다. 입력 \(x\)와 정답 \(y\)는 한 token만큼 어긋나 있으며, 각 위치에서 지금까지의 context로 바로 다음 token을 예측하도록 학습한다. batch_size는 이런 독립적인 sequence를 병렬로 몇 개 처리할지 정한다.

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)

def get_batch(split):
    source = train_data if split == 'train' else val_data
    ix = torch.randint(len(source) - block_size, (batch_size,))
    x = torch.stack([source[i:i + block_size] for i in ix])
    y = torch.stack([source[i + 1:i + block_size + 1] for i in ix])
    return x, y

2. Bigram model로 기준선 만들기

가장 단순한 모델은 현재 token 하나만 보고 다음 token을 예측한다. nn.Embedding(vocab_size, vocab_size)의 각 행은 현재 token에 대응하며, 각 열의 값은 다음 token 후보에 대한 logits다. Softmax를 적용하기 전이므로 이 값 자체가 확률은 아니다.

학습할 때는 logits의 \((B, T, C)\)를 \((B \times T, C)\)로, target의 \((B, T)\)를 \((B \times T)\)로 펼친 뒤 cross-entropy loss를 계산한다. 생성 단계에서는 마지막 timestep의 logits를 확률로 바꾸고 다음 token을 sampling해 기존 sequence 뒤에 붙인다.

class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        logits = self.token_embedding_table(idx)  # (B, T, C)
        if targets is None:
            return logits, None
        B, T, C = logits.shape
        loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
        return logits, loss

3. Self-Attention으로 context를 읽기

Self-Attention은 각 token이 같은 sequence 안의 다른 token을 얼마나 참고할지 계산한다. Query와 Key의 내적으로 관계 점수를 만들고, causal mask로 미래 token을 가린 뒤 softmax weight로 Value를 가중 합산한다.

왜 \(\sqrt{d_k}\)로 나누는가?

Query와 Key의 차원이 커질수록 내적의 분산도 커진다. 큰 score를 그대로 softmax에 넣으면 분포가 한 위치에 지나치게 몰려 gradient가 불안정해질 수 있으므로, \(\sqrt{d_k}\)로 나눠 scale을 맞춘다.

k = self.key(x)
q = self.query(x)
v = self.value(x)

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

4. Transformer Block 조립하기

Karpathy는 Transformer block을 token 사이의 communication을 담당하는 Multi-Head Self-Attention과, 각 token이 얻은 정보를 개별적으로 처리하는 computation 단계인 Feed-Forward Network의 조합으로 설명한다.

class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        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

5. 구현 결과는 Decoder-only Transformer

원 논문의 Transformer는 encoder와 decoder로 구성된다. Encoder는 입력 전체를 양방향으로 읽고, decoder는 causal mask를 사용하며 encoder output을 cross-attention으로 참고한다. 반면 이 강의에서 만든 언어 모델은 encoder와 cross-attention 없이 masked self-attention만 쌓은 decoder-only 구조다.

6. GPT 학습 단계와 연결하기

여기서 구현한 next-token predictor는 대규모 corpus로 base language model을 만드는 pretraining의 축소판으로 볼 수 있다.