Build a mini deep learning framework from scratch

Same XOR points from Chapter 2, same tiny network shape: 2 inputs, 4 hidden neurons, 1 output. Click train and watch the shaded decision region fold around the two "on" corners — except this time the optimizer behind every step is AdamW, not plain gradient descent.

AdamW is Adam plus one change: weight decay applied directly to the parameter, decoupled from the gradient-based moment estimates.

• $\hat m, \hat v$ — Chapter 9's bias-corrected running averages of the gradient and its square.
• $\theta$ — the parameter being updated: any single weight or bias in the network's Tensor.
• $\eta$ — the learning rate.
• $\lambda$ — the weight decay coefficient, shrinking every parameter a little on every step.

1. Everything except one term is Adam, unchanged

   $\hat m/(\sqrt{\hat v}+\epsilon)$ is exactly Chapter 9's Adam update — nothing new there.

2. The new piece: decoupled weight decay

   $\lambda\theta$ is added straight onto the update, applied to the raw parameter — not folded into the gradient the way plain L2 regularization would be. That decoupling is the "W" in AdamW.

3. One update rule, every parameter in the network

   The same formula runs once per weight and bias — 17 parameters total in this tiny network, billions in a real one — each keeping its own $m$, $v$, and $t$, all driven by the one gradient backprop computed for it.

Watch the four individual outputs. AdamW's adaptive step size means the two that need to climb toward $1$ and the two that need to fall toward $0$ move faster and more evenly than plain gradient descent managed on this exact same network in Chapter 2 — the per-parameter scaling matters even on a network this small.

From this network's fixed random start, before any training:

1. Pick one parameter: the weight from hidden neuron 1 to the output

   $\theta = 0.166168$. A full backward pass through all four XOR points gives its gradient: $g = -0.012081$.

2. At the very first step, the moving averages start from zero

   $t=1$, so:

   • $\hat m = m/(1-\beta_1) = g$ exactly.
   • $\hat v = v/(1-\beta_2) = g^2$ exactly.

   The bias correction perfectly cancels the very first $(1-\beta)$ factor on both averages.

3. That collapses the ratio to just the gradient's sign

   $\hat m/\sqrt{\hat v} = g/|g| = -1$ — same "first step ignores magnitude" behavior Chapter 9 found for RMSProp and Adam.

4. Add the weight decay term and take the step

   With $\eta=0.1$ and $\lambda=0.01$: update $= 0.1\,(-1 + 0.01\times0.166168) \approx -0.099834$, so $\theta \leftarrow 0.166168 - (-0.099834) = 0.266002$.