Let's build a tiny GPT, from scratch, that writes text

1 Goal & the build plan

Our goal for this project is concrete: build a decoder-only Transformer — a GPT-style model (the same family that powers ChatGPT) — small enough to train on a laptop, and have it generate text. "Decoder-only" just means it does one job: read the text so far and predict what comes next. It only ever looks backwards, never forwards (we'll enforce that with causal masking in Topic 3).

What "from scratch" means here

We will use PyTorch's tensors and autograd (the automatic differentiation engine from Session 15) and basic layers like nn.Linear and nn.Embedding, but we will not call any pre-built attention block. We write the attention math ourselves so you can see exactly what happens. This mirrors Andrej Karpathy's famous "Let's build GPT" walkthrough (see References), simplified for one sitting.

Tokens: characters, not words

Real LLMs split text into sub-word tokens (recall the tokenisation idea — it shows up across modern NLP). To keep our project tiny and self-contained, we use the simplest possible tokeniser: each character is one token. So our vocabulary is just the set of unique characters in our training file (maybe ~65 of them). This is enough to learn real structure — words, spacing, even punctuation rhythm — without any external library.

The components we'll build

# minigpt.py — a tiny decoder-only Transformer in PyTorch
import torch
import torch.nn as nn
from torch.nn import functional as F

# 1. config + data + tokeniser    (Topic 1 & 3)
# 2. Head, MultiHeadAttention, FeedForward, Block   (Topic 2)
# 3. MiniGPT model with embeddings + stacked blocks (Topic 2 & 3)
# 4. training loop                 (Topic 3)
# 5. generate() autoregressive loop (Topic 4)

2 Building the blocks

Every block here is a small PyTorch nn.Module. We'll define a tiny config first so the dimensions line up, then build each piece bottom-up.

Setup & hyperparameters

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

torch.manual_seed(1337)          # reproducible runs

# --- hyperparameters (small on purpose) ---
block_size = 128      # context length: how many chars the model sees at once
n_embd     = 128      # size of each token's vector (the "width" of the model)
n_head     = 4        # number of attention heads (128 / 4 = 32 dims per head)
n_layer    = 4        # how many Transformer blocks we stack
dropout    = 0.1      # regularisation: randomly zero 10% of activations while training
device     = 'cuda' if torch.cuda.is_available() else 'cpu'

Part A — Token & positional embeddings

An embedding turns a token's integer ID into a learnable vector of n_embd numbers that captures meaning. But a Transformer reads all positions in parallel, so on its own it has no sense of order ("dog bites man" would look identical to "man bites dog"). We fix that with a positional embedding: a second lookup table, one learnable vector per position 0..block_size-1. We add the two together so each token vector carries both "what I am" and "where I am."

# inside the model's __init__ (full model assembled in Topic 3):
self.token_embedding    = nn.Embedding(vocab_size, n_embd)   # what  (one row per char)
self.position_embedding = nn.Embedding(block_size, n_embd)   # where (one row per slot)

# inside forward(), idx is a (B, T) batch of token IDs:
B, T = idx.shape
tok = self.token_embedding(idx)                              # (B, T, n_embd)
pos = self.position_embedding(torch.arange(T, device=idx.device))  # (T, n_embd)
x   = tok + pos                                              # broadcast-add → (B, T, n_embd)

Part B — Scaled dot-product self-attention (one head)

This is the heart of the Transformer. For every token we build three vectors (recall the library analogy): a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what I'll hand over"). We compare each token's query against every token's key to get attention scores, scale them, mask out the future, softmax them into weights, and use those weights to average the values.

The "scaled dot-product" formula is just:

attention(Q, K, V) = softmax( (Q · Kᵀ) / √d_k ) · V

The √d_k (square root of the head size) keeps the scores from getting huge, which would make softmax too "peaky" and kill the gradients. Here's one head as code:

class Head(nn.Module):
    """One head of causal 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)
        # lower-triangular matrix of 1s; 'buffer' = saved but not a learned parameter
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        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) * k.shape[-1]**-0.5   # (B, T, T), scaled by 1/√d_k
        # causal mask: a token may NOT attend to future tokens (Topic 3)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)                       # (B, T, T), rows sum to 1
        wei = self.dropout(wei)
        v = self.value(x)                                  # (B, T, head_size)
        out = wei @ v                                      # (B, T, head_size)
        return out

Part C — Multi-head attention

One head learns one kind of relationship. Multi-head attention runs several heads in parallel — each with its own Q/K/V — so one head can track, say, subject–verb links while another tracks nearby punctuation. We concatenate their outputs and pass them through a final linear "projection" to mix the heads back together.

class MultiHeadAttention(nn.Module):
    """Several self-attention heads running in parallel."""
    def __init__(self, n_head, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(n_head)])
        self.proj  = nn.Linear(head_size * n_head, n_embd)   # recombine the heads
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)  # (B, T, head_size*n_head)
        out = self.dropout(self.proj(out))                  # (B, T, n_embd)
        return out

Part D — The MLP / feed-forward

Attention lets tokens communicate. The feed-forward network (an MLP) lets each token think on its own about what it just gathered. It's a simple two-layer network applied to every position independently. By convention the hidden layer is 4× wider than n_embd, with a non-linearity (ReLU here) in between.

class FeedForward(nn.Module):
    """Position-wise MLP: think after you've gathered context."""
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),   # expand
            nn.ReLU(),                       # non-linearity
            nn.Linear(4 * n_embd, n_embd),   # project back down
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)                  # (B, T, n_embd) → (B, T, n_embd)

