Train next-token logits

Learn it

Words to know

bigram model

The simplest language model: it predicts the next token from the current token alone.

cross-entropy

The usual loss for choosing between classes: minus the log of the probability the model gave to the right answer.

logit

A raw score a model gives to one choice, before the scores are turned into probabilities.

optimizer

The part of the training code that changes the parameters, using their gradients.

How it works

1. The model is one embedding table of shape [vocab, vocab], starting at zero. Row i holds the scores for what follows token i.
2. Looking up the input IDs gives one row of scores per input token.
3. torch.nn.functional.cross_entropy(logits, targets) applies softmax and -log in one call and averages over the tokens.
4. Run the four-call training step from chapter 4 steps times. Then compute the loss once more, without updating, and return it.

Worked example

Input
ids = [0]
targets = [1]
vocab = 2
rate = 1
steps = 1

Output
0.313262

The gradient of row 0 was [0.5, -0.5]. One step lowered the score of the wrong token and raised the right one, and the loss fell by more than half.

Why it matters

tiny_gpt.py trains with the same four calls, with the GPT in place of the table. It also limits the size of the gradients and uses the AdamW optimizer, not SGD. A real model should beat a bigram table's loss.

Your turn

Using PyTorch, train a bigram table to predict each target token from its input token, and report the loss afterwards.

ids

the input token IDs; at least one

targets

the token that follows each input, in the same order

vocab

the number of different tokens

rate

the learning rate for SGD

steps

how many training steps to run; 0 or more

Returns

The mean cross-entropy over the pairs, measured after the last update, as a Python float. The table is torch.nn.Embedding(vocab, vocab) with float64 weights set to zero.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
ids = [0]
targets = [1]
vocab = 2
rate = 1
steps = 0

Output
0.693147

Input
ids = [0]
targets = [1]
vocab = 2
rate = 1
steps = 1

Output
0.313262

Input
ids = [0]
targets = [1]
vocab = 2
rate = 1
steps = 2

Output
0.194609

Input
ids = [0, 1]
targets = [1, 0]
vocab = 2
rate = 0.5
steps = 1

Output
0.575939

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   The model is one torch.nn.Embedding. Set its weights to zero, then write the same loop as in the training-loop lesson.

2. Hint 2: The key step

   model(x) gives the scores, and torch.nn.functional.cross_entropy(model(x), y) gives the loss. After the loop, compute that loss one more time and return its .item().

3. Hint 3: The outline

   Create model = torch.nn.Embedding(vocab, vocab, dtype=torch.float64) and zero it with torch.nn.init.zeros_(model.weight). Make x and y tensors and torch.optim.SGD(model.parameters(), lr=rate). Loop: zero_grad, loss, backward, step. Return the loss computed after the loop.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. Row i holds the scores for what follows token i. An embedding table is used as the whole model. Zero weights make every next token equally likely at the start.
2. The usual four calls. cross_entropy takes raw scores, not probabilities. It does the softmax itself, in a numerically stable way.
3. Measure again, now that the last update has been applied. The loss inside the loop was computed before that step's update. The task asks for the loss after it.

The same numbered steps are the comments in the solution.

import torch


def solve(ids, targets, vocab, rate, steps):
    F = torch.nn.functional
    # 1. Row i holds the scores for what follows token i
    model = torch.nn.Embedding(vocab, vocab, dtype=torch.float64)
    torch.nn.init.zeros_(model.weight)
    x = torch.tensor(ids)
    y = torch.tensor(targets)
    optimizer = torch.optim.SGD(model.parameters(), lr=rate)
    for _ in range(steps):
        # 2. The usual four calls
        optimizer.zero_grad()
        loss = F.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()
    # 3. Measure again, now that the last update has been applied
    return F.cross_entropy(model(x), y).item()

Watch out

• Do not apply softmax before cross_entropy. The function expects raw scores, and softmax twice gives a wrong loss.
• Returning the loss from inside the loop reports the value from before the last update, and fails when steps is 0 because the variable does not exist.

Check yourself