Skip to content

← Optax step 11 of 25

Medium Primitives

Chain: Clip + SGD

Jump to the task

An Optax chain applies transformations in the order you give them. Here, first limit the gradient's L2 norm, then let plain SGD turn the clipped gradient into a parameter update.

Follow the values through the chain

For params = [1.0, 2.0], grads = [3.0, 4.0], max_norm = 1.0 and lr = 0.1:

  1. The gradient norm is 5.0, so clipping multiplies both elements by 1 / 5, giving [0.6, 0.8].
  2. SGD multiplies by -lr, giving updates [-0.06, -0.08].
  3. Adding these updates to the parameters gives [0.94, 1.92].

If the gradient norm is already at or below the threshold, keep its values unchanged before the SGD step. A zero gradient produces no change.

Why the order matters

Build the chain with optax.clip_by_global_norm(max_norm) first and optax.sgd(lr) second. The first transform limits the gradient; the second sets the direction and scale of the parameter change.

Reversing them limits the norm of the already scaled updates instead. In the example above, SGD first would produce [-0.3, -0.4], whose norm is 0.5. A threshold of 1.0 would leave it unchanged, giving parameters [0.7, 1.6]. That is a different rule from the one requested here.

For this plain SGD chain with a nonnegative learning rate, the final update norm is at most lr * max_norm, up to floating-point rounding. The threshold does not bound the parameter values themselves.

Inputs

Implement clip_then_sgd(params, grads, lr, max_norm) using optax.chain. Initialize its 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 nonnegative scalar learning rate.
  • max_norm: a positive scalar threshold for the gradient's L2 norm.

Use SGD without momentum. The gradients are supplied; you do not need to compute a loss or differentiate it.

Output

Return the updated JAX array with the same shape as params. Leave the input arrays unchanged. Use optax.apply_updates once: the chain's updates already contain the negative sign and learning-rate scaling.