Words as numbers

Your phone's keyboard guesses your next word. It has seen a lot of text, and it knows which words tend to follow which. A language model makes this 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 uses this process too, with much more training and a much larger model.

We'll build the smallest version of that guess. It reads one letter and guesses the next one, and it learns from 32,000 first names.

First, a problem. A computer can only do arithmetic, and a letter is not a number. So we give each letter a number, the way a takeaway menu numbers its dishes.

Cutting text into pieces

At a takeaway you can order a 12 and a 31 without spelling out the dishes. Text works the same way. Cut the text into pieces, give each piece its number from a menu, and the computer works with the numbers.

The computer reads each piece as a unit and looks up its number. That piece is a token. The list of all the pieces it knows is its vocabulary. The program that cuts the text and looks up those numbers is its tokenizer.

There is more than one way to cut. The demo below cuts one sentence three ways. Each row says how many items its menu has and how many pieces the sentence became. Every menu was learned from the words of this note, so a word this note never uses is not on the words menu. A digit or a mark, such as a comma, is on no menu at all. Type your own sentence to try it.

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 is already longer than the pieces menu. A new word is not 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 it to the menu. It does that 200 times, so the menu ends at 226 pieces. A word that is not 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 demo keeps only lowercase English letters, so its missing marks are a limitation of this toy.

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, and 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%.

Every count starts at one, even for a pair that no name has, so every letter keeps a small chance. That is why u gets 69.2% after q, not the 76% that 206 out of 272 would give. The 26 other 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 is measured. For now, smaller is better, and 2.45 is the number to beat.

For readers who know JavaScript, here is the code the demo runs. It is TypeScript, like the code in part 1. The tokenizer gives each letter its place in symbols, so the dot is 0 and z is 26. Anything else, such as a capital letter or a space, stops it with an error.

export const symbols = '.abcdefghijklmnopqrstuvwxyz';

export function encode(name: string) {
  const ids = [0];
  for (const letter of name) {
    const id = symbols.indexOf(letter);
    if (id < 1) throw new Error(`${JSON.stringify(letter)} is not a letter from a to z`);
    ids.push(id);
  }
  ids.push(0);
  return ids;
}

countPairs keeps all 729 counts in one long array, row after row, so the count for a pair sits at its row times 27 plus its column.

export const size = symbols.length;

export function countPairs(names: string[]) {
  const counts = new Uint32Array(size * size).fill(1);
  const dot = symbols.indexOf('.');
  counts[dot * size + dot] = 0;
  for (const name of names) {
    if (name === '') continue;
    const ids = encode(name);
    for (let i = 1; i < ids.length; i++) {
      counts[ids[i - 1] * size + ids[i]] += 1;
    }
  }
  return counts;
}

chances loads the die for one row at a temperature, and sampleName rolls it until the end dot comes up.

export function chances(counts: Uint32Array, row: number, temperature = 1) {
  const sides = Array.from(counts.subarray(row * size, (row + 1) * size), (count) =>
    Math.pow(count, 1 / temperature),
  );
  const total = sides.reduce((sum, side) => sum + side);
  return sides.map((side) => side / total);
}

export function sampleName(counts: Uint32Array, temperature: number, random = Math.random) {
  let name = '';
  let row = 0;
  while (true) {
    const odds = chances(counts, row, temperature);
    let roll = random();
    let next = 0;
    while (next < size - 1 && roll >= odds[next]) {
      roll -= odds[next];
      next += 1;
    }
    if (next === 0) return name;
    name += symbols[next];
    row = next;
  }
}

To run them, save the three blocks in order as names.mts, and save names.txt, the 32,033 names the demo counts, next to it. Then add these lines at the end of names.mts.

import fs from 'node:fs';

const names = fs.readFileSync('names.txt', 'utf8').split('\n');
const counts = countPairs(names);
for (let i = 0; i < 10; i++) console.log(sampleName(counts, 1));

Run node names.mts, and it prints ten new names, like the sample button does. That needs Node 22.18 or later.

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 is not 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 is 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 moves the scores a little downhill with adam, a method that uses recent slopes as well as the current one. Its updates 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 only one letter knows no more than the counts do, and its trained rows give nearly the same chances as the counting model's. In five trial runs, the slowest took 881 steps to come within 0.02 of 2.45, so training stops at twice that, 1,762 steps.

Here is the code. 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 overflowing. crossEntropy measures the surprise, and surpriseSlopes gives the slope of the surprise for each score. That slope is just the letter's chance, less 1 for the letter that really came next.

export function softmax(scores: Floats, chances: Float64Array) {
  let biggest = -Infinity;
  for (let i = 0; i < scores.length; i++) biggest = Math.max(biggest, scores[i]);
  let total = 0;
  for (let i = 0; i < scores.length; i++) {
    chances[i] = Math.exp(scores[i] - biggest);
    total += chances[i];
  }
  for (let i = 0; i < scores.length; i++) chances[i] /= total;
}

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

export function surpriseSlopes(chances: Float64Array, next: number, slopes: Float64Array) {
  for (let i = 0; i < chances.length; i++) slopes[i] = chances[i] - (i === next ? 1 : 0);
}

The one-letter model keeps its scores in p.tensors.scores. Its data holds the numbers and its grad holds their slopes, like part 1's Value, but for a whole table at once.

export const oneLetterKnobs = () => ({ scores: zeros([size, size]) });

