Core Foundation · Local CPULaptop CPULesson 12 of 12 in Phase 04

Phase 04 / Math & optimization

Adam, clipping, and learning-rate schedules

Revision v0.7.0 · 2026-10-20

Lesson · about 90 minutes · CPU practice

Explain Adam’s moments, inspect clipping and schedule order, and compare optimizers with a controlled update budget.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Optimize with Optax
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment

python phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded optimization-12.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: Two Adam updates match the manual moment calculation; the clipped SGD update is $(-0.06,-0.08)$ .

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 04.12 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Explain Adam’s moments, inspect clipping and schedule order, and compare optimizers with a controlled update budget.
Calculate the first two Adam steps
Clip a direction without rotating it
Make the schedule clock explicit
Compare decisions with evidence
Bring the phase together in a training audit

The problem

An optimizer can lower the loss while hiding several choices: how it averages gradients, rescales coordinates, clips large values, and changes the step size. We will expose those choices one at a time before running a controlled comparison. The goal is to explain a training decision, not to crown a universal winner.

The idea

Adam combines gradient history with coordinatewise scaling. Clipping limits a gradient or update according to a chosen rule, while a learning-rate schedule changes scale over training. These mechanisms interact, but we should inspect each separately before interpreting their combined loss curve.

Clipping changes a vector; a schedule changes a clocked scale

Use the analytic gradient $g=(3,4)$ , whose norm is $5$ . Global-norm clipping at $2.5$ scales the entire vector to $(1.5,2)$ , preserving its direction. Clipping each coordinate to the interval $[-2.5,2.5]$ instead produces $(2.5,2.5)$ , whose direction and norm differ.

Now consider four microbatches accumulated into one optimizer update. A schedule indexed by optimizer updates should advance once for that group, not four times. The schedule's clock is part of the algorithm and should be restored with optimizer state.

The existing optimizer plot uses its labeled learning rates and shows nonmonotone trajectories, including rebounds. It compares that specific fixture; it does not demonstrate clipping or scheduling. Use the separate clipping diagram to understand the geometry, then test each mechanism with its own controlled change.

A norm ball and a coordinate box

Predict: Does coordinate clipping at $2.5$ guarantee that the whole gradient norm is at most $2.5$ ?

$The blue arrow is $(3,4)$. Green scales it to norm $2.5$ along the same ray. Orange clips each coordinate to $2.5$, landing at the box corner with norm $2.5\sqrt{2}$. The circle and box are analytic constraints, not training measurements.$

Architecture and dataflow mechanism diagram.

The blue arrow is $(3,4)$ . Green scales it to norm $2.5$ along the same ray. Orange clips each coordinate to $2.5$ , landing at the box corner with norm $2.5\sqrt{2}$ . The circle and box are analytic constraints, not training measurements.

Open full-size diagram

Calculate the first two Adam steps

Let $g_t$ be the gradient at update $t$ . Adam tracks a first moment $m_t$ and a second raw moment $v_t$ , using decay rates $\beta_1$ and $\beta_2$ . Starting at zero makes the early averages small; divide by $1-\beta_1^t$ and $1-\beta_2^t$ to correct that bias. At the first step, corrected moments are $g_1$ and $g_1^2$ . With a small $\epsilon$ , each nonzero coordinate initially moves close to the learning rate in the opposite sign direction. This does not mean later steps ignore magnitude.

\begin{aligned}m_t&=\beta_1m_{t-1}+(1-\beta_1)g_t\\v_t&=\beta_2v_{t-1}+(1-\beta_2)g_t^2\\\hat m_t&=m_t/(1-\beta_1^t),\quad \hat v_t=v_t/(1-\beta_2^t)\\\Delta w_t&=-\eta_t\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}\end{aligned}

Clip a direction without rotating it

Global norm clipping rescales a gradient only when its length exceeds a threshold $c$ . For $g=(3,4)$ , the norm is $5$ , so clipping at $c=1$ gives $(0.6,0.8)$ . Coordinate-wise clipping to the same threshold instead gives $(1,1)$ , changing direction. Clipping before Adam limits the gradient entering its moments; it does not impose that same bound on the final parameter update. Clipping cannot repair NaNs, wrong labels or a wrong objective.

\tilde g=g\min\!\left(1,\frac{c}{\|g\|_2}\right)\quad\text{for }\|g\|_2>0

