Dominik Keller

Part two of two · The maths and the code

The same thing again, with the numbers

Part one described how a language model learns. This one does it, in about forty lines of code, with arithmetic you can check on your phone.

How to read this

Part one used pictures. This one uses numbers, and the numbers are small enough to verify by hand. Every value printed here came out of code that actually ran.

The setup

A model that knows four words

Real models have a vocabulary of around a hundred thousand words. Ours has four.

0 = what      1 = is      2 = two      3 = four

Words are stored as those numbers, so a sentence is a list of numbers.

"what is two four"   →   [0, 1, 2, 3]

The example is a question with an answer glued onto it, which is how instruction training data looks.

prompt:      what is two    ← a person types this
completion:  four           ← the model writes this

Finally, the model’s output. For each position in the sentence it gives a score to all four words, so you get a grid of sixteen numbers.

logits = [[0.0, 2.0, 1.0, 0.0],   # after reading "what"
          [0.0, 0.0, 1.0, 2.0],   # after reading "what is"
          [0.0, 0.0, 0.0, 3.0],   # after reading "what is two"
          [0.0, 0.0, 0.0, 0.0]]   # nothing left to predict

Those raw scores are called logits. I have made them up so we have something concrete to work with. In real training they come out of the model and change after every step.

Read the first row: after seeing what, the model scores is highest at 2.0, then two at 1.0, and is indifferent about the rest.

Step one

Throw away the last row

The bottom row is the model guessing what comes after four. Nothing does. The sentence ended.

There is no correct answer to compare that guess against, so it cannot be marked.

rows = logits[:-1]     # everything except the last

Four words give you three markable guesses, not four.

This sounds like housekeeping. It is one of the most common bugs in real training code, because getting it wrong produces no error message. Training simply learns the wrong thing, quietly, forever.

Step two

Turn scores into percentages

Take the first row: 0.0, 2.0, 1.0, 0.0. Two moves.

Make everything positive

Raise e to each score. e is a fixed number, roughly 2.718, and its useful property is that e to the power of anything is always positive.

e^0.0  =  1.000
e^2.0  =  7.389
e^1.0  =  2.718
e^0.0  =  1.000

Notice what happened to the gaps. A score of 2.0 against 1.0 was a small lead. After this it is 7.389 against 2.718, close to three times. Exponentiating stretches differences, which lets the model be decisive.

Share out the pot

Add them up, then divide each one by the total.

total = 1.000 + 7.389 + 2.718 + 1.000 = 12.107

what   1.000 / 12.107 =  8.3%
is     7.389 / 12.107 = 61.0%
two    2.718 / 12.107 = 22.5%
four   1.000 / 12.107 =  8.3%
                        ──────
                        100.0%

Those two moves together are softmax.

escore  /  sum of all escoresoftmax

One extra line, and it is not optional

Before exponentiating, subtract the largest score in the row from every score in that row.

0.0, 2.0, 1.0, 0.0   →   -2.0, 0.0, -1.0, -2.0

Why: e^710 is larger than a computer can store. It returns infinity. Then infinity divided by infinity is undefined, and that undefined value spreads into every number in the model, silently.

Subtracting the largest value first means the biggest exponent is always e^0, which is 1. Overflow becomes impossible.

And it changes nothing, because the same factor appears above and below the division line and cancels out. You can check this: run softmax on the shifted row and you get the identical percentages.

def softmax(x):
    biggest = np.max(x, axis=-1, keepdims=True)
    shifted = x - biggest
    e       = np.exp(shifted)
    return e / np.sum(e, axis=-1, keepdims=True)

Step three

Logarithms, and why they are here

This is the one piece of school maths in the article. If you already know logs, skip to the code block.

What a log is

A logarithm answers: how many times do I multiply 10 by itself to reach this number?

10 × 10 × 10 = 1,000        log(1,000) = 3
10 × 10      =   100        log(100)   = 2

Roughly speaking, the log of a number is how many zeros it has. A million becomes 6.

