Words as numbers

Your phone's keyboard guesses your next word. It has seen a lot of text, so it knows which words tend to follow which. A language model makes the same kind of guess, one piece of text at a time. To write, it picks a piece, adds it to the text so far, and guesses again. ChatGPT works this way too. It just has a much larger model and much more training behind it.

In this note we'll build the smallest version of that guess: a model that reads one letter and guesses the next, after learning from about 32,000 first names.

But first, a problem. A computer can only do arithmetic, and a letter isn't a number.

Cutting text into pieces

Think of a takeaway menu. You can order a 12 and a 31 without spelling out the dishes. Text works the same way. Cut it into pieces, give each piece its number from a menu, and let the computer work with the numbers.

Each piece is a token. The whole menu of pieces is the vocabulary. The program that does the cutting and the looking up is the tokenizer.

There's more than one way to cut, and the demo below cuts one sentence three ways. Type your own sentence and see which menu copes best.

Letters make a tiny menu but long orders. Whole words make short orders, but a words menu needs every word anyone might write. The one in the demo knows only the words of this note, and it's already longer than the pieces menu. A new word isn't on it at all.

Pieces of words sit in between. The pieces menu starts as the 26 letters. Then it finds the two pieces that sit side by side most often in this note, glues them into one new piece, and adds that piece to the menu. It does that 200 times, so the menu ends at 226 pieces. A word that isn't on the words menu still comes out as a few pieces that are. OpenAI's tiktoken tokenizer uses the same idea with a much bigger menu. It also handles spaces, punctuation, and other languages. This toy knows only lowercase English letters, which is why digits and commas get marked.

From here on, this note uses letters, because their menu fits on one screen. Names need the 26 letters and one more token, a dot that marks where a name starts and ends. That makes 27 tokens.

The counting model

Now the guess itself. Guessing what comes next is next-token prediction, and the simplest way to do it is to count.

Go through 32,033 names and count which letter follows which. The names come from Andrej Karpathy's makemore dataset, shared under the MIT license. In "emma", the dot is followed by e, e by m, m by m, m by a, and a by the dot. After "q" comes "u" 206 times out of 272, about three times out of four.

The grid has every count. Each row is a letter and each column is the letter after it, so each square is one pair, as big as its count. The dot row says which letters start names, and the dot column says which letters end them.

Type a letter and its row lights up. A list shows each next letter's share of the row, its count divided by the row's total. That share is the letter's probability, the model's chance that this letter comes next. The shares in a row add up to 100%.

One small cheat: every count starts at one, even for pairs no name has, so no letter's chance is ever zero. That's why u gets 69.2% after q in the demo, a bit less than the 76% you'd get from 206 out of 272. The other 26 choices share the rest. The one pair that stays at zero is the end dot right after the start dot, because a name has at least one letter.

To write a name, the model rolls a loaded die. Each letter gets sides in proportion to its share in the current row. Roll, write the letter down, move to that letter's row, and roll again until the end dot comes up. That is sampling. Press sample to write ten names.

You can change how often the die picks an unlikely letter without changing any counts. Turn the control down and the letters the model already likes get even more sides, so the names come out samey. Turn it up and the choices spread out, so the names get strange. This control is called temperature. It changes how we pick from the model, not what the model has learned.

The number under the grid is the counting model's loss on all 32,033 names, 2.45. A later section explains how it's measured. For now, smaller is better, and 2.45 is the number to beat.

Learning instead of counting

Counting works when the model reads one letter. It gets hard as soon as the model has to look further back, as the next section shows. A network like the one in part 1 can learn the same guess instead, and it can grow from there.

For the letter it just read, the network gives every possible next letter a score. A score can be any number, big or small, positive or negative. Those raw scores are the logits.

A score isn't a chance yet. Chances can't be negative, and they have to add up to 100%. To turn scores into chances, raise a fixed number, about 2.72, to the power of each score. That gives a positive result even for a negative score and makes the big scores count a lot more. Divide each result by their total, and the shares add up to 100%. This conversion is called softmax. Scores of -1, 0, and 1 become chances of 9%, 24%, and 67%.

