Gradient Clipping by Global Norm From Scratch in NumPy

Given a list grads of gradient arrays (any shapes) and max_norm, compute the total norm N = √(Σ over every element of every array of g²). If max_norm / (N + 10⁻⁶) < 1, multiply every array by that factor; otherwise leave them unchanged. Return [clipped, N].

Examples

Input grads = [[3], [4]], max_norm = 1

Output [[[0.6], [0.8]], 5]

Norm 5. Each gradient is scaled by about 1/5, so the new norm is 1.

Input grads = [[0.1, 0.2], [[0.1], [0.2]]], max_norm = 10

Output [[[0.1, 0.2], [[0.1], [0.2]]], 0.3162]

Already small: unchanged.

Input grads = [[[1, 1], [1, 1]], [2, 2, 2, 2]], max_norm = 3

Output [[[[0.6708, 0.6708], [0.6708, 0.6708]], [1.3416, 1.3416, 1.3416, 1.3416]], 4.4721]

Norm √20 ≈ 4.47 across both arrays; both are scaled by the same factor.

Constraints

1 ≤ number of arrays ≤ 1000
max_norm > 0
Checked to a relative tolerance of 10⁻¹²

Hints

Hint 1

Flatten conceptually: the total norm is the square root of the sum of every array's sum of squares.

Hint 2

The scaling factor is max_norm / (N + 1e-6). Apply it only if it is below 1.

Hint 3

Multiply every array by the same factor, and return arrays with their original shapes.

Where candidates go wrong

Clipping each array to max_norm separately, which changes the direction of the overall gradient.
Clipping element by element (clip by value), which is a different method.
Scaling up gradients whose norm is below the limit.
Dropping the 1e-6, which makes the result differ from PyTorch's when N is tiny.

The idea

Treat all the parameters' gradients as one vector g. Its norm is the square root of the summed squares. If it exceeds the limit, g is replaced by g · max_norm/(‖g‖ + ε): same direction, length capped. Because one factor multiplies every tensor, the relative sizes of the layers' gradients are untouched, which is the point.

Time O(total number of elements) · Space O(total number of elements)

Reference solution

import numpy as np

def clip_grad_norm(grads, max_norm):
    gs = [np.asarray(g, dtype=float) for g in grads]
    total = float(np.sqrt(sum((g ** 2).sum() for g in gs)))   # one norm over every element
    coef = max_norm / (total + 1e-6)
    if coef < 1:
        gs = [g * coef for g in gs]                          # same factor for all: direction kept
    return [[g.tolist() for g in gs], total]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Clipping is in every large training run. The question checks whether you clip the whole gradient as one vector, which preserves its direction, or each tensor separately, which does not, and whether you know the small constant PyTorch adds.

The theory behind it is in the lesson Dropout, gradient clipping and learning-rate schedules.

Keep going

Easy8 min

Gradient clipping by global norm