Transformers From Scratch Lab

A library gives you model.generate() in one line. That is convenient until the model attends to words that come later in the sentence, or a cached decode gives different text from an uncached one, and no message says why. This lab has eight exercises. Each is a short function you write yourself in numpy, then check with hidden tests that run in your browser. Each one starts with a plain-words setup and a trace by hand, so you know the right number before you type a line.

Recap of earlier lessons. Tokenization covers byte pair encoding. Positional Encoding covers the sine and cosine waves. Multi-Head Attention covers queries, keys, values and heads. The Transformer Architecture covers layer norm, residual connections and the feed-forward layer. GPT and Decoder-Only Models covers the causal mask. Attention Efficiency covers the KV cache. This lab does not teach those ideas again. It makes you build them.

Learning Objectives

After this lesson, you will be able to:

Write one byte-pair-encoding merge step: count neighbouring pairs, pick the most frequent, and merge it left to right
Write sinusoidal positional encoding with one broadcast, and show that moving k positions rotates each sine and cosine pair
Write scaled dot-product attention that survives huge scores, and say why the scores are divided by the square root of d_k
Write the causal mask and apply it before the softmax, so that no token can see a later one
Write the split and merge of attention heads, and keep the shapes and the column order straight
Write layer norm, then a full pre-norm transformer block, and prove that cached decoding gives exactly the numbers of full recomputation

Prerequisites

Before You Start

Review

#How This Lab Works

Every exercise has the same shape. First comes a short setup in plain words. Then a hand trace: one small calculation done on paper, with the answer your code should give. Then a card with the goal, a ladder of hints, and a code editor. The editor holds a starter with TODO marks. It has a Run button and a Check my work button, and under it sits a Show solution button. The hints appear one at a time, so you can take the smallest nudge that gets you moving.

Check my work runs your code, then runs a list of hidden tests, one at a time. It tells you how many passed and shows the message of the first one that failed. The starter is built to fail on purpose, so a fresh cell always says some checks failed. When you finish the TODO lines, all of them pass.

Three things about the tests are worth knowing before you start.

They are tolerance based. A result close enough to the right answer passes, so a harmless rounding difference never fails you.
They are deterministic. Any random numbers come from a fixed seed, so a test that fails today fails the same way tomorrow.
They check more than the one example in the hand trace. Some compare your answer with a second route to the same number, such as a slow loop version. Others test a property that must hold, such as "changing a later token must not change an earlier output". If your code only matches the trace, it will still fail those.

The cells use numpy only. The browser has no torch, so none of this needs a graphics card, and each cell runs in a moment.

Words you need8 terms

A piece of text the model reads as one unit. It may be a whole word, part of a word, or a single character.

Three lists of numbers made from each token. A query asks what to look for, a key says what a token offers, and a value is what gets passed on when the key matches.

For one token, a list of positive numbers that add up to 1, one per token it may look at. They say how much of each value to mix in.

A table of allowed and forbidden pairs that stops a token from looking at tokens after it. It is what lets a model learn to predict the next token.

One of several small attention computations that run side by side, each on its own slice of the columns.

Adding a layer's input back onto its output, so the layer only has to learn a change to the input.

Rescaling the numbers of one token so they have mean 0 and variance 1 across its own features, then applying a learned scale and shift.

The keys and values of the tokens already processed, kept so that generating the next token does not recompute them.

#The Conventions Used Here

Many mistakes in this area are not about maths. They are about which of two conventions a sentence uses. Fix them before writing code. This table is the single source for the whole lab, and the tests assume it.

