[GPT] 모델의 구조

https://levelup.gitconnected.com/lets-build-our-own-gpt-model-from-scratch-with-pytorch-236a65a1fb54

Let’s Build our own GPT Model from Scratch with PyTorch

이 블로그를 기반으로 GPT 모델이 어떻게 동작하고 어떤 코드로 되어 있는지 공부한 내용을 정리해 보았습니다

 

1. GPT의 정의와 구조
GPT(Generative Pre-trained Transformer)는 Auto-Regressive 모델로, 주어진 입력 데이터에 대해 다음 토큰을 예측하는 언어 모델입니다. 주로 디코더 블록으로 구성된 Transformer 아키텍처를 기반으로 동작합니다.

 

핵심 단계:

입력 데이터는 토큰화된 후 토큰 임베딩과 위치 임베딩을 통해 벡터화됩니다.
멀티-헤드 어텐션 및 피드포워드 네트워크를 통해 다음 단어의 확률 분포를 생성합니다.
Softmax를 사용해 예측 분포를 확률로 변환하고, 샘플링이나 Argmax를 통해 다음 단어를 결정합니다.
2. 주요 파라미터

batch_size: 한 번에 처리하는 데이터 샘플의 수.
block_size: 입력 시퀀스에서 모델이 고려하는 최대 문맥 길이.
num_embeddings: 단어 집합(Vocabulary)의 크기.
embedding_dim: 각 단어의 벡터 차원.
head_size: 한 어텐션 헤드에서 사용하는 벡터 차원.
num_heads: 어텐션 헤드의 개수.
num_layers: Transformer 블록의 개수.
3. Transformer 블록 구조
Transformer 디코더 블록은 다음과 같은 단계를 거칩니다:

입력 데이터 → LayerNorm
입력 데이터를 정규화합니다.
LayerNorm 출력 → Multi-Head Attention
쿼리(Q), 키(K), 값(V)를 통해 스케일드 내적 어텐션을 계산하고, 모든 헤드의 결과를 결합합니다.
어텐션 출력 → LayerNorm → Feed-Forward Network (FFN)
어텐션 출력에 대해 정규화를 적용하고, 비선형 활성화를 포함하는 FFN을 수행합니다.
예제 코드 (PyTorch)
아래는 GPT 모델의 간단한 구현 코드입니다.

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

# 하이퍼파라미터 설정
batch_size = 32
block_size = 128
num_embeddings = 50257  # 단어 집합 크기
embedding_dim = 768
head_size = embedding_dim // 12  # 각 헤드의 크기
num_heads = 12
num_layers = 12
dropout = 0.1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Multi-Head Attention 블록
class MultiHeadAttention(nn.Module):
    def __init__(self, embedding_dim, num_heads, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.head_size = embedding_dim // num_heads
        self.query = nn.Linear(embedding_dim, embedding_dim)
        self.key = nn.Linear(embedding_dim, embedding_dim)
        self.value = nn.Linear(embedding_dim, embedding_dim)
        self.proj = nn.Linear(embedding_dim, embedding_dim)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        q = self.query(x).view(B, T, self.num_heads, self.head_size).transpose(1, 2)
        k = self.key(x).view(B, T, self.num_heads, self.head_size).transpose(1, 2)
        v = self.value(x).view(B, T, self.num_heads, self.head_size).transpose(1, 2)

        # Scaled Dot-Product Attention
        scores = (q @ k.transpose(-2, -1)) / (self.head_size ** 0.5)
        scores = scores.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
        probs = F.softmax(scores, dim=-1)
        probs = self.dropout(probs)

        # Attention output
        out = (probs @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

# Transformer 블록
class TransformerBlock(nn.Module):
    def __init__(self, embedding_dim, num_heads, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(embedding_dim)
        self.mha = MultiHeadAttention(embedding_dim, num_heads, dropout)
        self.ln2 = nn.LayerNorm(embedding_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embedding_dim, 4 * embedding_dim),
            nn.GELU(),
            nn.Linear(4 * embedding_dim, embedding_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        x = x + self.mha(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

# GPT 모델
class GPT(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, num_heads, num_layers, block_size, dropout=0.1):
        super().__init__()
        self.token_embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.position_embedding = nn.Embedding(block_size, embedding_dim)
        self.blocks = nn.Sequential(*[TransformerBlock(embedding_dim, num_heads, dropout) for _ in range(num_layers)])
        self.ln = nn.LayerNorm(embedding_dim)
        self.head = nn.Linear(embedding_dim, num_embeddings)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        token_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = token_emb + pos_emb
        x = self.blocks(x)
        logits = self.head(self.ln(x))

        loss = None
        if targets is not None:
            B, T, C = logits.shape
            logits = logits.view(B * T, C)
            targets = targets.view(B * T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

# 모델 생성 및 테스트
model = GPT(num_embeddings, embedding_dim, num_heads, num_layers, block_size, dropout).to(device)
x = torch.randint(0, num_embeddings, (batch_size, block_size)).to(device)
logits, loss = model(x, targets=None)
print(f"Logits shape: {logits.shape}")

위치 임베딩:
pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
토큰의 순서를 나타내는 위치 임베딩을 추가합니다.

멀티-헤드 어텐션: 각 헤드는 독립적으로 어텐션을 계산하고, 결과를 결합합니다.

정규화와 FFN: LayerNorm과 피드포워드 네트워크를 사용해 모델의 성능을 안정화합니다.

최종 예측:
logits = self.head(self.ln(x))
마지막으로 선형 계층을 통해 다음 토큰의 확률 분포를 계산합니다.