Source: experiment.py

pip install numpy torch matplotlib
python3 experiment.py    # a few seconds on any CPU; writes captures/ and diagrams/

#!/usr/bin/env python3
"""
LLM Self-Attention by Hand -- the one script behind the workshop.

Everything numeric in the workshop (every table, every plot, every quoted
number) comes from ONE deterministic run of this file:

    python3 experiment.py

Outputs:
    ../captures/results.log     -- every measurement, in citable sections
    ../diagrams/*.png           -- every matplotlib figure

The script builds self-attention in the same order the workshop teaches it:

    [S1] a tiny hand-designed "toy world": 4 words, 4 named embedding axes
    [S2] the two ingredients on their own: dot product and softmax
    [S3] the simplest possible attention: raw embeddings, no learned parts
    [S4] a hand-wired query/key head: why Q and K exist
    [S5] why scores are divided by sqrt(d_k): a measurement, not a slogan
    [S6] the general head, checked against PyTorch; causal masking; limits

The only randomness is in [S5] (random vectors to measure dot-product
statistics) and [S6] (random projection matrices for the PyTorch check),
and both draw from one generator seeded with 0, so the log is reproduced
byte for byte on every run.
"""

import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

import matplotlib
matplotlib.use("Agg")  # render to files; no display needed
import matplotlib.pyplot as plt

# Where this script lives and where its outputs go. The script sits in
# experiments/; the log and figures land in sibling directories so the
# wiki pages can reference them.
HERE = Path(__file__).resolve().parent
CAPTURES = HERE.parent / "captures"
DIAGRAMS = HERE.parent / "diagrams"

SEED = 0

# Larger default font so every figure stays readable at phone width.
plt.rcParams.update({"font.size": 13})

# ---------------------------------------------------------------------------
# logging: everything printed is also collected and written to results.log,
# so the workshop can cite "line NN" of one canonical capture file.
# ---------------------------------------------------------------------------

_LOG_LINES = []


def log(msg=""):
    """Print a line and record it for captures/results.log."""
    print(msg)
    _LOG_LINES.append(msg)


def save_plot(fig, name):
    """Save a figure under diagrams/ and log where it went."""
    path = DIAGRAMS / name
    fig.savefig(path, dpi=110, bbox_inches="tight")
    plt.close(fig)
    log(f"  [plot saved: diagrams/{name}]")


# ---------------------------------------------------------------------------
# [S1] the toy world
#
# Four words, and embedding vectors with only FOUR dimensions, each with a
# human-readable meaning. Real LLMs learn thousands of unlabeled dimensions
# during training; we hand-design ours so that every dot product in the
# workshop can be checked with pencil and paper.
# ---------------------------------------------------------------------------

# The four named axes of our toy embedding space.
AXES = ["water", "money", "place", "filler"]

# One vector per word. Note that "bank" carries a little water-ness, a
# little money-ness (it is ambiguous between the two) and a lot of
# place-ness. "the" carries nothing but filler-ness.
EMBED = {
    "the":   np.array([0.0, 0.0, 0.0, 1.0]),
    "river": np.array([0.9, 0.0, 0.3, 0.0]),
    "money": np.array([0.0, 0.9, 0.1, 0.0]),
    "bank":  np.array([0.3, 0.3, 0.9, 0.0]),
}

# The two test sentences. Same last word, opposite intended meanings.
SENT_RIVER = ["the", "river", "bank"]
SENT_MONEY = ["the", "money", "bank"]


def embed(sentence):
    """Stack the embedding rows for a sentence into a matrix X.

    X has one ROW per word and one COLUMN per axis, so a 3-word sentence
    in our 4-dimensional space gives a 3 x 4 matrix.
    """
    return np.stack([EMBED[w] for w in sentence])


