Unit 3.1 · about 4 h
304 / 324The next-token objective: shifts, masks and perplexity
Every later number in this section, from optimizer steps to evaluation claims, reads off this loss. A shifted-by-one or unmasked loss trains the wrong objective while the curve still falls.
- derive
- implement
You can now
- Derive token-weighted cross-entropy and perplexity with an explicit valid-token mask.
- Explain the input-target shift and predict the failure when labels are off by one.
- Explain which positions a packed or padded batch must exclude from the loss and why.
Covers
- Autoregressive factorization and next-token likelihood
- Input-target shift and off-by-one failure modes
- Packed document, padding and valid-token loss masks
- Token-weighted cross-entropy and perplexity
Study, in this order
3 steps
- CS336 Assignment 1 handout: cross-entropy and perplexityStanford CS336 Assignment 1 - training and decoding contract
Section 4.1, Cross-entropy loss (pp. 28 to 29): equations 16 to 18 for shifted likelihood, the numerically stable cross-entropy requirements and perplexity
- LLMs from Scratch ch05: generation loss and train/val lossLLMs from Scratch companion - training and generation
Sections 5.1.2 (shifted inputs and targets, flattening
logitsforcross_entropy, perplexity as the exponential of the loss) and 5.1.3 (train and validation loss withcalc_loss_batchandcalc_loss_loader) - Modern LLM Notebook 10-training-loss: shift, hand loss and valid tokensModern LLM Notebook - configuration, loss and generation
Sections 1 (input and label shift on
[1, 3, 4, 5, 6]), 3 (hand-computed cross-entropy fromlogits[2.0, 1.0, 0.1, -1.0]), 10 (perplexity table) and 11 (shift in the data versus shift in the model), plus Homework 3 (mean over valid tokens only)
Focus on
- Position \(i\) predicts token \(i+1\) from context \(1\) to \(i\). Inputs and labels are the same sequence offset by one.
- The mean runs over valid tokens only. Padding and
-100label positions leave both the sum and the count. - Perplexity is exp of the mean cross-entropy, so it reads as the effective shortlist size at each step.
- One forward pass yields every position's distribution at once, which is why training parallelizes but generation does not.
Reproduce closed-book
F.cross_entropy on the flattened tensors.Traps
- Averaging over all positions including padding. The denominator changes, so short and padded batches get different objectives.
- Shifting twice: slicing tokens in the data loader and again inside the model. Every target then predicts two steps ahead.
- Comparing perplexity across tokenizers. Different vocabularies split the same text into different token counts, so the means are over different units.
Practice
CS336 Assignment 1: `cross_entropy`
Problem (cross_entropy) in Section 4.1; connect your function to run_cross_entropy in tests/adapters.py and pass test_cross_entropy in tests/test_nn_utils.py
The tests compare against F.cross_entropy on a fixed batch and on a 1000x scaled copy, so subtract the row max and cancel log with exp. Implement it in your own repository.
Self-test
-
Why does the Raschka code flatten
logitsto \((B \times T, V)\) before callingcross_entropy?The function expects a 2D \((N, C)\) input with one target index per row. Flattening keeps one row per position while targets flatten to \((N,)\), so each position is scored against its own shifted target. -
Logits are
[2.0, 1.0, 0.1, -1.0]and the correct token is 0. What is the loss?About 0.45. Softmax gives the correct token probability \(\exp(2.0)/(\exp(2.0)+\exp(1.0)+\exp(0.1)+\exp(-1.0))=0.64\), and \(-\log(0.64)=0.45\). -
Labels are not shifted, so position \(i\) supervises with token \(i\). What breaks?
The model is asked to output a token already in its context, so it learns copying. The loss still falls, but the trained objective is not next-token prediction and the perplexity is meaningless. -
Six positions have losses
[0.2, 0.4, 1.0, 0.3, 0.8, 0.5]with labels[10, 11, -100, 3, -100, 12]. What is the training loss?0.35. Only the four positions with real labels count: \((0.2+0.4+0.3+0.5)/4\). Including the masked positions changes both the sum and the denominator.
Completion criterion
Stop when
You hand-compute a masked loss and perplexity that match code, and your cross_entropy passes test_cross_entropy including the overflow case.
Leads to
- 3.2 Feeding the model: batches, initialization and first steps Feed correctly shifted and masked batches from a real data loader.
- 4.1 Likelihood and perplexity, and their limits Carry perplexity into evaluation and learn when it misleads.
Comes after
- 1.5 From documents to training examples: packing and loss masksUse the valid-target mask in the next-token loss.
- 2.7 Assembling and verifying a decoderTurn raw vocabulary logits into the shifted next-token loss.
Resources for this unit
- LLMs from Scratch companion - training and generation Use the ch05 training and generation code for the small pretraining loop. Add full checkpoint state, explicit loss masks and resume-equivalence tests that the base example does not supply.
- Stanford CS336 Assignment 1 - training and decoding contract Use cross_entropy, adamw, adamw_accounting, learning_rate_schedule, gradient_clipping, data_loading, checkpointing, training_together and decoding specifications plus public tests. Debug on CPU and implement independently in your own repository.
- Modern LLM Notebook - configuration, loss and generation Use English notebooks 09-model-config, 10-training-loss and 20-generation to visualize configuration ownership, shifted loss and decoding policies. Change edge cases and cross-check probability calculations against your own implementation.
Also in this section
- PyTorch saving, loading and reproducibility references Use the state-dict and reproducibility notes to enumerate all state needed for the supported resume guarantee. Test model, optimizer, schedule, scaler, data order and RNG state across an interruption.
- Hugging Face LLM Course - library translation Only after the blank-file loop passes, map tokenization, model configuration, data collation, generation configuration and stopping behavior into Transformers APIs. Keep the from-scratch tests as the authority.
- PyTorch Automatic Mixed Precision recipe Use the
autocast,GradScaler, unscale-before-clip, scaler checkpointing and inf-NaN troubleshooting sections to add mixed precision to the blank-file loop. Keep the CS336 clipping contract as the authority for the clipping rule.