Skip to content

Building a Transformer-Based Language Model from Scratch

This guide walks through the process of building a character-level Transformer language model, similar in architecture to GPT, trained on the complete works of Shakespeare. You'll understand the core components that power systems like ChatGPT.

Introduction

The Transformer architecture, introduced in the seminal 2017 paper "Attention Is All You Need", forms the foundation of modern large language models (LLMs) like GPT (Generative Pre-trained Transformer). While production systems like ChatGPT involve massive models trained on internet-scale data with complex fine-tuning stages, the core architecture can be implemented in a few hundred lines of code.

This guide implements a decoder-only Transformer - the same architecture used by GPT models - trained at the character level on a dataset of Shakespeare's complete works. You'll build the model piece by piece, understanding each component: token embeddings, positional encodings, multi-head self-attention, feed-forward networks, layer normalization, and residual connections.

Prerequisites

  • Proficiency in Python
  • Basic understanding of PyTorch
  • Familiarity with calculus and statistics fundamentals
  • A machine with Python 3.8+ and PyTorch installed (GPU recommended for training)

Project Setup

Install Dependencies

pip install torch numpy

Download the Dataset

The Tiny Shakespeare dataset is a concatenation of all Shakespeare's works in a single 1MB text file.

wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tiny-shakespeare/input.txt

Dataset Overview

# Read the dataset
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

print(f"Length of dataset in characters: {len(text)}")
# Output: Length of dataset in characters: ~1,000,000

Building the Tokenizer

A tokenizer converts raw text into sequences of integers that the model can process. For simplicity, we use a character-level tokenizer.

import torch

# Get all unique characters in the dataset
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size}")
# Output: Vocabulary size: 65

# Create mapping from characters to integers
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Encoder: string -> list of integers
encode = lambda s: [stoi[c] for c in s]

# Decoder: list of integers -> string
decode = lambda l: ''.join([itos[i] for i in l])

# Test the encoder/decoder
print(encode("hi there"))
# Output: [46, 47, 56, 1, 58, 46, 43, 56, 43]

print(decode(encode("hi there")))
# Output: hi there

[!NOTE] Production systems like GPT use subword tokenizers (like Byte-Pair Encoding or SentencePiece) with vocabularies of 50,000+ tokens. Character-level tokenization is simpler but produces longer sequences.

Data Preparation

Train/Validation Split

# Split into train and validation (90/10)
n = int(0.9 * len(text))
train_data = text[:n]
val_data = text[n:]

# Tokenize the entire dataset
train_data = torch.tensor(encode(train_data), dtype=torch.long)
val_data = torch.tensor(encode(val_data), dtype=torch.long)

Batching

We never feed entire texts to the Transformer at once. Instead, we sample small chunks (blocks) and process them in batches.

batch_size = 4      # How many independent sequences to process in parallel
block_size = 8      # Maximum context length for predictions

def get_batch(split):
    """Generate a batch of inputs and targets"""
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x, y

# Test the batching
xb, yb = get_batch('train')
print(f"Input shape: {xb.shape}")     # torch.Size([4, 8])
print(f"Target shape: {yb.shape}")    # torch.Size([4, 8])

# The batch contains 32 independent examples (4 batches × 8 positions)

[!NOTE] For each position in the input, the target is the next character. This means a single chunk of 8 characters actually contains 8 training examples (context lengths from 1 to 8).

Building the Bigram Language Model

Start with the simplest possible language model - a Bigram model that predicts the next character based only on the current character.

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

class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # Each token directly reads off the logits for the next token
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # idx and targets are (B, T) tensors
        logits = self.token_embedding_table(idx)  # (B, T, C) where C=vocab_size

        if targets is None:
            loss = None
        else:
            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

    def generate(self, idx, max_new_tokens):
        """Generate new tokens autoregressively"""
        for _ in range(max_new_tokens):
            # Get predictions
            logits, _ = self(idx)
            # Focus on the last time step
            logits = logits[:, -1, :]  # (B, C)
            # Apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1)
            # Sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
            # Append to the sequence
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
        return idx

