promptdojo_

Attention and Transformer blocks — step 1 of 7

Attention and the transformer block

The architecture behind every model in this course's AI half is the transformer, introduced in "Attention Is All You Need" (Vaswani et al., 2017). Its core operation — attention — fits in the editor above, and you should run it once before ever repeating the word.

What the code does

Every token carries three learned vectors: a query ("what am I looking for?"), a key ("what do I offer?"), and a value ("what do I contribute if chosen?"). For each token:

  1. Dot its query against every token's key — chapter 38's weighted-sum atom, used as a relevance score (scaled by √dim to keep scores tame).
  2. Softmax the scores into weights that are positive and sum to 1 — a differentiable "soft choice."
  3. Output the weighted mix of all tokens' values.

Read the printout: each token's output is literally a blend of every token's contribution, weighted by computed relevance. That's how "bank" can end up represented differently next to "river" than next to "money" — context flows in through the weights. And because every token scores against every other, cost grows with the square of sequence length — the mechanical reason long context is expensive (last lesson's promise, kept).

The block, and the stack

A transformer layer wraps attention with machinery that makes deep stacks trainable: multi-head attention (several attentions in parallel, each free to learn a different relationship), a small per-token feed-forward network, plus residual connections and normalization — the plumbing that keeps chapter 43's gradients flowing through dozens of layers. Stack N of these blocks, feed in token embeddings plus position information, and you have the architecture; scale N, the dims, and the data, and you have the models you've been calling through APIs since chapter 13.

LLMs of the chat kind are trained as decoders: each position may attend only to earlier positions (a causal mask), because the training task is chapter 00's next-token prediction, industrial edition.

Where AI specifically gets this wrong

  • Explaining attention with metaphors only. "The model focuses on relevant words" — you now have the actual mechanism: dot, softmax, mix. Prefer it.
  • Rolling custom attention when a standard block exists. Subtle masking and scaling bugs, in the one place they're hardest to notice — use the framework's implementation (nn.MultiheadAttention, or better, a whole prebuilt model).
  • Ignoring the quadratic. Generated "just increase max_length" advice meets a square law. Longer context costs compute and memory superlinearly — now you know why.