Lab 07 - Training Loop and Cross Entropy Loss

Overview

Lab 07: Training Loop and Cross Entropy Loss

Annotated code reading lab. Running code is optional.

Related handbook section

Training Loop and Cross Entropy Loss

This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Open handbook section Open starter file

Concept Goal

Training Loop and Cross Entropy Loss

Connect next-token prediction to shifted labels and cross entropy over the vocabulary.

Mental Model

Mechanism to keep in mind

`input_ids` and `labels` are offset by one token.
`logits` has one vocabulary distribution per position.
`loss` ignores positions that should not train, such as padding.

Annotated Code Preview

Starter preview

Excerpt from code/lab-07-training-loop-cross-entropy/training_loop.py. The linked starter file is the source of truth.

Open starter file

# Training Loop and Cross Entropy Loss
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Connect next-token prediction to shifted labels and cross entropy over the vocabulary.

tokens = [11, 25, 42, 9]
input_ids = tokens[:-1]    # model reads 11,25,42
labels = tokens[1:]        # model predicts 25,42,9
logits_shape = (len(input_ids), "vocab_size")
loss = "cross_entropy(logits, labels)"

# What to explain while reading:
# - input_ids and labels are offset by one token.
# - logits has one vocabulary distribution per position.
# - loss ignores positions that should not train, such as padding.
#
# Common traps:
# - Teacher forcing is not KV-cache decoding.
# - Cross entropy compares logits with ids, not generated strings.

Line-by-line Explanation

What each block is doing

Setup / contract: `input_ids` and `labels` are offset by one token.
Main transition: `logits` has one vocabulary distribution per position.
Interview hook: `loss` ignores positions that should not train, such as padding.

What to Notice

Reading checkpoints

Training is parallel over positions.
Inference is sequential over generated tokens.
Loss is computed before optimizer state updates.

Common Misunderstandings

What this lab prevents

Teacher forcing is not KV-cache decoding.
Cross entropy compares logits with ids, not generated strings.

Interview Explanation

How to say it out loud

Connect next-token prediction to shifted labels and cross entropy over the vocabulary. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.

External intuition notes

Additional intuition

Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.

InfraLens