“Know how to solve every problem that has been solved.” “What I cannot create, I do not understand.” — Richard Feynman

Project: A Transformer From Scratch

Projects

A decoder-only GPT with every internal hand-written, verified against PyTorch and trained to learn an algorithm. Stack: Python · PyTorch (tensors + autograd) · from-scratch architecture.

What this is

A transformer is not magic. It is a stack of two ideas — attention and a per-token MLP — wrapped in residual connections — each block's output is added to its input, x + f(x), so a layer learns a correction rather than a replacement and gradients keep an unobstructed path — and normalization. This project builds a decoder-only GPT with every one of those pieces written by hand on top of raw tensors (no torch.nn.Transformer, no attention library), then trains it to learn an actual algorithm: sorting. Two checks verify it — the attention mechanism is matched against PyTorch's fused kernel, and the trained model is judged on sequences held out by a hash split, so a high score means it learned to sort, not to memorize.

The architecture

Each block does two things, both inside the residual stream. First, tokens exchange information through causal self-attention: every position emits a query, a key, and a value, and attends to the earlier positions whose keys best match its query. Then each token is transformed independently by a small MLP. Pre-norm LayerNorm and the residual connections around both keep gradients flowing through the stack. Attention itself is one line of linear algebra:

Attention (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}} + M) V,

where $M$ is the causal mask — $0$ on and below the diagonal, $- \infty$ above it — so a token can never attend to its own future. The whole engine of a GPT is that expression, run in parallel across several heads and stacked in depth.

def sdpa(q, k, v, causal=True):
    """Scaled dot-product attention, by hand. q,k,v: (B, H, T, d)."""
    T, d = q.shape[-2], q.shape[-1]
    att = (q @ k.transpose(-2, -1)) / math.sqrt(d)          # (B, H, T, T)
    if causal:
        mask = torch.triu(torch.ones(T, T, device=q.device), 1).bool()
        att = att.masked_fill(mask, float("-inf"))          # no peeking ahead
    att = att.softmax(-1)
    return att @ v, att


class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.n_head, self.head_dim = n_head, n_embd // n_head
        self.qkv  = Linear(n_embd, 3 * n_embd)              # one projection, split 3 ways
        self.proj = Linear(n_embd, n_embd)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        split = lambda t: t.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        y, att = sdpa(split(q), split(k), split(v), causal=True)
        y = y.transpose(1, 2).reshape(B, T, C)             # recombine heads
        return self.proj(y)


class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # tokens mix, in the residual stream
        x = x + self.mlp(self.ln2(x))      # then each is transformed alone
        return x

Verification 1 — the attention mechanism

The one place a subtle bug hides forever is the attention kernel: a wrong scale factor or a mis-applied mask still trains, just to a worse model, with nothing to flag it. So the hand-written version is pinned against PyTorch's fused scaled_dot_product_attention on random inputs with a causal mask.

def mechanism_check():
    B, H, T, d = 2, 4, 7, 16
    q, k, v = (torch.randn(B, H, T, d, dtype=torch.float64) for _ in range(3))
    mine, _ = sdpa(q, k, v, causal=True)                        # hand-written
    ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # PyTorch fused
    return (mine - ref).abs().max().item()

They agree to $4.4 \times 1 0^{- 16}$ — floating-point rounding and nothing else. The mechanism is exactly right, so anything the model fails to learn later is the model's fault, not the kernel's.

Verification 2 — training it to sort

To test the whole model end to end, I train it to sort. Each training sequence is a run of digits followed by their sorted order; the model sees the digits and must generate the sorted tail one token at a time, scored only on that tail. The detail that makes the number mean something: the input space is split by a deterministic hash, one quarter held out. Training never touches a held-out sequence, so accuracy there measures generalization — sorting arrangements the model has genuinely never seen — rather than recall.

device: mps
[1] mechanism  max |my_attention - torch SDPA| : 4.44e-16
[2] training GPT to sort 6 digits (vocab 10, 1/4 held out by hash):
  step     0   loss 2.5828   test-acc 0.000
  step   200   loss 0.0036   test-acc 1.000
  step  1000   loss 0.1409   test-acc 0.702
  step  2000   loss 0.0009   test-acc 1.000
  step  4000   loss 0.0009   test-acc 0.998
    final held-out exact-match sort accuracy   : 0.9995

A 3-layer, 4-head, 64-dimensional model reaches 99.95% exact-match accuracy on the held-out split — essentially every unseen sequence sorted correctly — and it gets there almost immediately, within a couple hundred steps. (The occasional dips are the aggressive learning rate overshooting and recovering, not the model forgetting.)

Two panels. Left: cross-entropy loss falling from 2.6 to near zero within about 200 steps while held-out exact-match sort accuracy rises to 1.0, with two brief dips to 0.7 and 0.85 that recover. Right: an attention heatmap for one held-out example, strictly lower-triangular because of the causal mask, with the brightest weights on the diagonal and spread across the earlier input positions.

The right panel is the last block's attention on one held-out example. The strict lower-triangular shape is the causal mask doing its job — every token attends only to itself and the past — and as the model generates the sorted tail, each output position draws from the input block where the digits actually live. The mechanism verified in isolation above is here doing real work.

Extensions

Sorting is the "hello world." The same hand-built model reaches much further.

Addition. Train it to add two numbers. Carries force the model to route information across positions in a way sorting doesn't — a genuinely harder and more revealing task.
Grokking. On modular addition with weight decay, watch held-out accuracy stay at chance long after the training loss hits zero, then suddenly jump to 100%. One of the strangest phenomena in deep learning, reproducible on this exact model.
Language. Point it at a character-level corpus (tiny-shakespeare) and sample. The same code that sorts digits writes text.
A KV-cache. Generation recomputes attention over the whole prefix every step. Cache the keys and values and measure the speedup — the optimization every serving stack depends on.
Rotary positions (RoPE). Swap the learned positional table for RoPE and test length generalization: train on length 6, evaluate on length 8.
Ablations. Delete the LayerNorms, or the residual connections, and watch training collapse. The fastest way to understand why each piece is there is to remove it.