Gradient clipping does different jobs under SGD and Adam

Gradient clipping is usually introduced as a guard against rare explosions. Pascanu et al. (2013)¹ proposed rescaling the gradient whenever its norm passes a threshold $c$, so that a sudden cliff in the loss surface cannot throw the parameters far away:

A belief I had held, and had once written down, goes further. If the threshold is set so low that the clip fires on almost every step, then every update has the same length, and the learning rate times the threshold sets the step size while the gradient's magnitude carries no information. It is a tidy argument. I wanted to check it, because most of my training uses Adam, and Adam does something to gradient magnitudes before the argument's conclusion can arrive.

Under SGD the argument holds

With plain SGD and a clip that always fires, the update is

which is normalised gradient descent: a step of length exactly $\eta c$ in the direction of the gradient. The step can no longer shrink as the model approaches a minimum, because the gradient's size has been divided out.

I tested it on a small regression network: three fully connected layers learning a fixed nonlinear function of 32 inputs, with a little label noise, 4,000 steps, three seeds. With SGD at learning rate 0.05 I compared no clipping, a threshold of 10 that never fired, and a threshold of 0.05 that fired on every step.

Variation is the coefficient of variation of the update size after the first 500 steps. The always-firing clip made every update exactly $\eta c = 0.0025$ long, with no variation at all, and the test loss ended seven times worse. A model that can only take steps of one size circles its minimum at a distance set by that size. Its median gradient norm stayed near 0.3, against 0.11 without the clip: it never settled.

Under Adam it mostly doesn't

Adam's update divides a running average of the gradient by the root of a running average of its square:

Kingma and Ba (2015)² note that this makes it "invariant to diagonal rescaling of the gradients". Multiply every gradient by the same constant and both averages scale by it, so the update doesn't change. A clip that fires on every step multiplies each gradient by $c/\|g_t\|$. If the gradient norm changes slowly compared with the averaging windows, that factor is nearly constant over each window, and Adam divides it back out. My prediction was that an always-firing clip under Adam would barely matter.

The clip fired on 95% of steps and the test loss moved from 0.0031 to 0.0033. Update sizes varied exactly as much as without it. The steps were not fixed-length and the learning rate was not suddenly in charge. Under Adam the size of each step was already set by the learning rate and by how consistent the gradients were, not by how large they were, so the clip's rescaling was absorbed. The fraction of clipped steps, which under SGD says a lot, says very little here.

Where it matters for Adam: a rare spike

If an always-firing clip does nothing for Adam, what does clipping do? My guess was rare spikes, and I expected a particular mechanism. A huge gradient inflates the second-moment average $\hat v$, which decays slowly with $\beta_2 = 0.999$, so for hundreds of steps afterwards every update would be divided by an inflated denominator and training would stall.

To test it I made one batch in every 500 corrupt, with targets multiplied by 100, and trained Adam with no clip and with a clip at 1.0, set above the normal gradient norms of around 0.1 so it fires only on spikes. I tracked the update size and the loss on clean batches around each spike, relative to the 100 steps before it.

The stall I expected exists, but it is small and late: 300 steps after a spike the updates are 14% smaller than before. The real damage is the opposite and immediate. The spike step moved the weights 94 times further than a normal step, the next ten steps still about 50 times further, and the loss on clean data jumped 68-fold and took about 300 steps to recover. By the end, the unclipped runs had a test loss of 0.012 against 0.0030 for the clipped ones. The clipped runs were as good as runs with no spikes at all, 0.0031.

Why a spike makes Adam take its largest steps

The size of Adam's normal step surprised me once I looked at it. Per coordinate, $|\hat m_i| / \sqrt{\hat v_i}$ behaves like a signal-to-noise ratio. When a coordinate's gradient is noisy, positive and negative values cancel in the average $\hat m$ but not in the average of squares $\hat v$, so the ratio is well below 1 and the step is well below $\eta$. In ordinary training most coordinates are noisy, and the typical step is a small fraction of what Adam is allowed to take.

A spike is the opposite: one gradient so large that it dominates both averages at once. Ignoring the bias correction, which is close to 1 late in training,

on every coordinate the spike touches, all pointing the same way. Adam's step is bounded; a spike pushes every coordinate to that bound at once, in a coherent and wrong direction. Then $\hat m$ forgets the spike at the rate $\beta_1^k$, which takes tens of steps, while $\hat v$ forgets it at $\beta_2^k$, which takes hundreds. That accounts for both phases in the table: large steps for a few dozen steps, then slightly small ones for a few hundred.

Clipping at a threshold above normal gradients and below spikes removes the spike before either average sees it. That is the job clipping does under Adam.

Choosing the threshold, depending on the optimizer

Under SGD the threshold sets the step length whenever the clip fires, so a clip that fires often is a change of learning rate, and a clip that always fires is a different algorithm. Zhang et al. (2020)³ analyse clipped and normalised gradient descent and show they can converge faster than plain gradient descent when the smoothness of the loss grows with the gradient norm. That is a reason to use clipping deliberately, as part of the optimizer, and to know that is what it has become.

Under Adam the step length is already normalised, so the useful threshold is one that separates ordinary gradients from outliers. From these runs, that means:

• logging the pre-clip gradient norm every step and reading its median and its upper tail, not only the fraction clipped;
• setting the threshold above the bulk of that distribution, where it fires rarely;
• watching the update size per step relative to its running median; a step tens of times larger than usual marks a spike whether or not anything was clipped.

The bound of about 3.16 per coordinate reappears when a fresh optimizer takes its first step, in restarting Adam moves every weight by the full learning rate. When the spikes keep coming, taking one apart by loss term and by module is the next step, in taking a gradient spike apart.

What I still don't know

My spikes are blunt: whole batches with targets multiplied by 100. Real spikes come from rare hard examples, from numerical blow-ups and from sharp regions of the loss surface, and they are rarely that clean. I also used only one $\beta_2$. A smaller value such as 0.95, common in large-model training, should shorten the stall after a spike and might also make Adam react harder to the spike itself, and I haven't measured either. And all of this is global clipping over the whole network. Clipping each layer separately changes which spikes are caught, and I don't know whether it helps Adam or only SGD.

References

1. On the difficulty of training Recurrent Neural Networks — Razvan Pascanu, Tomas Mikolov and Yoshua Bengio, ICML 2013
2. Adam: A Method for Stochastic Optimization — Diederik P. Kingma and Jimmy Ba, ICLR 2015
3. Why gradient clipping accelerates training: A theoretical justification for adaptivity — Jingzhao Zhang, Tianxing He, Suvrit Sra and Ali Jadbabaie, ICLR 2020