Skip to content

← Optax step 3 of 25

Medium Primitives

Optax Adam Step

Jump to the task

Adam uses recent gradients to choose a separate update scale for each parameter. It keeps two buffers: m, an average of gradients, and v, an average of squared gradients. The second buffer measures gradient size; it is not a variance.

From buffers to an update

With the defaults used here, each step updates the buffers as follows:

m = 0.9 * m + 0.1 * grads
v = 0.999 * v + 0.001 * grads**2

Both start at zero. Adam corrects for that starting point before computing the parameter update. With t counting updates from 1:

m_hat = m / (1 - 0.9**t)
v_hat = v / (1 - 0.999**t)
new_params = params - lr * m_hat / (sqrt(v_hat) + 1e-8)

These operations apply independently to each element. The small 1e-8 constant prevents division by zero. It sits outside the square root; Optax's separate eps_root setting is zero by default.

What happens on the first step

This exercise performs one update from freshly initialized state. In exact arithmetic, bias correction gives m_hat = grads and v_hat = grads**2. With learning rate 0.01, the result is therefore:

new_params = params - 0.01 * grads / (abs(grads) + 1e-8)

For a gradient much larger in magnitude than 1e-8, the parameter moves by about 0.01 against the gradient's sign. A gradient of zero gives no change. For a gradient of 1e-8, the change is about -0.005, so replacing this rule with a sign function would be incorrect. These are first-step properties, not a fixed step-size rule for the rest of training.

Inputs

Implement adam_step_loss(params, grads) using optax.adam(learning_rate=0.01) with its other settings left at their defaults. Initialize state, calculate one update and apply it.

  • params: a 1-D floating-point JAX array of current parameters.
  • grads: a floating-point JAX array with the same shape.

The gradients are supplied. Despite the function's name, you do not compute or return a loss. In a full training loop, initialize only once and carry the returned state between steps.

Output

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