To learn, the network needs a loss, like the squares in part 1. For letters, we measure how surprised it was by the letter that really came next. Giving that letter a 50% chance is a small surprise, 0.69. Giving it a 1% chance is a big one, 4.6. This way of scoring its guesses is called cross-entropy. The counting model's 2.45 is this same surprise, averaged over every letter in the names. Training tries to lower the average surprise. An individual step can still make it rise. Part 1's playground scored its yes or no answers the same way, with two possible answers instead of 27.

The network in the demo is small. It's a table of 27 rows of 27 scores, one row for each letter it might read, and every score starts at 0. Each step, it reads 4,096 letter pairs picked at random from the names. It works out its average surprise and the slope of every score, then nudges every score a little downhill using Adam, the same method as part 1's playground. Those nudges need not follow the current downhill direction exactly. Press train.

At the start every score is 0, so every letter gets the same chance, 1 in 27, and the surprise is 3.30. The loss falls fast, then flattens just above the counting model's 2.45. It can't go much lower. A model that sees one letter knows no more than the counts do. Its trained rows end up giving nearly the same chances as the counting model's.

Letters on a map

To do better, the model has to see further back. How far back it sees is its context. The counting model's context is one letter. Counting three letters back would need a row for every run of three letters, 19,683 rows, and 13,997 of those runs never appear in the names.

A letter's menu number is just a label. Giving q number 17 doesn't mean it's similar to the letter with number 18. Instead, give each letter a list of three numbers that training can adjust, like coordinates on a map. Letters that behave alike can sit close together, so what the model learns about one can help with the others. This learned list is the letter's embedding.

Give the model the lists for the last three letters, nine numbers in all. A layer of 32 neurons reads them, as in part 1, each squashing its sum with tanh. Then 27 more neurons turn that into a score for each possible next letter.

You've already seen one of these maps. The cloud at the top of every page on this site is a hand-built toy of the same idea, where nothing is trained. This demo trains one for real.

The map draws each letter at its three numbers, one number on each axis, while the model trains. A letter too far out to fit is pulled closer to the center, with an arrow toward where it really is. Choose a letter to follow its darker square. Its label stays above the map. Turn the map and watch whether the vowels find each other.

Watch the vowels. At the start every letter sits in a small heap in the middle, because every number starts close to 0. As training runs, the five vowels drift toward each other until they're closer to one another than to most other letters, and y's nearest neighbor is i. The letter q, which is rare and comes before u three times out of four, ends up far from everything. That's why the map pulls it in and marks it with an arrow.

Meanwhile the loss beats the counting model's 2.45 after about 300 steps, dips below 2.3 after about 5,000, and settles near 2.28 after about 30,000. It wobbles along the way because each step practices on just 128 random letters, and some picks are harder than others.

Beyond three letters

The model now has a way to read text as numbers, learn from the letter that followed, and write by making one choice after another. The names can look real without being real. A likely next letter isn't a checked fact, and that stays true all the way up to ChatGPT.

Three letters isn't much. Your phone can use a word from much further back. How does a model know which earlier words to look at? That's part 3.

For JavaScript readers

For readers who know JavaScript, these are short versions of the demos. Each block has its own data and runs on its own. Press Run to see the result, or Edit to change it.

Cut a sentence three ways. This version starts with two already learned merges, t with h, then th with e. merge joins matching neighbors in order. The demo learns a longer merge list from this note. This small list makes the cuts easy to follow.

function merge(parts, [left, right]) {
  const joined = [];
  for (let i = 0; i < parts.length; i++) {
    if (parts[i] === left && parts[i + 1] === right) {
      joined.push(left + right);
      i++;
    } else {
      joined.push(parts[i]);
    }
  }
  return joined;
}

function tokenize(text, merges) {
  const lowercase = text.toLowerCase();
  const words = lowercase.match(/[a-z]+|[0-9]+|[^\s\w]/g);
  const letters = [];
  const pieces = [];
  for (const word of words ?? []) {
    letters.push(...word);
    let parts = [...word];
    for (const pair of merges) {
      parts = merge(parts, pair);
    }
    pieces.push(...parts);
  }
  return { letters, words: words ?? [], pieces };
}

const merges = [
  ['t', 'h'],
  ['th', 'e'],
];
const cuts = tokenize('the thin cat', merges);
console.log('Letters:', cuts.letters.join(' | '));
console.log('Words:', cuts.words.join(' | '));
console.log('Pieces:', cuts.pieces.join(' | '));

