Skip to content
wannacode

Predicting the next token

In short: A language model writes text by computing the probabilities of the next token and picking one; temperature and top-k decide how random the choice is.

Problem

One start, three answers

The support bot has to continue replies. A colleague runs the same start, "the order", three times with a different seed for the generator. Here are the three answers:

const table = countNext(TEXT);
for (const seed of [2026, 34567, 12345])
  console.log(`seed ${seed}: ${generate(table, "the order", 1, 100, new Rng(seed))}`);
Output
seed 2026: the order will arrive within two days
seed 34567: the order arrived at the order will arrive within five days
seed 12345: the order will come to two days

The first sounds reasonable. The second runs on into nonsense. The third is almost a sentence, but not quite. Every two neighbouring words in all three are from the shop text. And all three come from the same program, which does not look up an answer anywhere: it writes word by word.

How does it know which word to put? And why is it different every time?

The idea

The probability of the next word

A language model takes text and gives a probability for every possible next token. Nothing more: it does not know what an order is, it only knows what usually comes after what.

Text is written with a loop:

  1. The model gives the probabilities of the next token.
  2. One token is chosen by them.
  3. It is added to the text and everything repeats, until the end or the limit.

The choice is a separate step, not part of the model. That is why the same model can give different answers.

The model in this lesson is the smallest possible: a table of how many times each word came right after each word. It looks only at the last word and counts whole words, not tokens. Real models are neural networks that look at the whole text, but they give the same thing: the probabilities of the next token.

See it

The table: which word after which

The shop text is the same as in the lesson on tokens. For every sentence we count every pair of neighbouring words. The dot "." marks the end of the sentence, so that the model knows when to stop.

// how many times each word comes right after each word; "." ends a sentence
function countNext(text) {
  const table = new Map();
  for (const line of text) {
    const w = `${line} .`.split(" ");
    for (let i = 0; i + 1 < w.length; i++) {
      if (!table.has(w[i])) table.set(w[i], new Map());
      const counts = table.get(w[i]);
      counts.set(w[i + 1], (counts.get(w[i + 1]) ?? 0) + 1);
    }
  }
  return table;
}

The probability of the next word is its count divided by the sum. For now look only at t = 1 and "all words": then the weight is the count itself. The temperature t and k come further down.

// the k most frequent, weight count^(1/t) / sum
function nextWords(counts, t, k) {
  const top = [...counts.keys()]
    .sort((a, b) => counts.get(b) - counts.get(a))
    .slice(0, k);
  let total = 0;
  for (const w of top) total += counts.get(w) ** (1 / t);
  return top.map((w) => {
    const weight = counts.get(w) ** (1 / t);
    return { word: w, p: weight / total };
  });
}
const table = countNext(TEXT);
const counts = table.get("will");
console.log('after "will":');
for (const c of nextWords(counts, 1, 100))
  console.log(c.word.padEnd(11) + String(counts.get(c.word)).padStart(2) + "  " + (c.p * 100).toFixed(1).padStart(5) + "%");
Output
after "will":
arrive      4   40.0%
be          3   30.0%
come        1   10.0%
call        1   10.0%
get         1   10.0%

After "will" the text has five different words, 10 times in all. "arrive" is 4 of them: 40.0%. This is the language model: for "will" it says "arrive" with 40.0%, "be" with 30.0% and each of the others with 10.0%.

See it

The choice: a number from 0 to 1

How do you choose a word with probability 40.0%? The seeded generator gives a number r from 0 to 1. The bars of the words stand side by side: "arrive" from 0 to 0.4, "be" from 0.4 to 0.7, and so on. The word whose bar r falls into is the chosen one.

class Rng {
  constructor(seed) {
    this.state = seed;
  }


  // the next number between 0 and 1; one seed gives the same numbers
  next() {
    this.state = (this.state * 48271) % 2147483647;
    return this.state / 2147483647;
  }
}
// a number from 0 to 1 falls into one of the bars
function pick(choices, rng) {
  const r = rng.next();
  let sum = 0;
  for (const c of choices) {
    sum += c.p;
    if (r < sum) return c.word;
  }
  return choices[choices.length - 1].word;
}

