SGD, Momentum and RMSProp: how weights go down the loss surface

Why we need it

Training a neural network means searching for the minimum of a loss function $L(\theta)$ over millions or billions of weights $\theta$. There is no closed-form formula for the minimum: the loss depends on the weights in too complicated a way. All we can do is walk towards the minimum step by step, and the only guide at every step is the gradient $\nabla L$, the vector pointing in the direction of the fastest growth of the loss. We step the opposite way.

This page is about three steps of one method:

• SGD (stochastic gradient descent) — a step against the gradient computed on a small random batch of examples, a mini-batch;
• Momentum (the heavy-ball method) — a step against the accumulated sum of past gradients: inertia damps zigzags and speeds up descent along gentle directions;
• RMSProp — a step divided by the root-mean-square size of recent gradients: every weight gets its own step scale.

Combine the second and third steps, add bias correction, and you get Adam. And SGD with momentum still trains many convolutional networks for images, often generalizing better on them than adaptive methods.

Intuition: a step downhill

Imagine walking down a hill in thick fog. All you can see is the slope under your feet: which way and how steeply it goes down. A sensible strategy is to step in the direction of the steepest descent, look around and repeat. The length of the step is set by the learning rate $\eta$.

Short steps are safe but slow. A step that is too long jumps over the bottom of the pit onto the opposite slope — and if it lands higher there, the next step throws it even further. See how this looks on the simplest loss $f(\theta)=L\theta^2/2$: its gradient is $f'(\theta)=L\theta$ and its curvature (second derivative) is $L$.

The formula, bone by bone

$\theta_{t-1}$

the current weights, the vector of all model parameters;

$g_t$

the stochastic gradient: the mean of the loss gradients $\ell$ over a mini-batch $B_t$ of $B$ random examples;

$\eta$

the learning rate, a positive number, usually between $10^{-4}$ and $10^{-1}$;

$-\eta\,g_t$

the step: against the gradient, of length $\eta\|g_t\|$ — the steeper the slope, the longer the step.

Why it works. Expand the loss in a Taylor series around the current point: $L(\theta-\eta g)\approx L(\theta)-\eta\,\|g\|^2$. While the step is small the loss drops by about $\eta\|g\|^2$, and this is the fastest decrease among all steps of the same length: the dot product $\langle\nabla L,\,d\rangle$ for a fixed length $\|d\|$ is smallest when $d$ points against the gradient.

The condition number: why descent can be slow

If the curvatures differ between directions, the step has to adapt to the steepest one: $\eta<2/\lambda_{\max}$. Then along the gentlest direction a coordinate shrinks per step by only the fraction $\eta\lambda_{\min}<2/\kappa$, where $\kappa=\lambda_{\max}/\lambda_{\min}$ is the condition number. Even with the best constant step $\eta=2/(\lambda_{\max}+\lambda_{\min})$ the error shrinks per step only by the factor $(\kappa-1)/(\kappa+1)$:

At $\kappa=100$ the factor is $0.980$, and reducing the error a thousandfold takes about 345 steps. A rough rule: SGD spends on the order of $\kappa$ steps per digit of accuracy.

Stochasticity: mini-batches and noise

The full gradient is the mean over all $N$ examples of the training set. With millions of examples one such step would cost a pass over all the data. SGD takes a random mini-batch of $B$ examples. Its gradient is an unbiased estimate of the full one: it equals $\nabla L$ on average, but every time it deviates by random noise.

$\Sigma$

the covariance of per-example gradients at the current point (normalized by $N$);

$\frac{N-B}{N-1}$

the correction for sampling without replacement: at $B=N$ the noise vanishes.

The noise variance falls as $1/B$ while the cost of a step grows as $B$. So within one epoch (one pass over the data) small batches make many imprecise steps and large batches few precise ones. The figure shows linear regression $y\approx wx+b$ on 200 points with the loss $L=\frac1{2N}\sum(wx_i+b-y_i)^2$: a quadratic “bowl” over the plane of the weights $(w,\,b)$.

Momentum: the inertia of a heavy ball

Gradient descent is a massless point: at every moment it moves exactly where the slope points. Boris Polyak (1964) proposed to give it mass. A heavy ball rolls, builds up speed and does not react to every bump. In the formula the velocity is the vector $u$:

$u_t$

the accumulated direction (the momentum buffer), $u_0=0$;

$\beta$

the momentum coefficient, $0\le\beta<1$, usually $0.9$; at $\beta=0$ this is plain SGD.

