Stochastic Gradient Descent

Stochastic gradient descent (SGD) is gradient descent with a cheap, noisy gradient. Instead of averaging the loss gradient over every training image before each step, it averages over a small random mini-batch of images. Goodfellow, Bengio and Courville (2016, §8.3.1) call SGD and its variants probably the most used optimizers in deep learning; the gradients come from backpropagation.

Problem

Training a neural network minimizes the regularized empirical risk of the loss function article,

$$J(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell_i(\theta) + \frac{\lambda}{2}\lVert \theta \rVert^2, \qquad \ell_i(\theta) = \ell\big(f(\mathbf{x}_i; \theta), y_i\big),$$

where $\mathbf{x}_i$ is the $i$-th training image and $y_i$ its label. The gradient $\nabla J = \frac{1}{N}\sum_i \nabla \ell_i + \lambda\theta$ is an average over images, so each exact gradient needs forward and backward passes through all $N$ images, about 1.28 million for ImageNet. SGD replaces it with an estimate whose cost does not depend on $N$.

Inputs and Outputs

• Input: a differentiable model and per-image loss, $N$ labeled images, initial $\theta_0$, batch size $B$, learning rates $\eta_k$, and optionally momentum $\mu$ and weight decay $\lambda$.
• Output: parameters after a number of epochs (full passes over the training set), usually chosen by validation accuracy.

Intuition

The gradient of $J$ is a vote among all training images on which way the weights should move. A random sample of 32 images gives a noisy but fair estimate, and many quick, roughly right steps beat a few exact ones. Images are also redundant: thousands of handwritten 3s pull the weights in similar directions. Goodfellow et al. (§8.1.3) make both arguments: the standard error of a mean falls only as $1/\sqrt{B}$, so 100 times more images buys only 10 times less noise.

Algorithm

1. Initialize $\theta$ and, if momentum is used, a velocity $\mathbf{v} = \mathbf{0}$.
2. At the start of each epoch, shuffle the training images.
3. Take the next $B$ images $S_k$ from the shuffled order and compute the mini-batch gradient $\mathbf{g}_k = \frac{1}{B}\sum_{i \in S_k} \nabla \ell_i(\theta_k) + \lambda \theta_k$ by backpropagation.
4. Update: without momentum, $\theta_{k+1} = \theta_k - \eta_k \mathbf{g}_k$; with momentum, $\mathbf{v} \leftarrow \mu\mathbf{v} - \eta_k\mathbf{g}_k$ and $\theta_{k+1} = \theta_k + \mathbf{v}$.
5. Repeat steps 3–4 to the end of the epoch, then return to step 2.

$B = N$ recovers full-batch gradient descent; $B = 1$ is SGD in the strict sense.

Mathematical Formulation

An unbiased, noisy gradient

If $S_k$ is a uniformly random set of $B$ distinct indices, each image is equally likely to be included, so $\mathbb{E}[\mathbf{g}_k] = \nabla J(\theta_k)$: the mini-batch gradient is unbiased. Its noise is set by the spread of the per-image gradients. With $\Sigma = \frac{1}{N}\sum_i (\nabla\ell_i - \overline{\nabla\ell})(\nabla\ell_i - \overline{\nabla\ell})^\top$, the standard result for sampling without replacement gives

$$\operatorname{Cov}[\mathbf{g}_k] = \frac{1}{B}\,\frac{N - B}{N - 1}\,\Sigma \;\approx\; \frac{\Sigma}{B} \quad (B \ll N).$$

The variance falls as $1/B$, the factor Bottou, Curtis and Nocedal (2018, §4.2) use for mini-batches. $\Sigma$ does not vanish at a minimizer, where individual images still disagree although their average gradient is zero.

Momentum

Momentum, Polyak’s (1964) heavy-ball method, accumulates a velocity:

$$\mathbf{v}_{k+1} = \mu \mathbf{v}_k - \eta\, \mathbf{g}(\theta_k), \qquad \theta_{k+1} = \theta_k + \mathbf{v}_{k+1},$$

with $0 \le \mu < 1$ (notation of Sutskever et al., 2013). Eliminating $\mathbf{v}$ gives $\theta_{k+1} = \theta_k - \eta\,\mathbf{g}_k + \mu(\theta_k - \theta_{k-1})$: a gradient step plus a fraction of the previous step. Each step is an exponentially weighted sum of past gradients, so directions of persistent descent accumulate and oscillating directions cancel (Bottou et al., §7.1). Under a constant gradient the step approaches $\eta\mathbf{g}/(1 - \mu)$, ten times the plain step for $\mu = 0.9$ (Goodfellow et al., §8.3.2).

Nesterov momentum, the form of Nesterov’s accelerated gradient that Sutskever et al. (2013) introduced to deep learning, evaluates the gradient at the point the velocity is about to carry the parameters to:

$$\mathbf{v}_{k+1} = \mu \mathbf{v}_k - \eta\, \mathbf{g}(\theta_k + \mu\mathbf{v}_k), \qquad \theta_{k+1} = \theta_k + \mathbf{v}_{k+1}.$$

When the velocity carries the weights uphill, the look-ahead gradient corrects it sooner; Sutskever et al. argued that this makes it more stable than classical momentum, especially at large $\mu$.

Parameters

• Learning rate $\eta$. A crucial setting, usually decreased during training by a learning-rate schedule (Goodfellow et al., §8.3.1); Goyal et al. (2017) divided it by 10 at epochs 30, 60 and 80 of 90.
• Batch size $B$. Larger batches give less noisy gradients and use parallel hardware better, but memory grows with $B$. Goodfellow et al. (§8.1.3) give 32 to 256 as typical.
• Batch size and learning rate together. Goyal et al. (2017) proposed the linear scaling rule: multiply the learning rate by $k$ when the batch size is multiplied by $k$, because $k$ small steps of size $\eta$ match one step of size $k\eta$ if the gradient barely changes in between. That fails early in training, so they raised the learning rate linearly over the first 5 epochs (warmup). With batches of 8,192 images on 256 GPUs they trained ResNet-50 on ImageNet in one hour, with a top-1 validation error of 23.74% ± 0.09 against 23.60% ± 0.12 at batch 256; without warmup it was 24.84%, and at batch 16k with warmup, 24.79%.
• Momentum $\mu$. Common values are 0.5, 0.9 and 0.99 (Goodfellow et al., §8.3.2); Goyal et al. used Nesterov momentum with 0.9.
• Weight decay $\lambda$. The $\lambda\theta$ term, an L2 penalty that shrinks the weights every step, a form of regularization; Goyal et al. used $10^{-4}$.
• Shuffling. It matters when the stored order is meaningful, such as images grouped by class; never shuffling can seriously hurt (Goodfellow et al., §8.1.3).

Complexity

An update costs $B$ forward and backward passes and memory for $B$ images’ activations; an epoch costs $N$ per-image gradients for any $B$. The batch size sets how many updates an epoch buys and how well hardware is used: in the example below, a batch-1 epoch took about twenty times as long as a batch-32 epoch on our machine.

Implementation

The example trains a softmax classifier, a linear map from the 64 pixels of scikit-learn’s 8×8 digit images to 10 class scores, with cross-entropy loss and weight decay $\lambda = 10^{-2}$ on the weight matrix $W$ only: the $\lambda\theta$ terms above apply to $W$, not to the bias $\mathbf{b}$. The objective is convex with a well-defined minimum $J^\star$, so $J - J^\star$ measures optimization progress. The gradient was checked against centered finite differences (relative error below $10^{-9}$).

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 2

digits = load_digits()                       # 1797 grayscale 8x8 images, values 0..16
X_tr, X_te, y_tr, y_te = train_test_split(
    digits.data / 16.0, digits.target, test_size=0.25, random_state=0, stratify=digits.target)
N, D, K = X_tr.shape[0], X_tr.shape[1], 10   # 1347 training images, 64 pixels, 10 classes
Y_tr = np.eye(K)[y_tr]                       # one-hot labels
lam = 1e-2                                   # weight decay (L2 penalty on W)

def loss_grad(W, b, X, Y):
    """Softmax cross-entropy averaged over the images in X, plus lam/2 ||W||^2."""
    s = X @ W + b
    s -= s.max(axis=1, keepdims=True)
    P = np.exp(s)
    P /= P.sum(axis=1, keepdims=True)
    loss = -np.mean(np.log(P[Y > 0])) + 0.5 * lam * np.sum(W * W)
    G = (P - Y) / len(X)                     # d loss / d scores
    return loss, X.T @ G + lam * W, G.sum(axis=0)

def J(W, b):                                 # full training objective
    return loss_grad(W, b, X_tr, Y_tr)[0]

def accuracy(W, b):
    return np.mean((X_te @ W + b).argmax(axis=1) == y_te)

# Reference minimizer from many full-batch steps (the objective is convex).
W_star, b_star = np.zeros((D, K)), np.zeros(K)
for _ in range(3000):
    _, gW, gb = loss_grad(W_star, b_star, X_tr, Y_tr)
    W_star -= gW
    b_star -= gb
J_star = J(W_star, b_star)

# 1. Mini-batch gradients are noisy even at the minimizer, where the full gradient is ~0.
rng = np.random.default_rng(SEED)
per_image = np.stack([np.concatenate([g.ravel() for g in loss_grad(W_star, b_star, X_tr[i:i+1], Y_tr[i:i+1])[1:]])
                      for i in range(N)])
full = per_image.mean(axis=0)
var1 = np.mean(np.sum((per_image - full) ** 2, axis=1))      # per-image variance (trace)
print(f"|full gradient| at minimizer: {np.linalg.norm(full):.1e}")
for B in [1, 8, 32, 128, 512]:
    errs = [np.sum((per_image[rng.choice(N, B, replace=False)].mean(axis=0) - full) ** 2)
            for _ in range(2000)]
    predicted = var1 / B * (N - B) / (N - 1)                  # sampling without replacement
    print(f"B = {B:4d}: E|g_B - grad J|^2 = {np.mean(errs):.4f}  (predicted {predicted:.4f})")

# 2. Equal budget: 20 epochs (20 x 1347 per-image gradients) for every method.
def train(B, lr, momentum=0.0, epochs=20, seed=SEED, schedule=None, report=()):
    rng = np.random.default_rng(seed)
    W, b = np.zeros((D, K)), np.zeros(K)
    vW, vb = np.zeros_like(W), np.zeros_like(b)
    k, gaps = 0, []
    for epoch in range(1, epochs + 1):
        order = rng.permutation(N)           # reshuffle once per epoch
        for i in range(0, N, B):
            idx = order[i:i + B]
            _, gW, gb = loss_grad(W, b, X_tr[idx], Y_tr[idx])
            eta = lr if schedule is None else schedule(k)
            vW = momentum * vW - eta * gW    # heavy-ball momentum (Polyak)
            vb = momentum * vb - eta * gb
            W, b, k = W + vW, b + vb, k + 1
        if epoch in report:
            gaps.append(J(W, b) - J_star)
    return W, b, k, gaps

print(f"\nJ* = {J_star:.4f}; after 20 epochs:")
for name, B, lr, mom in [("full batch GD", N, 1.0, 0.0),
                         ("SGD, batch 1", 1, 0.01, 0.0),
                         ("mini-batch 32", 32, 0.1, 0.0),
                         ("mini-batch 32 + momentum 0.9", 32, 0.03, 0.9)]:
    W, b, k, _ = train(B, lr, mom)
    print(f"{name:29s} lr {lr:<5} {k:6d} updates  J - J* = {J(W, b) - J_star:.1e}  "
          f"test accuracy {accuracy(W, b):.3f}")

# 3. Constant steps stall at a noise floor; Robbins-Monro steps keep converging.
print("\nmini-batch 32, J - J* after 10 / 50 / 200 epochs:")
for name, lr, sched in [("constant 0.3", 0.3, None),
                        ("constant 0.1", 0.1, None),
                        ("1 / (1 + k/420)", None, lambda k: 1.0 / (1 + k / 420))]:
    *_, gaps = train(32, lr, epochs=200, schedule=sched, report=(10, 50, 200))
    print(f"{name:16s} " + "  ".join(f"{g:.1e}" for g in gaps))

Output:

|full gradient| at minimizer: 3.8e-05
B =    1: E|g_B - grad J|^2 = 2.4983  (predicted 2.5604)
B =    8: E|g_B - grad J|^2 = 0.3177  (predicted 0.3184)
B =   32: E|g_B - grad J|^2 = 0.0774  (predicted 0.0782)
B =  128: E|g_B - grad J|^2 = 0.0181  (predicted 0.0181)
B =  512: E|g_B - grad J|^2 = 0.0030  (predicted 0.0031)

J* = 0.7353; after 20 epochs:
full batch GD                 lr 1.0       20 updates  J - J* = 1.7e-01  test accuracy 0.911
SGD, batch 1                  lr 0.01   26940 updates  J - J* = 3.5e-03  test accuracy 0.942
mini-batch 32                 lr 0.1      860 updates  J - J* = 7.7e-03  test accuracy 0.940
mini-batch 32 + momentum 0.9  lr 0.03     860 updates  J - J* = 5.2e-03  test accuracy 0.940

mini-batch 32, J - J* after 10 / 50 / 200 epochs:
constant 0.3     2.1e-02  1.1e-02  1.4e-02
constant 0.1     4.0e-02  2.2e-03  1.6e-03
1 / (1 + k/420)  4.7e-02  3.1e-03  3.7e-04

At the minimizer the full gradient is essentially zero, yet one image’s gradient misses it by a squared norm of about 2.5; the error falls as $1/B$, as predicted.

With the same 20 epochs of per-image gradients, full-batch gradient descent has taken 20 steps and is still far from the minimum; the stochastic runs are 20 to 50 times closer and about 3 points more accurate on the 450 test images. Each learning rate is the best by training loss of a grid from 0.003 to 10. Over ten shuffling seeds the median gaps were $4.4 \times 10^{-3}$ (batch 1), $8.0 \times 10^{-3}$ (batch 32) and $4.8 \times 10^{-3}$ (momentum, lower than plain batch 32 in nine seeds), with test accuracies from 93.6% to 95.6%.

The last table shows the noise floor. The constant step of 0.3 leads at epoch 10, then stalls between $10^{-2}$ and $2 \times 10^{-2}$; the step of 0.1 levels off near $2 \times 10^{-3}$. The decaying step $\eta_k = 1/(1 + k/420)$ satisfies the Robbins–Monro conditions and keeps improving; the ordering at epoch 200 held in all ten seeds.

The PyTorch version uses torch.optim.SGD with Nesterov momentum (run with PyTorch 2.14.1, continuing the script above) and runs Sutskever et al.’s update in NumPy on the same mini-batches for comparison:

import torch
import torch.nn.functional as F

def run_torch(lr_at_epoch, mu=0.9, epochs=20, B=32):
    """torch.optim.SGD with Nesterov momentum next to Sutskever et al.'s form;
    returns the PyTorch weights and their largest difference from theta + mu v."""
    W = torch.zeros(D, K, dtype=torch.float64, requires_grad=True)
    b = torch.zeros(K, dtype=torch.float64, requires_grad=True)
    opt = torch.optim.SGD([W, b], lr=lr_at_epoch(0), momentum=mu, nesterov=True)
    Xt, yt = torch.tensor(X_tr), torch.tensor(y_tr)

    theta = [np.zeros((D, K)), np.zeros(K)]       # Sutskever: parameters theta
    v = [np.zeros((D, K)), np.zeros(K)]           # and velocity v
    rng = np.random.default_rng(SEED)
    for epoch in range(epochs):
        eps = lr_at_epoch(epoch)
        for group in opt.param_groups:
            group["lr"] = eps
        order = rng.permutation(N)
        for i in range(0, N, B):
            idx = order[i:i + B]
            opt.zero_grad()
            loss = F.cross_entropy(Xt[idx] @ W + b, yt[idx]) + 0.5 * lam * (W * W).sum()
            loss.backward()
            opt.step()
            # v <- mu v - eps grad J(theta + mu v);  theta <- theta + v
            look = [t + mu * u for t, u in zip(theta, v)]
            g = loss_grad(look[0], look[1], X_tr[idx], Y_tr[idx])[1:]
            v = [mu * u - eps * gi for u, gi in zip(v, g)]
            theta = [t + u for t, u in zip(theta, v)]
    W_t, b_t = W.detach().numpy(), b.detach().numpy()
    look = [t + mu * u for t, u in zip(theta, v)]
    diff = max(np.abs(W_t - look[0]).max(), np.abs(b_t - look[1]).max())
    return W_t, b_t, diff

print(f"torch {torch.__version__}")
W_t, b_t, diff = run_torch(lambda epoch: 0.03)
print(f"constant lr 0.03:     J - J* = {J(W_t, b_t) - J_star:.1e}, test accuracy "
      f"{np.mean((X_te @ W_t + b_t).argmax(1) == y_te):.3f}, "
      f"max |torch - (theta + mu v)| = {diff:.1e}")
W_t, b_t, diff = run_torch(lambda epoch: 0.03 if epoch < 10 else 0.003)
print(f"lr 0.03 -> 0.003 at epoch 10:  max |torch - (theta + mu v)| = {diff:.1e}")

Output:

torch 2.14.1
constant lr 0.03:     J - J* = 4.2e-03, test accuracy 0.940, max |torch - (theta + mu v)| = 1.6e-15
lr 0.03 -> 0.003 at epoch 10:  max |torch - (theta + mu v)| = 3.2e-02

As its documentation notes, PyTorch writes momentum differently: it keeps a buffer of raw gradients and applies the learning rate $\gamma$ afterwards,

$$\mathbf{b}_k = \mu\,\mathbf{b}_{k-1} + \mathbf{g}_k, \qquad \theta_k = \theta_{k-1} - \gamma\,\mathbf{b}_k \quad (\text{or } \theta_k = \theta_{k-1} - \gamma(\mathbf{g}_k + \mu\,\mathbf{b}_k) \text{ with Nesterov}),$$

where $\mathbf{g}_k$ is the mini-batch gradient at the stored parameters $\theta_{k-1}$ (plus $\lambda\theta_{k-1}$ if weight_decay is set) and $\mathbf{b}_1 = \mathbf{g}_1$. With a constant learning rate, classical momentum is the same in both forms under $\mathbf{v} = -\gamma\mathbf{b}$ (a check of nesterov=False against the heavy-ball loop above agreed to $3 \times 10^{-16}$). PyTorch’s Nesterov update matches Sutskever’s because its stored parameters are Sutskever’s look-ahead point $\boldsymbol{\phi}_k = \theta_k + \mu\mathbf{v}_k$. Substituting $\mathbf{v}_{k+1} - \mu\mathbf{v}_k = -\eta\,\mathbf{g}(\boldsymbol{\phi}_k)$ into $\boldsymbol{\phi}_{k+1} = \theta_{k+1} + \mu\mathbf{v}_{k+1}$ gives

$$\boldsymbol{\phi}_{k+1} = \boldsymbol{\phi}_k - \eta\,\mathbf{g}(\boldsymbol{\phi}_k) + \mu\,\mathbf{v}_{k+1},$$

which with $\mathbf{v} = -\eta\mathbf{b}$ is PyTorch’s update, as the first printed line confirms. The forms differ when the learning rate changes: PyTorch’s buffer holds no learning rate, so a new rate rescales the whole accumulated history, while Sutskever’s velocity keeps each past step at the rate it was taken with; after the tenfold drop the trajectories separate. Goyal et al. (2017) call the factor $\eta_{k+1}/\eta_k$ that reconciles the two forms a momentum correction, and found it especially important when the learning rate rises sharply.

Properties and Behavior

• Noise floor with a constant step. For strongly convex objectives, Bottou, Curtis and Nocedal (2018, Theorem 4.6) show that SGD with a small enough fixed step converges linearly in expectation, but only to a neighborhood of the optimum whose size is proportional to the step and the gradient variance; a smaller step lowers the floor but slows the approach.
• Decaying steps. With $\sum_k \eta_k = \infty$ and $\sum_k \eta_k^2 < \infty$, conditions that go back to Robbins and Monro (1951), for strongly convex objectives and $\eta_k = \beta/(\gamma + k)$ with suitable constants, the expected gap falls as $O(1/k)$ (Bottou et al., Theorem 4.7). Goodfellow et al. (§8.3.1) note that SGD’s $O(1/\sqrt{k})$ (convex) and $O(1/k)$ (strongly convex) rates cannot be improved without further assumptions.
• Nonconvex objectives. For network losses the guarantees concern gradients: with a fixed step the average expected squared gradient norm tends to below a bound proportional to the step, and Robbins–Monro steps drive $\liminf_k \mathbb{E}\lVert\nabla J(\theta_k)\rVert^2$ to zero (Bottou et al., Theorems 4.8–4.9).
• Momentum under noise. With noisy gradients, acceleration does not improve the asymptotic rate (Goodfellow et al., §8.3.3; Bottou et al., §7.2). Sutskever et al. (2013) argued that the early, transient phase matters more in deep learning.
• Noise and generalization. Keskar et al. (2017) compared batches of 256 with batches of 10% of the training set on networks for MNIST, CIFAR-10, CIFAR-100 and TIMIT: large batches reached similar training accuracy but lower test accuracy. They presented evidence that large batches converge to sharp minimizers and small batches to flat ones, supporting the view that gradient noise is responsible. Their main experiments used Adam for both batch sizes; they report that other optimizers for the large-batch runs, including SGD, gave similar results. Goyal et al.’s warmup results above closed the gap at batch 8,192. Our reading: the gap depends on how other hyperparameters are adapted, and why gradient noise helps is not established.

Limitations

• Tuning. Learning rate, schedule, batch size and momentum interact and depend on the model and data; Bottou et al. (§3.3) note that their efficient SGD run required trying several step sizes.
• Ill-conditioning. One learning rate for all parameters makes badly scaled problems slow, as for gradient descent.
• Theory and practice differ. The analyses above assume independently drawn mini-batches, whereas training usually sweeps through a shuffled dataset without replacement.

Variants

• Adaptive methods such as Adam scale the step separately for each parameter using running averages of the gradient and its square.
• Iterate averaging returns an average of past iterates, or for neural networks an exponentially decaying running average, instead of the last one, reducing the effect of gradient noise (Goodfellow et al., §8.7.3; Bottou et al., §5.4).

History

Robbins and Monro (1951) sought the root of $M(x) = \alpha$ when each experiment at level $x$ returns only a noisy response with mean $M(x)$. Their iteration $x_{n+1} - x_n = a_n(\alpha - y_n)$ uses steps with $\sum a_n^2 < \infty$ and a condition implying $\sum a_n = \infty$, such as $a_n = 1/n$; they proved convergence in mean square and allowed several observations at the same level to be averaged before each step, in our reading an early mini-batch. SGD applies the scheme to $\nabla J(\theta) = \mathbf{0}$.

Boris Polyak (1964) proposed the heavy-ball method; as Sutskever et al. (2013) summarize, he showed that with $\mu = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1)$ it needs about $\sqrt{\kappa}$ times fewer iterations than steepest descent near a minimum with condition number $\kappa$. Yurii Nesterov’s accelerated gradient followed in 1983. According to Sutskever et al., interest in momentum diminished after the 1990s because theory predicts that its asymptotic advantage is lost with noisy gradients. They revived it, reporting that with a good initialization and a slowly increasing momentum schedule, SGD with momentum matched results previously reached only with Hessian-free optimization on deep autoencoders and recurrent networks; PyTorch’s Nesterov option is based on their formula.