Paying attention

Read this sentence. "The bag didn't fit in the car because it was too small." What was too small? The car. Now change one word. "The bag didn't fit in the car because it was too big." Now the bag was too big.

You did that without thinking. When you reached "small", you looked back at the earlier words and picked the one that fit. A model has to do the same, and the way it does it is called attention.

Part 2's best model saw three letters back. Here the word that matters is far behind. Between "car" and "small" sit 20 other letters and spaces. The model needs a larger window and a way to choose which pieces inside it matter. Attention makes that choice. It still cannot read beyond its context window.

The model in this series is too small to solve that sentence. It learned from Shakespeare's plays, one letter at a time. So the sentence explains why attention exists, and the demos show what the tiny model really looks at.

A soft lookup

Think of a library search. You type what you are looking for. That is the query. Each book has a label on its spine, its key, and what is inside, its value. A normal search hands back the one book whose label matches best.

Attention reads a little of every book, and more of the ones whose labels match better. To score a match, multiply two lists of numbers pair by pair and add the results. This calculation is the dot product. For the lists 2, 1 and 3, 0 that is 2 × 3 + 1 × 0 = 6. Draw each list as an arrow, and the total is bigger when the arrows point the same way. Longer arrows count more.

Part 2's softmax turns the match scores into shares that add up to 100%. Those shares are the attention weights. Then the word adds up the values, each scaled by its weight, and that mix is what it takes away.

In the demo, four earlier words from the sentence are arrows, and each arrow is that word's key. The thick arrow is the query for "small". Drag it, or focus it and use the arrow keys. The bars show how much "small" reads from each word, and the chart under them shows the four values and the mix that "small" gets back.

Move the query toward a word's arrow and watch its share change. Both the direction and the length of each arrow affect the match. Point it between two words and they share. Pull it in toward the middle and the shares even out, because every score shrinks toward 0. At the very middle each word gets 25%.

These arrows are made up, with two numbers each so they fit on a flat chart. In a real model, each word's query, key, and value come from its own list of numbers through three sets of learned knobs, the same knobs at every position. The tiny model starts each letter as a list of 32 numbers. Its code stores the three sets of knobs side by side in one table.

No peeking

When the model writes, it writes one letter at a time, and the next letter does not exist yet. So a letter may read itself and the letters before it, never the ones after. That rule is the causal mask.

The triangle below is the trained model reading the sentence. Each row is a letter doing the reading, and each column is a letter being read. A square's area is how much the row's letter reads the column's letter, and each row adds up to 100%. Point at a row to read it, or focus the triangle and use the up and down arrow keys. Nothing sits above the diagonal, because no letter reads ahead.

The model has eight of these triangles, and the next section shows why. This one often puts its biggest square on a space. Read the row for any letter after the first in a word, and the biggest square is usually the space just before that word.

Attention on its own does not know which letter came first. Shuffle the earlier letters, and each one gets the same score it had before. So each position also gets a learned list of numbers added to its letter's list, its position embedding, like a seat number. The tiny model has 64 seats, so it reads at most the last 64 letters.

Several readers at once

One attention makes one mix of what it reads. A model runs several side by side, each with its own knobs for queries, keys, and values, so each can learn to look for something different. Each one is a head, and together they are multi-head attention. The tiny model repeats attention followed by a small network twice. Each repeat is a block, with four heads, so the model has eight heads in all.

The demo loads the tiny model, trained ahead of time with part 4's trainer. Point at a letter, or focus the demo and use the left and right arrow keys, to see where each head looks from it. Each arc runs from the letter back to a letter it reads, and the darker the arc, the more it reads.

Nobody told the heads what to look for. Training found it. In this sentence, the second head of the first block gives its largest weight to the previous character at 39 of the 53 positions that have one. The fourth head of the second block often reads the space before a word. Choose a character after a word's first letter, with a space before the word. At 20 of those 29 positions, this head gives that space its largest weight.

The heads of the first block read close by. Averaged over the sentence after its first letter, each head puts about 80% to 85% of its weight on the letter itself and the two before it. So the four overlap. The heads of the second block reach further, 3.6 to 6.6 letters back on average over the same positions.

For readers who know JavaScript, here is the attention the tiny model runs. lookups holds each letter's query, key, and value side by side, and each head reads its own 8 of each letter's 32 numbers. For each letter t and each head, it scores every letter from 0 to t with a dot product, turns the scores into weights with softmax, and adds up the values. The loop stops at t, and that is the causal mask. It also divides each score by Math.sqrt(size), about 2.83 here, so that longer lists do not make the scores too big for softmax.

function attend(lookups: Floats, { length, width, heads }: Sizes) {
  const size = width / heads;
  const scale = 1 / Math.sqrt(size);
  const rows = lookups.length / (3 * width * length);
  const found = like(lookups, rows * length * width);
  const weights = like(lookups, rows * heads * length * length);
  for (let r = 0; r < rows; r++) {
    for (let t = 0; t < length; t++) {
      for (let h = 0; h < heads; h++) {
        const query = (r * length + t) * 3 * width + h * size;
        const row = ((r * heads + h) * length + t) * length;
        let biggest = -Infinity;
        for (let s = 0; s <= t; s++) {
          const key = (r * length + s) * 3 * width + width + h * size;
          let dot = 0;
          for (let i = 0; i < size; i++) dot += lookups[query + i] * lookups[key + i];
          weights[row + s] = dot * scale;
          if (weights[row + s] > biggest) biggest = weights[row + s];
        }
        let total = 0;
        for (let s = 0; s <= t; s++) {
          weights[row + s] = Math.exp(weights[row + s] - biggest);
          total += weights[row + s];
        }
        const out = (r * length + t) * width + h * size;
        for (let s = 0; s <= t; s++) {
          weights[row + s] /= total;
          const value = (r * length + s) * 3 * width + 2 * width + h * size;
          for (let i = 0; i < size; i++) found[out + i] += weights[row + s] * lookups[value + i];
        }
      }
    }
  }
  return { found, weights };
}

Keeping what was already there

function block(knobs: BlockKnobs, notebooks: Floats, sizes: Sizes): BlockTape {
  const { attention, network } = knobs;
  const first = level(notebooks, attention.level);
  const lookups = multiply(first.out, attention.lookup);
  const { found, weights } = attend(lookups, sizes);
  const middle = add(notebooks, multiply(found, attention.write));
  const second = level(middle, network.level);
  const hidden = multiply(second.out, network.hidden);
  const fired = gelu(hidden);
  const out = add(middle, multiply(fired, network.write));
  return { notebooks, first, lookups, weights, found, middle, second, hidden, fired, out };
}