Part E — Residual connections + LayerNorm, packaged as a Block

Two more ingredients make deep Transformers trainable, both from Session 15:

• Residual connection — instead of x = layer(x), we do x = x + layer(x). The layer only has to learn a small adjustment; the original signal flows straight through. This gives gradients a "highway" and lets us stack many layers without them collapsing.
• LayerNorm — normalises each token's vector (mean 0, variance 1) so numbers stay in a sane range. We apply it before each sub-layer ("pre-norm"), which is the modern, more stable arrangement.

A Block is just: norm → attention (+residual), then norm → MLP (+residual).

class Block(nn.Module):
    """One Transformer block: communicate (attention), then think (MLP)."""
    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))     # residual around attention (pre-norm)
        x = x + self.ffwd(self.ln2(x))   # residual around MLP
        return x

3 Assembling & training

Step 1 — Load data & build the character tokeniser

Grab any plain-text file (a book, your favourite lyrics, the tiny Shakespeare file Karpathy uses). We map each unique character to an integer and back.

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

chars = sorted(list(set(text)))   # the vocabulary = unique chars
vocab_size = len(chars)

stoi = { ch: i for i, ch in enumerate(chars) }   # string → int
itos = { i: ch for i, ch in enumerate(chars) }   # int → string
encode = lambda s: [stoi[c] for c in s]            # "hi" → [40, 41]
decode = lambda l: ''.join([itos[i] for i in l])    # [40, 41] → "hi"

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))                # 90% train, 10% validation
train_data, val_data = data[:n], data[n:]

Step 2 — Batches of (input, target) pairs

For training we grab random chunks of length block_size. The clever trick: the target is the same chunk shifted one character to the right. So at every position the model learns "given everything up to here, what's the next char?" — that's next-token prediction, and it gives us block_size training signals per chunk.

batch_size = 32   # how many chunks per step

def get_batch(split):
    d = train_data if split == 'train' else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))  # random start points
    x = torch.stack([d[i      : i+block_size  ] for i in ix])  # (B, T) inputs
    y = torch.stack([d[i+1    : i+block_size+1] for i in ix])  # (B, T) targets (shifted by 1)
    return x.to(device), y.to(device)

Step 3 — Causal masking (why the model can't cheat)

Because each position is asked to predict its next character, the model must not be allowed to see that character (or any later one) when making the prediction — otherwise it would just copy the answer and learn nothing. Causal masking enforces this. Recall the tril (lower-triangular) buffer in our Head: we set every score for a future position to -inf before softmax, so after softmax those weights become exactly 0.

Step 4 — The full model

Now we assemble everything: the two embedding tables, a stack of n_layer Blocks, a final LayerNorm, and a "language-model head" — a linear layer mapping each token's vector to one logit (raw score) per vocabulary character.

class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding    = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f   = nn.LayerNorm(n_embd)              # final norm
        self.lm_head = nn.Linear(n_embd, vocab_size)    # → one score per char

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok = self.token_embedding(idx)                              # (B, T, C)
        pos = self.position_embedding(torch.arange(T, device=device))# (T, C)
        x = tok + pos                                                # (B, T, C)
        x = self.blocks(x)                                           # (B, T, C)
        x = self.ln_f(x)
        logits = self.lm_head(x)                                     # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, V = logits.shape
            # cross-entropy expects (N, V) logits and (N,) targets
            loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))
        return logits, loss

