Capstones · 1

Build a mini-transformer

A transformer reads a whole sequence at once with self-attention: each position pulls in information from the others, weighted by how relevant they are. A causal mask is what turns it into a language model — a token may attend only to its past. Here you assemble one causal-attention block from the pieces you have already met. Everything runs in this page, in both languages.

Self-attention: every position weighs the others

Score each pair of positions by their dot product (scaled by √D), then softmax each row into a distribution. Row i says how much token i draws from every token — itself included.

The causal mask: no peeking ahead

Add −1e9 to every score above the diagonal before the softmax, so those future positions get zero weight. The result is lower-triangular: token i attends only to tokens 0…i.

Assemble the block

Mix each position's value vectors by its attention weights, then add the input back — the residual connection. A sequence goes in and a sequence comes out, each token now carrying a summary of its own past. Stack this block a few times and you have the core of a GPT.

Your turn: make it causal

The block reports how much weight token 0 puts on the future (positions 1…T−1). For a language model that must be zero — the first token cannot see what comes after it. Replace the empty mask with the causal one.

The causal mask, drawn

The attention weights as a heatmap: a lower triangle of light and a dark upper triangle. Each row is a token; the dark cells are the future it is forbidden to see. This one picture is what makes a transformer a language model.

What to remember

Self-attention lets each position mix information from the others, weighted by relevance.
A causal mask (−∞ above the diagonal, before the softmax) makes it a language model: a token sees only its past.
A transformer block is attention + a feed-forward net, with residual connections and layer norm; stack them for a GPT.
Built here from mm, softmax and a triu mask — the arithmetic nn.MultiheadAttention wraps.

If a block above errors, that is worth knowing. Every example on this page runs against the same library the tests run against — nothing here is a screenshot. Press Reset to get the original code back.