Restarting Adam moves every weight by the full learning rate

Training runs get interrupted and resumed all the time: to change a learning rate, to add data, to continue a run that stopped early, or to fine-tune a model that has already converged. A common shortcut is to load the weights and build a new optimizer. I wanted to understand what that costs with Adam specifically, because its update has a property that makes the first step after a reset unusual.

The first step of a fresh Adam, on paper

Adam (Kingma and Ba, 2015¹) keeps running averages of each gradient coordinate and of its square, $m_t$ and $v_t$, both starting at zero, and corrects them for that start:

At $t = 1$, $m_1 = (1 - \beta_1)\, g$ and $v_1 = (1 - \beta_2)\, g^2$, and the corrections divide those factors out exactly: $\hat m_1 = g$ and $\hat v_1 = g^2$. So, ignoring $\epsilon$,

Every parameter moves by exactly $\eta$, in the direction of its gradient, whatever the size of that gradient. A parameter whose gradient is tiny noise moves as far as one with a strong, consistent signal. The length of the whole step is $\eta \sqrt{P}$ for $P$ parameters. In the middle of training, by contrast, most coordinates have noisy gradients that partly cancel in $\hat m$ but not in $\hat v$, so the typical step is a small fraction of $\eta$ per coordinate. A fresh optimizer throws that away, and its first steps behave like a sign-gradient method at full rate.

Checking it

The network is a small convolutional classifier with 77,140 parameters, trained on a 20-class synthetic task with Adam and a cosine schedule from $2 \times 10^{-3}$ over 4,000 steps, three seeds. At step 2,000, where the scheduled rate is exactly $10^{-3}$, I compared the uninterrupted run with one that replaces the optimizer with a fresh Adam, keeping the schedule's rate. I measured the length of each update.

Ranges are over the three seeds.

The predicted length of the first fresh step is $10^{-3} \times \sqrt{77140} = 0.278$. The measured value was 0.274 to 0.275 in all three seeds; the small shortfall is $\epsilon$ and a few parameters with zero gradient. That is 6 to 14 times the steps just before the resume. The following steps shrink as $\hat v$ accumulates a history, but they stay several times larger than normal for a while.

On this task the jolt cost little that a metric would notice. Validation loss rose from about 0.007 to about 0.013 over the next 50 steps, and the run recovered by the end of training. A model trained to near-perfect accuracy on an easy task has plenty of room to absorb a bad step. A fine-tuning run that starts with one has less.

A new full-rate window on a finished model

The other common restart is to take a model that has finished its schedule and start a new one from the full rate, to train it further or on new data. That combines the fresh optimizer with a large learning rate. I did it from the end of each run, with a new 2,000-step cosine schedule from $2 \times 10^{-3}$, and compared it with a gentle continuation at a twentieth of that rate.

The first step of the new window was 0.54 to 0.55 long, against a predicted $2 \times 10^{-3} \times \sqrt{77140} = 0.556$. Validation loss rose from about 0.0025 to between 0.009 and 0.016 within the first 100 to 250 steps, a five- to six-fold increase at its peak, and validation accuracy fell by a few tenths of a percent. By step 1,000 it had recovered, and after the full 2,000 steps it ended slightly lower than where it started, 0.0011 to 0.0015. The gentle continuation also ended lower, at 0.0019 to 0.0022, without the detour.

I had expected the new window to knock the model back and then return it to where it began: a fixed point, set by the architecture, the data and the learning rate. It didn't, because my model hadn't saturated, and 2,000 more steps of training simply helped. The knock-back itself was real and is what the formula predicts. Whether the model recovers to the same value, a better one or a worse one depends on how much room it had left, and on a converged model at capacity I would expect the detour to buy nothing.

What a resume should carry

From these runs, a resume that should behave like a continuation needs:

• the optimizer's state, deep-copied, with any change to the learning rate applied after loading;
• the scheduler's state, so the rate continues along the planned curve;
• if a fresh optimizer is unavoidable, a short warmup from a small rate, so that the sign-gradient first steps are small.

The last point is the most useful one. Adam's first steps after a reset are the reason a warmup helps even for a model that is already trained. The warmup isn't there to protect the weights from a bad initialisation; it protects them from the optimizer's own bias correction.

The same ratio of averages is what makes a single gradient spike so destructive under Adam, in gradient clipping does different jobs under SGD and Adam, and what lets round-off move weights that should never move, in softmax over one token has zero gradient. Restoring the optimizer properly has traps of its own, collected in a checkpoint is more than its state dict.

What I still don't know

I measured the size of the jolt closely, and its cost only on an easy task, where the cost was small. On a harder task, or one where the model sits near a sharp region of the loss, a full-rate sign step on every parameter could do much more damage, and I haven't measured how much. I also only tried Adam. Optimizers without bias correction, or with a different normalisation, will behave differently after a reset, and comparing them on the same restart would be a clean next experiment.

References

1. Adam: A Method for Stochastic Optimization — Diederik P. Kingma and Jimmy Ba, ICLR 2015