Skip to content

← Optax step 23 of 25

Hard Primitives

Train Step with Global-Norm Clipping

Jump to the task

Write the parameter-update part of a training step: clip the supplied gradients by their global L2 norm, then apply SGD. The loss and gradients have already been computed. This function does not run a model or calculate derivatives.

Limit the gradient before taking the step

Compute the gradient length as sqrt(sum(grads**2)). If it exceeds max_norm, multiply every gradient element by max_norm / norm. Otherwise keep the values unchanged, including a zero gradient vector. Then update the parameters with:

new_params = params - lr * clipped_grads

For params = [1.0, 2.0, 3.0], grads = [3.0, 4.0, 0.0], lr = 0.1 and max_norm = 1.0, the gradient norm is 5.0. Clipping gives [0.6, 0.8, 0.0], so the new parameters are [0.94, 1.92, 3.0].

The threshold bounds the gradient vector's length, not each element independently. Optax can also compute a global norm across a nested collection of gradient arrays; this task uses one 1-D array.

Compose the update with Optax

Place optax.clip_by_global_norm(max_norm) before optax.sgd(lr) in an optax.chain. Initialize its state from params, call update, then apply the returned updates with optax.apply_updates.

The order expresses the rule above. Clipping after SGD would instead limit the already scaled parameter updates. That is a different operation, not a meaningless one. Here the final update norm is at most lr * max_norm, up to floating-point rounding.

Clipping limits large gradient-driven steps. It does not guarantee that the loss decreases, and there is no threshold that is best for every model.

Inputs

Implement train_with_clip(params, grads, lr, max_norm) with the chain above. Use plain SGD without momentum and perform one update.

  • params: a 1-D floating-point JAX array of current parameters.
  • grads: a floating-point JAX array with the same shape.
  • lr: a nonnegative scalar learning rate.
  • max_norm: a positive scalar threshold for the gradient norm.

Output

Return the updated JAX array with the same shape as params. Leave the input arrays unchanged. Return the parameters, not the raw updates or optimizer state.