The property that matters

Logs turn multiplication into addition.

100 × 1,000 = 100,000
  2  +   3  =       5      ← the logs simply add

And division into subtraction.

1,000 ÷ 100 = 10
    3  -  2  =  1

Which log

Everything above counted in tens, because zeros are easy to count. The code uses a different base: e, the same 2.718 from the last step. That version has its own name and its own calculator button, ln.

Changing the base changes the numbers but not the property that matters. Multiplication still becomes addition.

Every logarithm from here on is ln. If you are checking the arithmetic yourself, use that button and not log, or your numbers will not match mine.

Why anyone would want that

To score a whole sentence you multiply the percentages for each word together. A thousand-word document means multiplying a thousand numbers, each around 0.01.

Try it in code and watch what happens.

0.01 multiplied by itself 1000 times   →   0.0
the same thing using logs              →   -4605.2

The true answer is a decimal point, then one thousand nine hundred and ninety-nine zeros, then a one. Computers cannot store numbers that small, so it becomes exactly zero. Not approximately. Zero.

Once the score is zero, every possible setting of the model scores zero, nothing can be compared to anything, and training does nothing at all. It would run for a week and produce a model exactly as bad as when it started.

Logs are not here to be clever. They are here because multiplying many small numbers destroys the answer and adding does not.

Log-softmax, in one move

We want the log of the percentage. The obvious approach is to compute the percentage and then take its log. That works, and it has a failure mode: if a percentage rounds to zero, its log is minus infinity, which poisons everything downstream.

So instead of dividing by the total and then logging, you take the log of the total and subtract it. Same answer, and the percentage never gets built.

shifted row:   -2.0    0.0   -1.0   -2.0
e^shifted:      0.135  1.000  0.368  0.135
total:          1.639
ln(total):      0.494

subtract 0.494 from every value in the shifted row:

log-probs:     -2.494 -0.494 -1.494 -2.494

Check the second one against the long way round: the percentage for is was 61.0%, and ln(0.610) is -0.494. Identical.

def log_softmax(x):
    biggest = np.max(x, axis=-1, keepdims=True)
    shifted = x - biggest
    total   = np.sum(np.exp(shifted), axis=-1, keepdims=True)
    return shifted - np.log(total)

Run that on our three rows and you get the full grid.

          what      is      two     four
slot 0  -2.494  -0.494  -1.494  -2.494
slot 1  -2.494  -2.494  -1.494  -0.494
slot 2  -3.139  -3.139  -3.139  -0.139

Every value is negative, because percentages are less than one and the log of anything under one is negative. Closer to zero means the model was more confident.

Step four

Building the answer key

You have twelve numbers. You need to know which one in each row is the correct word.

Position 0 has read what and should predict is. Position 1 should predict two. Position 2 should predict four.

Those answers are already sitting in the sentence, one place to the right. So chop the first element off and everything lines up.

labels = input_ids[1:]
input_ids:   what   is   two   four
              [0]   [1]   [2]   [3]

labels:              is   two   four
                    [0]   [1]   [2]

Now labels[0] is the answer for slot 0, labels[1] for slot 1, and so on. Same index in both, which is exactly what the next step needs.

You never wrote an answer key. One line of slicing produced a complete one.

This is the whole reason a model can be trained on text scraped from the internet without anybody labelling it. Every sentence already contains its own marking scheme.

Two things fall off the ends, and it is worth seeing why. The first word never appears in the labels, because nothing comes before it to predict it. The last slot gets dropped, because nothing comes after it to mark it against.

Step five

Keeping one number per row

Labels tell you which column is correct. Go to each row, take that one value, discard the other three.

picked = log_probs[np.arange(len(labels)), labels]

That line looks stranger than it is. It hands numpy two lists which get paired up into coordinates.

rows:     [0, 1, 2]        ← np.arange(3)
columns:  [1, 2, 3]        ← labels

pairs:    (0,1)  (1,2)  (2,3)

Every log-probability the model produced

