Learning to guess the next letter
A language model is trained by next-token prediction: it guesses the next token, sees the right one, and nudges its numbers to guess better. That one loop, repeated over a huge amount of text, is all the training there is.
Here you train a tiny model on one short line in a few seconds and watch its guesses improve.
By the end of this lesson
Note
I can train a tiny model to predict the next character and read its loss.
Key words
| Word | What it means | Example |
|---|---|---|
| Next-token prediction | guessing the token that comes next | |
| Target | the token that really comes next | e after t |
| Probability | the share a model gives one token, from 0 to 1 | |
| Loss | how surprised the model was by the target | |
| Training step | one small nudge to every score | |
| Learning rate | how big each nudge is | 2.0 |
The cells students fill in are left blank in this preview.
Every place in a text is a question
Take the line "tea team meat team tea". At each place, the question is: what comes next? The answer, the target, is the next character in the text.
So 22 characters give 21 questions, with no labelling by hand. That is why a model can learn from any text at all.
Worked example: the first six questions
| Place | This token | Target |
|---|---|---|
| 0 | t | e |
| 1 | e | a |
| 2 | a | space |
| 3 | space | t |
| 4 | t | e |
| 5 | e | a |
The cells students fill in are left blank in this preview.
Our model looks back one token
Our tiny model is a table of scores with one row for each token and one column for each possible next token. Softmax turns a row into probabilities. Five characters make a table of 25 scores.
This model only sees the token just before. A transformer sees every earlier token through masked attention. The way it is trained is the same.
The loss measures surprise
The loss for one question is minus the log of the probability the model gave the target. A sure, right guess costs almost nothing. A small probability on the target costs a lot.
At the start every score is 0, so each of the 5 characters gets 0.2. The loss is then 1.61 on every question.
Worked example: probability on the target and loss
| Probability on the target | Loss |
|---|---|
| 1.0 | 0 |
| 0.9 | 0.11 |
| 0.5 | 0.69 |
| 0.2 | 1.61 |
| 0.1 | 2.30 |
The cells students fill in are left blank in this preview.
A training step nudges the scores
For each question, look at the row of this token. Every score in it goes down by its probability. The target's score also goes up by 1.
Add these nudges over all 21 questions and divide by 21. Multiply by the learning rate and apply them. That is one training step. Repeat it, and the loss falls.
Train the tiny model for 50 steps
import math
text = "tea team meat team tea"
chars = sorted(set(text)) # the vocabulary: space, a, e, m, t
n = len(chars)
pairs = len(text) - 1 # 21 pairs of this token and the next
table = [[0.0] * n for i in range(n)] # every score starts at 0
rate = 2.0
for step in range(51):
nudge = [[0.0] * n for i in range(n)]
loss = 0
for k in range(pairs):
a = chars.index(text[k]) # this token
b = chars.index(text[k + 1]) # the right next token
powers = [math.exp(s) for s in table[a]]
total = sum(powers)
loss = loss - math.log(powers[b] / total)
for j in range(n):
nudge[a][j] += powers[j] / total - (1 if j == b else 0)
if step % 10 == 0:
print("step", step, "loss", round(loss / pairs, 3))
for a in range(n):
for j in range(n):
table[a][j] -= rate * nudge[a][j] / pairs
powers = [math.exp(s) for s in table[chars.index("t")]]
print("after t:", [(chars[j], round(powers[j] / sum(powers), 2)) for j in range(n)])
# Change rate to 0.5 and press Run. Does the loss fall faster or slower?
Worked example: the loss as training goes on
| Step | 0 | 10 | 20 | 30 | 40 | 50 |
|---|---|---|---|---|---|---|
| Loss | 1.609 | 0.779 | 0.660 | 0.615 | 0.591 | 0.576 |
The cells students fill in are left blank in this preview.
A common mistake
Warning
Many people think a model stores the sentences it read and looks them up. Ours stores 25 scores and no text. It learnt that t is usually followed by e. A model holds patterns in numbers, and that is also why it can be confidently wrong.
Write the training script into the lab
docker exec jupyter sh -c "printf '%s\n' 'import numpy as np' 'text = \"tea team meat team tea\"' 'chars = sorted(set(text))' 'ids = np.array([chars.index(c) for c in text])' 'x, y = ids[:-1], ids[1:]' 'W = np.random.default_rng(0).normal(0, 0.1, (5, 5))' 'for step in range(101):' ' p = np.exp(W[x]) / np.exp(W[x]).sum(axis=1, keepdims=True)' ' if step % 25 == 0:' ' print(\"step\", step, \"loss\", round(-np.log(p[np.arange(21), y]).mean(), 3))' ' p[np.arange(21), y] -= 1' ' np.add.at(W, x, -2.0 * p / 21)' 'p = np.exp(W[4]) / np.exp(W[4]).sum()' 'print(\"after t:\", dict(zip(chars, p.round(2).tolist())))' > steps/tiny.py"Shown here, not run. Create a free account to press Run and type it into your own lab’s terminal.
Train from random scores with a seed
docker exec jupyter python3 steps/tiny.pyShown here, not run. Create a free account to press Run and type it into your own lab’s terminal.