def s1_toy_world():
    """Print the embedding table that everything else builds on."""
    log("[S1] the toy world")
    log("  embedding axes (hand-named): " + ", ".join(AXES))
    header = "  word    " + "".join(f"{a:>8}" for a in AXES)
    log(header)
    for word, vec in EMBED.items():
        log(f"  {word:<8}" + "".join(f"{v:8.1f}" for v in vec))
    log(f'  test sentences: "{" ".join(SENT_RIVER)}" / "{" ".join(SENT_MONEY)}"')
    log("  note: 'bank' has exactly ONE row here, whichever sentence it is in")
    log("")


# ---------------------------------------------------------------------------
# [S2] the two ingredients on their own: dot product and softmax
# ---------------------------------------------------------------------------

def softmax(scores):
    """Turn a row (or rows) of scores into positive weights that sum to 1.

    Bigger score -> bigger weight, and every weight is between 0 and 1.
    Works on the LAST axis, so feeding it a matrix softmaxes each row
    independently.
    """
    # Subtracting the row maximum leaves the output unchanged (both the
    # top and bottom of the fraction get multiplied by the same constant)
    # but keeps exp() from overflowing on large scores.
    shifted = scores - np.max(scores, axis=-1, keepdims=True)

    exp = np.exp(shifted)
    return exp / np.sum(exp, axis=-1, keepdims=True)


def plot_dot_product_arrows():
    """Three vector pairs: aligned, perpendicular, opposite."""
    pairs = [
        (np.array([1.0, 0.4]), np.array([0.9, 0.5]), "similar direction"),
        (np.array([1.0, 0.0]), np.array([0.0, 1.0]), "unrelated (perpendicular)"),
        (np.array([1.0, 0.3]), np.array([-0.9, -0.4]), "opposite direction"),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.6))
    for ax, (a, b, title) in zip(axes, pairs):
        ax.quiver(0, 0, a[0], a[1], angles="xy", scale_units="xy", scale=1,
                  color="#1976d2", width=0.02, label="a")
        ax.quiver(0, 0, b[0], b[1], angles="xy", scale_units="xy", scale=1,
                  color="#ef6c00", width=0.02, label="b")
        ax.set_xlim(-1.3, 1.3)
        ax.set_ylim(-1.3, 1.3)
        ax.set_aspect("equal")
        ax.axhline(0, color="gray", lw=0.5)
        ax.axvline(0, color="gray", lw=0.5)
        ax.set_title(f"{title}\na . b = {np.dot(a, b):+.2f}", fontsize=12)
        ax.set_xticks([])
        ax.set_yticks([])
    fig.tight_layout()
    save_plot(fig, "dot_product_arrows.png")


def plot_softmax_in_action(scores, words):
    """Scores, their softmax, and the softmax of the scores times 5."""
    weights = softmax(scores)
    sharp = softmax(scores * 5.0)
    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.4))
    panels = [
        (scores, "raw scores", "#90a4ae"),
        (weights, "softmax(scores)", "#1976d2"),
        (sharp, "softmax(scores x 5)", "#c2185b"),
    ]
    for ax, (vals, title, color) in zip(axes, panels):
        ax.bar(words, vals, color=color)
        ax.set_title(title, fontsize=12)
        ax.set_ylim(0, max(1.05, vals.max() * 1.15))
        for i, v in enumerate(vals):
            ax.text(i, v + 0.02, f"{v:.3f}", ha="center", fontsize=11)
    fig.tight_layout()
    save_plot(fig, "softmax_in_action.png")


