Skip to content

← Optax step 2 of 25

Easy Primitives

Optax SGD with Momentum

Jump to the task

Momentum lets an SGD update depend on earlier gradients as well as the current one. A buffer called velocity stores that history. With ordinary momentum (not Nesterov momentum), each step uses:

new_velocity = momentum * velocity + grads
new_params = params - lr * new_velocity

Repeated gradients in the same direction build up in the buffer. Gradients in opposing directions can cancel some of that history. Whether this helps training depends on the problem and the learning rate.

The first step has no history

This exercise starts with zero velocity and performs just one step. Therefore new_velocity = momentum * 0 + grads = grads, and the returned parameters match plain SGD. Changing the momentum coefficient cannot change this first result. Use optax.sgd(lr, momentum=momentum) to practice the optimizer API; leave Nesterov momentum disabled, as it is by default.

Why carrying state matters

Consider one parameter starting at 1.0, with learning rate 0.1 and momentum 0.5. Supply gradient 2.0 on the first step and 0.0 on the next:

Step Gradient New velocity New parameter
1 2.0 0.5 * 0 + 2 = 2.0 1.0 - 0.1 * 2 = 0.8
2 0.0 0.5 * 2 + 0 = 1.0 0.8 - 0.1 * 1 = 0.7

The second step still moves because the buffer retains part of the previous gradient. Here is a separate two-step example:

import jax.numpy as jnp
import optax

params = jnp.array([1.0])
optimizer = optax.sgd(0.1, momentum=0.5)
state = optimizer.init(params)
for grads in (jnp.array([2.0]), jnp.array([0.0])):
    updates, state = optimizer.update(grads, state, params)
    params = optax.apply_updates(params, updates)
# params is approximately [0.7].

Initializing state inside that loop would erase the velocity. The second step would then leave the parameter at 0.8. The function you implement below deliberately covers only a fresh first step, not this full loop.

Inputs

Implement momentum_step(params, grads, lr, momentum) using Optax. Initialize the optimizer state, calculate one update, and apply it.

  • 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.
  • momentum: a scalar momentum coefficient, such as 0.9.

Output

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

The widget below illustrates how a buffer carries values across steps. Its rule is carry = 0.5 * carry + x, so it corresponds to velocity with momentum 0.5 and the displayed inputs as gradients. It does not apply parameter updates.

Loading visualization…