Gradient Clipping by Global Norm

The problem

One bad batch can produce a gradient hundreds of times larger than usual, and a single step with it can wreck a training run. Clipping is the guard: whenever the gradients are too large overall, shrink them all by the same factor, so that their direction stays the same and only their size changes.

A model has many parameter tensors (one weight matrix per layer, the biases, the embeddings), and each has its own gradient array. "Too large overall" is measured on all of them together.

Step 1: the global norm

The global norm is the length of all the gradients together, as if every array were flattened and laid end to end in one long vector: square every entry of every array, add everything up, take the square root.

where:

• $g_k$ is the gradient array of the $k$-th parameter tensor;
• $g_{k,i}$ is its $i$-th entry, and the double sum runs over every entry of every array.

Step 2: one scale factor for everything

where:

• $c$ is max_norm, the largest norm allowed;
• if the norm is above $c$, the factor is about $c / \lVert g \rVert$, which brings the norm down to $c$;
• if the norm is already below $c$, $\min(1, \cdot)$ keeps the factor at 1, so nothing changes;
• the $10^{-6}$ is what PyTorch's clip_grad_norm_ adds, to avoid a division by zero when all gradients are 0.

Example. Two parameter tensors with gradients [3.0] and [4.0, 0.0], and max_norm = 1.

1. Global norm: $\sqrt{3^2 + 4^2 + 0^2} = \sqrt{25} = 5$.
2. Factor: $\min(1, 1 / 5.000001) \approx 0.2$.
3. The arrays become about [0.6] and [0.8, 0.0], whose global norm is 1. The second entry is still $4/3$ of the first: same direction.

The function returns 5.0, the norm before clipping. With max_norm = 10 the factor would be 1 and nothing would change.

What to implement

clip_grad_norm(grads, max_norm):

• grads is a list of NumPy arrays of any shapes, one per parameter tensor.
• Scale every array in place, like PyTorch does: the caller's arrays must change.
• Return the global norm before clipping, as a float. Training loops log it to spot trouble.

Requirements

• One scale factor for all arrays. Clipping each array to max_norm separately would change the direction of the overall gradient: in the example, [3.0] would become [1.0] and [4.0, 0.0] would become [1.0, 0.0]. The two entries would end up equal, while the true gradient says the second is $4/3$ of the first.
• When the norm is below max_norm, the gradients must not change.
• No Python loops over elements (a loop over the list of arrays is fine).

Hint

"In place" means writing into the existing array: g *= scale or g[...] = .... Rebinding the name, g = g * scale, creates a new array that the caller never sees.