def s2_dot_and_softmax():
    """Measure the two ingredients before combining them."""
    log("[S2] dot product and softmax, on their own")

    # The dot product multiplies matching entries and adds the products.
    # Spell one out term by term so it can be checked by hand.
    b, r = EMBED["bank"], EMBED["river"]
    terms = " + ".join(f"{x:.1f}*{y:.1f}" for x, y in zip(b, r))
    log(f"  bank . river = {terms} = {np.dot(b, r):.3f}")
    for other in ["money", "the", "bank"]:
        log(f"  bank . {other:<5} = {np.dot(b, EMBED[other]):.3f}")

    # Softmax the row of scores that 'bank' produces against the words
    # of "the river bank".
    scores = np.array([np.dot(b, EMBED[w]) for w in SENT_RIVER])
    weights = softmax(scores)
    log(f"  softmax([{', '.join(f'{s:.3f}' for s in scores)}]) = "
        f"[{', '.join(f'{w:.3f}' for w in weights)}]  (sum = {weights.sum():.6f})")

    # The same scores, five times larger: softmax turns winner-take-all.
    sharp = softmax(scores * 5.0)
    log(f"  softmax(same scores x 5)       = "
        f"[{', '.join(f'{w:.3f}' for w in sharp)}]  (one clear winner)")

    plot_dot_product_arrows()
    plot_softmax_in_action(scores, SENT_RIVER)
    log("")


# ---------------------------------------------------------------------------
# [S3] the simplest possible attention: raw embeddings, no learned parts
# ---------------------------------------------------------------------------

def simple_attention(X):
    """Attention with no learned parts at all.

    Each word scores every word in the sentence with a plain dot product
    of raw embeddings, softmaxes its row of scores into weights, and
    replaces its own vector with the weighted average of ALL the vectors.
    """
    # scores[i, j] = how similar word i's embedding is to word j's.
    scores = X @ X.T

    # Each row of scores becomes a row of weights that sums to 1.
    weights = softmax(scores)

    # Each output row is a weights-blended average of the embedding rows.
    return weights @ X, weights


def log_matrix(title, M, row_labels, fmt="{:8.3f}"):
    """Log a small matrix with row and column labels."""
    log(f"  {title}")
    log("    " + " " * 7 + "".join(f"{w:>8}" for w in row_labels))
    for label, row in zip(row_labels, M):
        log(f"    {label:<7}" + "".join(fmt.format(v) for v in row))


def heatmap(ax, W, words, title):
    """Draw an annotated attention-weight heatmap on the given axes."""
    ax.imshow(W, cmap="Blues", vmin=0.0, vmax=1.0)
    ax.set_xticks(range(len(words)), words)
    ax.set_yticks(range(len(words)), words)
    ax.xaxis.set_ticks_position("top")
    for i in range(len(words)):
        for j in range(len(words)):
            v = W[i, j]
            ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=12,
                    color="white" if v > 0.55 else "black")
    ax.set_title(title, fontsize=12, pad=28)


def plot_raw_attention_heatmaps(W_river, W_money):
    fig, axes = plt.subplots(1, 2, figsize=(8.6, 4.2))
    heatmap(axes[0], W_river, SENT_RIVER, '"the river bank"')
    heatmap(axes[1], W_money, SENT_MONEY, '"the money bank"')
    fig.text(0.5, 0.02, "row = word doing the looking; columns sum to nothing, ROWS sum to 1",
             ha="center", fontsize=11)
    fig.tight_layout(rect=(0, 0.05, 1, 1))
    save_plot(fig, "raw_attention_heatmaps.png")


def plot_bank_on_water_money_plane(out_river, out_money):
    """Where 'bank' sits on the water/money plane, before and after."""
    fig, ax = plt.subplots(figsize=(6.4, 5.6))
    pts = {
        "river": (EMBED["river"][0], EMBED["river"][1], "#1976d2"),
        "money": (EMBED["money"][0], EMBED["money"][1], "#2e7d32"),
        "bank (embedding)": (EMBED["bank"][0], EMBED["bank"][1], "#616161"),
    }
    for label, (x, y, color) in pts.items():
        ax.scatter(x, y, s=90, color=color, zorder=3)
        ax.annotate(label, (x, y), textcoords="offset points", xytext=(8, 6),
                    fontsize=12)
    for out, label, color in [
        (out_river, 'bank after "the river bank"', "#c2185b"),
        (out_money, 'bank after "the money bank"', "#c2185b"),
    ]:
        ax.scatter(out[0], out[1], s=90, color=color, marker="D", zorder=3)
        ax.annotate(label, (out[0], out[1]), textcoords="offset points",
                    xytext=(8, 6), fontsize=12)
        ax.annotate("", xy=(out[0], out[1]),
                    xytext=(EMBED["bank"][0], EMBED["bank"][1]),
                    arrowprops=dict(arrowstyle="->", color="#c2185b", lw=1.6))
    ax.set_xlabel("water axis")
    ax.set_ylabel("money axis")
    ax.set_xlim(-0.05, 1.05)
    ax.set_ylim(-0.05, 1.05)
    ax.grid(alpha=0.3)
    ax.set_title("one embedding, two context-aware outputs", fontsize=12)
    fig.tight_layout()
    save_plot(fig, "bank_water_money_plane.png")


