The training loop·30 min

The Adam step

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage the-adam-step --watch

Gradients tell you the direction; Adam decides the step. We use plain Adam — β1 = 0.9, β2 = 0.999, ε = 1e-8, no weight decay — the same optimizer torch.optim.Adam runs. (If you came from train-gpt, that course used AdamW with decoupled decay and β2 = 0.95; here you set the decay to zero and β2 to 0.999. The moment buffers and bias correction are identical.)

one Adam update per parameter

m = β1·m + (1−β1)·g          v = β2·v + (1−β2)·g²
m̂ = m / (1 − β1ᵗ)            v̂ = v / (1 − β2ᵗ)        # bias correction
θ = θ − lr · m̂ / (√v̂ + ε)

What to build

Your task

Implement the Adam step

Implement adamStep over all five parameter tensors (E, W1, b1, W2, b2), maintaining per-parameter first/second-moment buffers and a global step count t. The graded adapter runAdamStepFlat(seed, gradSeed, lr, steps) inits params from seed, uses a fixed seeded gradient (gradSeed), applies steps updates, and returns the flattened params.

tests/adapters.ts

export function runAdamStepFlat(
  seed: number, gradSeed: number, lr: number, steps: number,
): number[] {
  // initParams(seed); grad = initParams(gradSeed); apply 'steps' Adam updates
  // beta1=0.9, beta2=0.999, eps=1e-8; return flattenParams(p)
}

How it's graded — 1-step exact

Adapter mode: from the seeded init, applying 1 step must match the reference weights with allclose, atol 1e-6, rtol 1e-5 — this pins the update rule exactly. (Many-step runs would compound float drift; this course grades the single step tightly and leaves multi-step training to the process-mode loss check.)

Global step, not per-row

Bias correction uses the global step count t, incremented once per adamStep call — not a per-parameter or per-row counter. Even though the table gradient is sparse, t advances for the whole update, or the bias correction 1 − β1ᵗ drifts.

Why finite differences?The batcher