A complete "build a mini-GPT from scratch" project — implementing a small decoder-only Transformer and training it as a character-level language model, the single most direct way to genuinely understand how modern LLMs actually work under the hood.
Problem Statement
Implement a small, decoder-only Transformer (a simplified GPT-style architecture) entirely from the components covered in the Transformers category, and train it on a text corpus for next-token prediction.
Dataset
Any moderately sized plain-text corpus — the same kind used in the Text Generation project works well here too, letting you directly compare the LSTM-based and Transformer-based approaches on identical data.
Architecture & Approach
A stack of decoder blocks, each containing masked multi-head self-attention and a feed-forward network, with residual connections and layer normalization throughout — directly assembling the pieces implemented individually in the Practice: Transformers exercises into one complete, trainable model.
Step-by-Step Build
import torch
import torch.nn as nn
import math
class DecoderBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
self.norm1 = nn.LayerNorm(d_model)
self.ff = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x, causal_mask):
attn_out, _ = self.attn(x, x, x, attn_mask=causal_mask)
x = self.norm1(x + attn_out) # residual connection + layer norm
ff_out = self.ff(x)
x = self.norm2(x + ff_out) # residual connection + layer norm
return x
class MiniGPT(nn.Module):
def __init__(self, vocab_size, d_model=128, num_heads=4, num_layers=4, d_ff=512, max_len=256):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.position_embedding = nn.Embedding(max_len, d_model)
self.blocks = nn.ModuleList([DecoderBlock(d_model, num_heads, d_ff) for _ in range(num_layers)])
self.final_norm = nn.LayerNorm(d_model)
self.output_head = nn.Linear(d_model, vocab_size)
def forward(self, x):
batch, seq_len = x.shape
positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
h = self.token_embedding(x) + self.position_embedding(positions)
causal_mask = torch.triu(torch.full((seq_len, seq_len), float('-inf')), diagonal=1)
for block in self.blocks:
h = block(h, causal_mask)
h = self.final_norm(h)
return self.output_head(h) # (batch, seq_len, vocab_size)
model = MiniGPT(vocab_size=len(chars)) # reusing the character vocab from the Text Generation project
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
loss_fn = nn.CrossEntropyLoss()
for step in range(3000):
x_batch, y_batch = get_batch(data) # same next-token batching as the LSTM project
optimizer.zero_grad()
logits = model(x_batch)
loss = loss_fn(logits.view(-1, len(chars)), y_batch.view(-1))
loss.backward()
optimizer.step()
if step % 500 == 0:
print(f"Step {step}: loss={loss.item():.4f}")
# Generation -- autoregressive, same idea as the LSTM project, but with a Transformer
@torch.no_grad()
def generate(model, start_text, length=200):
model.eval()
tokens = [char_to_idx[c] for c in start_text]
for _ in range(length):
x = torch.tensor([tokens[-256:]]) # respect the model's max context length
logits = model(x)
probs = torch.softmax(logits[0, -1], dim=0)
next_token = torch.multinomial(probs, 1).item()
tokens.append(next_token)
return ''.join(idx_to_char[t] for t in tokens)Expected Results
With a small model (a few million parameters) and a modest corpus, expect generation quality broadly comparable to, and often somewhat better than, the LSTM from the Text Generation project — the real value of this project isn't necessarily better final text quality at this small scale, but a genuine, from-scratch understanding of exactly how the architecture powering every modern LLM actually works.
Key Learnings & Extensions
- Compare training speed against the LSTM project on the same hardware and data — the Transformer's parallelizable attention computation should train noticeably faster per step, directly demonstrating the parallelization advantage covered in Transformer Motivation.
- Extension: Scale up — more layers, larger
d_model, more training steps — and observe how generation quality improves, a small-scale, hands-on taste of the scaling behavior that underlies real LLM development. - Extension: Switch to word-level or subword (BPE) tokenization instead of character-level, and compare training dynamics and generation quality.