def s3_raw_attention():
    """Run the no-learned-parts attention on both sentences."""
    log("[S3] attention with raw embeddings (no learned parts)")

    X_river = embed(SENT_RIVER)
    X_money = embed(SENT_MONEY)

    # First the input fact the whole workshop turns on: the embedding row
    # for 'bank' is the same array in both sentences.
    same = bool(np.array_equal(X_river[2], X_money[2]))
    log(f"  input check: 'bank' row identical in both sentences -> {same}")

    out_river, W_river = simple_attention(X_river)
    out_money, W_money = simple_attention(X_money)

    S_river = X_river @ X_river.T
    log_matrix('score matrix S = X @ X.T for "the river bank":',
               S_river, SENT_RIVER)
    log(f"  symmetric? {bool(np.allclose(S_river, S_river.T))} "
        "(S[i,j] always equals S[j,i] here)")

    log_matrix('attention weights (each row softmaxed) for "the river bank":',
               W_river, SENT_RIVER)
    log_matrix('attention weights for "the money bank":',
               W_money, SENT_MONEY)

    # The payoff: the same 'bank' row produced two different outputs.
    log("  bank's output vector (water, money, place, filler):")
    log("    in \"the river bank\": "
        + "  ".join(f"{a}={v:.3f}" for a, v in zip(AXES, out_river[2])))
    log("    in \"the money bank\": "
        + "  ".join(f"{a}={v:.3f}" for a, v in zip(AXES, out_money[2])))
    log(f"  water vs money, river sentence: {out_river[2][0]:.3f} vs {out_river[2][1]:.3f}")
    log(f"  water vs money, money sentence: {out_money[2][0]:.3f} vs {out_money[2][1]:.3f}")

    plot_raw_attention_heatmaps(W_river, W_money)
    plot_bank_on_water_money_plane(out_river[2], out_money[2])
    log("")
    return W_river


# ---------------------------------------------------------------------------
# [S4] a hand-wired query/key head: why Q and K exist
#
# Raw-embedding attention can only ask "who looks like me?". Real attention
# lets every word ask a QUESTION that is different from its own identity.
# Two small matrices do it: W_Q turns an embedding into a query ("what I am
# looking for") and W_K turns an embedding into a key ("what I advertise").
# We wire them by hand, with just two channels, so the effect is legible.
# ---------------------------------------------------------------------------

# Channel 1 is "topic": place words ASK for topic words on this channel
# (W_Q), and watery/money words ADVERTISE topic content on it (W_K).
# Channel 2 is "place": filler words ASK for place words, and place words
# ADVERTISE their place-ness. Rows are input axes (water, money, place,
# filler); columns are the two channels.
W_Q_WIRED = np.array([
    [0.0, 0.0],   # water-ness asks for nothing
    [0.0, 0.0],   # money-ness asks for nothing
    [2.0, 0.0],   # place words ask for TOPIC words   (channel 1)
    [0.0, 1.5],   # filler words ask for PLACE words  (channel 2)
])
W_K_WIRED = np.array([
    [1.0, 0.0],   # water-ness advertises "I carry topic"  (channel 1)
    [1.0, 0.0],   # money-ness advertises "I carry topic"  (channel 1)
    [0.0, 1.0],   # place-ness advertises "I am a place"   (channel 2)
    [0.0, 0.0],   # filler advertises nothing
])


