Module 0: The three tools

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

The run is deterministic, so your captures/results.log should match the captured log byte for byte. The capture on these pages was made with python 3.10.12, numpy 2.2.6, torch 2.12.1+cpu, and matplotlib 3.10.6, seed 0.

Words become lists of numbers

A language model is a program trained to continue text: given the words so far, it produces a guess about the next one. Its very first move is humble: every word it knows is looked up in a big table, and the table hands back a fixed list of numbers for that word. A list of numbers is a vector, each slot in the list is a dimension, and this particular vector is called the word's embedding.

Why lists of numbers? Because you can do arithmetic on them. The table is learned during training (a topic outside this workshop), and it ends up arranged so that words used in similar ways get vectors that point in similar directions. "Meaning" becomes geometry.

Real models use vectors with hundreds or thousands of dimensions, and none of the dimensions has a human-readable label. For this workshop we build a toy world instead: four words, and embeddings with only four dimensions that we design by hand and name. That choice is what makes every later computation checkable with a pencil.

Here is the entire toy world, exactly as the experiment script defines it near the top:

# 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 first thing the script does when it runs is print this table, so the record starts with it (cross-check: results.log lines 7-15):

[S1] the toy world
  embedding axes (hand-named): water, money, place, filler
  word       water   money   place  filler
  the          0.0     0.0     0.0     1.0
  river        0.9     0.0     0.3     0.0
  money        0.0     0.9     0.1     0.0
  bank         0.3     0.3     0.9     0.0
  test sentences: "the river bank" / "the money bank"
  note: 'bank' has exactly ONE row here, whichever sentence it is in

Let's go deeper on what this table encodes. "river" is strongly watery (0.9) and somewhat a place (0.3). "money" is strongly money-ish. "the" carries no content at all, just a 1.0 on the filler axis, our stand-in for grammatical glue words. And "bank" is the interesting one: mostly a place (0.9), with a little water-ness (0.3) and a little money-ness (0.3), because the word alone, out of context, is genuinely ambiguous between riverbank and financial bank. The last log line states the fact the whole workshop turns on: this table has exactly one row for "bank". Whatever sentence it appears in, the lookup returns the same four numbers. Fixing that is the job of everything that follows.

Tool 1: the dot product measures agreement

We will constantly need to ask "how similar are these two words?". The tool for that is the dot product: multiply the two vectors slot by slot, then add the products up. For vectors $a$ and $b$ with four dimensions:

$$a·b=a_1{b}_1+a_2{b}_2+a_3{b}_3+a_4{b}_4$$

Why does this measure agreement? Each term $a_i{b}_i$ is large only when both vectors have a large value in the same slot. Overlap on an axis is rewarded, and axes where either vector is zero contribute nothing. Vectors pointing in similar directions therefore get a large positive dot product; vectors with nothing in common get zero; vectors pointing in opposite directions go negative:

The script computes these four dot products with np.dot, spelling the first one out term by term so you can check it by hand (cross-check: results.log lines 18-21):

  bank . river = 0.3*0.9 + 0.3*0.0 + 0.9*0.3 + 0.0*0.0 = 0.540
  bank . money = 0.360
  bank . the   = 0.000
  bank . bank  = 0.990

Let's go deeper, because two of these four numbers carry the rest of the workshop:

• bank . river = 0.540. Read the expanded form: the overlap comes from the water axis (0.3 x 0.9 = 0.27) and the place axis (0.9 x 0.3 = 0.27). The two words agree on "wet" and on "location", and the dot product noticed.
• bank . money = 0.360 is positive but smaller: they overlap on the money axis (0.27) and only barely on place (0.09).
• bank . the = 0.000. No shared axes at all: "the" lives entirely on the filler axis, where "bank" has a zero. The dot product calls them unrelated.
• bank . bank = 0.990 is the largest of all, and this is a general fact worth pausing on: a vector dotted with itself sums the squares of its entries, so every axis matches perfectly. Every word is maximally similar to itself. Keep this in your pocket; it becomes a real problem in Module 2.

One more piece of notation you will see everywhere: matrix multiplication, NumPy's @, is just many dot products done at once. If X is a matrix with one word-vector per row, then X @ X.T (where .T flips rows and columns) produces a table whose entry at row i, column j is the dot product of word i with word j: every pairwise similarity in one expression. That is the only way this workshop ever uses @.

Tool 2: softmax turns scores into shares

The plan for Module 1 is to build each word's new vector as a blend of all the words in the sentence, where more similar words contribute more. A blend needs mixing proportions: numbers that are all positive and add up to 1, like percentages. Raw dot products are neither (bank's scores against "the river bank" are 0.000, 0.540, and 0.990, which add up to 1.53). The converter from raw scores to proper shares is the softmax function:

$$\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(s{)}_i=\frac{e^{s_i}}{{\sum }_j{e}^{s_j}}$$

In words: raise $e$ (Euler's number, about 2.718) to each score, then divide each result by the sum of all of them. The exponential does two jobs at once. First, $e^x$ is always positive, even for negative or zero scores, so every word gets a nonzero share. Second, dividing by the total forces the shares to sum to exactly 1. Bigger scores always get bigger shares, and, because the exponential grows so fast, gaps between scores get amplified.

Here is the function exactly as the script defines it (this same six-line definition is reused by every experiment in the workshop):

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)

The script then feeds it bank's three scores, twice: once as measured, and once multiplied by 5 to expose a property we will need later:

    # 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)

and then, a few lines later (the log(...) calls in between write the capture below):

    # The same scores, five times larger: softmax turns winner-take-all.
    sharp = softmax(scores * 5.0)

Here is what comes out (cross-check: results.log lines 22-23):

  softmax([0.000, 0.540, 0.990]) = [0.185, 0.317, 0.498]  (sum = 1.000000)
  softmax(same scores x 5)       = [0.006, 0.095, 0.899]  (one clear winner)

Let's go deeper on both lines, because you can verify the first one by hand in under a minute. The three exponentials are $e^0=1$, $e^{0.54}\approx 1.716$, and $e^{0.99}\approx 2.691$; their sum is about 5.407; and dividing each by the sum gives 0.185, 0.317, and 0.498. Notice that "the", with a score of exactly zero, still receives an 18.5% share: $e^0$ is 1, not 0. Softmax never fully silences anyone (a fact Module 3 turns into a puzzle: what if we need a word silenced?).

The second line is the same three scores multiplied by 5, and the shares change character completely: 0.006, 0.095, 0.899. The order is the same, but the top word now takes almost everything. Softmax cares about the size of the gaps between scores, and scaling all the scores up widens every gap. File this away: in Module 2 this exact effect comes back at realistic model sizes, uninvited, and we will have to measure it and cancel it.

What you now have

Three tools, each verified on the toy world: embeddings turn words into vectors (one fixed row per word), the dot product turns two vectors into a similarity score (bank . river = 0.540, checkable term by term), and softmax turns a row of scores into a row of shares that sum to 1 (with a sharpness that grows as scores grow). Module 1 snaps the three together into a working attention mechanism.

Next: Module 1: The simplest attention that works