%matplotlib inline
from d2l import torch as d2l
import math
import random
import torch9.3 Stochastic Gradient Descent
Stochastic gradient descent replaces the exact gradient with an estimate computed from sampled examples. Building on Section 9.2, this section examines how the estimate changes the optimization trajectory, why convergence requires a decreasing learning rate, and how batch size controls its variance.
%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
import math9.3.1 Stochastic Gradient Updates
In deep learning, the objective function is an average of per-example losses. Given a training dataset of \(n\) examples with loss \(f_i(\mathbf{x})\) on the example of index \(i\), where \(\mathbf{x}\) is the parameter vector, the objective is
\[f(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n f_i(\mathbf{x}), \tag{9.3.1}\]
with gradient
\[\nabla f(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n \nabla f_i(\mathbf{x}). \tag{9.3.2}\]
Gradient descent therefore costs \(\mathcal{O}(n)\) per update—a full pass over the dataset to move the parameters once. Stochastic gradient descent (SGD) instead samples an index \(i\in\{1,\ldots, n\}\) uniformly at random and updates using that single example’s gradient:
\[\mathbf{x} \leftarrow \mathbf{x} - \eta \nabla f_i(\mathbf{x}), \tag{9.3.3}\]
where \(\eta\) is the learning rate. The cost per iteration drops from \(\mathcal{O}(n)\) to \(\mathcal{O}(1)\), and the estimate is unbiased: because \(i\) is uniform,
\[\mathbb{E}_i \nabla f_i(\mathbf{x}) = \frac{1}{n} \sum_{i = 1}^n \nabla f_i(\mathbf{x}) = \nabla f(\mathbf{x}). \tag{9.3.4}\]
On average the stochastic gradient points the right way; any single draw may point almost anywhere. To see what that does to the trajectory we revisit the quadratic \(f(x_1, x_2) = x_1^2 + 2x_2^2\) from Section 9.2 and simulate the sampling noise by adding zero-mean, unit-variance Gaussian noise to each gradient evaluation (the exercises show this simulation is faithful to sampling from an actual dataset).
def f(x1, x2): # Objective function
return x1 ** 2 + 2 * x2 ** 2
def f_grad(x1, x2): # Gradient of the objective function
return 2 * x1, 4 * x2def sgd(x1, x2, s1, s2, f_grad):
g1, g2 = f_grad(x1, x2)
# Simulate noisy gradient (Python's random.gauss avoids a GPU sync per
# step that a framework-tensor .item() would force in this 1000-step
# demo; the noise is scalar so a framework tensor buys nothing).
g1 += random.gauss(0, 1)
g2 += random.gauss(0, 1)
eta_t = eta * lr()
return (x1 - eta_t * g1, x2 - eta_t * g2, 0, 0)def sgd(x1, x2, s1, s2, f_grad):
global key
g1, g2 = f_grad(x1, x2)
# Simulate noisy gradient: split off a fresh subkey per step, the JAX
# idiom for drawing a stream of random numbers
key, subkey = jax.random.split(key)
n1, n2 = jax.random.normal(subkey, (2,))
eta_t = eta * lr()
return (x1 - eta_t * (g1 + n1), x2 - eta_t * (g2 + n2), 0, 0)
key = jax.random.PRNGKey(42)def constant_lr():
return 1
eta = 0.1
lr = constant_lr # Constant learning rate
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=50, f_grad=f_grad))epoch 50, x1: -0.052828, x2: -0.065022
epoch 50, x1: -0.141533, x2: -0.069805
The trajectory is far noisier than the gradient-descent trace in Section 9.2, and, more troubling, it does not settle: run it for more steps and the quality stops improving (we encourage you to check). The reason is visible in the update rule. Near the minimum the true gradient vanishes, but the noise does not, so each step still moves the parameters by \(\eta\) times pure noise. With a constant learning rate the iterates end up executing a random walk around the optimum instead of converging to it. Shrinking \(\eta\) tames the walk but slows the approach; the way out of this bind is to change \(\eta\) during optimization, which is why the step function above threads through a learning rate function lr, dormant so far.
9.3.2 Dynamic Learning Rate
How close does that random walk get? For a quadratic bowl the answer is sharp enough to guide practice. Model one coordinate of the update as \(x_{t+1} = x_t - \eta\, (\lambda x_t + \xi_t)\), with curvature \(\lambda\) and zero-mean noise \(\xi_t\) of variance \(\sigma^2\). Far from the optimum the contraction term dominates and SGD makes gradient-descent-like progress. Close to it, contraction and noise injection balance at
\[\mathbb{E}\big[x_\infty^2\big] \approx \frac{\eta\, \sigma^2}{2\lambda}. \tag{9.3.5}\]
The iterates fluctuate within a noise ball whose squared radius grows linearly with the learning rate. Section 2.1 first introduced this phenomenon. Section 26.1.5 supplies the derivation and compares GD with SGD on the same quadratic objective. Equation Equation 9.3.5 says exactly what the demo showed: a constant \(\eta\) stalls at a noise floor proportional to \(\eta\). Halving \(\eta\) halves the floor — and also halves the speed of the initial approach. A time-dependent learning rate \(\eta(t)\) resolves this tradeoff: large early steps make rapid progress, while smaller later steps reduce the noise floor.
Choosing how fast \(\eta(t)\) decays becomes the new problem. Too fast, and optimization stops prematurely; too slow, and we waste time bouncing in the ball. A few basic strategies (more refined ones follow in Section 9.8):
\[ \begin{aligned} \eta(t) & = \eta_i \textrm{ if } t_i \leq t \leq t_{i+1} && \textrm{piecewise constant} \\ \eta(t) & = \eta_0 \cdot e^{-\lambda t} && \textrm{exponential decay} \\ \eta(t) & = \eta_0 \cdot (\beta t + 1)^{-\alpha} && \textrm{polynomial decay} \end{aligned} \]
In the piecewise constant scenario we drop the learning rate whenever progress stalls, lowering the noise floor after the iterates reach it. Exponential decay is more aggressive. Polynomial decay with \(\alpha = 0.5\) is a standard well-behaved choice for convex problems. The next experiment shows exponential decay in practice.
def exponential_lr():
# Global variable that is defined outside this function and updated inside
global t
t += 1
return math.exp(-0.1 * t)
t = 1
lr = exponential_lr
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=1000, f_grad=f_grad))epoch 1000, x1: -0.898006, x2: -0.098858
epoch 1000, x1: -0.810909, x2: -0.028332
The parameter variance is substantially smaller, but the iterates never reach the optimum \(\mathbf{x} = (0, 0)\), even after 1000 steps. The schedule quenches the noise too eagerly: its total budget \(\sum_t \eta(t)\) is finite, so the iterate can only ever travel a bounded distance from where it started. It stops not because it has arrived but because it has run out of learning rate. Polynomial decay with \(\alpha = 0.5\) avoids that trap, and convergence improves after only 50 steps.
def polynomial_lr():
# Global variable that is defined outside this function and updated inside
global t
t += 1
return (1 + 0.1 * t) ** (-0.5)
t = 1
lr = polynomial_lr
d2l.show_trace_2d(f, d2l.train_2d(sgd, steps=50, f_grad=f_grad))epoch 50, x1: 0.171984, x2: 0.004072
epoch 50, x1: 0.037837, x2: 0.002780
The two experiments are instances of a classical theorem, due to the paper that founded stochastic approximation. Robbins and Monro (1951) showed that convergence requires exactly the two properties the demos probe: the learning rates must sum to infinity, \(\sum_t \eta(t) = \infty\), so the iterate can travel arbitrarily far (the condition exponential decay just violated), while their squares must remain controlled, \(\sum_t \eta(t)^2 < \infty\), so the injected noise is eventually quenched. Polynomial decay with \(\alpha \in (1/2, 1]\) satisfies both; our \(\alpha = 1/2\) sits at the edge of the window, where running for a fixed budget of \(T\) steps with averaged iterates gives the classical \(\mathcal{O}(1/\sqrt{T})\) guarantee for convex objectives. We state these results rather than prove them: the convex rates, and the traps hidden in their constants, are worked out in Section 26.1.5. None of the noise-ball reasoning needs convexity, either — for smooth nonconvex objectives SGD still drives the expected gradient norm to zero at a noise-controlled rate, the theorem of Ghadimi and Lan (2013) proved in Section 26.2.1.
9.3.3 Gradient Variance and Batch Size
So far the noise was ours: we chose \(\sigma = 1\). In real training the noise comes from which examples land in the minibatch, and there we hold a dial — the batch size \(b\). A minibatch gradient averages \(b\) independent draws, so its variance is \(1/b\) times the single-example variance (Section 26.1.5 states and proves this precisely). We test the \(1/b\) prediction on a real network: a small two-layer MLP on the airfoil-noise regression dataset that serves as this chapter’s workhorse from Section 9.4 on. We freeze the parameters at a random initialization, take the full-dataset gradient as ground truth, and measure the mean squared deviation of minibatch gradients from it, for batch sizes spanning nearly three orders of magnitude — sampling with replacement, matching the theory.
data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
X = torch.cat([Xb for Xb, yb in data_iter])
y = torch.cat([yb for Xb, yb in data_iter])
torch.manual_seed(1)
W1, b1 = torch.randn(feature_dim, 64) * 0.1, torch.zeros(64)
W2, b2 = torch.randn(64, 1) * 0.1, torch.zeros(1)
params = [W1, b1, W2, b2]
for p in params:
p.requires_grad_(True)
def batch_grad(idx): # Flattened loss gradient on the minibatch X[idx]
h = torch.relu(X[idx] @ W1 + b1)
loss = ((h @ W2 + b2).squeeze() - y[idx]).pow(2).mean() / 2
return torch.cat([g.reshape(-1) for g in torch.autograd.grad(loss, params)])
g_full = batch_grad(torch.arange(len(y)))
batch_sizes = [1, 8, 64, 512]
var = []
for b in batch_sizes:
idx = torch.randint(0, len(y), (200, b))
var.append(torch.stack([((batch_grad(i) - g_full) ** 2).sum()
for i in idx]).mean().item())
d2l.plot(batch_sizes, [var, [var[0] / b for b in batch_sizes]],
'batch size', 'gradient variance', xscale='log', yscale='log',
legend=['measured', '1/b'])
data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
X = jnp.concatenate([jnp.asarray(Xb) for Xb, yb in data_iter])
y = jnp.concatenate([jnp.asarray(yb) for Xb, yb in data_iter])
k1, k2 = jax.random.split(jax.random.PRNGKey(1))
params = dict(W1=0.1 * jax.random.normal(k1, (feature_dim, 64)),
b1=jnp.zeros(64),
W2=0.1 * jax.random.normal(k2, (64, 1)), b2=jnp.zeros(1))
def batch_loss(params, idx):
h = jax.nn.relu(X[idx] @ params['W1'] + params['b1'])
return jnp.mean(((h @ params['W2'] + params['b2']).squeeze()
- y[idx]) ** 2) / 2
def batch_grad(idx): # Flattened loss gradient on the minibatch X[idx]
grads = jax.grad(batch_loss)(params, idx)
return jnp.concatenate([g.reshape(-1) for g in jax.tree.leaves(grads)])
g_full = batch_grad(jnp.arange(len(y)))
batch_sizes = [1, 8, 64, 512]
var, key = [], jax.random.PRNGKey(0)
for b in batch_sizes:
key, subkey = jax.random.split(key)
idx = jax.random.randint(subkey, (200, b), 0, len(y))
err = jax.vmap(lambda i: ((batch_grad(i) - g_full) ** 2).sum())(idx)
var.append(float(err.mean()))
d2l.plot(batch_sizes, [var, [var[0] / b for b in batch_sizes]],
'batch size', 'gradient variance', xscale='log', yscale='log',
legend=['measured', '1/b'])
The measured points fall on the \(1/b\) reference line across the whole range: a factor of 500 in batch size reduces the variance by the same factor. Note what the log scale conceals. Variance falling like \(1/b\) means noise amplitude falls like \(1/\sqrt{b}\), so spending \(100\times\) more compute per step reduces the noise amplitude by only a factor of \(10\). Batch size is thus a genuine second dial next to the learning rate, but one with diminishing returns. Section 9.4 takes up how to spend a compute budget between the two, and how batching interacts with the hardware that made it cheap; Section 9.10 asks when the \(1/b\) variance reduction stops improving data efficiency at the scale of modern language models.
One loose end. Our theory and the measurement above both sampled indices with replacement: each draw is independent, which is what makes the analysis clean. Practice does something else. Drawing \(n\) times with replacement touches only a fraction \(1 - (1-1/n)^n \approx 1 - e^{-1} \approx 0.63\) of the dataset per pass while picking other examples twice or more — lower data efficiency and higher variance than shuffling the dataset and traversing it exactly once per epoch. So that is what every training loop in this book does, reshuffling before each epoch. The cost is theoretical: within an epoch, successive gradients are no longer independent, and the clean proofs no longer apply verbatim. The gap is real but benign in practice, and the exercises take it up.
9.3.4 Summary
SGD trades exactness for speed: an unbiased \(\mathcal{O}(1)\)-per-step gradient estimate in place of an \(\mathcal{O}(n)\) exact one. This introduces variance, and this section met both of the dials that control it. A constant learning rate leaves the iterates rattling in a noise ball of squared radius proportional to \(\eta\); decaying learning rates converge, provided the decay permits unbounded cumulative movement while making the accumulated noise finite. Batch size is the other dial: minibatch gradient variance falls like \(1/b\), measured on a real network over a factor of 500 in batch size, though at a linear cost in compute per step. The proofs owed here live in Section 26.1.5 and Section 26.2.1: the convex rates and the nonconvex Ghadimi–Lan theorem.
9.3.5 Exercises
- Experiment with different learning rate schedules for stochastic gradient descent and with different numbers of iterations. In particular, plot the distance from the optimal solution \((0, 0)\) as a function of the number of iterations.
- Prove that for the function \(f(x_1, x_2) = x_1^2 + 2 x_2^2\) adding normal noise to the gradient is equivalent to minimizing a loss function \(f(\mathbf{x}, \mathbf{w}) = (x_1 - w_1)^2 + 2 (x_2 - w_2)^2\) where \(\mathbf{x}\) is drawn from a normal distribution.
- Compare convergence of stochastic gradient descent when you sample from \(\{(x_1, y_1), \ldots, (x_n, y_n)\}\) with replacement and when you sample without replacement.
- How would you change the stochastic gradient descent solver if some gradient (or rather some coordinate associated with it) was consistently larger than all the other gradients?
- Assume that \(f(x) = x^2 (1 + \sin x)\). How many local minima does \(f\) have? Can you change \(f\) in such a way that to minimize it one needs to evaluate all the local minima?
- Repeat the gradient-variance measurement of this section at the parameters reached after training the network (use any optimizer from this chapter). Does the \(1/b\) dependence still hold? What changes — the slope or the level — and why?