For names, symbols gives each letter its number, with dot at 0 and z at 26. countPairs keeps a row of counts for every letter. It adds a start and end dot to each name, then counts each neighboring pair. Try adding a name to the list.

const symbols = '.abcdefghijklmnopqrstuvwxyz';

function countPairs(names) {
  const counts = [];
  for (let i = 0; i < symbols.length; i++) {
    counts.push(Array(symbols.length).fill(1));
  }
  counts[0][0] = 0;
  for (const name of names) {
    if (!name) continue;
    const letters = '.' + name + '.';
    for (let i = 1; i < letters.length; i++) {
      const before = symbols.indexOf(letters[i - 1]);
      const next = symbols.indexOf(letters[i]);
      counts[before][next] += 1;
    }
  }
  return counts;
}

const counts = countPairs(['emma', 'ella', 'anna']);
console.log('Start then e:', counts[0][5]);
console.log('m then m:', counts[13][13]);
console.log('a then end:', counts[1][0]);

chances loads the die for one row at a temperature, and sampleName rolls it until the end dot comes up. These counts come from just one name, "anna". The rolls are written out so Run gives the same result each time. Change one to a number from 0 up to, but not including, 1 to try a different roll.

const symbols = '.abcdefghijklmnopqrstuvwxyz';

function chances(counts, row, temperature = 1) {
  const sides = [];
  let total = 0;
  for (const count of counts[row]) {
    const share = count ** (1 / temperature);
    sides.push(share);
    total += share;
  }
  return sides.map((share) => share / total);
}

function sampleName(counts, temperature, random) {
  let name = '';
  let row = 0;
  while (true) {
    const shares = chances(counts, row, temperature);
    let roll = random();
    let next = 0;
    while (
      next < symbols.length - 1 &&
      roll >= shares[next]
    ) {
      roll -= shares[next];
      next++;
    }
    if (next === 0) return name;
    name += symbols[next];
    row = next;
  }
}

const counts = [];
for (let i = 0; i < symbols.length; i++) {
  counts.push(Array(symbols.length).fill(1));
}
counts[0][0] = 0;
const letters = '.anna.';
for (let i = 1; i < letters.length; i++) {
  const before = symbols.indexOf(letters[i - 1]);
  const next = symbols.indexOf(letters[i]);
  counts[before][next] += 1;
}
const rolls = [0.02, 0.5, 0.5, 0.05, 0];
let turn = 0;
const name = sampleName(counts, 1, () => rolls[turn++]);
const chance = chances(counts, 1)[14];
console.log('Name:', name);
console.log('Chance of n after a:', chance.toFixed(2));

This block counts three names and writes one new name, using the counts as the shares at temperature 1. Its three fixed rolls pick two letters and an end dot. The demo uses all 32,033 names and draws fresh rolls.

const symbols = '.abcdefghijklmnopqrstuvwxyz';

function countPairs(names) {
  const counts = [];
  for (let i = 0; i < symbols.length; i++) {
    counts.push(Array(symbols.length).fill(1));
  }
  counts[0][0] = 0;
  for (const name of names) {
    if (!name) continue;
    const letters = '.' + name + '.';
    for (let i = 1; i < letters.length; i++) {
      const before = symbols.indexOf(letters[i - 1]);
      const next = symbols.indexOf(letters[i]);
      counts[before][next] += 1;
    }
  }
  return counts;
}

function sampleName(counts, random) {
  let name = '';
  let row = 0;
  while (true) {
    let total = 0;
    for (const count of counts[row]) total += count;
    let roll = random();
    let next = 0;
    while (next < 26 && roll >= counts[row][next] / total) {
      const chance = counts[row][next] / total;
      roll -= chance;
      next++;
    }
    if (next === 0) return name;
    name += symbols[next];
    row = next;
  }
}

const counts = countPairs(['emma', 'ella', 'anna']);
const rolls = [0.1, 0.2, 0];
let turn = 0;
const name = sampleName(counts, () => rolls[turn++]);
console.log('Name:', name);

softmax turns scores into chances. It first takes the biggest score away from every score, which gives the same chances and keeps Math.exp from making numbers too big for the computer. crossEntropy measures the surprise, and surpriseSlopes gives the slope of the surprise for each score. That slope is the letter's chance, less 1 for the letter that really came next.