whatistwofourWhat we keep
slot 0-2.494-0.494-1.494-2.494-0.494
slot 1-2.494-2.494-1.494-0.494-1.494
slot 2-3.139-3.139-3.139-0.139-0.139

Twelve numbers become three. The highlighted cell steps diagonally, which is the shift from step four, made visible.

Only one number per row survives: the one the labels pointed at.

Look at the middle row. The model gave two only -1.494, which is 22.5%, while putting 61% on four. It got that one wrong, and now you know exactly where and by how much.

Step six

Deciding what counts

Three numbers, one per slot. But two of them are for words inside the question, which the model will never have to write.

slot         0        1        2
predicts    is      two     four
part      prompt  prompt  answer
mask         0        0        1

Multiply and the zeros delete.

picked:   -0.494  -1.494  -0.139
mask:          0       0       1
                       ─────────
product:       0       0  -0.139

Note which number disappeared. Slot 1 was the model’s worst performance by a wide margin, contributing -1.494. And it does not count, because that mistake was inside the question.

Without the mask, most of the training signal would come from an error that does not matter.

The mask is not computed from anything. A person decides where the question ends when the training data is built, and that decision gets stored alongside the words. Nothing in what is two four reveals where the answer starts.

Step seven

Add up, flip the sign

0 + 0 + (-0.139) = -0.139

One number. Then flip its sign to get 0.139.

Why flip: log-probabilities are negative and closer to zero is better, so bigger is better. But the method that improves the model works by making numbers smaller. Flip the sign and now smaller is better, and zero is perfect.

Without the flip, training would drive the model steadily toward being wrong. It would be doing its job correctly, aimed at the wrong target.

The whole loss, assembled

def loss_of(w):
    rows      = w[:-1]
    biggest   = np.max(rows, axis=-1, keepdims=True)
    shifted   = rows - biggest
    total     = np.sum(np.exp(shifted), axis=-1, keepdims=True)
    log_probs = shifted - np.log(total)
    labels    = input_ids[1:]
    picked    = log_probs[np.arange(len(labels)), labels]
    return -np.sum(picked * mask)

Eight lines. That is the loss function used to train every large language model. Real ones handle a hundred thousand words instead of four and thousands of positions instead of three, but the steps are these steps.

What the number means

with the mask   0.139   grading only the answer
without it      2.127   grading everything

And a sense of scale for the loss itself, so the numbers are not abstract.

correct word got 99%  →  loss 0.01
                 87%  →  loss 0.14
                 50%  →  loss 0.69
                  1%  →  loss 4.61

Being unsure costs 0.69. Being confidently wrong costs 4.61. That ratio comes straight out of the logarithm, and it encodes a real judgement: a model that shrugs is doing less damage than one that is certain about something false.

Step eight

Making the number smaller

Now the other half. You have a score for how wrong the model is. How do you improve it?

For this article the model is those sixteen logits. Start them all at zero, which means the model gives every word 25% everywhere. It knows nothing.

starting loss:  1.386

Poke one number and watch

Take the entry at row 2, column 3. That is the score for four at the slot that has to predict four. Nudge it up by a hair, then down by a hair, computing the loss both times.

nudged up by 0.00001     →   loss 1.38628686
nudged down by 0.00001   →   loss 1.38630186

The up version is smaller. So raising that number improves the model. Obvious in hindsight, but nothing told the code that in advance. It found out by trying.

Turn the observation into a number

1.38628686 - 1.38630186 = -0.000015

Negative means up was better. But the value is tiny, because the poke was tiny. Scale it back up by dividing by how far apart the two pokes were.

-0.000015 / 0.00002 = -0.75

This is exactly how a road gradient sign works. Rise divided by run. A hill that climbs 8 metres over 100 metres is an 8% gradient, and nobody quotes it as 0.0008 metres over 0.01 metres even though that is the same slope.

( loss with it up − loss with it down ) / ( 2 × poke size )the gradient for one number

Move