Make the schedule clock explicit

A schedule maps an optimizer update count to a learning rate. Our linear example goes from $0.1$ at count $0$ to $0.01$ at count $4$ , then stays there. Count $0$ is the first update, not an epoch. Gradient accumulation changes how many optimizer updates occur per data pass. Restoring only model weights loses the moments and schedule position; retain optimizer state to continue the same run.

Compare decisions with evidence

Use the same objective, initialization, data and number of updates when comparing SGD, momentum and Adam. Report the loss history and chosen rates; equal updates are not equal wall-clock time, and one rate per method is not a thorough tuning study. An adaptive method does not make scaling or validation irrelevant. Adding an $L_2$ penalty to an Adam loss also differs from decoupled weight decay in AdamW: the penalty gradient is transformed by Adam’s moments, while decoupled decay acts separately on the weights.

Bring the phase together in a training audit

For the regression audit project, first document feature and target shapes, loss reduction and dtype. Use a direct least-squares fit as a reference when its assumptions match. Check gradients at more than one point, inspect feature scales, and preserve optimizer state for replay. If you introduce a penalty or change the batch size, state how that changes the objective or noise. Finish with separate training and held-out losses, an unstable run, and a short explanation of your choice. The existing project checks a subset of these skills; keep the new lesson evidence alongside it rather than treating a project pass as assessment of all twelve lessons.

Compare optimizer paths under a fixed update budget

Predict before looking: Will equal update counts produce equal loss histories?

Read the figure

Start with the axes. The horizontal axis counts completed parameter updates; $0$ is the shared starting point. The vertical axis is the objective $L(w)=\tfrac12(w_0^2+10w_1^2)$ , so lower is better. It is logarithmic: falling from $10^{-2}$ to $10^{-3}$ means a tenfold reduction, not a decrease of one unit. Zero cannot appear on this scale.

All three methods begin at $w=(1,1)$ , with loss $5.5$ . The solid SGD curve drops smoothly and then develops a slower tail. With learning rate $0.05$ , each update multiplies the two coordinates by $0.95$ and $0.5$ . The steep direction shrinks quickly; the flatter direction accounts for the slow remaining progress.

Connect it to the computation

Now follow the dashed momentum curve. Its loss falls to about $0.568$ at update $2$ , then rises to about $3.714$ at update $4$ . Momentum retains past gradients, so a favorable current position does not guarantee a favorable next step: the accumulated update can carry the parameters past the minimum. The repeated dips and rebounds occur under a generally declining envelope. They are deterministic optimizer behavior in this example, not noise from shuffled minibatches.

The dash-dot Adam curve makes the same caution especially clear. It reaches about $1.45\times10^{-4}$ at update $11$ , then rises to about $0.410$ by update $19$ . Adam adapts coordinate-wise update sizes using gradient history, but that history can still move parameters away after a close approach. A deep dip means the current parameters are near the minimum; it does not mean subsequent updates will stay there. The scalar loss alone does not tell us which side of the minimum each coordinate occupies.

At update $50$ , the losses are approximately $1.28\times10^{-4}$ for Adam, $2.96\times10^{-3}$ for SGD, and $2.49\times10^{-2}$ for momentum. Adam finishes lowest under these settings, although it was not lowest at every earlier update. Compare the same update budget, read the endpoint as well as the dips, and remember that equal update counts do not imply equal runtime. This figure uses fixed rates without clipping or a schedule; the later experiments examine those separate choices.

Open full-size figure Download data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Where to put & run this code: Work in phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py (or downloaded exercises/optimization-12.py / exercises/optimization-12.ipynb) and execute from the repository root with python3 phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py.

1 Build a transparent Adam reference

Create main.py. Use an explicit two-gradient sequence to inspect the optimizer independently of a training loop.

main.py · append build step 1

# Step 1 — Build a transparent Adam reference: The supplied gradients isolate optimizer arithmetic; they are not...
import jax
import jax.numpy as jnp
import optax
params = jnp.zeros(2)
# Configure or step the Optax optimizer state (`adam`).
adam = optax.adam(0.1, b1=0.9, b2=0.99, eps=1e-08)
state = adam.init(params)
m = jnp.zeros(2)
v = jnp.zeros(2)
gradients = [jnp.array([2.0, -4.0]), jnp.array([1.0, 3.0])]

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
optax.adam(lr) / optax.apply_updates(params, updates) — Optax gradient transformations and numerically stable loss functions over parameter PyTrees.