Generation repeats this word by word from the start, until "." comes or 12 words:

// continue the text word by word until "." or 12 words
function generate(table, start, t, k, rng) {
  const text = start.split(" ");
  let word = text[text.length - 1];
  while (text.length < 12) {
    word = pick(nextWords(table.get(word), t, k), rng);
    if (word === ".") break;
    text.push(word);
  }
  return text.join(" ");
}

If you always take the most likely word (k = 1), the generator has nothing to choose and the answer is the same for every seed:

const table = countNext(TEXT);
for (const seed of [2026, 34567, 12345])
  console.log(`seed ${seed}: ${generate(table, "the order", 1, 1, new Rng(seed))}`);
Output
seed 2026: the order will arrive within three days
seed 34567: the order will arrive within three days
seed 12345: the order will arrive within three days

With a random choice the answer depends on the seed. But the same seed gives the same numbers, so the same answer too:

const table = countNext(TEXT);
for (const seed of [2026, 2026, 34567])
  console.log(`seed ${seed}: ${generate(table, "the order", 1, 100, new Rng(seed))}`);
Output
seed 2026: the order will arrive within two days
seed 2026: the order will arrive within two days
seed 34567: the order arrived at the order will arrive within five days

That is why the three answers at the start were different: each seed gives other numbers, and so other words.

Try it

Predict the output

After "delivery" the text has five continuations, one of them the end ".". What does the program print with k = 1?

const table = countNext(TEXT);
for (const c of nextWords(table.get("delivery"), 1, 1))
  console.log(`${c.word} ${(c.p * 100).toFixed(1)}%`);
Answer

The idea

Temperature and top-k

The choice has two knobs, and real models have them too:

  • temperature t: the weight is count^(1/t). With t below 1 the frequent words get even more likely and the answer is more predictable. With t above 1 the probabilities even out and rare words come up more often.
  • top-k: only the k most frequent words stay, and the others cannot be chosen. The sum is taken over them only, so it is 100% again.
const table = countNext(TEXT);
const counts = table.get("will");
const temps = [0.5, 1, 2];
const rows = temps.map((t) => nextWords(counts, t, 100));
console.log(" ".repeat(11) + temps.map((t) => `t=${t}`.padStart(7)).join(""));
for (let i = 0; i < counts.size; i++)
  console.log(rows[0][i].word.padEnd(11) + rows.map((row) => ((row[i].p * 100).toFixed(1) + "%").padStart(7)).join(""));
Output
             t=0.5    t=1    t=2
arrive       57.1%  40.0%  29.7%
be           32.1%  30.0%  25.7%
come          3.6%  10.0%  14.9%
call          3.6%  10.0%  14.9%
get           3.6%  10.0%  14.9%

At t = 0.5 "arrive" is 57.1%, and at t = 2 only 29.7%. The words seen once are 3.6%, 10.0% and 14.9%: the higher the temperature, the more often the model chooses them.

Here are the bars after "will" at the three temperatures, then with k = 3 and the choice with one number from the generator:

// the k most frequent, weight count^(1/t) / sum
function nextWords(counts, t, k) {
  const top = [...counts.keys()]
    .sort((a, b) => counts.get(b) - counts.get(a))
    .slice(0, k);
  let total = 0;
  for (const w of top) total += counts.get(w) ** (1 / t);
  return top.map((w) => {
    const weight = counts.get(w) ** (1 / t);
    return { word: w, p: weight / total };
  });
}


// a number from 0 to 1 falls into one of the bars
function pick(choices, rng) {
  const r = rng.next();
  let sum = 0;
  for (const c of choices) {
    sum += c.p;
    if (r < sum) return c.word;
  }
  return choices[choices.length - 1].word;
}
Step by step, in words
  1. After "will" the text has 5 words. A bar is the count divided by all 10: "arrive" is 4 of 10, 40.0%.
  2. t = 0.5: the weight is count². The frequent word gets even more likely: "arrive" 57.1%.
  3. t = 2: the weight is √count. The bars even out: "arrive" 29.7%, "get" 14.9%.
  4. t = 1, k = 3: only the three most frequent stay, and together they are 100% again: 50.0%, 37.5%, 12.5%.
  5. The three bars stand side by side from 0 to 1. The generator with seed 31337 gives 0.7044.
  6. 0.7044 falls into the bar of "be" (from 50.0% to 87.5%): the next word is "be", not the most frequent one.
