Global Norm Gradient Rescaling

In deep learning optimization, gradients can explode during backpropagation—growing to extremely large magnitudes that destabilize training and cause parameter updates to overshoot optimal values. This phenomenon, known as the exploding gradient problem, is particularly prevalent in recurrent neural networks (RNNs), deep transformers, and networks with many layers.

Global Norm Gradient Rescaling is a powerful stabilization technique that addresses this issue by constraining the total magnitude of all gradients across an entire model. Unlike per-parameter clipping methods, this approach maintains the relative proportions between different gradient components while ensuring the combined magnitude stays within safe bounds.

The Algorithm

Given a collection of gradient arrays (G = {g_1, g_2, ..., g_k}) representing gradients for different parameters and a maximum norm threshold (\tau), the algorithm works as follows:

Step 1: Compute the Global L2 Norm

Calculate the global norm by treating all gradient values across all arrays as a single flattened vector:

$$|G|2 = \sqrt{\sum{i=1}^{k} \sum_{j} g_{i,j}^2}$$

This represents the total Euclidean magnitude of all gradients combined.

Step 2: Determine the Scaling Factor

If the global norm exceeds the threshold, compute a scaling factor to bring it exactly to the threshold:

$$\alpha = \frac{\tau}{|G|_2} \quad \text{if } |G|_2 > \tau$$

If the global norm is already within bounds ((|G|_2 \leq \tau)), set (\alpha = 1) (no scaling needed).

Step 3: Apply Uniform Rescaling

Multiply every gradient value by the scaling factor:

$$g'{i,j} = \alpha \cdot g{i,j}$$

This ensures all gradients are scaled proportionally, preserving their relative directions and magnitudes while constraining the total norm.

Your Task

Write a function that takes a list of gradient arrays and a maximum norm threshold, then returns the rescaled gradients. The function should preserve the original structure of the input (same number of arrays, same dimensions per array).