JAX Pathways

Course · Math & optimizationPhase 04 · Study guide

Can you explain an optimizer’s behavior from the objective and data?

Core Foundation · Local CPULaptop CPU12 lessons (~15.5h)

Revision v0.7.0 ·

Use a model small enough to solve independently. Connect residuals, gradients, curvature and update state before comparing optimizer names. A lower training loss is useful evidence, but it does not by itself establish better held-out predictions.

Preparation: State, randomness & control flow

Runtime: CPU

Start the phaseSee the project

Check your starting point

A model predicts one value for each observation. Targets accidentally have an extra trailing singleton axis. Why might mean squared error still run?

Compare your reasoning

Broadcasting can construct a pairwise residual matrix instead of one residual per observation. Assert identical prediction and target shapes, then compute one residual by hand before reducing.

If this is unfamiliar, work through Linear algebra and loss intuition before attempting the project.

Build in stages

1. Audit the objective and derivative

Work through vectors, residuals, rank and directional derivatives. Use a linear solve and a finite-difference sweep as independent checks.

2. Connect dynamics to evaluation

Predict a stable versus divergent quadratic update, inspect stable likelihood arithmetic and separate fitting, validation selection and final test reporting.

3. Make optimizer comparisons fair

Count-weight unequal batches, preserve optimizer moments and schedule position, then change feature scale while retaining the prediction contract.

Try a changed condition

Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.

Rescale one feature by a factor of ten. Is the same learning rate automatically a fair comparison?

Compare an approach

No. The parameterization changes directional curvature. Map coefficients to preserve predictions, verify the objective and compare trajectories under stated rates. Fit preprocessing on training data only; never use test performance to choose the rate.

Diagnose a likely failure

Symptom: Gradient descent diverges while all values initially remain finite.

Check next: Check the update sign, loss reduction, feature scales and largest quadratic curvature before trying a different seed.

Decide what is ready

Complete regression-audit, then optimizer-audit. Keep an analytic reference, stable/unstable trajectories, a resumed optimizer comparison and an untouched final evaluation split.

Where to put your project code

Write your code in: projects/optimizer-audit/my_optimization.py (copy projects/optimizer-audit/starter/optimization.py first: cp projects/optimizer-audit/starter/optimization.py projects/optimizer-audit/my_optimization.py).

Verify your implementation: python3 projects/optimizer-audit/tests/check.py --stage all --impl projects/optimizer-audit/my_optimization.py

Run the staged project checks on your implementation in projects/optimizer-audit/my_optimization.py and record your transfer experiment results in your portfolio.

Open A fitted regression model

What still needs further work

Constrained, nonsmooth and large-scale stochastic optimization are not covered comprehensively. Current ranking evidence is limited to declared synthetic objectives.