Gradient Descent

Let $L:\mathbb{R}^p\to\mathbb{R}$ be a loss function and let $\theta_t\in\mathbb{R}^p$ be the current parameter vector. Under the standard Euclidean inner product, the gradient

points in the direction of steepest local increase. Gradient descent uses the update

where $\eta > 0$ is the learning rate. Cauchy (1847) proposed this method of steepest descent for solving systems of equations.

If $L$ is differentiable near $\theta$ and $\nabla L(\theta)\neq 0$, the first-order Taylor approximation says

That is why stepping against the gradient lowers the loss, for a small enough step. The learning rate says how far to trust this local approximation: too small and progress is slow, too large and the step jumps across the valley or even raises the loss. How large is too large depends on curvature.

At a stationary point, $\nabla L(\theta)=0$, so this first-order decrease argument disappears. In a nonconvex loss, such a point might be a minimum, a saddle, or a maximum.

For the quadratic used below,

with $A$ symmetric positive definite, the gradient is $\nabla L(\theta)=A\theta$. In an eigen-direction of $A$ with curvature $\lambda_i$, the update becomes

For convergence to zero from every starting point, fixed-step descent needs $|1-\eta\lambda_i|<1$ for every direction, so

At equality, the steepest direction's factor is $-1$ and that coordinate alternates without decaying. Above the bound it grows. Below it, a positive factor shrinks without changing sign, a zero factor lands in one step, and a negative factor alternates while shrinking.

On the figure's bowl, with $\eta = 0.1$ exactly at the edge and start $(4.4, 2.6)$, the recurrence gives $x_t=4.4(0.9)^t$ and $y_t=2.6(-1)^t$. So $L_t=67.6+9.68(0.9)^{2t}$ falls at every step, toward $67.6$ rather than zero. A falling loss alone does not show convergence to the minimum.

In deep learning, the exact gradient over the whole dataset is usually too expensive. For a training set of $n$ examples, define the empirical risk

For a mini-batch $B_t$, we instead compute

With uniform sampling, $g_t$ is an unbiased estimate of $\nabla L_{\mathrm{emp}}(\theta_t)$. Using it turns gradient descent into stochastic gradient descent. The update has the same form, but the direction is noisy:

Replay math

For the continuation example, all quantities are dimensionless scalars. At the end of update $t$, retain parameter $\theta_t$, momentum $v_t$, generator state $s_t$, next-row cursor $c_t$, and next-rate index $j_t$. Cursors start at zero. The ordered targets are $(y_0,y_1,y_2)=(1,-1,2)$, with row labels r1, r2, r3. Initially $(\theta_0,v_0,s_0,c_0,j_0)=(0,0,1,0,0)$.

Before each update, read its row and rate. Advance the generator once, then derive a perturbation from its new state:

The update-specific loss is $L_t(\theta)=\tfrac12(\theta-y_{c_t}-\xi_t)^2$. Different updates can have different effective targets; this is not the fixed quadratic above. Differentiating, then carrying half the previous gradient history, gives

Only after these calculations, set $c_{t+1}=(c_t+1)\bmod3$ and $j_{t+1}=j_t+1$, and increment the completed-update count. Save after the whole update, not between the parameter and cursor changes. After update 2, the checkpoint is $(5/64,15/16,15,2,2)$, with completed count 2.

For update 3, read target 2 and rate $1/8$. The generator advances from 15 to 12, giving perturbation $-1/4$. Thus $g_2=-107/64$, the momentum is $v_3=15/32-107/64=-77/64$, and $\theta_3=5/64-(1/8)(-77/64)=117/512$. Resetting only momentum leaves the same gradient but gives $\theta_3=147/512$. The current parameter alone does not determine the next step.

For this recipe, saved progress $n$ reconstructs $c=n\bmod3$, $j=n$ and the generator state after exactly $n$ advances from its seed. Restoring $\theta$ and $v$ while reconstructing these positions is sufficient. If extra calls occurred whose number or order is not recoverable, progress alone is no longer sufficient. Saving each field separately is one method, not a universal rule.

This generator is a small deterministic rule chosen for the example, not independent noise or uniform mini-batch sampling, and an exact continuation here says nothing about convergence or about restarting training spread over many machines. The PyTorch checkpoint tutorial explains why resuming real training requires the optimizer's state as well as the parameters, and PyTorch's reproducibility notes explain why the same seed need not give identical results on different hardware or software versions.

import numpy as np

# Minimize L(theta) = 0.5 * theta^T A theta.
# The second coordinate has much higher curvature.
# Shapes: A is (2, 2), theta is (2,), grad(theta) is (2,).
A = np.diag([1.0, 20.0])
theta = np.array([5.0, 5.0])
lr = 0.08

def loss(theta):
    return 0.5 * theta @ A @ theta

def grad(theta):
    return A @ theta

for step in range(20):
    theta = theta - lr * grad(theta)
    if step in [0, 1, 2, 5, 19]:
        print(step + 1, "theta=", np.round(theta, 3), "loss=", round(loss(theta), 3))

