11 · H · Train

The training loop

Batch, forward, loss, backward, accumulate, clip, step. Train loss is not validation loss.

What actually runs

Repeat millions of steps:

Sample a batch of sequences. Shape: $B \times T$ token IDs.
Forward: embeddings → $n$ blocks → logits $Z \in R^{B \times T \times V}$ .
Compute $L$ on next-token targets (ignore pad / masked positions).
Backward: reverse-mode AD gives $g = \nabla_{θ} L$ .
Optionally accumulate g over several micro-batches to fake a larger batch.
Clip g, AdamW update, zero grads.
Periodically evaluate validation loss on held-out data (no parameter update).

g_{accum} = \frac{1}{K} k = 1 \sum K \nabla_{θ} L^{(k)}

(11.1)

$K$: number of micro-batches
$L^{(k)}$: mean token loss on micro-batch k

Initialization

Embeddings and most projections start as small Gaussians. Residual-output projections are often scaled so the untrained skip path dominates at depth (GPT-2):

W_{out} \sim N (0, \frac{σ _{0}^{2}}{2 n}), σ_{0} \approx 0.02

(11.2)

Why it mattersThe skip Jacobian still starts at I. Scaling the branch keeps H^{(ℓ)} from exploding before the first useful gradient. μP (Yang et al.) is the width-aware version: change d without retuning η.

Train loss is not the target

Batch loss is the mean of per-token losses. Mixed precision (BF16/FP16 forward, FP32 master weights and loss) is an implementation detail; the math above is the same. The loss itself is computed in FP32 so the log-sum-exp is not wrecked by underflow.

Training loss can keep falling by memorization. Validation loss estimates generalization. The entropy floor applies to both; validation is the one that tells you the model is not just reciting.