Unroll the recursion: $u_t=g_t+\beta g_{t-1}+\beta^2 g_{t-2}+\dots$ — a weighted sum of all past gradients with exponentially decaying weights. If gradients keep pointing the same way from step to step, they add up and the step grows to $\eta/(1-\beta)$, ten times larger at $\beta=0.9$. If the sign flips (a zigzag across a ravine), the terms cancel each other.

Why momentum accelerates: theory on a quadratic

Take one direction with curvature $L$ and write $h=\eta L$. Substitute $g_t=L\theta_{t-1}$ and eliminate $u$ (from the second formula $\eta u_{t-1}=\theta_{t-2}-\theta_{t-1}$). We get a second-order linear recurrence:

Its behaviour is set by the roots of the characteristic equation $\lambda^2-(1+\beta-h)\lambda+\beta=0$. The convergence rate is the largest root modulus $\rho$: the error decays roughly like $\rho^t$.

Nesterov: the gradient where momentum is taking you

Yurii Nesterov (1983) proposed to compute the gradient not at the current point but where momentum will carry the ball anyway. In PyTorch’s form (nesterov=True) this is a correction to the step:

The characteristic equation becomes $\lambda^2-(1+\beta)(1-h)\lambda+\beta(1-h)=0$. Along gentle directions oscillations decay faster ($\rho=\sqrt{\beta(1-h)}$ instead of $\sqrt\beta$), but the stability region is narrower: $\eta L<2(1+\beta)/(1+2\beta)$, only $1.36$ at $\beta=0.9$. Switch Figure 4 to “Nesterov” to see both features.

RMSProp: a step of its own for every weight

SGD and Momentum use one rate $\eta$ for all weights. But the gradients of different weights can differ by orders of magnitude: embeddings of rare words get rare large kicks, the biases of the last layer get constant small ones. A shared $\eta$ is too large for some and too small for others. RMSProp (Geoffrey Hinton, Coursera lectures, 2012) divides the step of each weight by its typical gradient size:

$g_t^2$

an element-wise square: every weight has its own $v$;

$v_t$

the exponential moving average of the squared gradient, an estimate of $\mathbb E[g^2]$ over the last $\sim1/(1-\beta)$ steps, $v_0=0$;

$\sqrt{v_t}$

the root mean square (RMS) of the gradient — hence the name;

$\epsilon$

a small constant ($10^{-8}$) against division by zero.

Dividing by $\sqrt{v_t}$ makes the step dimensionless: multiply all gradients of a weight by 1000 and $\sqrt v$ grows 1000 times too, so the step does not change. A typical step becomes of order $\eta$ — the rate now sets the step length directly in the units of the weights.

Trajectories in 2D and 3D

Let us compare the three methods on one problem — an elongated quadratic valley $f=\tfrac12(a^2+\kappa b^2)$, where $a$ and $b$ are the coordinates along and across the ravine: the curvature across is $\kappa$ times larger than along. The start is at $a=-3.4$, $b=0.9$, where $f=46.3$.

The second classic is the Rosenbrock function $f=(1-x)^2+100(y-x^2)^2$: a narrow curved “banana” with its minimum at $(1,\,1)$. Here the curvature changes from point to point, and no constant step is ideal.

In code: from scratch and in PyTorch

All three methods fit into two small classes. Note that the state ($u$ or $v$) is an array of the same size as the weights, so Momentum and RMSProp double the memory for parameters and Adam triples it.

pythonimport numpy as np

class SGD:
    """SGD with momentum and the Nesterov variant, as in torch.optim.SGD."""
    def __init__(self, lr=0.01, momentum=0.0, nesterov=False):
        self.lr, self.beta, self.nesterov, self.u = lr, momentum, nesterov, None

    def step(self, theta, g):
        if self.beta == 0:
            return theta - self.lr * g
        self.u = g.copy() if self.u is None else self.beta * self.u + g
        d = g + self.beta * self.u if self.nesterov else self.u
        return theta - self.lr * d

class RMSProp:
    """RMSProp: every weight gets its own step scale."""
    def __init__(self, lr=0.01, beta=0.99, eps=1e-8):
        self.lr, self.beta, self.eps, self.v = lr, beta, eps, None

    def step(self, theta, g):
        if self.v is None:
            self.v = np.zeros_like(theta)
        self.v = self.beta * self.v + (1 - self.beta) * g * g
        return theta - self.lr * g / (np.sqrt(self.v) + self.eps)

