Your Gradient Clipping Code Might Not Be Clipping the Gradient Norm

Y

ym p

Guest
Two short pieces of code can both be called “gradient clipping” while doing different things. One caps each coordinate. The other measures the length of the entire gradient vector and, when necessary, rescales every coordinate by the same factor.



That distinction is easy to miss when reviewing generated code. A loop full of min and max calls looks reasonable, and its outputs may even be smaller than its inputs. Neither observation establishes that it implements global norm clipping.



Here is a complete numerical example and a small Python implementation you can run without a machine-learning framework.

Start with a vector you can inspect​


Take a gradient [3, 4] and a maximum L2 norm of 2. Its original length is 5. For global L2 norm clipping, multiply both coordinates by 2 / 5. The result is [1.2, 1.6], whose length is 2. The vector retains its direction because both coordinates receive the same positive multiplier.



Clamping each coordinate to the interval [-2, 2] instead produces [2, 2]. Its length is approximately 2.828, so it does not satisfy the requested global bound. Its direction has also changed.



The geometry explains the mismatch: an L2 bound defines a ball, whereas a coordinate-wise bound defines a box. A point can lie inside that box while remaining outside the ball.

What does “global” include?​


A model usually stores gradients in separate parameter tensors. For a global L2 norm, imagine concatenating all the included gradients into one vector before calculating its length. There is one scale factor for the whole collection.



Splitting [3, 4] into blocks [[3], [4]] must not change the output. Clipping each block separately would produce [[2], [2]], reproducing the same failure as coordinate clamping in this example.



That gives us a useful test: repartitioning the same coordinates should not change a global clipping result. The test checks a behavioral property instead of merely repeating the implementation's loop structure.

A complete teaching implementation​


This function accepts lists of gradient blocks and returns new lists. max_norm must be finite and positive. Each coordinate must be a finite float, and the combined norm must be representable as a finite float. This is a teaching contract, not a general tensor API; input validation, sparse gradients and optimizer state are outside its scope.

Code:
import math


def clip_global_l2(blocks, max_norm):
    flat = [value for block in blocks for value in block]
    norm = math.hypot(*flat)
    scale = 1.0 if norm <= max_norm else max_norm / norm
    return [[value * scale for value in block] for block in blocks]



math.hypot avoids the naive intermediate squaring used by sqrt(sum(x*x)). It still has the limits of floating-point arithmetic, which is why the contract explicitly limits the representable combined norm.

The zero vector takes the unchanged branch, so there is no division by zero. A vector already below the bound is also unchanged. Flattening and rebuilding the lists take linear time and linear extra space in the number of coordinates.

Tests that expose the wrong operation​


Run this block after the implementation above:

Code:
def close_list(actual, expected):
    assert len(actual) == len(expected)
    assert all(abs(a - b) < 1e-12 for a, b in zip(actual, expected))


original = [[3.0, 4.0]]
clipped = clip_global_l2(original, 2.0)
close_list(clipped[0], [1.2, 1.6])
assert original == [[3.0, 4.0]]  # No input mutation.
assert abs(math.hypot(*clipped[0]) - 2.0) < 1e-12

# Repartitioning coordinates must preserve the result.
split = clip_global_l2([[3.0], [4.0]], 2.0)
close_list([split[0][0], split[1][0]], clipped[0])

assert clip_global_l2([[0.0, 0.0]], 2.0) == [[0.0, 0.0]]
assert clip_global_l2([[0.3, 0.4]], 2.0) == [[0.3, 0.4]]
assert clip_global_l2([[0.0, -5.0]], 2.0) == [[0.0, -2.0]]
assert clip_global_l2([[3.0, 4.0]], 5.0) == [[3.0, 4.0]]

wrong = [min(2.0, max(-2.0, x)) for x in original[0]]
assert math.hypot(*wrong) > 2.0
assert wrong[0] * 4.0 != wrong[1] * 3.0  # Direction changed.



The wrong implementation fails both the global bound and the direction check. A test that only checks whether each absolute coordinate is at most 2 would incorrectly approve it.

Compare framework APIs carefully​


PyTorch exposes separate functions for norm clipping and value clipping. The documentation for clip_grad_norm_ describes computing the norm over the collection as if the gradients were concatenated, and modifying gradients in place. Our list-based example returns a new value instead.



Do not expect every framework's floating-point output to match this teaching function bit for bit. Implementations may include a small denominator constant and support different norm types, data types, and execution paths. Check the actual API contract before translating an equation into a test.



Clipping also does not by itself fix an incorrect loss, invalid data, or a non-finite gradient. Here we are verifying one transformation, not making a claim that training will converge.

Try changing the example​


Change the gradient to [6, 8], keep the bound at 2, and predict the result before running the function. Then split the coordinates into two blocks and check that the answer stays the same.



I use this distinction in the global norm clipping walkthrough on TensorDrill, the AI coding practice project I maintain. It provides an interactive comparison alongside the exercise solution. The implementation and tests needed for this article are all above.

References​

 

Thread statistics

Created
ym p,
Replies
0
Views
0
Back
Top