[Deep Learning] MinGPT

프로젝트: br4c3/MinGPT

MinGPT란?

MinGPT는 PyTorch로 GPT 구조를 직접 구현한 작은 LLM project이다. 이미 학습된 model을 가져와 사용하는 project가 아니라, GPT의 핵심 구성 요소를 code로 직접 만들고 학습 loop와 text generation까지 연결한다.

repository 구조는 단순하다.

MinGPT/
  GPT.py
  main.py
  requirements.txt
  img/

핵심 파일은 두 개이다.

즉 이 project는 다음 흐름을 한 파일 단위로 따라가기 좋게 구성되어 있다.

text
  -> token
  -> input/target sequence
  -> GPT model
  -> cross entropy loss
  -> parameter update
  -> next token generation

목적

LLM을 API로만 사용하면 내부에서 어떤 계산이 일어나는지 보기 어렵다. MinGPT는 GPT를 구성하는 부품을 직접 구현하면서 다음 질문에 답하는 project이다.

token은 어떻게 model input이 되는가?
self-attention은 어떤 shape으로 계산되는가?
causal mask는 왜 필요한가?
transformer block은 어떤 순서로 쌓이는가?
학습된 model은 다음 token을 어떻게 고르는가?

특히 MultiHeadAttention, LayerNorm, FeedForward, TransformerBlock, GPTModel이 class 단위로 나뉘어 있어서 GPT 구조를 code로 추적하기 쉽다.

의존성

필요한 package는 requirements.txt에 정리되어 있다.

tiktoken
torch
torchvision
numpy

설치는 다음처럼 할 수 있다.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

tiktoken은 GPT-2 tokenizer를 사용하기 위해 필요하고, torch는 model 구현과 학습에 사용된다.

Dataset 구성

GPT.py의 GPTDatasetV1은 text를 token으로 바꾼 뒤 input sequence와 target sequence를 만든다.

예를 들어 token sequence가 다음과 같다고 하자.

[10, 20, 30, 40, 50]

language model은 현재 token들을 보고 다음 token을 예측해야 한다. 그래서 input과 target은 한 칸씩 밀린 형태가 된다.

input:  [10, 20, 30, 40]
target: [20, 30, 40, 50]

긴 text에서는 sliding window 방식으로 여러 sequence를 만든다.

for i in range(0, len(token_ids) - max_length, stride):
    input_chunk = token_ids[i:i + max_length]
    target_chunk = token_ids[i + 1: i + max_length + 1]

이 구조는 autoregressive language model의 기본 학습 방식이다. model은 각 위치에서 다음 token을 맞히도록 학습된다.

Tokenizer

tokenizer는 tiktoken의 GPT-2 encoding을 사용한다.

tokenizer = tiktoken.get_encoding("gpt2")

문장을 그대로 neural network에 넣을 수는 없다. 먼저 text를 token id의 list로 바꿔야 한다.

"Can you explain BFS?"
  -> [token_id_1, token_id_2, token_id_3, ...]

그리고 token id는 embedding layer를 거쳐 vector가 된다.

token id
  -> token embedding
  -> position embedding과 더함
  -> transformer block 입력

Multi-Head Attention

MultiHeadAttention은 GPT에서 가장 중요한 부분이다. 입력 x를 query, key, value로 projection한 뒤 attention score를 계산한다.

keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)

그리고 head 수에 맞게 tensor shape을 나눈다.

(batch, tokens, emb_dim)
  -> (batch, heads, tokens, head_dim)

attention score는 query와 key의 내적으로 계산한다.

attn_scores = queries @ keys.transpose(2, 3)

GPT는 미래 token을 보면 안 된다. 예를 들어 세 번째 token을 예측할 때 네 번째 token을 보면 정답을 미리 보는 것이 된다.

그래서 causal mask를 사용한다.

self.register_buffer(
    "mask",
    torch.triu(torch.ones(context_length, context_length), diagonal=1)
)

mask가 적용된 위치는 -inf로 바뀌고, softmax 이후 attention weight가 0에 가까워진다.

attn_scores.masked_fill_(mask_bool, -torch.inf)
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

결과적으로 각 token은 자기 자신과 과거 token만 참고한다. 이것이 GPT가 text를 왼쪽에서 오른쪽으로 생성할 수 있는 이유이다.

Layer Normalization

LayerNorm은 embedding dimension 방향으로 평균과 분산을 계산해 값을 정규화한다.

mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)

그다음 학습 가능한 scale과 shift를 적용한다.

return self.scale * norm_x + self.shift

Transformer에서는 layer normalization이 학습 안정성에 중요하다. MinGPT에서는 attention과 feed forward 앞에 normalization을 적용하는 pre-norm 형태를 사용한다.

x -> LayerNorm -> Attention -> Residual
x -> LayerNorm -> FeedForward -> Residual

Feed Forward Network

FeedForward는 각 token position에 독립적으로 적용되는 MLP이다. 구조는 다음과 같다.

nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"])
GELU()
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"])

embedding dimension을 4배로 키웠다가 다시 원래 크기로 줄인다. Transformer block 안에서 attention이 token 간 관계를 섞는 역할을 한다면, feed forward는 각 token representation을 비선형 변환하는 역할을 한다.

activation function은 GELU를 직접 구현해서 사용한다.

0.5 * x * (1 + torch.tanh(
    torch.sqrt(torch.tensor(2.0 / torch.pi)) *
    (x + 0.044715 * torch.pow(x, 3))
))

Transformer Block

TransformerBlock은 attention, feed forward, layer normalization, residual connection을 묶은 단위이다.

구조는 다음과 같다.

input
  -> LayerNorm
  -> MultiHeadAttention
  -> Dropout
  -> Residual Add
  -> LayerNorm
  -> FeedForward
  -> Dropout
  -> Residual Add
  -> output

code에서는 shortcut 변수를 사용해 residual connection을 구현한다.

shortcut = x
x = self.norm1(x)
x = self.att(x)
x = self.drop_shortcut(x)
x = x + shortcut

residual connection은 깊은 network에서 gradient 흐름을 유지하는 데 중요하다. GPT는 이 block을 여러 층 쌓아서 token representation을 점점 풍부하게 만든다.

GPTModel

GPTModel은 전체 GPT 구조를 정의한다.