export function oneLetterLoss(p: OneLetter, batch: Examples) {
  const { scores } = p.tensors;
  const chances = new Float64Array(size);
  const slopes = new Float64Array(size);
  const share = 1 / batch.count;
  p.grad.fill(0);
  let total = 0;
  for (let e = 0; e < batch.count; e++) {
    const row = batch.before[e] * size;
    softmax(scores.data.subarray(row, row + size), chances);
    total += crossEntropy(chances, batch.next[e]);
    surpriseSlopes(chances, batch.next[e], slopes);
    for (let j = 0; j < size; j++) scores.grad[row + j] += slopes[j] * share;
  }
  return total * share;
}

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 does not mean it is 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. The model reads the places of the three letters, nine numbers in all. A layer reads those nine numbers, as in part 1, but with 32 neurons instead of eight. Each neuron squashes its sum with tanh. Then 27 more neurons turn those outputs into a score for each next letter. Letters that turn up in the same kinds of places often drift close together, because the model can treat them alike. Then what it learns about one letter helps with the others.

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. Drag the map or use the arrow keys to turn it. Watch whether the vowels find each other.

At the start every letter sits in a small heap in the middle, because every number starts close to 0. After about 300 steps the loss passes the counting model's 2.45. From there it falls slowly and wobbles, because each step practices on 128 letters picked at random, and some picks are harder than others. It starts to dip below 2.3 after about 5,000 steps, and after about 30,000 it stays near 2.28. The vowels find each other. The five of them end up closer to one another than to most other letters, and y's nearest letter is i. The letter q, which is rare and comes before u three times out of four, ends up far from everything, so the map pulls it closer and marks it with an arrow.

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.

Here is the three-letter model. threeLetterScores reads the three places and works out the 27 scores. threeLetterLoss then walks back from the scores to every number, the way part 1's backward walk goes back up the receipt. On the way it passes back through tanh, whose slope is 1 minus its output squared.

export const threeLetterKnobs = ({ hidden }: { hidden: number }) => ({
  places: random([size, 3], 0.1),
  hiddenWeights: random([9, hidden], 1 / 3),
  hiddenBias: zeros([hidden]),
  scoreWeights: random([hidden, size], 0.1 / Math.sqrt(hidden)),
  scoreBias: zeros([size]),
});

export function threeLetterScores(
  p: ThreeLetter,
  before: Uint8Array,
  input: Float64Array,
  hidden: Float64Array,
  scores: Float64Array,
) {
  const { places, hiddenWeights, hiddenBias, scoreWeights, scoreBias } = p.tensors;
  const width = hidden.length;
  for (let letter = 0; letter < 3; letter++) {
    for (let k = 0; k < 3; k++) input[letter * 3 + k] = places.data[before[letter] * 3 + k];
  }
  for (let h = 0; h < width; h++) {
    let sum = hiddenBias.data[h];
    for (let i = 0; i < 9; i++) sum += input[i] * hiddenWeights.data[i * width + h];
    hidden[h] = Math.tanh(sum);
  }
  for (let j = 0; j < size; j++) {
    let sum = scoreBias.data[j];
    for (let h = 0; h < width; h++) sum += hidden[h] * scoreWeights.data[h * size + j];
    scores[j] = sum;
  }
}

export function threeLetterLoss(p: ThreeLetter, batch: Examples) {
  const { places, hiddenWeights, hiddenBias, scoreWeights, scoreBias } = p.tensors;
  const width = hiddenBias.data.length;
  const input = new Float64Array(9);
  const hidden = new Float64Array(width);
  const scores = new Float64Array(size);
  const chances = new Float64Array(size);
  const slopes = new Float64Array(size);
  const hiddenSlopes = new Float64Array(width);
  const inputSlopes = new Float64Array(9);
  const share = 1 / batch.count;
  p.grad.fill(0);
  let total = 0;
  for (let e = 0; e < batch.count; e++) {
    const before = batch.before.subarray(e * 3, e * 3 + 3);
    threeLetterScores(p, before, input, hidden, scores);
    softmax(scores, chances);
    total += crossEntropy(chances, batch.next[e]);
    surpriseSlopes(chances, batch.next[e], slopes);
    for (let j = 0; j < size; j++) scoreBias.grad[j] += slopes[j] * share;
    for (let h = 0; h < width; h++) {
      let sum = 0;
      for (let j = 0; j < size; j++) {
        scoreWeights.grad[h * size + j] += hidden[h] * slopes[j] * share;
        sum += scoreWeights.data[h * size + j] * slopes[j];
      }
      hiddenSlopes[h] = sum * (1 - hidden[h] * hidden[h]);
      hiddenBias.grad[h] += hiddenSlopes[h] * share;
    }
    for (let i = 0; i < 9; i++) {
      let sum = 0;
      for (let h = 0; h < width; h++) {
        hiddenWeights.grad[i * width + h] += input[i] * hiddenSlopes[h] * share;
        sum += hiddenWeights.data[i * width + h] * hiddenSlopes[h];
      }
      inputSlopes[i] = sum;
    }
    for (let letter = 0; letter < 3; letter++) {
      for (let k = 0; k < 3; k++) {
        places.grad[before[letter] * 3 + k] += inputSlopes[letter * 3 + k] * share;
      }
    }
  }
  return total * share;
}

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 choices can look like names without being real names. A likely continuation is not a checked fact.

Three letters still leave out most of a sentence. Your phone can use a word from much further back to make its next suggestion. Part 3 shows how a model learns which earlier pieces to read when it makes the next guess.