# Test generation with random model
m = BigramLanguageModel(vocab_size)
idx = torch.zeros((1, 1), dtype=torch.long)
print(decode(m.generate(idx, max_new_tokens=100)[0].tolist()))

Training the Bigram Model

# Create optimizer
optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)

# Training loop
for steps in range(10000):
    # Sample batch
    xb, yb = get_batch('train')

    # Evaluate loss
    logits, loss = m(xb, yb)

    # Backward pass
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

    if steps % 1000 == 0:
        print(f"Step {steps}: loss = {loss.item():.4f}")

# Output (approximate):
# Step 0: loss = 4.8732
# Step 1000: loss = 3.5176
# Step 9000: loss = 2.5018

The Mathematical Trick Behind Efficient Attention

Before implementing self-attention, understand the key mathematical operation that makes it efficient: weighted aggregation via matrix multiplication.

# Toy example: B=4, T=8, C=2
B, T, C = 4, 8, 2
x = torch.randn(B, T, C)

# Version 1: Inefficient loop-based averaging
xbow = torch.zeros((B, T, C))
for b in range(B):
    for t in range(T):
        xprev = x[b, :t+1]  # (t, C)
        xbow[b, t] = torch.mean(xprev, 0)

# Version 2: Efficient matrix multiplication approach
wei = torch.tril(torch.ones(T, T))
wei = wei / wei.sum(1, keepdim=True)  # Normalize rows
xbow2 = wei @ x  # (B, T, T) @ (B, T, C) -> (B, T, C)

