Kahibaro
Discord Login Register

6.5 Gradient Clipping

What gradient clipping is and when you need it

Gradient clipping is a training technique that limits how large parameter gradients are allowed to become during backpropagation. It is most useful when training becomes unstable because some gradients occasionally explode in magnitude, which can cause the loss to spike, produce NaNs, or make optimization diverge.

You most often see gradient clipping in sequence models such as RNNs, LSTMs, and GRUs, but it can also help with Transformers, very deep networks, high learning rates, and noisy objectives. It is not a substitute for choosing a reasonable learning rate or fixing data and loss issues, but it is a practical guardrail that can turn rare catastrophic updates into harmless ones.

Gradient clipping must be applied after loss.backward() computes gradients and before optimizer.step() updates parameters.

Two common types: clip by value and clip by norm

Clipping by value clamps each individual gradient element into a fixed range. Conceptually, for each scalar gradient component $g_i$ you apply $g_i \leftarrow \mathrm{clip}(g_i, -c, c)$. This is simple, but it can distort the gradient direction heavily, especially when many components are clamped.

Clipping by norm rescales the entire gradient tensor set so the overall norm does not exceed a threshold. For parameters with gradients flattened into a single vector $g$, define the total norm as $||g||$. If $||g|| > c$, you rescale all gradients by $c / (||g|| + \epsilon)$ so the update direction is preserved but the step is bounded.

In practice, prefer clipping by global norm for deep learning training stability because it limits step size while keeping the gradient direction.

How to do gradient clipping in PyTorch

PyTorch provides utilities in torch.nn.utils that operate on an iterable of parameters. The most common is clip_grad_norm_, which clips by global norm in place.

A typical training step with clipping looks like this:

python
import torch
from torch.nn.utils import clip_grad_norm_
model.train()
optimizer.zero_grad(set_to_none=True)
y_pred = model(x)
loss = loss_fn(y_pred, y)
loss.backward()
max_norm = 1.0
total_norm = clip_grad_norm_(model.parameters(), max_norm=max_norm)
optimizer.step()

clip_grad_norm_ returns the total norm before clipping, which you can log to see how often clipping activates and how severe the spikes are.

If you specifically want elementwise clamping, you can use clip_grad_value_:

python
from torch.nn.utils import clip_grad_value_
loss.backward()
clip_grad_value_(model.parameters(), clip_value=0.5)
optimizer.step()

Do not clip gradients before backward(), there is nothing to clip yet. Do not clip after optimizer.step(), the update already happened.

Choosing a clipping threshold

There is no universal best value, but there are common starting points. For global norm clipping, values like 0.5, 1.0, or 5.0 are typical. Smaller values clip more aggressively and can slow learning if they activate constantly. Larger values clip only rare extreme spikes.

A practical approach is to log the returned total_norm for a few hundred steps. If you see occasional very large norms that correlate with loss spikes, choose max_norm just below the spike region. If the norm is almost always above max_norm, you are likely clipping too hard or your learning rate is too high.

If clipping happens nearly every step, treat it as a warning sign. Often the learning rate is too large, the loss scale is off, or the input normalization is broken.

Interaction with mixed precision and gradient accumulation

With mixed precision training, gradients may be scaled internally. If you use torch.cuda.amp.GradScaler, you should unscale gradients before clipping so that the norm is computed in the true scale.

The order is:

python
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()

With gradient accumulation, clipping should typically happen right before the optimizer step, after all micro batch gradients have been accumulated. If you clip after every micro batch, you are changing the effective optimization behavior because you are repeatedly shrinking partial gradients.

Common mistakes and quick diagnostics

A frequent mistake is clipping the gradients of only some parameters by passing the wrong parameter list. Using model.parameters() is simplest, but if you have multiple parameter groups or optimizers, clip the parameters that will be stepped by that optimizer.

Another mistake is thinking clipping fixes NaNs from invalid operations. If NaNs appear in the forward pass due to invalid inputs or unstable losses, clipping will not help because the gradients are already NaN. Clipping also cannot correct systematic divergence from an excessively high learning rate, it can only bound individual updates.

A quick diagnostic is to log both the loss and the gradient norm. If the loss occasionally jumps sharply and the gradient norm jumps at the same time, clipping is a good fit. If the loss steadily increases without gradient spikes, clipping is unlikely to be the main fix.

Clipping is a stability tool, not a performance trick. Use it to prevent rare destructive updates, and still tune learning rate, normalization, and model configuration for healthy training.

Views: 62

Comments

Please login to add a comment.

Don't have an account? Register now!