def wired_attention(X):
    """Attention with the hand-wired W_Q and W_K (values stay raw).

    Keeping the values equal to the raw embeddings (formally, W_V is the
    identity matrix) is a hand-wiring choice: it keeps every output
    readable on our four named axes while Q and K do the interesting work.
    """
    # Each word's question, and each word's advertisement.
    Q = X @ W_Q_WIRED
    K = X @ W_K_WIRED

    # scores[i, j] = how well word j's advertisement answers word i's
    # question. No longer symmetric: asking and advertising differ.
    scores = Q @ K.T

    # Scale into softmax's responsive range (measured in [S5]).
    d_k = W_K_WIRED.shape[1]
    weights = softmax(scores / np.sqrt(d_k))

    return weights @ X, weights, scores


def plot_qk_vs_raw_heatmaps(W_raw, W_wired):
    fig, axes = plt.subplots(1, 2, figsize=(8.6, 4.2))
    heatmap(axes[0], W_raw, SENT_RIVER, "raw embeddings\n(who looks like me?)")
    heatmap(axes[1], W_wired, SENT_RIVER, "wired Q/K head\n(who answers my question?)")
    fig.tight_layout()
    save_plot(fig, "qk_vs_raw_heatmaps.png")


def s4_wired_qk(W_raw):
    """Compare raw-embedding attention with the hand-wired Q/K head."""
    log("[S4] a hand-wired query/key head")

    X = embed(SENT_RIVER)
    Q = X @ W_Q_WIRED
    K = X @ W_K_WIRED

    log('  queries and keys for "the river bank" (2 channels: topic, place):')
    for word, q, k in zip(SENT_RIVER, Q, K):
        log(f"    {word:<7} q = [{q[0]:5.2f}, {q[1]:5.2f}]   "
            f"k = [{k[0]:5.2f}, {k[1]:5.2f}]")

    out, W_wired, scores = wired_attention(X)

    log_matrix("score matrix Q @ K.T:", scores, SENT_RIVER)
    log(f"  symmetric? {bool(np.allclose(scores, scores.T))}  "
        f"(bank->river = {scores[2][1]:.3f}, river->bank = {scores[1][2]:.3f})")

    log_matrix("weights after dividing by sqrt(2) and softmaxing each row:",
               W_wired, SENT_RIVER)

    log("  bank's row,  raw vs wired: "
        f"[{', '.join(f'{v:.3f}' for v in W_raw[2])}] -> "
        f"[{', '.join(f'{v:.3f}' for v in W_wired[2])}]")
    log("  'the's row,  raw vs wired: "
        f"[{', '.join(f'{v:.3f}' for v in W_raw[0])}] -> "
        f"[{', '.join(f'{v:.3f}' for v in W_wired[0])}]")
    log("  bank now takes half its update from 'river'; 'the' now looks at its noun")

    log("  bank's output vector: "
        + "  ".join(f"{a}={v:.3f}" for a, v in zip(AXES, out[2])))

    plot_qk_vs_raw_heatmaps(W_raw, W_wired)
    log("")


# ---------------------------------------------------------------------------
# [S5] why scores are divided by sqrt(d_k): a measurement, not a slogan
#
# Claim to verify: with high-dimensional queries and keys, raw dot-product
# scores grow so large that softmax collapses to winner-take-all, and
# dividing by sqrt(d_k) prevents it. We MEASURE both halves.
# ---------------------------------------------------------------------------