The supplied gradients isolate optimizer arithmetic; they are not presented as a model’s measured training gradients.

2 Verify both moment updates

Append the manual recurrences and compare each update with Optax.

main.py · append build step 2

# Step 2 — Verify both moment updates: The first update is near (-0.1,0.1); the next depends on both...
for t, g in enumerate(gradients, start=1):
    m = 0.9 * m + 0.1 * g
    v = 0.99 * v + 0.01 * g * g
    expected = -0.1 * (m / (1 - 0.9 ** t)) / (jnp.sqrt(v / (1 - 0.99 ** t)) + 1e-08)
    update, state = adam.update(g, state, params)
    assert jnp.allclose(update, expected, atol=2e-06, rtol=2e-06)
    # Apply the computed gradient updates to update the model parameters.
    params = optax.apply_updates(params, update)
    print('step / update:', t, update)

Code building blocks in this step

jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

The first update is near $(-0.1,0.1)$ ; the next depends on both gradients.

3 Inspect one clipped update

Append a global clipping plus SGD chain and run python main.py.

main.py · append build step 3

# Step 3 — Inspect one clipped update: For this SGD chain, the update norm is 0.1.
g = jnp.array([3.0, 4.0])
# Configure or step the Optax optimizer state (`clip`).
clip = optax.clip_by_global_norm(1.0)
clipped, _ = clip.update(g, clip.init(jnp.zeros(2)))
assert jnp.allclose(clipped, jnp.array([0.6, 0.8]), atol=1e-06)
# Configure or step the Optax optimizer state (`chain`).
chain = optax.chain(optax.clip_by_global_norm(1.0), optax.sgd(0.1))
update, _ = chain.update(g, chain.init(jnp.zeros(2)))
assert jnp.allclose(update, jnp.array([-0.06, -0.08]), atol=1e-06)

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
optax.adam(lr) / optax.apply_updates(params, updates) — Optax gradient transformations and numerically stable loss functions over parameter PyTrees.

For this SGD chain, the update norm is $0.1$ . That bound does not automatically apply to a chain containing Adam.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — Build a transparent Adam reference: The supplied gradients isolate optimizer arithmetic; they are not...
import jax
import jax.numpy as jnp
import optax
params = jnp.zeros(2)
# Configure or step the Optax optimizer state (`adam`).
adam = optax.adam(0.1, b1=0.9, b2=0.99, eps=1e-08)
state = adam.init(params)
m = jnp.zeros(2)
v = jnp.zeros(2)
gradients = [jnp.array([2.0, -4.0]), jnp.array([1.0, 3.0])]

# Step 2 — Verify both moment updates: The first update is near (-0.1,0.1); the next depends on both...
for t, g in enumerate(gradients, start=1):
    m = 0.9 * m + 0.1 * g
    v = 0.99 * v + 0.01 * g * g
    expected = -0.1 * (m / (1 - 0.9 ** t)) / (jnp.sqrt(v / (1 - 0.99 ** t)) + 1e-08)
    update, state = adam.update(g, state, params)
    assert jnp.allclose(update, expected, atol=2e-06, rtol=2e-06)
    # Apply the computed gradient updates to update the model parameters.
    params = optax.apply_updates(params, update)
    print('step / update:', t, update)

# Step 3 — Inspect one clipped update: For this SGD chain, the update norm is 0.1.
g = jnp.array([3.0, 4.0])
# Configure or step the Optax optimizer state (`clip`).
clip = optax.clip_by_global_norm(1.0)
clipped, _ = clip.update(g, clip.init(jnp.zeros(2)))
assert jnp.allclose(clipped, jnp.array([0.6, 0.8]), atol=1e-06)
# Configure or step the Optax optimizer state (`chain`).
chain = optax.chain(optax.clip_by_global_norm(1.0), optax.sgd(0.1))
update, _ = chain.update(g, chain.init(jnp.zeros(2)))
assert jnp.allclose(update, jnp.array([-0.06, -0.08]), atol=1e-06)

Run the example

python3 phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py

About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example

main.py · complete Python example

