Taking a gradient spike apart, by loss term and by module

When a training run with several losses starts throwing gradient spikes, the usual first moves are to lower the learning rate, add clipping, or skip the batch that seemed to cause it. All three treat the spike as one number. I wanted to try the opposite: take a single spike and look inside it. Which loss term produced it, which part of the network it landed in, and whether the batch that triggered it was the cause or only the occasion.

That needs two tools. One is an exact replay of the step where the spike happened, so that every experiment on it starts from the same state. The other is a way to split the gradient. Both turned out to have traps of their own.

The model and the spike

The model is small on purpose: a shared trunk of two hidden layers and two heads. One head is trained with squared error on a smooth function. The other predicts a mean and a log standard deviation for a second target and is trained with the Gaussian negative log-likelihood. That loss is known to produce spiky gradients, because it divides each error by the predicted variance. The second target is nearly noiseless for most inputs and very noisy for a few, so some predicted variances get small. I trained with Adam at learning rate $10^{-2}$ on the CPU, single-threaded, with PyTorch's deterministic algorithms on.

I defined the onset as the first step whose gradient norm exceeded 20 times the median of the previous 200 steps. It came at step 446: the norm jumped from a median of about 8 to 187. Over the next 100 steps the regression head's loss rose from 0.060 to 0.090, even though nothing about its own target had changed.

Replaying one step exactly

The plan was simple: save the model, the optimizer and the data generator's state every 50 steps, restore the last checkpoint before the onset, run forward to the onset, and confirm that the gradient norm matches the logged one exactly, not approximately.

My first attempt didn't match, and two runs of the same script disagreed with each other. I was seeding the data generator but not the global random state that initialises the weights, so each run started from different weights. Every source of randomness has to be part of the saved state: the weight initialisation, the data order and any augmentation, dropout masks, and the library's own generators.

The second problem was harder to see. After fixing the seed, replaying once from a checkpoint worked, matching to every printed digit (186.9585…). But a second replay from the same in-memory checkpoint went somewhere else. When a PyTorch optimizer loads a saved state, it doesn't copy the moment tensors if they're already on the right device with the right type. The restored optimizer shared memory with the checkpoint, so the first replay's training rewrote the checkpoint's stored Adam averages in place. I measured it: one replay changed the stored first-moment estimate by up to 0.011. The model's weights are copied on load and were untouched, which is what made it easy to miss. A deep copy of the checkpoint before every load fixed it.

The third trap appeared when I tried a lower learning rate from the checkpoint and got results identical, to every digit, to the original rate. Restoring the optimizer also restores the learning rate saved with it, silently overriding the one the optimizer was constructed with. Setting the rate after loading fixed that.

With those fixed, the replay was exact, and every experiment below starts from the same bits.

Splitting the gradient by term

The total gradient at the onset is the sum of the two terms' gradients, $g = g_{\text{reg}} + g_{\text{lik}}$, so each can be computed on its own by backpropagating one term at a time from the same graph. Two numbers describe each part: its norm, and its cosine with the total, which says how much of the total's direction it supplies.

The spike is the likelihood term. Its gradient has essentially the total's length and direction. The regression term contributes a gradient about 500 times smaller, nearly orthogonal to the total.

Multi-task methods use the same quantities in a different role. Yu et al. (2020)¹ watch for task gradients whose cosine is negative, conflicting gradients, and project them apart. Chen et al. (2018)² rescale each task's gradient so that their norms stay balanced. Both start from the per-task gradient as a measurement, and so does this kind of forensics.

Splitting it by module

The same gradient vector can also be cut by where its parameters live.

Most of the spike stays inside the likelihood head, which only affects the likelihood term. But 35 units of it, about four times the whole network's typical gradient norm, land on the shared trunk. That is the path by which a problem in one loss term reaches the other head, and it explains why the regression loss rose after the onset. Adam then normalises each parameter's step, so a large gradient on the trunk turns into a full-size step in a direction the regression head never asked for.

What the obvious fixes did

From the exact state before the onset I ran the next 300 steps five ways and recorded the largest gradient norm afterwards and the regression loss over the last 50 steps.

Skipping the batch made the later spikes larger. The batch at the onset was only the occasion: the conditions for a spike were already in the state, small predicted variances behind a large learning rate, and the next batches found them. Dropping the term for one step did no better. A lower learning rate halved the later spikes. A floor under the predicted standard deviation, $\sigma = 0.05 + \mathrm{softplus}(\cdot)$, removes the mechanism that produced them, a weight $1/\sigma^2$ with no upper bound. The largest later norm fell from 225 to 59, and the regression head did as well as with the lower learning rate.

What I would do with the next spike

The order I would follow, from this experiment:

1. Make the run replayable first: seed everything, save the optimizer and generator states with the weights, deep-copy checkpoints before loading them, and set the learning rate explicitly after loading.
2. Check the replay against the logged norm exactly. An approximate match means something is not captured.
3. Split the onset gradient by loss term and by module, and read the cosines as well as the norms.
4. Test fixes from the exact pre-onset state, including skipping the batch, so that "the batch was bad" is checked instead of assumed.
5. Prefer the fix that changes the structure of the offending term over one that only changes how hard the optimizer pushes.

The term that carried the spike divides each error by a predicted variance; why that loss produces spiky gradients in the first place is in a model that learns its own uncertainty chooses what to fit. The replay traps, and a few more, are collected in a checkpoint is more than its state dict.

What I still don't know

This is a single onset in a single run. I chose the floor because I knew which mechanism to suspect; in a larger model the term that carries the spike may be less obvious, and its mechanism less clean. I also replayed on a CPU with one thread, where exact determinism is easy. On a GPU, several operations are non-deterministic by default, and I haven't tried to make a replay exact there. And I measured each fix over only 300 steps. A fix that calms the next 300 steps can still lose over a full run, and comparing them properly would take many onsets and many seeds.

References

1. Gradient Surgery for Multi-Task Learning — Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman and Chelsea Finn, NeurIPS 2020
2. GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks — Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee and Andrew Rabinovich, ICML 2018