Visual walkthrough — Saving and loading models (checkpoints)

Everything below is drawn on a chalkboard. Follow the coloured arrows — the pictures carry the argument; the words just point at them.

Step 1 — What a "parameter" and a "gradient" even are

WHAT. A neural network is just a big pile of numbers we can tune. Call any one of those tunable numbers $\theta$ (the Greek letter "theta" — read it as "the knob"). There are millions of them, but to see the idea we track just one.

Attached to the network is a loss $\mathcal{L}$ — a single number that says "how wrong are you right now?" Big $\mathcal{L}$ = very wrong. Our whole job is to turn the knob $\theta$ until $\mathcal{L}$ is small.

WHY a gradient? We need to know which way to turn the knob. The tool that answers "if I nudge $\theta$ up a tiny bit, does the loss go up or down, and how steeply?" is the derivative, written

$g={\nabla }_{\theta }\mathcal{L}$

• ${\nabla }_{\theta }$ (the "nabla" symbol with a little $\theta$) means slope of the loss as this one knob $\theta$ changes.
• $g$ is just a shorter name for that slope, the gradient.
• If $g$ is positive, loss rises when $\theta$ rises → so we should decrease $\theta$.
• If $g$ is negative, we should increase $\theta$.

We pick the derivative and not, say, the raw loss value, because the loss alone tells us "how bad" but never "which direction fixes it." Only slope encodes direction.

PICTURE. The board shows a valley (the loss curve). A ball sits on the slope; the arrow shows the downhill direction the gradient points to.

Step 2 — Plain gradient descent, and its weakness

WHAT. Repeat the rule from Step 1 over and over:

${\theta }_t={\theta }_{t-1}-\alpha g_t$

• $t$ is the step counter (step 1, 2, 3, …).
• ${\theta }_{t-1}$ is the knob's old value; ${\theta }_t$ is the new value.
• $g_t$ is the slope measured right now, at this step.
• $\alpha g_t$ is the size and direction of the nudge.

WHY look closer? In a nice round bowl this works fine. But real loss surfaces are long narrow valleys — steep across, nearly flat along the bottom. Plain descent zig-zags: it bounces off the two steep walls and crawls along the useful direction. Each step's direction is decided only by the current slope, with no memory of where we were heading.

PICTURE. A narrow valley seen from above. The plain-descent path (pink) bounces wall to wall, making painfully slow progress toward the goal.

Step 3 — Momentum: giving the ball memory

WHAT. Instead of stepping along the current slope, we step along a running average of past slopes. We invent a new quantity ${\mathbf{v}}_t$, the velocity, and update it each step:

${\mathbf{v}}_t=\beta {\mathbf{v}}_{t-1}+(1-\beta )g_t$ ${\theta }_t={\theta }_{t-1}-\alpha {\mathbf{v}}_t$

Term by term:

• ${\mathbf{v}}_{t-1}$ — the velocity we already had (accumulated history).
• $\beta$ (beta, the momentum coefficient, usually $0.9$) — how much of the old velocity we keep. Close to 1 = long memory.
• $(1-\beta )g_t$ — how much of the fresh slope we mix in. With $\beta =0.9$ that's only $0.1$ of the new gradient each step.
• Second line: we now step along ${\mathbf{v}}_t$ (the smoothed direction), not the raw $g_t$.

WHY this tool? A running average cancels the sideways bouncing (those flips average to near zero) while adding up the consistent downhill direction (same sign every step → grows). It's exactly a ball with inertia rolling down: it ignores jitter and builds speed along the true slope.

PICTURE. Same narrow valley. The momentum path (blue) barely bounces — it flows smoothly down the trough, reaching the goal in far fewer steps than the pink zig-zag.

Step 4 — The moment of the crash: what a "fresh optimizer" throws away

WHAT. You save a checkpoint but store only $\theta$. Training crashes. You reload $\theta$, but the optimizer is brand new, so its velocity resets:

${\mathbf{v}}_t⟶0$

The very next update, with ${\mathbf{v}}_{t-1}=0$, collapses back to plain gradient descent:

${\mathbf{v}}_{\text{new}}=\beta \cdot 0+(1-\beta )g=(1-\beta )g\Rightarrow {\theta }_{\text{next}}=\theta -\alpha (1-\beta )g$

Notice the step is now only $(1-\beta )$ as long — for $\beta =0.9$ that's a tiny $0.1\times$ step, in the raw jittery direction, with no inherited momentum.

WHY it hurts. The ball was cruising fast down the valley; you froze its position but set its speed to zero. It has to laboriously build speed again from a standstill — and while it does, it re-enters the sideways zig-zag it had already escaped.

PICTURE. The blue momentum path is cruising; a red "X" marks the crash. On reload the ball restarts from rest (dashed red) — slow, hesitant, bouncing again, before it eventually rebuilds the smooth blue flow.

Step 5 — Exactly how many steps you lose

WHAT. After reset, velocity is rebuilt by the recursion in Step 3. Unrolling it, the velocity is a weighted sum of the last many gradients with weights $(1-\beta ){\beta }^k$. The effective memory length — how many past steps meaningfully contribute — is

$N=\frac{1}{1-\beta }$

• For $\beta =0.9$: $N=\frac{1}{1-0.9}=10$ steps.
• For $\beta =0.99$: $N=\frac{1}{1-0.99}=100$ steps.

WHY this formula? The old contributions decay like ${\beta }^k$. The velocity reaches a fixed fraction (about $63\%$, i.e. $1-1/e$) of its full value after roughly $N=1/(1-\beta )$ steps — this is the momentum "warm-up time." Those are the steps you waste re-accelerating every time you resume without saved state.

PICTURE. A bar chart of how much each past gradient counts: tall bar for the newest, shrinking by $\times \beta$ each step back. The shaded region up to step $N=10$ holds the bulk of the weight — that's the memory you erased.

Step 6 — Edge cases: when does the reset not hurt?

Cover every scenario so you never hit a surprise:

Case A — $\beta =0$ (no momentum at all). Then ${\mathbf{v}}_t=g_t$; there is no history to lose. $N=1/(1-0)=1$. Saving weights alone is perfectly fine. Plain SGD, plain safety.

Case B — Adam / RMSProp. These carry two running averages: $m_t$ (mean of gradients) and $v_t$ (mean of squared gradients, used to scale each knob's step). Losing $v_t$ is worse than losing plain momentum — the per-parameter step sizes are wrong until it rebuilds, and Adam's bias-correction counter resets too. Always save optimizer.state_dict() for Adam.

Case C — You've reached the valley floor already ($g\approx 0$). Then ${\mathbf{v}}_t$ has already decayed toward zero on its own. Resetting it costs almost nothing — this is why inference-only exports (Example 2 in the parent) safely drop the optimizer: training is done, there's no motion left to preserve.

Case D — Fine-tuning / changing the learning-rate schedule. If you're deliberately restarting the schedule (see Learning Rate Schedules) or doing Transfer Learning on a new task, you sometimes want a fresh optimizer. Resetting is a choice, not an accident.

PICTURE. A 2×2 board: top-left "β=0 → safe", top-right "Adam → save both m and v", bottom-left "at minimum → harmless", bottom-right "new task → reset on purpose". A check or cross on each.

The one-picture summary

Everything on one board: position vs. motion. Save both, and the ball resumes cruising. Save only position, and it restarts from rest, costing you $N=1/(1-\beta )$ warm-up steps.

Practice recall: