The optimization step

One step forward, measure the loss, compute gradients, nudge every weight by the learning rate.

An optimization step is the atomic unit of training. It has four phases that happen in strict order: a forward pass produces predictions, cross-entropy loss measures how wrong they are, a backward pass computes how much each weight contributed to that error, and then every weight is nudged to reduce that error. Each step performs this once over the whole corpus.

The update rule is gradient descent: w ← w − lr × ∂L/∂w. The gradient ∂L/∂w is the slope of the loss with respect to each weight — it points in the direction that would make the loss increase. Subtracting it therefore moves the weight in the direction that reduces the loss. The learning rate lr is a small scalar (often around 0.001 to 0.01) that controls how large each step is. Too large and the updates overshoot the minimum and the loss bounces or diverges; too small and training converges uselessly slowly. There is no formula for the right value — it is a hyperparameter tuned by watching the loss curve.

In the playground, one click of "Train live" runs one optimization step over the entire corpus: a single forward pass scores every position at once, the loss is averaged across all the next-token predictions, one backward pass produces the gradients, and the optimizer applies the update. The loss curve plots the loss before the update; the weight chips below the diagram backbone light up after the update by how much each weight matrix changed. Heavier highlights mean that layer contributed more to the loss and received a larger gradient — early layers tend to update less than later ones because their gradients travel further through the chain rule and can shrink (vanishing gradient).

The optimizer used here is Adam (β₁ = 0.9, β₂ = 0.999, ε = 1e-8, with bias correction on both moments). Adam keeps a running mean and variance of each weight's gradient and normalises every update by them, so each weight effectively gets its own adaptive step size — which is why it converges reliably even on this tiny model. The w ← w − lr × ∂L/∂w rule above is the skeleton; Adam refines the raw gradient into that smoothed, per-weight-scaled step. Larger models use the same Adam (often as AdamW, which decouples weight decay).