new value = old - learning_rate × gradient
          =   0 - 0.5 × (-0.75)
          = +0.375

Subtracting a negative moves it up, which is what the measurement asked for. The 0.5 is there so you take a modest step rather than leaping the full distance and overshooting.

Three pokes, three different answers

row 2, col 3   four, the right word   gradient -0.75
row 2, col 0   what, a wrong word     gradient +0.25
row 0, col 1   inside the prompt      gradient  0.00

A negative gradient sends the number up, a positive one sends it down, and zero leaves it exactly where it is.

That last line is worth pausing on. Row 0 feeds slot 0, and slot 0 gets multiplied by a mask value of zero. So poking it changes the loss by precisely nothing, so its gradient is zero, so it never moves.

Nobody wrote a rule saying to skip the prompt. The mask made those numbers irrelevant, and the arithmetic did the rest.

Step nine

The training loop

Do the poke for all sixteen numbers, storing each answer. Then move all sixteen at once. Then repeat.

def train(logits, lr=0.5, steps=300, h=1e-5):

    w = logits.astype(float).copy()

    for _ in range(steps):

        grad = np.zeros_like(w)

        for i in range(w.shape[0]):
            for j in range(w.shape[1]):

                up = w.copy()
                up[i, j] += h

                dn = w.copy()
                dn[i, j] -= h

                grad[i, j] = (loss_of(up) - loss_of(dn)) / (2 * h)

        w = w - lr * grad

    return w

Two details in there matter.

Copies, not edits. up and dn are duplicates. The real model is never touched during a measurement, so all sixteen readings describe the same starting point.

Measure everything, then move once. The step is outside the inner loops. Move a number halfway through and the remaining measurements would be describing a model that no longer exists.

Watching it run

at the start     loss 1.386294
after step 1     loss 1.036592
after step 2     loss 0.780880
after step 3     loss 0.601321
after step 4     loss 0.476263
after step 5     loss 0.387892
      ...
after step 300   loss 0.005098

And here is what it learned.

[[ 0.000  0.000  0.000  0.000]   ← untouched
 [ 0.000  0.000  0.000  0.000]   ← untouched
 [-1.594 -1.594 -1.594  4.781]   ← trained
 [ 0.000  0.000  0.000  0.000]]  ← dropped in step 1

Only the row that mattered moved. Within it, the correct word was pushed up to 4.781 and the three wrong ones pushed down to -1.594.

before   what 25%    is 25%    two 25%    four 25%
after    what  0.2%  is  0.2%  two  0.2%  four 99.5%

A model that knew nothing now knows one thing. Poke, step, repeat, three hundred times.

Reality check

What real code does differently

The code above works and is honest, but no production system is written that way. Three differences.

Nobody pokes

Sixteen numbers means thirty-two loss computations per step, which is nothing. A model with eight billion numbers would need sixteen billion, and at roughly a second each that is about five hundred years for a single step. You need millions of steps.

Instead there is a method called backpropagation that computes every direction at once, in roughly one extra pass. It works by knowing what the model is built from, rather than treating it as a sealed box and prodding it.

The poking version is still used, though. When somebody implements a new piece of a model, they compute its gradients both ways and check the answers match. Slow and honest, so it catches bugs in the fast path.

Steps are smarter

Plain w - lr × grad treats every number identically. Real optimisers add two refinements.

Real training also works through thousands of examples at a time rather than one sentence, and each such batch gives a slightly different answer about which way to move. So optimisers step on a running average of recent measurements, which stops the noise from any single batch sending things zigzagging. And they give each number its own step size, because a model’s numbers live on wildly different scales and one setting cannot suit them all.

The standard combination is called AdamW. You will see that name in nearly every training script.

The loss is one function call

In PyTorch, the middle of loss_of — log-softmax, pick the right one, flip the sign — collapses to a single call.

loss = F.cross_entropy(logits, labels)

