Adam + Weight Decay (Chain)

This exercise adds an L2 penalty's gradient before applying Adam. The added term is proportional to the current parameter values:

effective_grad = grads + wd * params

This is the gradient you would get by adding 0.5 * wd * sum(params**2) to the original loss. Adam then processes the whole effective gradient, including the penalty term.

Build the chain in this order

Use optax.add_decayed_weights(wd) first, followed by optax.adam(lr). The first transform adds wd * params to its input. The second updates Adam's buffers and produces the signed, learning-rate-scaled updates. Finally, optax.apply_updates adds those updates to the parameters.

The first transform needs the parameter values, so pass params to the chain's update method along with the gradients and state.

How this differs from AdamW

In this chain, the penalty contributes to Adam's gradient averages and adaptive scaling. AdamW keeps weight decay outside those calculations and adds a separate update of -lr * wd * params. The two rules generally produce different results.

For example, take one parameter 2.0, gradient -0.1, wd = 0.1 and lr = 0.01. This chain gives an effective gradient of -0.1 + 0.1 * 2.0 = 0.1. On a fresh first Adam step, the parameter therefore decreases to about 1.99.

AdamW instead uses the negative gradient for its Adam contribution, which increases the parameter by about 0.01, then subtracts 0.002 for decay. Its result is about 2.008. The penalty can change even the direction of this chain's update relative to AdamW.