Global-Norm Gradient Clipping

Global-norm clipping limits the length of a gradient vector. If that length exceeds a chosen threshold, it scales every element by the same factor. This preserves the vector's direction while reducing its magnitude. Smaller gradients pass through unchanged.

Clipping can limit unusually large gradient-driven steps. It does not fix an unsuitable learning rate or guarantee stable training.

One scale factor for the whole vector

For a vector grads, its L2 norm is:

norm = sqrt(sum(grads**2))

If norm > max_norm, return grads * (max_norm / norm). Otherwise, return the original values. In particular, a zero vector stays zero; it does not need to be divided by its norm.

For grads = [3.0, 4.0], the norm is 5.0. With max_norm = 1.0, multiply both elements by 1 / 5 to get [0.6, 0.8]. Its norm is 1.0, up to floating-point rounding.

This differs from clipping each element to the interval [-1, 1], which would give [1.0, 1.0]. Elementwise clipping can change direction, and its result can still have a vector norm larger than 1.0.

Using the Optax transform

Create optax.clip_by_global_norm(max_norm), initialize its state and call its update method. The returned pair contains the clipped gradients and new state. This transform does not keep a history between calls.

Optax also supports nested collections of arrays, where "global" includes all elements across the collection. This exercise uses just one 1-D array.