Optax SGD Step

SGD updates each parameter using its loss gradient:

new_params = params - lr * grads

For a positive learning rate, a positive gradient decreases the parameter; a negative gradient increases it. The gradient gives a local direction, not a guarantee that any chosen step size will reduce the loss.

For params = [1.0, 2.0, 3.0], grads = [0.1, 0.2, 0.3] and lr = 0.5, the updates are [-0.05, -0.10, -0.15]. Adding them to the parameters gives [0.95, 1.90, 2.85].

Updates and parameters

Optax separates calculating an update from applying it. Create an optimizer with optax.sgd(lr), initialize its state with optimizer.init(params), then call optimizer.update(grads, state, params). This returns two values:

• updates: the changes to add to the parameters, already scaled by the learning rate and signed for gradient descent.
• new_state: the optimizer state to use on the next step.

Use optax.apply_updates(params, updates) to produce the updated parameters. Do not subtract the updates or multiply them by the learning rate again. The supplied gradients are already computed; this function does not need to differentiate a loss.

What state means here

With a fixed learning rate and no momentum, plain SGD has no gradient history to remember. Other optimizers use state for quantities such as a momentum buffer or a step counter. In a training loop, initialize once and carry the returned state into each subsequent update.

Optax exposes state explicitly rather than changing it inside an optimizer object. Passing params to update is optional for plain SGD, but needed by some other transformations, such as parameter-dependent weight decay.