def s5_scaling(rng):
    """Measure dot-product spread and softmax sharpness as d grows."""
    log("[S5] why scores are divided by sqrt(d_k)")

    ROW = 8        # keys per softmax row (an 8-word sentence)
    TRIALS = 2000  # rows sampled per dimension size
    dims = [4, 16, 64, 256]

    log(f"  random q, k entries drawn from N(0,1); {ROW} keys per row; "
        f"{TRIALS} rows per size")

    stds, max_unscaled, max_scaled = [], [], []
    for d in dims:
        dots, mx_raw, mx_scl = [], [], []
        for _ in range(TRIALS):
            # One query meeting a row of 8 keys, exactly like one row of
            # the score matrix in a real head of dimension d.
            q = rng.standard_normal(d)
            keys = rng.standard_normal((ROW, d))
            scores = keys @ q

            dots.extend(scores.tolist())
            mx_raw.append(softmax(scores).max())
            mx_scl.append(softmax(scores / np.sqrt(d)).max())

        stds.append(np.std(dots))
        max_unscaled.append(np.mean(mx_raw))
        max_scaled.append(np.mean(mx_scl))
        log(f"  d = {d:3d}: dot-product std = {stds[-1]:6.2f}  "
            f"(sqrt(d) = {np.sqrt(d):5.2f}) | avg top weight: "
            f"unscaled = {max_unscaled[-1]:.3f}, scaled = {max_scaled[-1]:.3f}")

    log("  unscaled top weight heads for 1.0 (winner-take-all);"
        " scaled stays flat")

    fig, ax = plt.subplots(figsize=(7.2, 4.6))
    ax.plot(dims, max_unscaled, "o-", color="#c2185b",
            label="softmax(Q K^T), unscaled")
    ax.plot(dims, max_scaled, "s-", color="#1976d2",
            label="softmax(Q K^T / sqrt(d_k))")
    ax.axhline(1.0 / ROW, color="gray", ls=":",
               label=f"perfectly even (1/{ROW})")
    ax.set_xscale("log", base=2)
    ax.set_xticks(dims, [str(d) for d in dims])
    ax.set_xlabel("query/key dimension d_k")
    ax.set_ylabel("average largest weight in a row")
    ax.set_ylim(0, 1.05)
    ax.legend(fontsize=11)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    save_plot(fig, "scaling_max_weight.png")
    log("")


# ---------------------------------------------------------------------------
# [S6] the general head, checked against PyTorch; masking; limits
# ---------------------------------------------------------------------------

def attention(X, W_q, W_k, W_v, causal=False):
    """One full head of self-attention: the real thing, in seven lines.

    X has one row per word. W_q, W_k, W_v are the head's three learned
    matrices (here: supplied by the caller). Returns the new row-per-word
    matrix and the attention weights.
    """
    # Project each word into its three roles.
    Q = X @ W_q   # what each word is looking for
    K = X @ W_k   # what each word advertises
    V = X @ W_v   # what each word hands over if attended to

    # Every query meets every key: scores[i, j] = Q[i] . K[j].
    scores = Q @ K.T / np.sqrt(W_k.shape[1])

    # A causal mask hides the future: word i may only look at words j <= i.
    # Setting a score to -infinity makes its softmax weight exactly 0.
    if causal:
        future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
        scores = np.where(future, -np.inf, scores)

    # Each row of scores becomes weights; each output row is the
    # weights-blended average of the value rows.
    weights = softmax(scores)
    return weights @ V, weights


def plot_causal_mask(W_masked):
    fig, ax = plt.subplots(figsize=(4.6, 4.2))
    heatmap(ax, W_masked, SENT_RIVER, "wired head + causal mask")
    fig.tight_layout()
    save_plot(fig, "causal_mask.png")


