Skip to content

← Optax step 4 of 25

Medium Primitives

AdamW with Decoupled Weight Decay

Jump to the task

AdamW combines Adam's gradient-based update with weight decay. Weight decay adds a separate change proportional to each parameter's current value. The decay term does not pass through Adam's gradient buffers or adaptive scaling; this separation is what "decoupled" means.

Two contributions to one update

Let m_hat and v_hat be Adam's bias-corrected averages of gradients and squared gradients. With the default epsilon settings, Optax's AdamW update is:

direction = m_hat / (sqrt(v_hat) + 1e-8)
new_params = params - lr * direction - lr * wd * params

The decay term uses the original params, not the parameters after the Adam update. Optax includes both contributions in the returned updates. Do not apply weight decay a second time.

Adding an L2 penalty to a loss works differently. A penalty of 0.5 * wd * sum(params**2) adds wd * params to the loss gradient. If that gradient goes into Adam, the penalty affects its buffers and adaptive scaling. It generally produces a different update from AdamW.

A first-step example

This exercise initializes fresh state and performs one update. With zero gradients, Adam's contribution is zero, leaving:

new_params = (1 - lr * wd) * params

For params = [2.0, -4.0], lr = 0.1 and wd = 0.5, the result is [1.9, -3.8]. Each magnitude shrinks by 5 percent, not by a fixed amount of 0.05. This shrinkage interpretation assumes 0 < lr * wd < 1.

With state retained from earlier steps, zero current gradients can still produce an Adam update because the buffers may be nonzero. The decay-only example above relies on fresh state. Setting wd = 0 instead gives the same update as Adam with matching settings and state.

Inputs

Implement adamw_step(params, grads, lr, wd) using optax.adamw(lr, weight_decay=wd). Leave the other settings at their defaults, initialize state and perform one update.

  • params: a 1-D floating-point JAX array of parameter values.
  • grads: a floating-point JAX array with the same shape.
  • lr: a scalar learning rate.
  • wd: a scalar weight decay coefficient.

Pass params to optimizer.update: AdamW needs the current values to calculate decay. In a full training loop, keep the returned optimizer state for the next step instead of initializing again.

Output

Return the updated JAX array with the same shape as params. Leave the input arrays unchanged. Do not return the optimizer state.