Gradient clipping

One huge gradient can undo many good steps

In gradient descent, each step moves the weights by the learning rate times the gradient. Most steps are small. Now and then, one batch of examples produces a gradient a hundred times larger than usual: an odd example, a steep cliff in the loss, or a numerical accident. The step is then a hundred times larger too. The weights land somewhere far away, the loss jumps up, and training may take thousands of steps to recover, or never recover.

Clipping shrinks a long gradient to a fixed length, keeping its direction

Think of the gradient as one long vector, with an entry for every weight in the model. Its length, called its norm, is the square root of the sum of the squares of all entries. That is the distance formula, with as many coordinates as the model has weights.

Choose a maximum length c. Then, before each step:

• if the norm is c or less, leave the gradient as it is;
• if the norm is larger, multiply every entry by c / norm.

Multiplying every entry by the same number keeps the direction and changes only the length. The clipped length is exactly c.

Try it. The figure clips a gradient to length c = 1, shown as the circle. Predict first: what is the clipped length when the gradient's length is 2.5? Then move the slider and find the length at which the two arrows stop matching.

The norm covers all the weights together

The usual method, clipping by the global norm, computes one norm over every weight in the model and scales all gradients by the same factor. Suppose one layer's gradient is (1, 2) and another's is (2). The global norm is √(1 + 4 + 4) = 3. With c = 1.5, every entry is halved.

A cruder method, clipping by value, caps each entry on its own. It changes the direction: capping (3, 4) at 1 gives (1, 1), which points somewhere else.

In PyTorch, clip after loss.backward() has filled the gradients and before optimizer.step() uses them. The function returns the norm from before clipping. That number is worth printing or logging: a norm that keeps rising warns that training is going wrong.

clip_torch.py

import torch

a = torch.zeros(2, requires_grad=True)   # weights of one layer
b = torch.zeros(1, requires_grad=True)   # weights of another
a.grad = torch.tensor([1.0, 2.0])
b.grad = torch.tensor([2.0])

total = torch.nn.utils.clip_grad_norm_([a, b], max_norm=1.5)
print(total)          # length before clipping: sqrt(1 + 4 + 4) = 3
print(a.grad, b.grad) # every entry scaled by 1.5 / 3 = 0.5

tensor(3.)
tensor([0.5000, 1.0000]) tensor([1.0000])

In a training loop, the three lines go in this order (an excerpt, not a whole program):

loss.backward()
# model.parameters() lists all the model's weights, like [a, b] above.
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()