Kiru Lab  /  Language Models  /  Anatomy of a Transformer

Self-Attention, Built From Scratch

Queries, keys, values — a differentiable lookup table. You already know every operation involved.

Hands-On Lab  ·  about 50 minutes

Each position produces three vectors from its current representation: a query (what am I looking for?), a key (what do I offer?), and a value (what will I pass along if selected?). Every query is compared against every key by dot product — the alignment measure from track two — and the resulting scores, softmaxed into weights, determine a weighted sum of values.

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V
Attention, complete
import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(X, Wq, Wk, Wv, causal=True):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv        # each (seq, d_k)
    scores = Q @ K.T / np.sqrt(Q.shape[-1])  # (seq, seq)
    if causal:
        mask = np.triu(np.ones_like(scores), k=1).astype(bool)
        scores[mask] = -np.inf               # a position may not read the future
    weights = softmax(scores)                # (seq, seq), rows sum to 1
    return weights @ V, weights

Three details that matter

  • The sqrt(d_k) division keeps dot products from growing with dimension and driving the softmax into saturation, where gradients vanish.
  • The causal mask is what makes the model autoregressive. Remove it and you have a bidirectional encoder instead.
  • Multiple heads run this in parallel with different projections, so different heads can specialize — one tracking syntax, another tracking a repeated name.
The standing warning

Attention weights are widely treated as explanations. They are not. A weight tells you where information was read from, not what was computed with it or whether it mattered to the output. The interpretability track replaces this intuition with causal tests that can actually distinguish the two.

Hold on to

  • Attention is a differentiable, content-addressed lookup
  • The causal mask is what makes generation autoregressive
  • Attention weights show routing, not causation

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Implement `attention` and verify each row of the weight matrix sums to 1 and respects the mask.
    Hint

    Softmax rows sum to 1; the causal mask should leave zeros strictly above the diagonal.

    Solution

    Each row sums to 1 because softmax normalizes across the key axis. Above the diagonal every entry is exactly 0, because -inf maps to exp(-inf) = 0 — that is the mask working. Row i therefore distributes all its weight across positions 0 through i, which is precisely what "may not read the future" means mechanically.

    import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(6, 32))
Wq, Wk, Wv = (rng.normal(size=(32, 32)) for _ in range(3))

out, weights = attention(X, Wq, Wk, Wv, causal=True)
assert np.allclose(weights.sum(axis=-1), 1.0)
assert np.allclose(np.triu(weights, k=1), 0.0)
    Check your work

    Paste this after your own code. If it runs without raising, you have it. 

    import numpy as np
assert np.allclose(weights.sum(axis=-1), 1.0)
assert np.allclose(np.triu(weights, k=1), 0.0)
assert out.shape == (6, 32)
print("ok")
  2. Run it on a sentence with a repeated noun and visualize the weight matrix. Describe what the pattern does and does not tell you.
    Hint

    Look at the row for the second occurrence of the repeated noun.

    Solution

    You typically see the second occurrence attending strongly to the first, and most positions attending heavily to the first token — a well-documented "attention sink" where models park unused attention mass. What the map does tell you: where information was read from. What it does not tell you: what was computed with it, whether it affected the output, or whether an ablation would change anything. Distinguishing those two lists is the entire point of the exercise.

In the Bio Mirror

Attention Networks

Sign in to track your progress through the lab.

← Anatomy of a Transformer The Residual Stream and the Transformer Block →