self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.trf_blocks = nn.Sequential(
    *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

forward 흐름은 다음과 같다.

input token ids
  -> token embedding
  -> positional embedding
  -> dropout
  -> transformer blocks
  -> final layer norm
  -> vocabulary logits

출력은 각 위치마다 vocabulary 전체에 대한 logit이다.

(batch, sequence_length, vocab_size)

이 logit을 target token과 비교해서 cross entropy loss를 계산한다.

Model 설정

main.py에는 GPT-2 124M class와 비슷한 크기의 설정이 들어 있다.

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

각 값의 의미는 다음과 같다.

주의할 점은 이 설정이 GPT-2 124M과 유사한 구조를 만든다는 뜻이지, GPT-2 pretrained weight를 불러온다는 뜻은 아니다. 이 project의 model은 random initialization에서 시작해 직접 학습된다.

학습 데이터 형식

main.py의 예제 데이터는 instruction과 output으로 구성되어 있다.

example_data = {
    "instruction": [
        "Can you explain the BFS algorithm?",
    ],
    "output": [
        "... BFS code ..."
    ]
}

이 데이터는 generate_prompt()에서 instruction-response 형식의 prompt로 바뀐다.

### Instruction: Can you explain the BFS algorithm?

### Response: ...
<eos>

<eos>는 응답이 끝나는 위치를 표시하기 위해 사용된다.

이 방식은 instruction tuning data 형식을 단순화한 것이다. 다만 예제 데이터가 매우 작기 때문에 일반적인 LLM처럼 넓은 지식을 학습하는 목적이라기보다, pipeline이 어떻게 연결되는지 확인하는 데 가깝다.

학습 과정

학습 함수는 일반적인 PyTorch training loop 형태이다.

for batch in dataloader:
    optimizer.zero_grad()
    input_ids, target_ids = batch
    outputs = model(input_ids)
    loss = criterion(outputs.view(-1, outputs.size(-1)), target_ids.view(-1))
    loss.backward()
    optimizer.step()

loss는 CrossEntropyLoss를 사용한다. model 출력은 vocabulary 전체에 대한 logit이고, target은 실제 다음 token id이다.

prediction: next token distribution
target: real next token id

optimizer는 AdamW를 사용한다.

optimizer = optim.AdamW(model.parameters(), lr=5e-5)

epoch는 150으로 설정되어 있다.

num_epochs = 150

학습 과정에서 loss가 출력되므로 model이 예제 prompt를 얼마나 잘 맞추는지 확인할 수 있다.

Text Generation

text generation은 generate_text()에서 수행된다.

핵심 흐름은 다음과 같다.

prompt를 token으로 encode
현재 sequence를 model에 입력
마지막 위치의 logits만 선택
argmax로 다음 token 선택
선택한 token을 sequence 뒤에 붙임
반복

code로 보면 다음 부분이 중요하다.

idx_cond = idx[:, -context_size:]
logits = model(idx_cond)
logits = logits[:, -1, :]
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
idx = torch.cat((idx, idx_next), dim=1)

idx[:, -context_size:]는 context length를 넘지 않도록 최근 token만 잘라서 사용한다. GPT는 context window 안의 token만 참고할 수 있기 때문이다.

현재 구현은 sampling이 아니라 greedy decoding을 사용한다. 즉 매 step에서 확률이 가장 높은 token 하나를 선택한다.

greedy decoding: 가장 높은 logit의 token 선택
sampling: 확률 분포에서 token을 샘플링

greedy decoding은 구현이 단순하고 결과가 deterministic하다. 반면 다양한 문장을 생성하려면 temperature, top-k, top-p 같은 sampling 기법을 추가할 수 있다.

실행 흐름

main.py를 실행하면 먼저 model을 만든다.

model = GPTModel(GPT_CONFIG_124M)

GPU가 가능하면 CUDA를 사용한다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

그다음 예제 데이터로 학습한다.

train(model, device, example_data)

학습 후에는 interactive loop에서 사용자 입력을 받는다.

while True:
    user_input = input("입력을 하고 싶으신 내용을 입력해주세요 (종료하려면 'x' 입력): ")

사용자 입력은 instruction prompt 형식으로 바뀐다.

formatted_prompt = f"### Instruction: {PROMPT}\n\n### Response:"

그리고 model이 이어질 token을 생성한다.

output_text = get_text(model, device, formatted_prompt)

이 project에서 볼 수 있는 핵심 개념

MinGPT는 큰 framework 없이 GPT의 핵심을 추적하기 좋다.

특히 다음 개념을 code로 확인할 수 있다.

전체 구조가 작기 때문에 논문이나 큰 codebase를 보기 전에 GPT의 뼈대를 이해하는 용도로 적합하다.

한계

이 project는 교육용 구현에 가깝다. 그래서 실제 LLM application과 비교하면 몇 가지 한계가 있다.

pretrained weight를 사용하지 않는다.
학습 데이터가 매우 작다.
checkpoint 저장과 불러오기가 없다.
validation loop가 없다.
generation이 greedy decoding 중심이다.
대규모 학습 최적화가 없다.

특히 GPT-2 124M과 비슷한 model 크기를 random initialization부터 학습하려면 매우 많은 data와 compute가 필요하다. 작은 예제 데이터만으로는 일반적인 질의응답 능력을 얻기 어렵다.

따라서 이 project는 "실용적인 chatbot"보다는 "GPT 내부 구조를 직접 구현하고 관찰하는 project"로 보는 것이 적절하다.

개선 아이디어

project를 더 발전시킨다면 다음 기능을 추가할 수 있다.

training checkpoint 저장
validation loss 측정
temperature sampling
top-k sampling
top-p sampling
gradient clipping
learning rate scheduler
dataset file loader
pretrained GPT-2 weight loading

예를 들어 generation에 temperature를 추가하면 logit 분포의 날카로움을 조절할 수 있다.

temperature 낮음 -> 보수적이고 반복적인 출력
temperature 높음 -> 다양하지만 불안정한 출력

top-k나 top-p sampling을 넣으면 greedy decoding보다 다양한 문장을 생성할 수 있다.

정리

MinGPT는 GPT 구조를 PyTorch로 직접 구현한 project이다.

핵심은 다음 흐름이다.

GPT-2 tokenizer로 text를 token화
input/target sequence 생성
causal multi-head attention 적용
transformer block 반복
vocabulary logits 출력
cross entropy로 다음 token 학습
학습된 model로 autoregressive generation 수행

이 project의 장점은 GPT를 구성하는 요소가 작은 codebase 안에 명확하게 드러난다는 점이다. MultiHeadAttention부터 GPTModel, train_model, generate_text까지 한 흐름으로 읽으면 LLM의 기본 구조를 이해하는 데 도움이 된다.

다만 pretrained model이 아니고 학습 데이터도 작기 때문에, 결과 품질보다 구조 학습에 초점을 맞춰 보는 것이 좋다.