# Step 1 — Build a transparent Adam reference: The supplied gradients isolate optimizer arithmetic; they are not...
import jax
import jax.numpy as jnp
import optax
params = jnp.zeros(2)
# Configure or step the Optax optimizer state (`adam`).
adam = optax.adam(0.1, b1=0.9, b2=0.99, eps=1e-08)
state = adam.init(params)
m = jnp.zeros(2)
v = jnp.zeros(2)
gradients = [jnp.array([2.0, -4.0]), jnp.array([1.0, 3.0])]

# Step 2 — Verify both moment updates: The first update is near (-0.1,0.1); the next depends on both...
for t, g in enumerate(gradients, start=1):
    m = 0.9 * m + 0.1 * g
    v = 0.99 * v + 0.01 * g * g
    expected = -0.1 * (m / (1 - 0.9 ** t)) / (jnp.sqrt(v / (1 - 0.99 ** t)) + 1e-08)
    update, state = adam.update(g, state, params)
    assert jnp.allclose(update, expected, atol=2e-06, rtol=2e-06)
    # Apply the computed gradient updates to update the model parameters.
    params = optax.apply_updates(params, update)
    print('step / update:', t, update)

# Step 3 — Inspect one clipped update: For this SGD chain, the update norm is 0.1.
g = jnp.array([3.0, 4.0])
# Configure or step the Optax optimizer state (`clip`).
clip = optax.clip_by_global_norm(1.0)
clipped, _ = clip.update(g, clip.init(jnp.zeros(2)))
assert jnp.allclose(clipped, jnp.array([0.6, 0.8]), atol=1e-06)
# Configure or step the Optax optimizer state (`chain`).
chain = optax.chain(optax.clip_by_global_norm(1.0), optax.sgd(0.1))
update, _ = chain.update(g, chain.init(jnp.zeros(2)))
assert jnp.allclose(update, jnp.array([-0.06, -0.08]), atol=1e-06)

Expected result: Two Adam updates match the manual moment calculation; the clipped SGD update is $(-0.06,-0.08)$ .

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:46:04.755758+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

step / update: 1 [-0.09999993  0.09999993]
step / update: 2 [-0.09334469  0.00893815]
step / update: 1 [-0.09999993  0.09999993]
step / update: 2 [-0.09334469  0.00893815]
schedule rates: [0.1    0.0775 0.055  0.0325 0.01   0.01  ]
sgd initial / final loss: 5.5 0.0029602647
momentum initial / final loss: 5.5 0.024945451
adam initial / final loss: 5.5 0.00012754407
PASS: optimization-12

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Inspect the schedule boundary

Inspect the schedule boundary · Python experiment

# Experiment — Inspect the schedule boundary: A schedule is indexed by optimizer updates.
# Configure or step the Optax optimizer state (`schedule`).
schedule = optax.linear_schedule(init_value=0.1, end_value=0.01, transition_steps=4)
rates = jnp.array([schedule(i) for i in range(6)])
assert jnp.allclose(rates, jnp.array([0.1, 0.0775, 0.055, 0.0325, 0.01, 0.01]), atol=1e-06)
# Configure or step the Optax optimizer state (`tx`).
tx = optax.sgd(schedule)
p = jnp.array(0.0)
s = tx.init(p)
for _ in range(5):
    delta, s = tx.update(jnp.array(1.0), s, p)
    # Configure or step the Optax optimizer state (`p`).
    p = optax.apply_updates(p, delta)
assert jnp.allclose(p, -0.275, atol=1e-06)
print('schedule rates:', rates)

Expected: The fifth update uses count $4$ and rate $0.01$ .

A schedule is indexed by optimizer updates. Recording only the number of epochs leaves its behavior ambiguous.

Compare a fixed update budget

Compare a fixed update budget · Python experiment

# Experiment — Compare a fixed update budget: This is a controlled demonstration at stated rates and update...
def bowl(w):
    return 0.5 * (w[0] ** 2 + 10 * w[1] ** 2)
start = jnp.array([1.0, 1.0])
# Configure or step the Optax optimizer state (`optimizers`).
optimizers = {'sgd': optax.sgd(0.05), 'momentum': optax.sgd(0.05, momentum=0.9), 'adam': optax.adam(0.1)}
traces = {}
for name, tx in optimizers.items():

    # Define `step(carry, _)` to evaluate the objective and its automatic derivatives:
    def step(carry, _):
        p, s = carry
        # Differentiate the objective to obtain `(delta, s)` via automatic differentiation.
        delta, s = tx.update(jax.grad(bowl)(p), s, p)
        # Configure or step the Optax optimizer state (`p`).
        p = optax.apply_updates(p, delta)
        return ((p, s), bowl(p))
    # Run compiled structured control flow via `jax.lax` (`((_, _), history)`).
    (_, _), history = jax.lax.scan(step, (start, tx.init(start)), None, length=50)
    # Confirm that all computed values remain finite (no NaN or Inf).
    assert jnp.all(jnp.isfinite(history))
    assert history[-1] < bowl(start)
    traces[name] = history
    print(name, 'initial / final loss:', bowl(start), history[-1])