Step 5 — The training loop

This is the same pattern from Session 15: forward pass → compute loss → backward pass (autograd) → optimiser step. We use AdamW, the go-to optimiser for Transformers.

model = MiniGPT().to(device)
print(sum(p.numel() for p in model.parameters())/1e6, 'M parameters')

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
max_iters, eval_interval = 5000, 500

@torch.no_grad()
def estimate_loss():           # average loss over a few batches, no gradients
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(50)
        for k in range(50):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

for step in range(max_iters):
    if step % eval_interval == 0:
        l = estimate_loss()
        print(f"step {step}: train {l['train']:.4f}, val {l['val']:.4f}")

    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)        # forward
    optimizer.zero_grad(set_to_none=True)
    loss.backward()                     # backward (autograd computes all gradients)
    optimizer.step()                    # nudge every parameter

0.21 M parameters
step 0:    train 4.36, val 4.36     # ≈ ln(65), basically random
step 500:  train 2.42, val 2.45
step 1000: train 2.05, val 2.08
step 2500: train 1.70, val 1.74
step 5000: train 1.52, val 1.57     # now it has learned real structure

4 Generating & understanding

Autoregressive sampling

Autoregressive generation means the model feeds its own output back in: predict a token, append it, predict the next using the longer context, and repeat. Crucially, we only need the logits for the last position (the prediction for "what comes next"), and we must crop the context to the last block_size tokens since our positional table only goes that far.

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]          # crop to context window
            logits, _ = self(idx_cond)               # (B, T, vocab_size)
            logits = logits[:, -1, :] / temperature  # keep last step; apply temperature
            if top_k is not None:                  # optional: keep only top-k options
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)         # (B, vocab_size)
            idx_next = torch.multinomial(probs, num_samples=1)  # sample one token
            idx = torch.cat((idx, idx_next), dim=1)   # append → loop
        return idx

Add that method inside the MiniGPT class, then generate:

context = torch.zeros((1, 1), dtype=torch.long, device=device)  # start token (index 0)
out = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=40)
print(decode(out[0].tolist()))

Inspecting attention — looking inside the box

One of the joys of building it yourself: you can peek at what the model attends to. The attention weights wei are literally a probability grid. We can expose and visualise them:

# quick hack: have Head.forward stash its weights, then read them after a forward pass
# in Head.forward, before returning:  self.last_attn = wei.detach()

model.eval()
xb, _ = get_batch('val')
_ = model(xb)
attn = model.blocks[0].sa.heads[0].last_attn   # (B, T, T) for layer 0, head 0
print(attn[0].shape)                            # (T, T) — one row per query position

# row i shows how token i distributed its attention over tokens 0..i
import matplotlib.pyplot as plt
plt.imshow(attn[0].cpu(), cmap='viridis'); plt.xlabel('key'); plt.ylabel('query'); plt.show()

What scales up to real LLMs

Our 0.2M-parameter character model and a 200-billion-parameter frontier model are the same architecture. The differences are quantitative, not structural:

Production models also add efficiency tricks (e.g. faster positional schemes like RoPE, attention kernels like FlashAttention, a key/value cache for fast generation), but the core Block you wrote is exactly what's inside them. This is the bridge to the GenAI subject: everything there — prompting, tokens-as-money, context windows, temperature/top-p, RAG, agents — sits on top of the very machine you just built. You now know what's under the hood.

★ Putting it all together

You went from a pile of components to a working baby-GPT. Here's the one-paragraph story that connects the whole project:

Quick self-check

📚 References & Further Reading

Class material

• SST Deep Learning handout (Session 19) — your course handout for this project session, covering the mini-Transformer build.

Papers, docs & deep dives

• SST Deep Learning — Neural Nets (course site) — the companion site for this course's Neural Nets track.
• Andrej Karpathy — "Let's build GPT: from scratch, in code, spelled out" — the definitive line-by-line video this project is modelled on; watch it alongside the code here.
• karpathy/minGPT (GitHub) — a clean, minimal reference implementation of a GPT in PyTorch; great to read once your own version runs.
• "The Annotated Transformer" (Harvard NLP) — the original Transformer paper re-implemented with code beside every equation.
• Vaswani et al., "Attention Is All You Need" (2017) — the paper that introduced the Transformer and scaled dot-product attention.
• PyTorch documentation — official reference for nn.Module, nn.Embedding, autograd, and the optimisers used here.