SGD, momentum, AdamW, and schedules
The trained state is parameters plus optimizer history, schedule position, data order, and randomness.
A nonlinear sensor classifier defeats every classical baseline.
This chapter moves from establish an sgd reference to test interruption and recovery. We will make the idea concrete in a small code example before placing it back inside a production system with a declared latency, cost, and operator-capacity budget.
Resume from a checkpoint containing parameters, optimizer moments, scheduler position, RNG state, and data cursor; compare with an uninterrupted run.
Verify the behavior illustrated in adamw_schedule.py, then measure decision utility, slice reliability, and system latency.
An excessive learning rate causes NaNs; dead ReLUs and a BatchNorm train/eval mismatch hide the cause. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
An optimizer turns noisy batch gradients into a parameter trajectory. Its state, scale, and learning-rate schedule often decide whether the same architecture converges smoothly, oscillates, or diverges.
Plain SGD follows the current gradient; momentum averages a velocity that damps zigzags and carries progress through shallow directions. Adam rescales updates with running first and second moments, which often accelerates early training. AdamW applies weight decay separately from the adaptive gradient update, making regularization behavior easier to reason about.
The learning rate is usually the dominant knob. Warmup protects unstable early steps, while cosine or step decay reduces motion near convergence. Compare optimizers at matched batch size, training budget, initialization, and tuned learning rates, and checkpoint optimizer and scheduler state for faithful resume.
- chain rule
- initialization
- AdamW
- dropout
Follow the reasoning before touching the code.
- 01
Establish an SGD reference
Start with a stable learning-rate range and record loss, gradient norm, update norm, and validation behavior before introducing adaptive state.
- 02
Tune the update bundle
Treat optimizer, base learning rate, momentum or betas, weight decay, batch size, clipping, and warmup as one interacting configuration.
- 03
Test interruption and recovery
Resume from a checkpoint containing parameters, optimizer moments, scheduler position, RNG state, and data cursor; compare with an uninterrupted run.
import torch
weight = torch.nn.Parameter(torch.tensor([5.0]))
optimizer = torch.optim.AdamW([weight], lr=0.2, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=6)
for step in range(6):
optimizer.zero_grad()
loss = (weight - 2).pow(2).mean()
loss.backward()
optimizer.step()
scheduler.step()
print(step, round(weight.item(), 3), optimizer.param_groups[0]['lr'])Optimize one parameter while decaying the learning rate.
What to notice
- AdamW stores moving moments in addition to the parameter, so saving weight alone cannot reproduce a resumed trajectory.
- The cosine schedule changes the size of later updates even when the raw gradient remains nonzero.
YOUR TURN Pause and predict before revealing the answer.
Resume at step three with a fresh optimizer and scheduler but the same weight. Will the remaining path match?
Need a hint?
List the hidden state accumulated by AdamW and the schedule.
Check your reasoning
No. The first and second moments and scheduler phase are missing, so subsequent updates differ despite an identical parameter value.
Lessons 1, 2, 3 in this course.
Connect the chapter to the course’s quantitative spine.
The trained state is parameters plus optimizer history, schedule position, data order, and randomness. The backpropagated error below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Backpropagated error
| Symbol | Meaning / shape / unit |
|---|---|
δ_l | loss gradient at layer l |
W | weight matrix |
φ′ | activation derivative |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: δ_l, W, φ′.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c09(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
Optimizer stability lab
Sweep learning rate over a deterministic convex-plus-curvature loss.
Assumption: Twenty gradient steps from w=4 on L(w)=w²/2.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| step 1 | 3.5200 | 6.19520 |
| step 5 | 2.1109 | 2.22801 |
| step 10 | 1.1140 | 0.62050 |
| step 15 | 0.5879 | 0.17281 |
| step 20 | 0.3103 | 0.04813 |
Trace the complete operating path.
Carry the chapter’s central idea into the system: The trained state is parameters plus optimizer history, schedule position, data order, and randomness. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Establish an SGD reference” into a written invariant and a failing test.
- 02
Promote adamw_schedule.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Test interruption and recovery” as the release checkpoint: Resume from a checkpoint containing parameters, optimizer moments, scheduler position, RNG state, and data cursor; compare with an uninterrupted run.
- 04
Serve or materialize behind explicit a declared latency, cost, and operator-capacity budget.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
An excessive learning rate causes NaNs; dead ReLUs and a BatchNorm train/eval mismatch hide the cause. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Network from first principles
Specify a NumPy MLP, finite-difference checks, PyTorch parity, calibrated inference, and debugging report.
- Forward/backward derivation
- Gradient check plan
- Optimizer/regularization ablations
- Export and batch-one parity
Read primary material with a purpose.
Return to the opening failure.
Package preprocessing, vocabularies, weights, calibration, output semantics, and batch-one parity as one artifact.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce a small network built from forward pass through backpropagation and optimization.