assert not jnp.allclose(traces['sgd'], traces['adam'])

Expected: All three reduce the loss for these specified settings; the trajectories differ.

This is a controlled demonstration at stated rates and update counts, not evidence that one optimizer universally wins.

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py (or inside the Make it yours cell of exercises/optimization-12.ipynb) and run python3 phases/04-optimization/12-adam-clipping-and-learning-rate-schedules/code/main.py to verify your prediction before opening the reference solution.

Compare global norm clipping with coordinate-wise clipping on $(3,4)$ . Check which preserves the ratio of the coordinates.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Compare global norm clipping with coordinate-wise clipping on (3,4).
coordinate = jnp.clip(...)  # TODO: compute coordinate
assert jnp.allclose(coordinate, jnp.array([1.0, 1.0]))  # TODO: complete assertion check
assert jnp.allclose(clipped[0] / clipped[1], g[0] / g[1])  # TODO: complete assertion check
assert not jnp.allclose(coordinate[0] / coordinate[1], g[0] / g[1])  # TODO: complete assertion check

Reveal the reference solution

Reference solution

# Exercise solution: Compare global norm clipping with coordinate-wise clipping on (3,4).
coordinate = jnp.clip(g, -1.0, 1.0)
assert jnp.allclose(coordinate, jnp.array([1.0, 1.0]))
assert jnp.allclose(clipped[0] / clipped[1], g[0] / g[1])
assert not jnp.allclose(coordinate[0] / coordinate[1], g[0] / g[1])

Catch schedule restart

Transfer / diagnosis

After five scheduled SGD steps with constant gradient $1$ , compare a retained optimizer state with a freshly initialized state for the next update.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
optax.adam(lr) / optax.apply_updates(params, updates) — Optax gradient transformations and numerically stable loss functions over parameter PyTrees.

2. Step-by-step implementation plan:

Configure or step the Optax optimizer state (tx).
Configure or step the Optax optimizer state (p).

Starter code scaffold · copy into your file & fill in the TODOs

# Catch schedule restart (Transfer / diagnosis): Restoring weights without optimizer state can silently...
# Configure or step the Optax optimizer state (`tx`).
tx = optax.sgd(...)  # TODO: compute tx
p = jnp.array(...)  # TODO: compute p
s = tx.init(...)  # TODO: compute s
for _ in range(5):
    delta, s = tx.update(...)  # TODO: compute delta, s
    # Configure or step the Optax optimizer state (`p`).
    p = optax.apply_updates(...)  # TODO: compute p
continued, _ = tx.update(...)  # TODO: compute continued, _
restarted, _ = tx.update(...)  # TODO: compute restarted, _
assert jnp.allclose(continued, -0.01, atol=1e-06)  # TODO: complete assertion check
assert jnp.allclose(restarted, -0.1, atol=1e-06)  # TODO: complete assertion check

Show a hint

The retained count has reached the final rate; a reset count starts at $0$ .

Reveal reference solution and reasoning

Catch schedule restart · Python solution

# Catch schedule restart (Transfer / diagnosis): Restoring weights without optimizer state can silently...
# Configure or step the Optax optimizer state (`tx`).
tx = optax.sgd(schedule)
p = jnp.array(0.0)
s = tx.init(p)
for _ in range(5):
    delta, s = tx.update(jnp.array(1.0), s, p)
    # Configure or step the Optax optimizer state (`p`).
    p = optax.apply_updates(p, delta)
continued, _ = tx.update(jnp.array(1.0), s, p)
restarted, _ = tx.update(jnp.array(1.0), tx.init(p), p)
assert jnp.allclose(continued, -0.01, atol=1e-06)
assert jnp.allclose(restarted, -0.1, atol=1e-06)

Restoring weights without optimizer state can silently restart the learning-rate schedule.