const table = countNext(TEXT);
const choices = nextWords(table.get("will"), 1, 3);
console.log("k=3:");
for (const c of choices) console.log(c.word.padEnd(11) + (c.p * 100).toFixed(1).padStart(5) + "%");
const rng = new Rng(31337);
const copy = new Rng(31337);
console.log(`number from the generator: ${copy.next().toFixed(4)}`);
console.log(`chosen: ${pick(choices, rng)}`);
Output
k=3:
arrive      50.0%
be          37.5%
come        12.5%
number from the generator: 0.7044
chosen: be

With k = 3 "arrive", "be" and "come" stay, and their probabilities are 4, 3 and 1 out of 8. The generator with seed 31337 gives 0.7044, and that falls into the bar of "be", from 50.0% to 87.5%. The choice follows the probabilities, but it is not always the most likely word.

Try it

Fill in

What goes in place of ___ in the weight of a word?

for (const w of top) total += counts.get(w) ** (1 / );

Try it

Find the bug

A colleague wants top-k with the three most likely words, but their model writes only rare words. Which line is the bug on?

Pick the line with the bug.

The idea

When it hurts

  • Probable is not true. The model writes "within three days" as easily as "within two days": both are in the text. It does not know when this particular order arrives. Smooth text sounds confident even when it is wrong. That is what the lesson on the limits of the model is about.
  • Every pair is right, the sentence is not. "the order arrived at the order will arrive within five days" is made of neighbouring words that exist. Our model looks only one word back. Real ones look at the whole text and write much better, but they can still glue together something smooth and meaningless.
  • Higher temperature, more mistakes. Rare continuations come up more often. For answers to customers, where accuracy matters, the temperature is low. For ideas and variants it can be higher.
  • One seed is not a promise. In the lesson one seed gives one answer. With a real model behind an API the same request sometimes gives a different answer even with the same settings, so do not rely on getting exactly the same text: check the answer.
More precisely, if you are curious

Softmax and temperature. A real model is a neural network: for every token in the vocabulary it gives a number (a logit), and softmax turns them into probabilities. Temperature divides the numbers by t before softmax. The weight count^(1/t) in the lesson is exactly that: softmax of log(count) / t.

top-p. Besides top-k there is top-p: the most frequent tokens stay until their sum reaches p (for example 0.9). So when the model is sure, few tokens stay, and when it is unsure, more.

Tokens, not words. A real model chooses tokens from its vocabulary, and a word comes out of several choices in a row. Temperature, top-k and the seed are set in the request to the model (the lesson on the call).

How far back it looks. Our model remembers one word. A real one sees the whole conversation up to its limit, and that is the context window.

Exercise

How likely is the sentence

Solve it in your editor.

The probability of a whole sentence is the product of the probabilities of each next word, including the "." at the end. Write a function that computes it with the table from the lesson, and print it in percent for three sentences. If a pair of words has never been seen, its probability is 0.

Output
the parcel will arrive within two days: 0.27%
the parcel will arrive in the original box: 0.04%
the parcel arrived tomorrow: 0.00%

The first sentence is 0.27% and the second 0.04%, although neither is in the text word for word: all the pairs in them have been seen. The third is 0.00%: "parcel" and "arrived" are both in the text, but never one after the other. This is how the model says how "natural" a text looks to it.

More tasks

  • The most likely sentence. Start from "the parcel" and always take the most likely word. What is the probability of the sentence you get?
  • Two words back. Count which word comes after a pair of words, not after one. How do the answers from the start of the lesson change?

Test

Check yourself

1. What does a language model give you when you pass it text?
2. You run the generation twice with seed 2026. What do you see?
3. How does temperature t = 0.5 change the probabilities after "will"?
4. The bot writes "the order will arrive within three days", but delivery really takes up to two days. Why?

Made by Ivo Gergov as part of Evo Tech's campaign for accessible programming education ·