Train a Bigram Language Model

The problem

Time to put the whole loop together: sample a batch, compute the loss and its gradient, update the parameters, repeat. To keep it fast in the browser, the model is the simplest neural language model: a bigram model whose only parameter is a logit table W of shape $(V, V)$. The logits for input token $i$ are row W[i]: the scores of every possible next token after $i$. For a batch of inputs x (any shape), the logits are W[x].

Part 1: loss and gradient

The loss is the mean cross-entropy of the logits W[x] against the targets y, over all the tokens of the batch (section 2).

The gradient. Two steps:

1. The gradient with respect to the logits of each token is its softmax probabilities minus the one-hot target, divided by $n$, the number of tokens in the batch (section 2).
2. Each logit row is a copy of a row of W, so its gradient goes back to that row, and repeated input tokens add up, exactly like the embedding of section 4.

Example. $V = 3$, W all zeros, x = [0, 1, 0], y = [1, 2, 2], so $n = 3$.

• Loss: with zero logits every token has probability $1/3$, so every term is $-\log(1/3)$, and the mean is $\log 3 \approx 1.099$.
• Logit gradients, $(1/3, 1/3, 1/3)$ minus the one-hot target, divided by 3: token 0 (target 1) gets $(0.111, -0.222, 0.111)$; token 1 (target 2) gets $(0.111, 0.111, -0.222)$; token 2 (input 0, target 2) gets $(0.111, 0.111, -0.222)$.
• Back to W: input 0 appears twice, so row 0 of the gradient is the sum of the first and third, $(0.222, -0.111, -0.111)$. Row 1 is $(0.111, 0.111, -0.222)$. Row 2 is 0: token 2 never appears as an input.

Part 2: the loop

Start from W = np.zeros((V, V)). At each step:

1. sample a batch with the given get_batch (previous exercise);
2. compute the loss and the gradient;
3. update W -= lr * grad (plain gradient descent);
4. record the loss.

Example. On the repeating text 0 1 2 0 1 2 ... with $V = 3$, the loss starts at $\log 3 \approx 1.099$ (step 0: all zeros, every token equally likely) and falls towards 0, because in this text the next token is always certain. On real text it never gets to 0, because the next token is never certain.

What to implement

• loss_and_grad(W, x, y): return (loss, grad): the mean cross-entropy, a float, and its gradient with respect to W, shape $(V, V)$.
• train(data, V, steps, batch_size, block_size, lr, rng): run the loop for steps steps and return (W, losses), where losses is the list of the training loss at every step.

Requirements

• Stable softmax.
• Scatter the logit gradients back into W with np.add.at: an input token appears many times in a batch.
• Keep it fast: the tests train for a few hundred steps in your browser.

Hint

Flatten first, x.reshape(-1) and y.reshape(-1): then the logits W[x] are a plain table with one row per token, as in the cross-entropy exercise.