# An elongated valley f = (a² + 100·b²)/2, as in Figure 6
H = np.diag([1.0, 100.0])
f = lambda th: 0.5 * th @ H @ th
grad = lambda th: H @ th

for name, opt in [("SGD", SGD(lr=0.018)),
                  ("Momentum", SGD(lr=0.018, momentum=0.9)),
                  ("Nesterov", SGD(lr=0.018, momentum=0.9, nesterov=True)),
                  ("RMSProp", RMSProp(lr=0.1, beta=0.9))]:
    th = np.array([-3.4, 0.9])
    for t in range(100):
        th = opt.step(th, grad(th))
    print(f"{name:9s} f = {f(th):.2e}")
# SGD       f = 1.53e-01
# Momentum  f = 2.23e-04
# Nesterov  f = 3.14e+57   ← diverges: ηλ_max = 1.8 > 2(1+β)/(1+2β) ≈ 1.36
# RMSProp   f = 2.60e-02

The numbers in the comments match Figure 6, and torch.optim.SGD and torch.optim.RMSprop give the same values on the same problem. Nesterov with the same $\eta$ diverges — exactly as the stability condition from the momentum section predicts.

pythonimport torch
import torch.nn.functional as F

model = torch.nn.Linear(10, 1)

# SGD with momentum and weight decay, the classic for convolutional networks
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9,
                      nesterov=True, weight_decay=1e-4)
# RMSProp: PyTorch calls the smoothing factor alpha (default 0.99)
# opt = torch.optim.RMSprop(model.parameters(), lr=1e-3, alpha=0.99, eps=1e-8)

for x, y in loader:                  # mini-batches from a DataLoader
    opt.zero_grad()                  # gradients accumulate, so reset them
    loss = F.mse_loss(model(x), y)
    loss.backward()                  # g_t lands in param.grad
    opt.step()                       # θ ← θ − η·(…)

Hyperparameters and practice

• Learning-rate schedule. Step decay (multiply by $0.1$ at epochs 30 and 60 out of 90 — the ResNet classic), cosine or one-cycle. A constant $\eta$ leaves a noise cloud around the minimum.
• Warmup. Over the first hundreds to thousands of steps $\eta$ grows linearly from a small value — especially important with a large batch.
• Gradient clipping by norm protects against rare huge gradients in RNNs and transformers.
• Learning-rate search (the LR range test, Leslie Smith): a short run with an exponentially growing $\eta$; pick a value well below the point where the loss starts to rise.
• Normalization of inputs, batch normalization and sensible initialization lower the condition number and allow a larger $\eta$.

Common mistakes

1. A forgotten zero_grad(). In PyTorch gradients accumulate in .grad; without resetting, $g_t$ becomes the sum of all past gradients — not momentum, but a bug.
2. A learning rate that is too large: the loss grows or turns into NaN. By Figure 1 you need $\eta<2/\lambda_{\max}$; first of all reduce $\eta$ 3–10 times.
3. Carrying $\eta$ over between optimizers. With momentum $\beta=0.9$ the steady-state step is 10 times larger than SGD’s with the same $\eta$. For RMSProp and Adam $\eta$ is in different units altogether: $10^{-3}$ is normal for them and usually too small for SGD.
4. Changing the batch without changing the rate. The batch grew 8 times and $\eta$ stayed — the network learns more slowly per epoch; the batch shrank — the noise may destabilize training.
5. $\epsilon=10^{-8}$ in low precision. In float16 this number rounds to zero, and when $v\approx0$ the RMSProp step explodes; with mixed precision $v$ is kept in float32 or $\epsilon$ is made larger.
6. A constant rate until the end of training. Without lowering $\eta$, SGD jitters around the minimum and the last epochs achieve almost nothing.

Links to other formulas

• Adam: the numerator is EMA momentum, the denominator is $\sqrt{v}$ from RMSProp, plus bias correction.
• MSE: its gradient $\frac2n\sum_i(\hat y_i-y_i)\,x_i$ is, up to a factor of 2, the $g_t$ of the linear regression in Figure 2 (the loss there is half the MSE).
• BCE and cross-entropy: the gradient with respect to the logits, $\hat y-y$, is the simplest descent signal.
• The sigmoid: saturation produces tiny gradients — SGD almost stands still on them, while RMSProp stretches the step.
• Huber: caps the gradient of outliers — a relative of gradient clipping.

Summary

Check yourself