def s6_pytorch_and_mask(rng):
    """Verify our NumPy head against PyTorch, then add the causal mask."""
    log("[S6] the general head, checked against PyTorch")

    X = embed(SENT_RIVER)

    # A general head: random learned-looking projections, d_model 4 -> 3.
    # Real models get these matrices from training; random ones are enough
    # to check that the two implementations compute the same function.
    W_q = rng.standard_normal((4, 3))
    W_k = rng.standard_normal((4, 3))
    W_v = rng.standard_normal((4, 3))

    ours, _ = attention(X, W_q, W_k, W_v)

    # PyTorch's fused kernel takes the already-projected Q, K, V.
    tQ = torch.from_numpy(X @ W_q)
    tK = torch.from_numpy(X @ W_k)
    tV = torch.from_numpy(X @ W_v)
    theirs = F.scaled_dot_product_attention(tQ, tK, tV).numpy()
    log(f"  ours vs torch.nn.functional.scaled_dot_product_attention:")
    log(f"    max |difference| = {np.abs(ours - theirs).max():.2e}")

    ours_causal, _ = attention(X, W_q, W_k, W_v, causal=True)
    theirs_causal = F.scaled_dot_product_attention(tQ, tK, tV,
                                                   is_causal=True).numpy()
    log(f"  same comparison with the causal mask on both sides:")
    log(f"    max |difference| = {np.abs(ours_causal - theirs_causal).max():.2e}")

    # The causal mask on the legible wired head from [S4].
    Xw = embed(SENT_RIVER)
    Q = Xw @ W_Q_WIRED
    K = Xw @ W_K_WIRED
    scores = Q @ K.T / np.sqrt(2)
    future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
    W_masked = softmax(np.where(future, -np.inf, scores))
    log_matrix("wired head with a causal mask (upper right forced to 0):",
               W_masked, SENT_RIVER)
    log(f"  every row still sums to 1: "
        f"{', '.join(f'{s:.6f}' for s in W_masked.sum(axis=1))}")

    # Limitation worth measuring: attention alone ignores word ORDER.
    # Reversing the sentence permutes the output rows but changes nothing
    # about what each word computes.
    X_rev = embed(list(reversed(SENT_RIVER)))
    out_fwd, _ = simple_attention(embed(SENT_RIVER))
    out_rev, _ = simple_attention(X_rev)
    same = bool(np.allclose(out_fwd[2], out_rev[0]))
    log(f"  order-blindness: bank's output in \"the river bank\" equals its")
    log(f"  output in \"bank river the\" -> {same} "
        "(real LLMs add position information)")

    # Scale check: the same seven lines at GPT-2's size.
    d_model, d_k, heads, layers = 768, 64, 12, 12
    per_head = 3 * d_model * d_k
    per_layer = heads * per_head + d_model * d_model
    log(f"  one GPT-2 sized head: d_model = {d_model}, d_k = {d_k} -> "
        f"W_q + W_k + W_v = {per_head:,} numbers")
    log(f"  {heads} heads x {layers} layers (+ output projections) -> "
        f"{layers * per_layer:,} attention parameters of GPT-2's ~124M")

    plot_causal_mask(W_masked)
    log("")


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    CAPTURES.mkdir(exist_ok=True)
    DIAGRAMS.mkdir(exist_ok=True)

    # One generator, one seed: [S5] and [S6] draw from it in a fixed
    # order, so every run reproduces the log byte for byte.
    rng = np.random.default_rng(SEED)

    log("=" * 72)
    log("LLM Self-Attention by Hand -- experiment log")
    log(f"torch {torch.__version__} | numpy {np.__version__} | "
        f"matplotlib {matplotlib.__version__} | "
        f"python {sys.version.split()[0]} | seed {SEED}")
    log("toy world: 4-word vocabulary, 4 hand-named embedding axes")
    log("=" * 72)
    log("")

    s1_toy_world()
    s2_dot_and_softmax()
    W_raw = s3_raw_attention()
    s4_wired_qk(W_raw)
    s5_scaling(rng)
    s6_pytorch_and_mask(rng)

    out = CAPTURES / "results.log"
    out.write_text("\n".join(_LOG_LINES) + "\n", encoding="utf-8")
    print(f"wrote {out}")


if __name__ == "__main__":
    main()