# Version 3: Using softmax (more flexible)
tril = torch.tril(torch.ones(T, T))
wei = torch.zeros((T, T))
wei = wei.masked_fill(tril == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
xbow3 = wei @ x

# All three versions produce identical results
print(torch.allclose(xbow, xbow2))  # True
print(torch.allclose(xbow, xbow3))  # True

Implementing Single-Head Self-Attention

Self-attention allows tokens to communicate with each other in a data-dependent manner.

class Head(nn.Module):
    """One head of self-attention"""

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape

        # Compute query, key, value for all tokens
        k = self.key(x)   # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)

        # Compute attention scores (affinities)
        wei = q @ k.transpose(-2, -1)  # (B, T, T)
        wei = wei * head_size**-0.5     # Scale by sqrt(head_size) for stable gradients

        # Mask out future tokens (decoder block)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)    # (B, T, T)

        # Weighted aggregation of values
        v = self.value(x)               # (B, T, head_size)
        out = wei @ v                   # (B, T, head_size)
        return out

[!WARNING] The scaling factor head_size**-0.5 is crucial. Without it, the softmax inputs become too large, causing attention to collapse to near one-hot vectors, especially at initialization.

Multi-Head Self-Attention

Multiple attention heads run in parallel, allowing the model to attend to different types of relationships simultaneously.

class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel"""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)  # Project back to residual pathway

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.proj(out)
        return out

Feed-Forward Network

After communication via attention, each token independently processes the information it gathered.

class FeedForward(nn.Module):
    """A simple linear layer followed by a non-linearity"""

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),  # Inner layer is 4x wider
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),  # Project back
        )

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

Transformer Block

Combine communication (multi-head attention) with computation (feed-forward), wrapped with residual connections and layer normalization.

class Block(nn.Module):
    """Transformer block: communication + computation with skip connections"""

    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):
        # Residual connection with pre-norm formulation
        x = x + self.sa(self.ln1(x))   # Communication
        x = x + self.ffwd(self.ln2(x)) # Computation
        return x

Complete GPT Model

Assemble all components into the final Transformer model.

class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Token and position embeddings
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)

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

        # Final layer norm and language model head
        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

        # Token and position embeddings
        tok_emb = self.token_embedding_table(idx)          # (B, T, C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, C)
        x = tok_emb + pos_emb                               # (B, T, C)

        # Apply transformer blocks
        x = self.blocks(x)                                  # (B, T, C)
        x = self.ln_f(x)                                    # (B, T, C)
        logits = self.lm_head(x)                            # (B, T, vocab_size)

        # Compute loss if targets provided
        if targets is None:
            loss = None
        else:
            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

    def generate(self, idx, max_new_tokens):
        """Generate new tokens"""
        for _ in range(max_new_tokens):
            # Crop idx to the last block_size tokens
            idx_cond = idx[:, -block_size:]

            # Get predictions
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]  # (B, C)
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Training Configuration

Set hyperparameters for training. These are much smaller than production models but demonstrate the same architecture.

# Hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
eval_iters = 200
n_embd = 384
n_head = 6
n_layer = 6
dropout = 0.2

# Device
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Initialize model
model = GPTLanguageModel()
model = model.to(device)

# Print number of parameters
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
# Output: Model parameters: ~10 million

Training Loop

# Loss estimation function
@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, Y = X.to(device), Y.to(device)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# Create optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# Training loop
for iter in range(max_iters):
    # Evaluate loss periodically
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"Step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # Sample batch
    xb, yb = get_batch('train')
    xb, yb = xb.to(device), yb.to(device)

    # Forward and backward pass
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

Generating Shakespeare-Like Text

After training, generate new text by sampling from the model.

# Generate from the model
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=2000)
print(decode(generated[0].tolist()))

The output will look something like:

verily my Lord the sites have left the again
the king coming with my curses with precious pale
and then tranos say something else

Understanding the Architecture

Decoder-Only vs. Encoder-Decoder

The model implemented here is a decoder-only Transformer, which uses a triangular mask to prevent tokens from attending to future positions. This is the architecture used by GPT models.

The original "Attention Is All You Need" paper used an encoder-decoder architecture for machine translation: - Encoder: Processes input sequence (e.g., French sentence) with full bidirectional attention - Decoder: Generates output sequence (e.g., English translation) with masked self-attention and cross-attention to encoder outputs

Key Components Summary

  1. Token Embeddings: Convert input integers to dense vectors
  2. Positional Embeddings: Add position information (since attention is permutation-invariant)
  3. Self-Attention: Communication mechanism where tokens aggregate information from other tokens
  4. Multi-Head Attention: Multiple attention channels in parallel
  5. Feed-Forward Networks: Computation layer where each token processes gathered information
  6. Layer Normalization: Stabilizes training by normalizing activations
  7. Residual Connections: Enable training of deep networks by providing gradient shortcuts

Scaling to Production Models

The architecture implemented here is fundamentally identical to what powers GPT-3 and ChatGPT. The differences are in scale:

Aspect Our Model GPT-3 (175B)
Parameters 10 million 175 billion
Training data 1 million chars 300+ billion tokens
Layers 6 96
Embedding dim 384 12,288
Heads 6 96
Batch size 64 3.2 million tokens
GPUs 1 (optional) Thousands

The Two Stages of Training ChatGPT

  1. Pre-training: Train on large internet corpus to predict next token (document completion)
  2. Fine-tuning (RLHF): Align the model to be helpful and safe:
  3. Supervised fine-tuning on human demonstrations
  4. Train a reward model on human preferences
  5. Optimize using Proximal Policy Optimization (PPO)

Repository: nanoGPT

The complete codebase is available as nanoGPT, which includes: - model.py: The GPT model definition (very similar to what we built) - train.py: Training loop with checkpointing, learning rate scheduling, and distributed training support

Next Steps

  • Experiment with different datasets (e.g., Python code, Wikipedia articles)
  • Implement the encoder-decoder architecture for translation tasks
  • Explore fine-tuning techniques (supervised fine-tuning, RLHF)
  • Scale up the model and use larger batch sizes with multiple GPUs