Here $\lambda_{\max}=20$, so convergence from every start needs $0<\eta<2/\lambda_{\max}=0.1$. Try lr = 0.005, 0.12 and 0.1, and read the coordinates as well as the loss. At lr = 0.1 the second coordinate never decays; this code starts at [5,5], not the demo's [4.4,2.6], so it bounces with amplitude $5$, not $2.6$.

The continuation example in exact arithmetic

This second program uses only Python's standard library. It runs the scalar recipe with exact fractions and an algebraically expanded update, so its results can be compared with the calculation shown in the page. theta, velocity, rng, cursor and scheduleIndex correspond to $\theta,v,s,c,j$. The seed is the generator's initial state; completed counts whole updates. The program allows at most eight updates.

from fractions import Fraction as F

ROWS = (("r1", 1), ("r2", -1), ("r3", 2))
POLICIES = ("restore", "velocity", "rng", "cursor", "schedule",
            "weights-seed", "reconstruct")

def initial(seed=1):
    return dict(theta=F(0), velocity=F(0), rng=seed,
                cursor=0, scheduleIndex=0, completed=0)

def step(state, rows=ROWS):
    t, v = state["theta"], state["velocity"]
    row_id, target = rows[state["cursor"]]
    rate = F(1, 4) if state["scheduleIndex"] < 2 else F(1, 8)
    rng = (5 * state["rng"] + 1) % 16
    jitter = F(1, 4) if rng < 8 else F(-1, 4)
    effective_target = target + jitter
    gradient = t - effective_target
    # Expanded affine form, equivalent to v_next = v/2 + gradient.
    v_next = t + v / 2 - effective_target
    t_next = (1 - rate) * t - rate * v / 2 + rate * effective_target
    after = dict(theta=t_next, velocity=v_next, rng=rng,
                 cursor=(state["cursor"] + 1) % 3,
                 scheduleIndex=state["scheduleIndex"] + 1,
                 completed=state["completed"] + 1)
    return dict(before=dict(state), rowId=row_id, target=target,
                rngAfter=rng, jitter=jitter, gradient=gradient,
                rate=rate, after=after)

def trace(state, count, rows=ROWS):
    if not 0 <= count <= 8 - state["completed"]:
        raise ValueError("At most eight whole updates")
    result = []
    for _ in range(count):
        update = step(state, rows)
        result.append(update)
        state = update["after"]
    return result

def checkpoint(completed=2, seed=1, rows=ROWS):
    start = initial(seed)
    updates = trace(start, completed, rows)
    return dict(updates[-1]["after"] if updates else start)

def restore(saved, policy, seed=1):
    if policy not in POLICIES:
        raise ValueError("Unknown restoration policy")
    state = dict(saved)
    if policy in ("velocity", "weights-seed"):
        state["velocity"] = F(0)
    if policy in ("rng", "weights-seed"):
        state["rng"] = seed
    if policy in ("cursor", "weights-seed"):
        state["cursor"] = 0
    if policy in ("schedule", "weights-seed"):
        state["scheduleIndex"] = 0
    if policy == "reconstruct":
        n = state["completed"]
        state["cursor"] = n % 3
        state["scheduleIndex"] = n
        # Closed form for n advances of s <- 5*s + 1 (mod 16).
        state["rng"] = (5**n * seed + sum(5**k for k in range(n))) % 16
    # Otherwise completed labels the comparison; it repairs no missing state.
    return state

if __name__ == "__main__":
    saved = checkpoint()
    reference = trace(saved, 2)
    for policy in POLICIES:
        resumed = trace(restore(saved, policy), 2)
        print(policy, "whole trace agrees:", resumed == reference)
        for u in resumed:
            print(u["after"]["completed"], u["rowId"],
                  "rng", u["before"]["rng"], "->", u["rngAfter"],
                  "xi", u["jitter"], "g", u["gradient"], "rate", u["rate"],
                  "v", u["after"]["velocity"], "theta", u["after"]["theta"])

The full restore and the reconstruction from progress both print True; the other five policies print False. Correct continuation gives $117/512$ then $1511/4096$. If you change ROWS, the seed, the checkpoint or the recipe, compare against an uninterrupted run with the same inputs. The program assumes the small inputs declared above; it is not a loader for untrusted checkpoints.

On a quadratic bowl nothing is left open: the multipliers decide the path before a single step is drawn. Neural-network losses are not quadratic, and their curvature changes as training moves. Cohen et al. (2021) found that full-batch gradient descent on neural networks typically pushes the sharpest curvature up until it reaches about $2/\eta$, the edge this page derives, and then hovers there, the edge of stability. The loss keeps falling over the long run while jumping around in the short run, which the quadratic picture says should not happen, and why it works is still being worked out. A wider question is open too: why the minima that gradient descent and its mini-batch form find in huge networks tend to generalize, when the same networks can fit random labels just as easily (Zhang et al., 2017).