Quantity	Meaning in this lab
Tokens for BPE	A list of strings. A pair is a tuple `(left, right)`. Counts include overlapping pairs, so `aaa` holds the pair `(a, a)` twice. A merge scans left to right and never reuses a token. A tie goes to the pair that appears first in the list. A pair seen fewer than `min_count` (2) times is not merged.
Sequence	X has shape (n, d_model): n tokens as rows, d_model numbers per token as columns.
Attention	Q is (n_q, d_k), K is (n_k, d_k), V is (n_k, d_v). The scores are `Q @ K.T / sqrt(d_k)`. Softmax runs across the keys (axis 1), so each row of the weights adds up to 1.
Mask	`True` means "allowed". Forbidden scores are set to minus infinity BEFORE the softmax, never after. Query i may see keys 0 to i.
Heads	Shape (n_heads, n, d_head) with d_head = d_model / n_heads. Head h holds columns h * d_head up to (h + 1) * d_head.
Positional encoding	Column 2i is sin(pos / 10000^(2i / d_model)) and column 2i + 1 is the cosine of the same angle. i counts pairs, and d_model is even.
Layer norm	Mean and variance over the last axis. The variance divides by d. eps (1e-5) is added INSIDE the square root. gamma scales and beta shifts, after normalizing.
Block	Pre-norm: `x = x + attention(norm(x))`, then `x = x + feed_forward(norm(x))`. The feed-forward layer is `relu(h @ W1 + b1) @ W2 + b2`. Real models often use a different activation, and the lab uses ReLU so that the numbers stay simple.
Cache	A dict with exactly two entries, "K" and "V", each with one row per token seen so far. A step returns a NEW cache and never edits the one it was given.

When a convention differs from the book you learned from, the tests follow this table. Neither is wrong, and a small change of convention gives different numbers, so name yours before you compare results.

#Exercise 1: One Byte-Pair-Encoding Merge Step

In plain words. A model cannot read raw text. It reads tokens, and byte pair encoding (BPE) is a way to decide what the tokens are. Start with single characters. Find the pair of neighbours that occurs most often, glue it into one new token, and repeat. Common pairs soon become single tokens, and rare words stay split into pieces. Each step has three parts: count the pairs, pick the best one, and merge it everywhere.

Two details decide whether your steps match the tests. When you count pairs, overlaps count: in aaa the pair (a, a) occurs at positions 0-1 and at positions 1-2. When you merge, you scan left to right and jump over both tokens you just joined, so aaa becomes aa followed by a.

Hand trace. Take the characters of aaabdaaabac. The pairs are aa, aa, ab, bd, da, aa, aa, ab, ba, ac. So (a, a) occurs 4 times, (a, b) twice, and the rest once each. Merging (a, a) gives aa a b d aa a b a c, which is 9 tokens. Now (aa, a) and (a, b) both occur twice, and (aa, a) appears first, so it wins the tie and the tokens become aaa b d aaa b a c. One more step merges (aaa, b) and leaves aaab d aaab a c.

Mini project

One BPE merge step

~6 min

Write count_pairs, merge_pair and bpe_merge_step. A step finds the most frequent pair (the first one on a tie), merges it from left to right, and does nothing when no pair occurs at least min_count times. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

import numpy as np

def count_pairs(tokens):
  """Count every pair of neighbouring tokens. Return a dict {(left, right): count}."""
  counts = {}
  # TODO: walk the list once; for each position i add 1 to the count of (tokens[i], tokens[i + 1])
  return counts

def merge_pair(tokens, pair):
  """Join every neighbouring occurrence of pair into one token, scanning left to right.
  Return a NEW list and leave tokens alone."""
  out = []
  # TODO: walk with an index i. If tokens[i] and tokens[i + 1] are the pair, append their
  # joined text and jump i by 2. Otherwise append tokens[i] and jump i by 1.
  return list(tokens)

def bpe_merge_step(tokens, min_count=2):
  """One BPE step. Find the most frequent pair and merge it. Return (new_tokens, pair).
  If two pairs tie, take the one that appears first in the list.
  If the best pair occurs fewer than min_count times, change nothing and return (list(tokens), None)."""
  # TODO: count the pairs, pick the best one, then merge it
  return list(tokens), None

text = list("aaabdaaabac")
print(count_pairs(text))
tokens, pair = bpe_merge_step(text)
print(pair, tokens)

The test that most often catches a wrong answer is the overlap one. A version that skips ahead after matching a pair counts (a, a) twice instead of four times in the trace, and the test says so. A random test then runs many steps on strings of a, b and c, and checks that the joined text never changes and the list always gets shorter.

#Exercise 2: Sinusoidal Positional Encoding

In plain words. Attention looks at every token at once and has no idea which came first. To give it order, each position gets a fingerprint: a list of numbers added to the token's embedding. The fingerprint is made of sine and cosine waves of different speeds. The first pair of columns changes quickly from one position to the next. The last pair changes very slowly. Together they are like the hands of a clock, and no two positions look alike.

The recipe, for position pos and pair number i: the angle is pos / 10000^(2i / d_model), column 2i holds its sine, and column 2i + 1 holds its cosine. The sine and cosine of one pair must share one angle. That sharing is what lets the model turn "k steps later" into a rotation.

Hand trace. Take d_model = 4. Position 0 has angle 0 everywhere, so its row is [0, 1, 0, 1]. For position 1 the first pair has angle 1 / 10000^0 = 1, giving sin 1 = 0.8415 and cos 1 = 0.5403. The second pair has angle 1 / 10000^(2/4) = 1 / 100 = 0.01, giving 0.0100 and 1.0000. So row 1 is [0.8415, 0.5403, 0.0100, 1.0000].

Mini project

Sinusoidal positional encoding

~5 min

Write positional_encoding(n_positions, d_model) so that it returns an (n_positions, d_model) array with sin in the even columns and cos in the odd columns, using the angle pos / 10000^(2i / d_model). Raise ValueError when d_model is odd. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

The test that does the most work is the rotation one. It takes the pair at position p, rotates it by the angle k times that pair's frequency, and demands the pair at position p + k. That holds only when sine sits in the even column, cosine in the odd column, and both use the same angle. Putting all the sines first and all the cosines second looks equally plausible, and fails it.

#Exercise 3: Scaled Dot-Product Attention

In plain words. Each token makes a query, a key and a value. A query is compared with every key by a dot product: a big number means "this token has what I want". The scores for one query are turned into weights that are positive and add up to 1, using softmax. Then the output for that query is the sum of the values, each multiplied by its weight. Attention is a soft lookup.

Two rules keep it healthy. The scores are divided by the square root of d_k. Without that, dot products of long vectors grow with d_k, softmax becomes close to a hard pick, and learning slows. And the softmax subtracts the row maximum first, because exp(1000) overflows a 64-bit float. This is the same trick as the stable softmax in the Deep Learning From Scratch Lab.

Hand trace. Take one query [1, 0], two keys [1, 0] and [0, 1], and values [10, 0] and [0, 10]. The dot products are 1 and 0. Dividing by sqrt(2) gives scores 0.7071 and 0. Then e^0.7071 = 2.0281 and e^0 = 1, so the weights are 2.0281 / 3.0281 = 0.6698 and 0.3302. The output is 0.6698 x [10, 0] + 0.3302 x [0, 10] = [6.6976, 3.3024].

Mini project

Scaled dot-product attention

~6 min

Write softmax_rows (safe for big scores) and scaled_dot_product_attention(Q, K, V), which returns the output and the weights. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

One test adds the same vector to every key. That adds the same number, Q dotted with that vector, to every score in a row, so the weights must not change. A softmax that is not shift-safe, or that runs along the wrong axis, fails it. Another test uses queries and keys a thousand times too large, with floating point errors turned into failures, to catch an overflow.

#Exercise 4: The Causal Mask

In plain words. A language model learns by predicting the next token. If token 3 could look at token 4 while predicting it, it would just copy the answer, and nothing would be learned. The causal mask forbids it: token i may look at tokens 0 to i, itself included, and never later ones.

The mask is a table of True and False, with True where looking is allowed. It is applied to the scores, before the softmax. A forbidden score becomes minus infinity, and exp(-inf) is exactly 0, so the weight is exactly 0 and the remaining weights still add up to 1. Zeroing the weights after the softmax is a classic slip: the rows then add up to less than 1.

Hand trace. For 3 tokens the mask is [[T,F,F],[T,T,F],[T,T,T]]. If every score is equal, token 0 puts all its weight on itself, [1, 0, 0]. Token 1 splits evenly over two tokens, [0.5, 0.5, 0]. Token 2 splits over three, [0.333, 0.333, 0.333].

Mini project

The causal mask

~6 min

Write causal_mask(n), a lower-triangular array of True, and finish causal_attention so that forbidden scores become minus infinity before the softmax. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

import numpy as np

def softmax_rows(scores):
  """Stable softmax along each row. This one is finished for you."""
  scores = np.asarray(scores, dtype=float)
  shifted = scores - np.max(scores, axis=-1, keepdims=True)
  e = np.exp(shifted)
  return e / np.sum(e, axis=-1, keepdims=True)

def causal_mask(n):
  """Return an (n, n) array of True and False. Entry [i, j] is True when query i may look at key j,
  which is when j <= i (a token sees itself and the past, never the future)."""
  # TODO: a lower triangle of True, including the diagonal
  return np.zeros((n, n), dtype=bool)

def causal_attention(Q, K, V):
  """Self-attention where each position sees only itself and earlier positions.
  Q, K, V all have n rows. Return (output, weights)."""
  n, d_k = Q.shape
  scores = Q @ K.T / np.sqrt(d_k)
  # TODO: replace the scores of forbidden (future) keys by -infinity, BEFORE the softmax
  weights = softmax_rows(scores)
  return weights @ V, weights

print(causal_mask(3).astype(int))
Q = np.zeros((3, 2))
out, w = causal_attention(Q, Q, np.eye(3)[:, :2])
print(np.round(w, 4))

The test that matters most changes the last two tokens completely and demands that the outputs of the first three do not move. That is the property the mask exists for, stated without mentioning the mask. A leak of any kind, such as a transposed mask or a mask applied after the softmax, breaks it.

#Exercise 5: Splitting and Merging Heads

In plain words. One attention computation can only learn one kind of lookup at a time. Multi-head attention runs several in parallel. The d_model columns of Q, K and V are cut into n_heads equal slices of d_head columns. Each head attends using only its own slice, and then the heads' outputs are glued back side by side and passed through one more weight matrix, Wo.

The work is all in the shapes. Splitting turns (n, d_model) into (n_heads, n, d_head). Merging does the exact opposite. The trap is the reshape: a bare reshape(n_heads, n, d_head) has the right shape and the wrong contents, because it cuts the flat memory without respecting which numbers belong to which token. You must first reshape to (n, n_heads, d_head), which keeps every token's columns together, and then swap the first two axes.

Hand trace. Take two tokens with four numbers each: [[1, 2, 3, 4], [5, 6, 7, 8]], and 2 heads. Head 0 gets the first two columns of every token: [[1, 2], [5, 6]]. Head 1 gets the last two: [[3, 4], [7, 8]]. The result has shape (2, 2, 2), and merging puts the columns back in the original order.

Mini project

Split and merge attention heads

~6 min

Write split_heads, merge_heads and multi_head_attention. Head h must hold exactly columns h * d_head to (h + 1) * d_head, and merge_heads(split_heads(x)) must give x back. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

import numpy as np

def attention(Q, K, V):
  """One head of scaled dot-product attention. Finished for you.
  Q (n, d_head), K (n, d_head), V (n, d_head) -> (n, d_head)."""
  scores = Q @ K.T / np.sqrt(Q.shape[-1])
  scores = scores - np.max(scores, axis=-1, keepdims=True)
  w = np.exp(scores)
  w = w / np.sum(w, axis=-1, keepdims=True)
  return w @ V

def split_heads(x, n_heads):
  """(n, d_model) -> (n_heads, n, d_head), where d_head = d_model // n_heads.
  Head h must hold columns h * d_head up to (h + 1) * d_head of x.
  If d_model is not divisible by n_heads, raise ValueError."""
  n, d_model = x.shape
  # TODO: check the divisibility, then reshape and reorder the axes
  return x[None, :, :]

def merge_heads(h):
  """(n_heads, n, d_head) -> (n, n_heads * d_head). The exact opposite of split_heads."""
  n_heads, n, d_head = h.shape
  # TODO: put the axes back in the order (n, n_heads, d_head), then reshape
  return h[0]

def multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads):
  """X is (n, d_model); all four weight matrices are (d_model, d_model). Return (n, d_model)."""
  # TODO: project, split, run attention on each head, merge, then apply Wo
  return X

x = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
h = split_heads(x, 2)
print(h.shape)
print(h)

The reshape trap has its own test: head 1 of the 2 by 4 example must be [[3, 4], [7, 8]], and the bare reshape gives something else. Another test shuffles the input tokens and demands that the output rows are shuffled in the same way. Attention without a mask has no sense of order, so that equality must hold, and a bug that mixes tokens across the head axis breaks it.

#Exercise 6: Layer Norm

In plain words. Numbers inside a deep network can drift to very different sizes from layer to layer, which makes training unstable. Layer norm fixes the scale of each token separately. For one token, take its d numbers, subtract their mean, and divide by their standard deviation. They now have mean 0 and variance 1. Then multiply by a learned gamma and add a learned beta, one pair per feature, so that the network can choose a different scale if it wants.