function softmax(scores) {
  const biggest = Math.max(...scores);
  const sides = [];
  let total = 0;
  for (const score of scores) {
    const share = Math.exp(score - biggest);
    sides.push(share);
    total += share;
  }
  return sides.map((share) => share / total);
}

const chances = softmax([-1, 0, 1]);
const shifted = softmax([999, 1000, 1001]);
const percentages = chances.map((chance) =>
  Math.round(chance * 100),
);
const same = shifted.map((chance) =>
  Math.round(chance * 100),
);
console.log('Chances in %:', percentages.join(', '));
console.log('After adding 1000:', same.join(', '));

function crossEntropy(chances, next) {
  return -Math.log(chances[next]);
}

function surpriseSlopes(chances, next) {
  const slopes = [...chances];
  slopes[next] -= 1;
  return slopes;
}

const chances = [0.1, 0.3, 0.6];
const surprise = crossEntropy(chances, 2);
const slopes = surpriseSlopes(chances, 2);
console.log('Surprise:', surprise.toFixed(2));
console.log('Score slopes:', slopes.join(', '));

The one-letter model keeps a row of scores for every letter. oneLetterLoss reads the pairs in one name, measures their average surprise, and adds up the slopes in grad. This short version takes ordinary downhill steps so you can see the update directly. The demo trains on a handful of names at a time and uses Adam.

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

function oneLetterLoss(scores, pairs) {
  const grad = [];
  for (const row of scores)
    grad.push(Array(row.length).fill(0));
  let loss = 0;
  for (const [before, next] of pairs) {
    const chances = softmax(scores[before]);
    loss -= Math.log(chances[next]) / pairs.length;
    for (let j = 0; j < chances.length; j++) {
      let slope = chances[j];
      if (j === next) slope -= 1;
      grad[before][j] += slope / pairs.length;
    }
  }
  return { loss, grad };
}

const scores = [];
for (let i = 0; i < 27; i++) scores.push(Array(27).fill(0));
// The pairs in ".anna.": dot = 0, a = 1, n = 14.
const pairs = [
  [0, 1],
  [1, 14],
  [14, 14],
  [14, 1],
  [1, 0],
];
for (let i = 0; i <= 20; i++) {
  const { loss, grad } = oneLetterLoss(scores, pairs);
  if (i % 10 === 0) {
    console.log('Step', i, 'loss:', loss.toFixed(2));
  }
  for (let t = 0; t < scores.length; t++) {
    for (let j = 0; j < scores[t].length; j++) {
      scores[t][j] -= grad[t][j];
    }
  }
}

threeLetterScores looks up the last three letters' lists, joins their nine numbers, and passes them through two layers. This version has two neurons in the middle layer and fixed knobs, so it prints scores before any training. The demo has 32 neurons in the middle layer and learns the knobs, including the letter lists, by walking backward from the loss.

function layer(input, weights, bias) {
  const scores = [];
  for (let j = 0; j < bias.length; j++) {
    let sum = bias[j];
    for (let i = 0; i < input.length; i++) {
      sum += input[i] * weights[i][j];
    }
    scores.push(sum);
  }
  return scores;
}

function threeLetterScores(before, knobs) {
  const input = [];
  for (const letter of before) {
    input.push(...knobs.places[letter]);
  }
  const sums = layer(
    input,
    knobs.hiddenWeights,
    knobs.hiddenBias,
  );
  const hidden = sums.map(Math.tanh);
  return layer(hidden, knobs.scoreWeights, knobs.scoreBias);
}

const places = [];
const hiddenWeights = [];
const scoreWeights = [Array(27).fill(0.3), []];
for (let i = 0; i < 27; i++) {
  places.push([i / 27, 0.1, -i / 27]);
  scoreWeights[1].push(i / 27);
}
for (let i = 0; i < 9; i++) {
  hiddenWeights.push([(i + 1) / 10, -0.1]);
}
const knobs = {
  places,
  hiddenWeights,
  hiddenBias: [0, 0.1],
  scoreWeights,
  scoreBias: Array(27).fill(0),
};
const scores = threeLetterScores([1, 14, 14], knobs);
const letterScore = scores[1].toFixed(3);
const endScore = scores[0].toFixed(3);
console.log('After ann, score for a:', letterScore);
console.log('After ann, score for end:', endScore);