Final Project: a Small Transformer

Time to put everything together. In this chapter we’ll build a small decoder-only transformer, a GPT in miniature, and train it end to end on a toy task that it can solve perfectly: copying. Each sequence is 8 random tokens followed by the same 8 again. The first half can’t be predicted, but once the model has seen it, every token of the second half can be, if the model learns to look back 8 positions. That’s exactly what attention is for.

The whole thing trains on a laptop CPU in a few seconds.

The plan

Three files in one folder:

• copier.ss: the sizes, the model, the data record, and the forward pass and loss;
• train.ss: training with the standard library’s fit;
• run_copier.py: the dataset, in Python, and the program that runs it all.

Sizes

devices cpu

## Vocabulary size: the tokens are 0 to 15.
dim V   = 16
## Length of each sequence: 8 random tokens, then the same 8 again.
dim T   = 16
## Model width.
dim D   = 64
## Number of attention heads.
dim H   = 4
## Width of each head.
dim Dh  = D / H
## Width of the MLP's hidden layer.
dim Dff = 4 * D
## Number of decoder blocks.
dim L   = 2

Everything that follows is checked against these. Change H to 5 and Dh is an error at once, since 5 doesn’t divide 64.

The model

A GPT’s parameters are a token embedding, a position embedding, and for each block, two layer norms, the attention projections and an MLP. We stack the blocks’ weights along a first axis of size L, so the model is one record:

## A small decoder-only transformer.
model Copier:
    param wte:   [V D]     = normal(std=0.02)
    param wpe:   [T D]     = normal(std=0.02)
    param ln1_w: [L D]     = ones
    param ln1_b: [L D]     = zeros
    param w_qkv: [L 3*D D] = normal(std=0.02)
    param w_o:   [L D D]   = normal(std=0.02)
    param ln2_w: [L D]     = ones
    param ln2_b: [L D]     = zeros
    param w1:    [L Dff D] = normal(std=0.02)
    param w2:    [L D Dff] = normal(std=0.02)
    param lnf_w: [D]       = ones
    param lnf_b: [D]       = zeros

linear(x, w) multiplies by the transpose of w, as in PyTorch, which is why w_qkv is [3*D D], output first.

The data

One example is a sequence and, for each position, the token that comes next:

## One training example: a sequence, and each position's next token.
data Example:
    ids:  [T] Index[V]
    next: [T] Index[V]

Attention

This is the function from the Introduction, with linear for the projections:

## Self-attention in which each token sees itself and the tokens before it.
fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(x, w_qkv)
        |> split(b n (3 H Dh) -> b H n Dh)
    att = einsum(q, k: b H n Dh, b H m Dh -> b H n m)
        |> scale(1 / sqrt(Dh))
        |> masked_fill(!causal[n m], -inf)
        |> softmax(over=m)
    return einsum(att, v: b H n m, b H m Dh -> b H n Dh)
        |> rearrange(b H n Dh -> b n (H Dh))
        |> linear(w_o)

Every line is checked: that 3 * H * Dh is 3 * D, that causal[n m] is a [n m] mask that broadcasts by name to [b H n m], that softmax normalizes over the keys m, and that the result is [b n D] again.

A block, and the forward pass

A pre-norm decoder block adds attention and then an MLP to its input. m.w_qkv[l] picks block l’s weights:

## Decoder block `l`: attention, then an MLP, each added to its input.
fn block(m: Copier, x: [b n D], l: int) -> [b n D]:
    h = x + (x
        |> layer_norm(over=D, weight=m.ln1_w[l], bias=m.ln1_b[l])
        |> attention(m.w_qkv[l], m.w_o[l]))
    return h + (h
        |> layer_norm(over=D, weight=m.ln2_w[l], bias=m.ln2_b[l])
        |> linear(m.w1[l])
        |> gelu
        |> linear(m.w2[l]))

## The logits of each position's next token.
fn forward(m: Copier, ids: [b T] Index[V]) -> [b T V] Logit over V:
    x = embedding(ids, m.wte) + m.wpe
    for l in 0..L:
        x = block(m, x, l)
    return x
        |> layer_norm(over=D, weight=m.lnf_w, bias=m.lnf_b)
        |> linear(m.wte)

embedding(ids, m.wte) + m.wpe adds a [T D] to a [b T D], which lines up by name: no unsqueeze. The loop runs the L blocks; x keeps its type each time round, as a loop requires. The output layer reuses the token embedding, linear(m.wte), a common trick called weight tying.