The standard deviation is the square root of the variance plus a tiny eps. The eps sits inside the root, so a token whose numbers are all equal, which has variance 0, does not cause a division by zero. This is the same idea as batch norm in the Deep Learning From Scratch Lab, with one change that matters a great deal: batch norm measures each feature down the batch, and layer norm measures each token across its own features. A language model often runs with a batch of one and with sequences of any length, where batch statistics are unreliable, so it uses the second.

Hand trace. Take the token [1, 2, 3, 4]. The mean is 2.5. The squared distances from the mean are 2.25, 0.25, 0.25, 2.25, so the variance is their average, 1.25, and the standard deviation is 1.118. Dividing the distances by that gives [-1.3416, -0.4472, 0.4472, 1.3416] (eps changes the fifth decimal). With gamma 1 and beta 0 that is the output.

Mini project

Layer norm

~5 min

Write layer_norm(x, gamma, beta, eps) for an array of any shape whose last axis holds the features. Normalize each token over its own features, then scale by gamma and shift by beta. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

Two tests tell batch norm and layer norm apart. After your function, every row should have mean 0, and the columns should not. And a token normalized alone must equal the same token normalized inside a batch. Choosing the wrong axis fails both, and still returns an array of the right shape.

#Exercise 7: A Pre-Norm Transformer Block

In plain words. A transformer is a stack of identical blocks. A block has two parts, attention and a feed-forward layer, and each part is wrapped the same way: normalize the input, run the part, and add the result back onto the input. Adding back is the residual connection. It gives the signal a straight path through the whole stack, so that each part only has to learn a change.

In a pre-norm block the layer norm comes before the part and the residual adds the un-normalized input:

x = x + attention(norm(x)), then x = x + feed_forward(norm(x)).

In a post-norm block the norm comes after the addition. Both orders exist, and many recent language models use pre-norm. They give different numbers, and the test knows which one you were asked for. In this exercise the layer norm, the softmax and the head split and merge are already written for you, and the weights are small fixed numbers. You write the causal attention step and the block that wires it up.

Hand trace. The weights here are fixed, so one trace is a property rather than a number: if the attention output matrix Wo and the feed-forward output weights W2 and b2 are all zero, both parts add nothing, and the block must return x exactly. That is a direct test of the two residual connections. For the fixed weights in the starter, and the input whose three rows are [0, 0.1, 0.2, 0.3], [0.4, 0.5, 0.6, 0.7] and [0.8, 0.9, 1.0, 1.1], the first output row should be [-0.5825, 0.02, 1.402, -0.8463].

Mini project

A pre-norm transformer block

~9 min

Write causal_self_attention(h, p), which runs multi-head attention with the causal mask and ends with Wo, and transformer_block(x, p), which wraps attention and a ReLU feed-forward layer each in a pre-norm residual step. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

import numpy as np

def layer_norm(x, gamma, beta, eps=1e-5):
  """Layer norm over the last axis. Finished for you (exercise 6)."""
  mean = x.mean(axis=-1, keepdims=True)
  var = x.var(axis=-1, keepdims=True)
  return gamma * (x - mean) / np.sqrt(var + eps) + beta

def softmax_rows(scores):
  """Stable softmax along each row. Finished for you (exercise 3)."""
  shifted = scores - np.max(scores, axis=-1, keepdims=True)
  e = np.exp(shifted)
  return e / np.sum(e, axis=-1, keepdims=True)

