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 it's the bag.

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 thing, and the way it does it is called attention.

Part 2's best model could only see three letters back. Here the word that matters is way behind: between "car" and "small" sit 20 other letters and spaces. So the model needs two things, a bigger window and a way to pick what inside it matters. Attention does the picking. (It still can't see past the 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 is here to show why attention exists, and the demos show what the tiny model really looks at. From here on it works on letters, not words, but the idea is the same.

A soft lookup

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

Attention is softer. It 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 up the results. That's the dot product. For the lists (2, 1) and (3, 0), it's 2 × 3 + 1 × 0 = 6. Draw each list as an arrow, and the score is bigger when the arrows point the same way. Longer arrows count more.

Part 2's softmax turns those scores into shares that add up to 100%. Those shares are the attention weights. The word then takes a mix of all the values, each scaled by its weight, and that mix is what it walks away with.

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 the thick arrow. 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.

No peeking

When the model writes, it writes one letter at a time, and the next letter doesn't 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. 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 doesn't 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 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. The second head of the first block mostly watches the character just before. In this sentence, that's where it looks hardest at 39 of the 53 positions that have one. The fourth head of the second block has learned to find the start of the current 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 a short version of one attention head. The example has two positions, each with two numbers in its query, key, and value. It uses the queries as keys too. For each position t, it scores the keys from 0 to t with a dot product, turns those scores into weights with part 2's softmax, and mixes the values. The loop stops at j <= t, so it never reads ahead. Dividing the dot product by the square root of the list's length keeps longer lists from making the scores too big for softmax. The tiny model's code stores the three sets of knobs side by side in one table. Press Run to see the last position's weights and mix, or Edit to try other vectors.

function dot(a, b) {
  let sum = 0;
  for (let i = 0; i < a.length; i++) {
    sum += a[i] * b[i];
  }
  return sum;
}

function softmax(scores) {
  const biggest = Math.max(...scores);
  const sides = scores.map((s) => Math.exp(s - biggest));
  let total = 0;
  for (const side of sides) total += side;
  return sides.map((side) => side / total);
}

function attention(queries, keys, values) {
  const weights = [];
  const found = [];
  for (let t = 0; t < queries.length; t++) {
    const size = Math.sqrt(queries[t].length);
    const scores = [];
    for (let j = 0; j <= t; j++) {
      scores.push(dot(queries[t], keys[j]) / size);
    }
    const chances = softmax(scores);
    const mix = Array(values[0].length).fill(0);
    for (let j = 0; j <= t; j++) {
      for (let i = 0; i < mix.length; i++) {
        mix[i] += chances[j] * values[j][i];
      }
    }
    weights.push(chances);
    found.push(mix);
  }
  return { weights, found };
}

const queries = [
  [1, 0],
  [1, 1],
];
const values = [
  [2, 0],
  [1, 1],
];
const keys = queries;
const { weights, found } = attention(queries, keys, values);
const shares = weights[1].map((share) => share.toFixed(2));
const mix = found[1].map((number) => number.toFixed(2));
console.log('Last weights:', shares.join(', '));
console.log('Last mix:', mix.join(', '));

Attention is the part that chooses what to read. The rest of the block keeps that working when you stack more than one.

Keeping what was already there

function level(notebook) {
  let total = 0;
  for (const number of notebook) total += number;
  const mean = total / notebook.length;
  let spread = 0;
  for (const number of notebook) {
    spread += (number - mean) ** 2;
  }
  const variance = spread / notebook.length;
  const scale = Math.sqrt(variance + 1e-5);
  return notebook.map((number) => (number - mean) / scale);
}

function gelu(number) {
  const bend = number + 0.044715 * number ** 3;
  const squash = Math.tanh(Math.sqrt(2 / Math.PI) * bend);
  return 0.5 * number * (1 + squash);
}

function block(notebook) {
  const copy = level(notebook);
  // The first token reads only itself, with weight 1.
  const value = copy.map((number) => number * 0.25);
  const middle = [];
  for (let i = 0; i < notebook.length; i++) {
    const change = value[i] * 0.5;
    middle.push(notebook[i] + change);
  }
  const leveled = level(middle);
  const out = [];
  for (let i = 0; i < middle.length; i++) {
    const weighted = leveled[i] * 0.7;
    const hidden = gelu(weighted);
    const change = hidden * 0.5;
    out.push(middle[i] + change);
  }
  return { middle, out };
}

const result = block([1, 2]);
const middle = result.middle.map((number) =>
  number.toFixed(2),
);
const out = result.out.map((number) => number.toFixed(2));
console.log('After attention:', middle.join(', '));
console.log('After network:', out.join(', '));