Two things do not vanish, and both trip people up. The mask changes shape rather than disappearing: you set the labels you do not want graded to -100, and cross_entropy skips those positions. And it averages across positions where ours added them up, which changes the number without changing what it measures.

Which raises a reasonable question: why write it out at all? Because when that line produces a strange number, or when someone mentions masking a completion, or when you read a paper that talks about log-probabilities, you will know precisely what is inside it.

Run it yourself

Everything assembled. Roughly forty lines, numpy only.

import numpy as np

VOCAB     = ["what", "is", "two", "four"]
input_ids = np.array([0, 1, 2, 3])
mask      = np.array([0.0, 0.0, 1.0])


def softmax(x):
    biggest = np.max(x, axis=-1, keepdims=True)
    e       = np.exp(x - biggest)
    return e / np.sum(e, axis=-1, keepdims=True)


def loss_of(w):
    rows      = w[:-1]
    biggest   = np.max(rows, axis=-1, keepdims=True)
    shifted   = rows - biggest
    total     = np.sum(np.exp(shifted), axis=-1, keepdims=True)
    log_probs = shifted - np.log(total)
    labels    = input_ids[1:]
    picked    = log_probs[np.arange(len(labels)), labels]
    return -np.sum(picked * mask)


def train(logits, lr=0.5, steps=300, h=1e-5):
    w = logits.astype(float).copy()
    for step in range(steps):
        grad = np.zeros_like(w)
        for i in range(w.shape[0]):
            for j in range(w.shape[1]):
                up = w.copy(); up[i, j] += h
                dn = w.copy(); dn[i, j] -= h
                grad[i, j] = (loss_of(up) - loss_of(dn)) / (2 * h)
        w = w - lr * grad
        if (step + 1) % 50 == 0:
            print(f"after step {step + 1:3}   loss {loss_of(w):.6f}")
    return w


trained = train(np.zeros((4, 4)))

print()
print(np.round(trained, 3))
print()
for word, p in zip(VOCAB, softmax(trained[2])):
    print(f"{word:6} {p:6.1%}")

Save it, run it with python3, and watch the loss fall. It takes a few seconds because of all that poking.

Things worth breaking

The fastest way to understand a piece of code is to make it wrong on purpose.

Set the mask to all ones. Now every slot is graded, which is what pretraining does. All three rows should train instead of one.

Remove the max subtraction and feed in logits around 800. Watch it produce nan and stop working.

Delete the minus sign in the last line of loss_of. The model will train enthusiastically toward being as wrong as possible.

Change the learning rate to 5. In most models a step that large overshoots, and the loss starts bouncing instead of settling. Here it will not. The loss simply falls faster, and that is worth understanding rather than shrugging at. The steepness of this loss can never exceed 1, and it flattens as the model gets the answer right, so a big step lands where the ground has already levelled off. Bouncing needs a surface that grows steeper the further you are from the bottom. This one grows gentler.

What you now know

You have the complete supervised training loop for a language model, at a scale you can check by hand. Every number in this article came out of code that runs.

Three things are worth carrying.

Logs are load bearing. Not decoration, not notation. Multiply enough probabilities together and the answer vanishes, so everything is done in logs where multiplying becomes adding.

The answer key is free. One slice of the sentence produces it. This is the fact that makes training on the open internet possible.

The masking behaviour is emergent. There is no instruction anywhere saying to skip the prompt. Multiplying by zero made those numbers irrelevant, so their gradients came out zero, so they never moved.

What comes after this is where it gets harder, and where the research is. Training on human preferences, where there is no correct next word to compare against, just a person saying they liked one response better. And training against a score, where the model writes forty words, gets back a single number, and something has to work out which of the forty deserve the credit.

But those all sit on top of this loop. Not next to it.

Two simplifications carried over from part one. The model reads fragments of words rather than whole words, so a long word might arrive as three pieces. And the numbers being trained here are the predictions themselves, whereas a real model trains the machinery that produces the predictions. The loop is identical either way.

The four-word vocabulary is not a toy version of the maths. It is the same maths, sized so you can check it.