Loss and predictions

fn loss(m: Copier, batch: [b] Example, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.ids), batch.next, over=V)

## The predicted next tokens.
fn predict(m: Copier, ids: [b T] Index[V]) -> [b T] Index[V]:
    return argmax(forward(m, ids), over=V)

loss takes a Mode, because that’s the signature the standard library’s fit expects. This model has no dropout, so it doesn’t use it.

Check the file:

shapescript check copier.ss

and hover Copier in the editor: it draws the architecture, with the embedding, the transformer block × L = 2 with its attention and MLP, and the tied output layer.

Training

train.ss imports the model and hands it to the standard library’s generic fit, which uses our loss:

import copier
import shapescript.std.fit as std

## Trains the copier with the standard library's loop: AdamW, batches of
## 64, a checkpoint after every epoch.
fn train(mut m: copier.Copier, data: Dataset[copier.Example], sink: LogSink, epochs: int):
    std.fit(mut m, data, sink, epochs, 64, 3e-3, Eval)

The program

The dataset is a few lines of PyTorch, in run_copier.py:

import sys

import torch

import shapescript
from copier import Copier, predict
from train import train


class CopyTask:
    """`n` sequences of 8 random tokens followed by the same 8 again."""

    def __init__(self, n, seed):
        g = torch.Generator().manual_seed(seed)
        half = torch.randint(0, 16, (n, 8), generator=g)
        seqs = torch.cat([half, half], dim=1)              # (n, 16)
        self.ids = seqs
        self.next = torch.roll(seqs, -1, dims=1)           # each position's next token

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, i):
        return {"ids": self.ids[i], "next": self.next[i]}


epochs = int(sys.argv[1]) if len(sys.argv) > 1 else 10
shapescript.loader_workers(0)        # the dataset is defined in this script
model = Copier.init(seed=0)
sink = shapescript.sinks.csv("log.csv", every=None)  # a row per epoch
train(model, CopyTask(4096, seed=0), sink, epochs)
sink.close()

test = CopyTask(256, seed=1)
with torch.no_grad():
    guess = predict(model, test.ids)
# Positions 8 to 14 predict tokens 9 to 15, copies of tokens 1 to 7.
right = (guess[:, 8:15] == test.next[:, 8:15]).float().mean().item()
print(f"second-half accuracy: {right:.1%}")

Run it:

$ shapescript run run_copier.py 10
epoch 0 done
epoch 1 done
epoch 2 done
epoch 3 done
epoch 4 done
epoch 5 done
epoch 6 done
epoch 7 done
epoch 8 done
epoch 9 done
second-half accuracy: 100.0%
$ cat log.csv
loss,grad_norm
1.9936360102146864,0.924981951713562
1.2352662440389395,0.38128986954689026
1.223402040079236,0.35159793496131897
1.220885556191206,0.31604573130607605
1.2193415965884924,0.31824299693107605
1.2186482679098845,0.6085124611854553
1.2172029335051775,0.6178174614906311
1.2268130257725716,1.7973766326904297
1.2175495941191912,0.659481406211853
1.2164913676679134,0.2807769179344177

The model copies perfectly. Why does the loss stop at about 1.22 rather than 0? Of the 16 next tokens in each sequence, 7 are random (the second to eighth tokens), and no model can predict them better than chance, ln 16 ≈ 2.77 each. The other 9 are copies. So the best possible mean loss is 7/16 × 2.77 ≈ 1.21, which is where it ended.

fit also left a checkpoint, ckpt.pt, which std.resume can continue from. To train on a GPU, change devices cpu to devices cuda[0], or run it on Modal as it is, with devices cuda[0]: shapescript run --modal T4 run_copier.py 10.

Where to go from here

You’ve now used nearly all of ShapeScript. The language reference has the exact rules, and a section for every diagnostic; examples/ in the repository has more programs, such as a ViT block and a GPT. The exercises below extend the project.

Exercises

1. Add dropout after the MLP’s gelu, with p = 0.1. What has to change in block, forward and loss, and what does train pass now?
2. Change the task to reversal: the second half is the first half reversed. Which file changes? Will the model still reach 100%?
3. Write accuracy(m: Copier, batch: [b] Example) -> [] in ShapeScript: the fraction of positions 8 to 14 predicted correctly.
4. Export predict with shapescript export. What has to change first?