Check transformation order

Transfer / diagnosis

For gradient $(3,4)$ , compare clipping before SGD at rate $0.1$ with clipping after it. Write a decision memo explaining which quantity you intended to bound.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
optax.adam(lr) / optax.apply_updates(params, updates) — Optax gradient transformations and numerically stable loss functions over parameter PyTrees.

2. Step-by-step implementation plan:

Configure or step the Optax optimizer state (before).
Configure or step the Optax optimizer state (after).

Starter code scaffold · copy into your file & fill in the TODOs

# Check transformation order (Transfer / diagnosis): Transformation order changes the algorithm.
# Configure or step the Optax optimizer state (`before`).
before = optax.chain(...)  # TODO: compute before
# Configure or step the Optax optimizer state (`after`).
after = optax.chain(...)  # TODO: compute after
u_before, _ = before.update(...)  # TODO: compute u_before, _
u_after, _ = after.update(...)  # TODO: compute u_after, _
assert jnp.allclose(jnp.linalg.norm(u_before), 0.1, atol=1e-06)  # TODO: complete assertion check
assert jnp.allclose(jnp.linalg.norm(u_after), 0.5, atol=1e-06)  # TODO: complete assertion check

Show a hint

The un-clipped SGD update already has norm $0.5$ , below the threshold $1$ .

Reveal reference solution and reasoning

Check transformation order · Python solution

# Check transformation order (Transfer / diagnosis): Transformation order changes the algorithm.
# Configure or step the Optax optimizer state (`before`).
before = optax.chain(optax.clip_by_global_norm(1.0), optax.sgd(0.1))
# Configure or step the Optax optimizer state (`after`).
after = optax.chain(optax.sgd(0.1), optax.clip_by_global_norm(1.0))
u_before, _ = before.update(g, before.init(jnp.zeros(2)))
u_after, _ = after.update(g, after.init(jnp.zeros(2)))
assert jnp.allclose(jnp.linalg.norm(u_before), 0.1, atol=1e-06)
assert jnp.allclose(jnp.linalg.norm(u_after), 0.5, atol=1e-06)

Transformation order changes the algorithm. State whether the bound applies to the raw gradient or the parameter update, and retain the observed evidence.

Diagnose the result

If a restart changes the trajectory, compare optimizer moments and schedule count. If a clipped run still jumps, measure both gradient and update norms and inspect transformation order. Record rates, reductions, dtype and validation behavior before recommending a method.

Carry forward

A schedule is indexed by optimizer updates. Recording only the number of epochs leaves its behavior ambiguous.
This is a controlled demonstration at stated rates and update counts, not evidence that one optimizer universally wins.

Keep your evidence

Download evidence template

Keep two hand-checked Adam steps, clipping-direction checks, schedule values and boundary, equal-budget loss traces, and the optimization decision memo.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references

Phase 04 Integration Project

Bring this phase together: A fitted regression model

Checkpoint to prove: Explain prediction geometry and loss assumptions; verify derivatives; diagnose conditioning, curvature and gradient noise; justify an optimizer using held-out evidence.

Complete the regression audit after the four core lessons. This deeper project follows all twelve optimization lessons and connects curvature, conditioning, gradient noise, optimizer state and regularization. Copy projects/optimizer-audit/starter/optimization.py to projects/optimizer-audit/my_optimization.py and write your code in projects/optimizer-audit/my_optimization.py. Run python3 projects/optimizer-audit/tests/check.py --implementation projects/optimizer-audit/my_optimization.py --stage 1 from the top-level folder to verify stages 1, 2, 3, 4, 5.

Open project workspace · stages 1, 2, 3, 4, 5 Read project guide & figures

Adam, clipping, and learning-rate schedules

What you will be able to do

The problem

The idea

Clipping changes a vector; a schedule changes a clocked scale

A norm ball and a coordinate box

Calculate the first two Adam steps

Clip a direction without rotating it

Make the schedule clock explicit

Compare decisions with evidence

Bring the phase together in a training audit

Compare optimizer paths under a fixed update budget

Read the figure

Connect it to the computation

1 Build a transparent Adam reference

2 Verify both moment updates

3 Inspect one clipped update

Run the example

Recorded reference execution

Inspect the schedule boundary

Compare a fixed update budget

Make it yours

Catch schedule restart

Check transformation order

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references

Bring this phase together: A fitted regression model