def split_heads(x, n_heads):
  """(n, d_model) -> (n_heads, n, d_head). Finished for you (exercise 5)."""
  n, d_model = x.shape
  return x.reshape(n, n_heads, d_model // n_heads).transpose(1, 0, 2)

def merge_heads(h):
  """(n_heads, n, d_head) -> (n, d_model). Finished for you (exercise 5)."""
  n_heads, n, d_head = h.shape
  return h.transpose(1, 0, 2).reshape(n, n_heads * d_head)

def make_params():
  """Tiny fixed weights: d_model = 4, 2 heads, feed-forward width 8. The same numbers every time."""
  rng = np.random.default_rng(7)
  d, d_ff = 4, 8
  return {
      "n_heads": 2,
      "Wq": rng.standard_normal((d, d)) * 0.5,
      "Wk": rng.standard_normal((d, d)) * 0.5,
      "Wv": rng.standard_normal((d, d)) * 0.5,
      "Wo": rng.standard_normal((d, d)) * 0.5,
      "W1": rng.standard_normal((d, d_ff)) * 0.5,
      "b1": rng.standard_normal(d_ff) * 0.1,
      "W2": rng.standard_normal((d_ff, d)) * 0.5,
      "b2": rng.standard_normal(d) * 0.1,
      "g1": 1.0 + 0.1 * rng.standard_normal(d),
      "be1": 0.1 * rng.standard_normal(d),
      "g2": 1.0 + 0.1 * rng.standard_normal(d),
      "be2": 0.1 * rng.standard_normal(d),
  }

def causal_self_attention(h, p):
  """Multi-head causal self-attention on h of shape (n, d_model). No residual here.
  Use p["Wq"], p["Wk"], p["Wv"], p["Wo"] and p["n_heads"]. Return shape (n, d_model)."""
  n = h.shape[0]
  # TODO: project to Q, K, V and split into heads. For each head: scores = Q K^T / sqrt(d_head),
  # forbid the future with -inf, softmax, then weights @ V. Merge the heads and apply Wo.
  return h

def transformer_block(x, p):
  """One pre-norm transformer block on x of shape (n, d_model). Return shape (n, d_model).
      x = x + attention(layer_norm(x))
      x = x + feed_forward(layer_norm(x))
  where feed_forward(h) = relu(h @ W1 + b1) @ W2 + b2.
  First layer norm uses p["g1"], p["be1"]; the second uses p["g2"], p["be2"]."""
  # TODO: write the two residual steps, normalizing BEFORE each sub-layer
  return x

p = make_params()
x = np.arange(12, dtype=float).reshape(3, 4) / 10.0
print(np.round(transformer_block(x, p), 4))

The reference in the tests is written a different way, with explicit loops over heads and rows, so a bug cannot hide by being the same bug twice. One test compares your output with the post-norm block and demands that it differs. Another changes the last token and demands that the earlier rows do not move, which checks the mask end to end.

#Exercise 8: KV-Cache Decoding That Equals Full Recompute

In plain words. A decoder produces text one token at a time. To produce token t it needs the keys and values of all earlier tokens. The slow way recomputes every key and value of the whole sequence for every new token. But the keys and values of old tokens never change, because the causal mask means they never see the new one. The fast way keeps them in a cache, and each step computes only the new token's query, key and value, adds the key and value to the cache, and attends over everything cached.

A cache is only useful if it is exact. It is an optimization, so the output must be identical to full recompute, to rounding error. Three details are easy to get wrong. The new token attends to itself, so its key and value must be added BEFORE it attends. Only keys and values are cached, never queries, since a query is used once. And a step must return a new cache and not edit the old one, or a caller that keeps the old cache gets silently wrong results.

Hand trace. Count the work. To produce 4 tokens without a cache, step 1 projects 1 token, step 2 projects 2, step 3 projects 3 and step 4 projects 4, so 1 + 2 + 3 + 4 = 10 key projections in total. With a cache each token is projected once, 4 in total. In general it is n(n + 1) / 2 against n. The starter prints the largest difference between the two methods, and when you are done it is about 1e-16, which is rounding error.

Mini project

KV-cache incremental decoding

~8 min

Write empty_cache, decode_step, decode_all and projection_counts. decode_all must give the same outputs as the provided attend_full, which recomputes everything. Press Check my work to run the hidden tests.

Your codeEdits save in this browser

pythonrunnable cell

import numpy as np

def softmax_rows(scores):
  """Stable softmax along each row. Finished for you."""
  shifted = scores - np.max(scores, axis=-1, keepdims=True)
  e = np.exp(shifted)
  return e / np.sum(e, axis=-1, keepdims=True)

def make_weights():
  """Fixed projections for a single attention head: d_model = 4, d_k = 3, d_v = 5."""
  rng = np.random.default_rng(3)
  return rng.standard_normal((4, 3)), rng.standard_normal((4, 3)), rng.standard_normal((4, 5))

def attend_full(X, Wq, Wk, Wv):
  """The slow way, finished for you: recompute everything for all n tokens, with a causal mask.
  X is (n, d_model). Returns (n, d_v)."""
  n = X.shape[0]
  Q, K, V = X @ Wq, X @ Wk, X @ Wv
  scores = Q @ K.T / np.sqrt(Q.shape[1])
  scores = np.where(np.tril(np.ones((n, n), dtype=bool)), scores, -np.inf)
  return softmax_rows(scores) @ V

def empty_cache(d_k, d_v):
  """A cache that holds no tokens yet: keys of shape (0, d_k) and values of shape (0, d_v)."""
  # TODO: return a dict {"K": ..., "V": ...} of empty float arrays with those shapes
  return {}

def decode_step(x_t, Wq, Wk, Wv, cache):
  """Process ONE new token x_t (shape (d_model,)) using the cache of earlier keys and values.
  Return (output_t, new_cache). output_t has shape (d_v,). Do not edit the cache you were given."""
  # TODO: compute q, k, v for this one token only; add k and v as the last row of the cached keys and values;
  # score q against ALL cached keys (including the new one), scale by sqrt(d_k), softmax, mix the cached values
  return np.zeros(Wv.shape[1]), cache

def decode_all(X, Wq, Wk, Wv):
  """Feed the tokens of X one at a time through decode_step, starting from an empty cache.
  Return (outputs of shape (n, d_v), the final cache)."""
  # TODO: loop over the rows of X, collecting each output
  return np.zeros((X.shape[0], Wv.shape[1])), empty_cache(Wq.shape[1], Wv.shape[1])

def projection_counts(n):
  """How many key (and value) projections of ONE token are computed to produce n outputs?
  Return (without_cache, with_cache). Without a cache, step t recomputes all t tokens so far."""
  # TODO: add up 1 + 2 + ... + n for the first number; the second number is n
  return 0, 0

Wq, Wk, Wv = make_weights()
X = np.arange(20, dtype=float).reshape(5, 4) / 10.0 - 1.0
full = attend_full(X, Wq, Wk, Wv)
out, cache = decode_all(X, Wq, Wk, Wv)
print("max difference:", float(np.max(np.abs(full - out))))
print("cache shapes  :", {name: arr.shape for name, arr in cache.items()})
print("projections   :", projection_counts(4))

The central test feeds six tokens through your decoder and demands the same six rows as attend_full. Further tests resume from the cache of the first three tokens and demand row 3, which is the real use of a cache, and run again with different weight shapes. A version that forgets to add the new key before attending, or that caches the raw inputs instead of the projected keys, gives answers that look plausible and differ by a large amount, and the test reports the size of the difference.

#Where the Traps Are

Every exercise has one common wrong answer that looks right. They fall into four families.

Shape and axis. A reshape can have the right shape and the wrong contents, and a sum or softmax along the wrong axis still returns numbers. Softmax goes across the keys (axis 1). Layer norm goes across the features of one token (the last axis). A bare reshape to (n_heads, n, d_head) mixes tokens. Use a non-square matrix and an input with a different number of queries and keys in your own tests, so that a transpose cannot hide.

Order of operations. The mask goes on the scores, before the softmax. The norm goes before the part in a pre-norm block, and the residual adds the un-normalized input. In decoding, the new key and value go into the cache before the new token attends. Swapping any two of these steps still runs.

Numerical stability. exp of a large number overflows. Subtract the row maximum before exp, use minus infinity for forbidden scores, and keep eps inside the square root of layer norm so a constant token does not divide by zero.

Equivalence. Several pieces are optimizations or rearrangements of something simpler, and the right test is equality with the simple version. Cached decoding must equal full recompute. Merging the heads must undo the split. A token's layer norm must not depend on its neighbours. A test that states what must stay the same catches bugs that a single example misses.

Symptom	Likely cause
Generated text looks fine but a model scores unreasonably well in training	The causal mask is missing, transposed, or applied after the softmax, so tokens see their own answers
Attention weights in each row add up to less than 1	The weights were masked after the softmax instead of the scores before it
nan or inf in the attention weights	A softmax without the row maximum shifted out, or a row where every score is minus infinity
Heads seem to attend to meaningless mixtures of tokens	A bare reshape was used in split_heads instead of reshape then transpose
The same token gets different outputs in different batches	Layer norm was taken down the batch instead of across the features
A deep stack trains badly or the signal vanishes	A residual connection is missing, or the skip path goes through the norm
Cached decoding drifts away from uncached decoding	The new key and value were added after attending, the raw inputs were cached, or the cache was edited in place

#Capstone Checklist

Use this list to check that you are ready to read and write real transformer code. Tick each item only if you can do it without looking.

I can write one BPE merge step and say why overlapping pairs count and why a merge scans left to right
I can write sinusoidal positional encoding with one broadcast and say why sine and cosine share an angle
I can write scaled dot-product attention that survives scores of 1000 and say why the scores are divided by the square root of d_k
I can write the causal mask and say why it goes on the scores before the softmax
I can write the split and merge of heads and say what a bare reshape gets wrong
I can write layer norm and say how it differs from batch norm
I can write a pre-norm block and say where each residual connection takes its input
I can write cached decoding and test it against full recompute
Given an attention implementation I have never seen, I can find the axis, order, stability and equivalence choices in under a minute

If you can tick every line, you can read the model code of a library and spot a convention mismatch. If you cannot, go back to the exercise for that line, open the solution, and write the function again without it.

Neural Machine Translation of Rare Words with Subword Units

Rico Sennrich, Barry Haddow, Alexandra Birch (2016)

Byte pair encoding applied to word segmentation, as in exercise 1.

Attention Is All You Need

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, Illia Polosukhin (2017)

Sinusoidal positional encoding, scaled dot-product attention, multi-head attention and the block, as in exercises 2 to 5 and 7.

Layer Normalization

Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey Hinton (2016)

Normalizing across the features of one example, as in exercise 6.

On Layer Normalization in the Transformer Architecture

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, Tie-Yan Liu (2020)

A comparison of the two placements of the norm, the pre-norm block of exercise 7 and the post-norm one.

#Key Takeaways

A BPE step counts every neighbouring pair including overlaps, merges the most frequent one left to right, and breaks ties by first appearance
Sinusoidal positional encoding puts sin(pos / 10000^(2i / d_model)) in column 2i and cos of the same angle in column 2i + 1, so a shift of k positions is a rotation of each pair
Attention is softmax(Q K^T / sqrt(d_k)) V, with the row maximum subtracted so that large scores cannot overflow
The causal mask sets forbidden scores to minus infinity before the softmax, so later tokens get exactly zero weight and each row still adds up to 1
Split heads with a reshape to (n, n_heads, d_head) and then a transpose. A bare reshape to (n_heads, n, d_head) has the right shape and the wrong contents
Layer norm takes the mean and variance over the features of one token, with eps inside the square root, and a token's result does not depend on its neighbours
A pre-norm block is x + attention(norm(x)) then x + feed_forward(norm(x)), and a KV cache must give exactly the numbers of full recompute, adding the new key and value before the new token attends

#Quick Check

Quick Check1 / 4

Why is the causal mask applied to the scores before the softmax, with minus infinity, rather than to the weights after it?

This lab closes the Transformers & LLMs track. To carry the same probabilistic objectives into images, video and audio, begin the next track with Generative AI: Start Here. To see how the cache you wrote is managed at scale, revisit Attention Efficiency.

Transformers From Scratch Lab

Learning Objectives

After this lesson, you will be able to:

Write one byte-pair-encoding merge step: count neighbouring pairs, pick the most frequent, and merge it left to right
Write sinusoidal positional encoding with one broadcast, and show that moving k positions rotates each sine and cosine pair
Write scaled dot-product attention that survives huge scores, and say why the scores are divided by the square root of d_k
Write the causal mask and apply it before the softmax, so that no token can see a later one
Write the split and merge of attention heads, and keep the shapes and the column order straight
Write layer norm, then a full pre-norm transformer block, and prove that cached decoding gives exactly the numbers of full recomputation

Prerequisites

Before You Start

Review

#How This Lab Works

Three things about the tests are worth knowing before you start.

They are tolerance based. A result close enough to the right answer passes, so a harmless rounding difference never fails you.
They are deterministic. Any random numbers come from a fixed seed, so a test that fails today fails the same way tomorrow.
They check more than the one example in the hand trace. Some compare your answer with a second route to the same number, such as a slow loop version. Others test a property that must hold, such as "changing a later token must not change an earlier output". If your code only matches the trace, it will still fail those.