Stochastic Gradient Descent and Mini-Batches, machine learning lesson

A training table with ten million rows makes every step of gradient descent expensive, because the exact gradient of the training objective averages one gradient per row. Stochastic gradient descent takes a step after reading only a few rows. Can such a cheap step still lead to a minimizer, and how would we notice if it does not? This lesson shows in what sense a mini-batch gradient is a correct estimate, what epochs and learning-rate schedules control, why feature scaling decides which step sizes work, and which evidence separates ordinary noise from a failed training run.

Key terms

Full-batch gradient descent

Definition Let the training objective be the empirical risk

J(\boldsymbol{\theta}) = \frac{1}{n} \sum_{i=1}^{n} L_i(\boldsymbol{\theta}),

the average over the $n$ training examples of the loss $L_i(\boldsymbol{\theta})$ of the model with parameters $\boldsymbol{\theta}$ on example $i$ , where every $L_i$ is differentiable. Full-batch gradient descent starts at a vector $\boldsymbol{\theta}_0$ and repeats

\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \alpha \nabla J(\boldsymbol{\theta}_t)

with a learning rate $\alpha > 0$ , so each update evaluates all $n$ example gradients at the current parameters. For a linear model with half squared loss $L_i(\mathbf{w}) = \tfrac12 (\mathbf{x}_i^{\top}\mathbf{w} - y_i)^2$ , the gradient of one example is its residual, prediction minus target, times its feature vector,

\nabla L_i(\mathbf{w}) = (\mathbf{x}_i^{\top}\mathbf{w} - y_i)\,\mathbf{x}_i ,

and $J$ is a quadratic with Hessian $H = X^{\top}X/n$ . When $H$ is positive definite, gradient descent with a constant rate converges to the unique minimizer from every starting point exactly when $0 < \alpha < 2/h_{\max}$ , where $h_{\max}$ is the largest eigenvalue of $H$ .

In plain words Each step uses the exact slope of the training objective, so the path is deterministic, and each step costs a full pass over the data. The rate has to suit the most sharply curved direction. If one feature is recorded on a much larger scale than another, the objective is sharply curved in one direction and nearly flat in another, and a rate small enough to stay stable in the sharp direction crawls in the flat one. Standardizing the features balances the curvature. The means and standard deviations it uses are computed from the training rows only and reused unchanged on validation, test and future rows. A training column with standard deviation zero cannot be divided by it, so the pipeline needs a stated policy, such as using a scale of $1$ for that column, which leaves a centered column of zeros, or dropping the column.

Example Four constructed training rows have features $(1, 100)$ , $(1, 300)$ , $(3, 100)$ and $(3, 300)$ and targets $4$ , $6$ , $6$ and $8$ . The training means of the features are $2$ and $200$ , and their standard deviations with divisor $n$ are $1$ and $100$ . Take the intercept to be the training mean $6$ of the targets and fit two coefficients to the centered targets $-2$ , $0$ , $0$ and $2$ , starting from $\mathbf{w} = \mathbf{0}$ , where $J = 1$ . With centered but unscaled features, $H$ is diagonal with entries $1$ and $10000$ , the gradient at $\mathbf{0}$ is $(-1, -100)$ , and one step with $\alpha = 0.1$ gives $\mathbf{w} = (0.1, 10)$ and $J = 499000.905$ . Stability would need $\alpha < 0.0002$ . After standardization the four rows are $(-1, -1)$ , $(-1, 1)$ , $(1, -1)$ and $(1, 1)$ , the Hessian is the identity, the gradient at $\mathbf{0}$ is $(-1, -1)$ and the same step gives $\mathbf{w} = (0.1, 0.1)$ and $J = 0.81$ . With this rate every step multiplies $J$ by $0.81$ , so twenty steps leave about $0.0148$ . Both coordinate systems describe the same prediction rules, and the exact fit is $\hat{y} = 2 + x_1 + 0.01x_2$ in the original units.

Stochastic gradient

Definition A stochastic gradient is an estimate of $\nabla J(\boldsymbol{\theta})$ computed from training examples chosen at random. The simplest one fixes $\boldsymbol{\theta}$ , draws an index $I$ uniformly from $\{1, \dots, n\}$ and uses $\mathbf{g} = \nabla L_I(\boldsymbol{\theta})$ . Each index has probability $1/n$ , so

\begin{aligned} E[\mathbf{g}] &= \sum_{i=1}^{n} \frac{1}{n} \nabla L_i(\boldsymbol{\theta}) \\[0.4em] &= \nabla J(\boldsymbol{\theta}), \end{aligned}

and the estimate is unbiased. Inside an algorithm the same argument gives $E[\mathbf{g}_t] = \nabla J(\boldsymbol{\theta}_t)$ conditionally on everything that happened before update $t$ , provided each index is drawn afresh and independently of all earlier draws. It does not hold for every batch of an epoch that walks through one shuffled order, because given the earlier batches of that epoch, the later ones can contain only examples not yet used.

In plain words One example’s gradient is a noisy guess of the full gradient. Averaged over the random choice it is exactly right, but a single draw can point the wrong way, and the step it produces can then raise the training objective.

Example A constant predictor $w$ is fitted to the targets $-2$ , $1$ and $7$ with $L_i(w) = \tfrac12 (w - y_i)^2$ , so $J$ is smallest at the mean $2$ . At $w = 0$ the example gradients $w - y_i$ are $2$ , $-1$ and $-7$ , and the full gradient is their mean, $-2$ . With $\alpha = 0.25$ the full-batch step moves $w$ to $0.5$ and lowers $J$ from $9$ to $8.125$ . If the first example is drawn, the gradient is $2$ , the step moves $w$ to $-0.5$ , and $J$ rises to $10.125$ . The three possible one-example steps end at $-0.5$ , $0.25$ and $1.75$ , and their average is the full-batch result $0.5$ .

Mini-batch

Definition A mini-batch is a selection $B$ of $m$ training examples with $1 \le m \le n$ , and the mini-batch gradient is

\mathbf{g}_B(\boldsymbol{\theta}) = \frac{1}{\lvert B \rvert} \sum_{i \in B} \nabla L_i(\boldsymbol{\theta}),

averaged over the actual number of examples in $B$ . With $m = 1$ it is a one-example stochastic gradient, and with all $n$ examples it is the full gradient. Let $\Sigma(\boldsymbol{\theta})$ be the covariance matrix of a one-example stochastic gradient, the average of $(\nabla L_i - \nabla J)(\nabla L_i - \nabla J)^{\top}$ over the $n$ examples at $\boldsymbol{\theta}$ . When the $m$ indices are drawn independently and uniformly with replacement, the mini-batch gradient is unbiased with covariance $\Sigma(\boldsymbol{\theta})/m$ . When $B$ is a uniformly random set of $m$ distinct examples and $n \ge 2$ , it is still unbiased, and its covariance is

\operatorname{Cov}(\mathbf{g}_B) = \frac{n - m}{n - 1} \cdot \frac{\Sigma(\boldsymbol{\theta})}{m},

which is zero when $m = n$ .

In plain words Averaging several gradients cancels part of their noise. Under sampling with replacement, the variance of each coordinate falls in proportion to $1/m$ and its standard deviation in proportion to $1/\sqrt{m}$ , so four times the batch size halves the typical size of the noise while each update costs about four times as much computation. When $m$ does not divide $n$ , the last batch of an epoch is smaller, and its summed gradient must be divided by its own size.

Example For the targets $-2$ , $1$ and $7$ at $w = 0$ , the one-example gradients $2$ , $-1$ and $-7$ have mean $-2$ and variance $(16 + 1 + 25)/3 = 14$ . A batch of two indices drawn with replacement has variance $14/2 = 7$ . The three sets of two distinct examples have batch gradients $0.5$ , $-2.5$ and $-4$ , with mean $-2$ and variance $3.5$ , which equals $\tfrac12 \cdot 14/2$ as the formula predicts. The set of the first two examples still gives a gradient of the wrong sign.

Stochastic gradient descent

Definition Stochastic gradient descent (SGD) starts at $\boldsymbol{\theta}_0$ and repeats

\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \alpha_t \mathbf{g}_t ,

where $\mathbf{g}_t$ is the gradient of a mini-batch $B_t$ evaluated at the current parameters $\boldsymbol{\theta}_t$ and $\alpha_t > 0$ is the learning rate of update $t$ . The name covers the one-example version and the mini-batch versions. In the version whose gradients are unbiased given the past, each $B_t$ is drawn afresh, uniformly and independently of earlier batches. In common practice the training rows are shuffled at the start of every epoch and taken in consecutive batches. An epoch is one pass through the $n$ training examples. When an epoch walks through a shuffled order with batch size $m$ , it has $\lceil n/m \rceil$ updates, while full-batch gradient descent makes one update per epoch.

In plain words SGD trades exact steps for cheap ones and makes many updates in the time full-batch descent makes one. Runs that differ only in the random seed follow different paths and can stop at different parameters. That spread is randomness of the optimizer on fixed training data. It is not uncertainty about future outcomes, and it does not measure how much the fitted model would change with a new training sample. Fixing the seed makes the shuffles repeatable on the same software, but library versions, hardware and parallel execution can still change the results.

Example Take the standardized four-row example with batch size $2$ , $\alpha = 0.1$ and $\mathbf{w} = \mathbf{0}$ . Suppose the first shuffle orders the rows as 2, 3, 1, 4. Rows 2 and 3 have centered target $0$ and prediction $0$ , so the first batch gradient is $\mathbf{0}$ and nothing moves. Rows 1 and 4 have residuals $2$ and $-2$ , so the second batch gradient is $\tfrac12 \bigl[ 2(-1, -1) - 2(1, 1) \bigr] = (-2, -2)$ , the step gives $\mathbf{w} = (0.2, 0.2)$ , and $J$ falls to $0.64$ after the first epoch, against $0.81$ after one full-batch step. In one seeded run of twenty epochs the mini-batch version reaches $J \approx 0.000162$ with $40$ updates and the full-batch version $J \approx 0.0148$ with $20$ updates, although both evaluate $80$ example gradients. That ranking belongs to this tiny constructed problem and is not a general guarantee.

Learning-rate schedule

Definition A learning-rate schedule is a rule that sets the rate $\alpha_t$ for each update $t = 0, 1, 2, \dots$ . A constant schedule keeps $\alpha_t = \alpha$ . A common decreasing schedule is

\alpha_t = \frac{\alpha_0}{1 + ct}

with $\alpha_0 > 0$ and $c > 0$ , and step decay, which multiplies the rate by a fixed factor after a fixed number of updates or epochs, is another. The counter has to be stated. A schedule written in epochs gives a different sequence of rates per update when the batch size changes, because an epoch then contains a different number of updates. The rate also depends on how the objective is normalized, since the gradient of a sum of $n$ losses is $n$ times the gradient of their average.

In plain words Large early steps make fast progress far from a minimizer, and small later steps damp the noise of stochastic gradients near one. With a constant rate, SGD cannot settle on a minimizer unless every example gradient is zero there, and with a small constant rate it keeps moving around near one.

Example With $\alpha_0 = 0.1$ and $c = 0.01$ , the rate is $0.1$ at update $0$ , $0.05$ at update $100$ and about $0.0091$ at update $1000$ . On $n = 1000$ training rows one epoch has $100$ updates with batch size $10$ but only $10$ updates with batch size $100$ . After one epoch the first run has reached the rate $0.05$ , while the second is still at about $0.091$ .

Convergence

Definition A sequence of parameter vectors $\boldsymbol{\theta}_t$ converges to $\boldsymbol{\theta}^*$ if $\lVert \boldsymbol{\theta}_t - \boldsymbol{\theta}^* \rVert_2 \to 0$ as $t \to \infty$ . SGD produces random iterates, so its convergence results are probability statements, such as convergence with probability one or $E[J(\boldsymbol{\theta}_t)] - J(\boldsymbol{\theta}^*) \to 0$ , and each holds only under stated assumptions. For full-batch descent on a quadratic with positive definite Hessian, the condition $0 < \alpha < 2/h_{\max}$ suffices. Classical stochastic approximation results for SGD use rates with

\begin{gathered} \sum_{t=0}^{\infty} \alpha_t = \infty, \\[0.4em] \sum_{t=0}^{\infty} \alpha_t^2 < \infty, \end{gathered}

together with conditions on the objective and the gradient noise, for example a strongly convex objective with Lipschitz continuous gradient and stochastic gradients that are unbiased given the past and have suitably bounded second moments. The schedule $\alpha_0/(1 + ct)$ meets both step conditions, and a constant rate fails the second. For nonconvex objectives, such as the losses of neural networks, results of this kind concern stationary points rather than a global minimum. In practice “the run converged” often means only that a stopping rule was met, such as a small change of the full training objective over several checkpoints.

In plain words Convergence is a claim about a limit, and the step conditions alone do not establish it. To monitor a run, evaluate the full training objective at fixed checkpoints instead of comparing the loss of one batch with the loss of a different batch, since an easy batch can look better while the model has become worse. Signs of failure are nonfinite values, a full training objective that keeps growing across checkpoints, and oscillations whose size grows. A noisy trace of batch losses is not by itself a failure, and a smooth falling training curve does not show that the model will do well on new data. Batch size and schedule are development choices, so compare them on validation data and keep the test set for the final evaluation.

Example For the targets $-2$ , $1$ and $7$ , one-example SGD with the constant rate $\alpha = 0.1$ gives $w_{t+1} - 2 = 0.9(w_t - 2) + 0.1(y_{I_t} - 2)$ . The sampled target has variance $14$ , so the variance $V_t$ of $w_t$ obeys $V_{t+1} = 0.81V_t + 0.14$ and approaches $0.14/0.19 \approx 0.74$ , a standard deviation of about $0.86$ around the minimizer $2$ that never dies out. The schedule $\alpha_t = 1/(t + 1)$ , which is $\alpha_0/(1 + ct)$ with $\alpha_0 = c = 1$ , makes $w_t$ the average of the first $t$ sampled targets for every $t \ge 1$ . Its variance is $14/t$ , and by the law of large numbers it converges to $2$ with probability one. The four-row regression converged with a constant rate because its model fits every training row exactly, which makes every example gradient zero at the minimizer.

Check yourself

A training set has $n = 10$ rows and is reshuffled every epoch into batches of size $4$ , with a smaller last batch. How many updates does one epoch contain, how many do five epochs contain, and by what number should the summed gradient of the last batch be divided?

Show the answer

The batches have sizes $4$ , $4$ and $2$ , so one epoch has $\lceil 10/4 \rceil = 3$ updates and five epochs have $15$ . The summed gradient of the last batch is divided by $2$ , its actual size. Dividing it by $4$ would make that step half as large as the averaged objective calls for. Full-batch gradient descent would make only $5$ updates in five epochs, each reading all $10$ rows.

A constant predictor $w$ has the targets $2$ and $6$ and the objective $J(w) = \tfrac14 \bigl[(w - 2)^2 + (w - 6)^2\bigr]$ . Start at $w = 0$ with $\alpha = 0.25$ . Find the full-batch update, the two possible one-example updates when each example is drawn with probability $\tfrac12$ , and their average. Does each one-example step lower $J$ ?

Show the answer

The objective is the average of $\tfrac12 (w - y_i)^2$ , so the example gradients at $w = 0$ are $-2$ and $-6$ and the full gradient is $-4$ . The full-batch update is $0 - 0.25 \cdot (-4) = 1$ . The one-example updates are $0.5$ and $1.5$ , and their average is $1$ , the full-batch update, as unbiasedness at a fixed starting point requires. Since $J(0) = 10$ , $J(0.5) = 8.125$ and $J(1.5) = 5.125$ , both steps lower $J$ here, because both example gradients have the same sign as the full gradient. With the targets $-2$ , $1$ and $7$ one possible step raised $J$ . The equality of averages concerns this single step from a fixed start. After it, the stochastic path does not follow the full-batch path.

In the four-row example, a run on centered but unscaled features with $\alpha = 0.1$ reports $J$ values of $1$ , $499000.905$ and about $5.0 \times 10^{11}$ after zero, one and two updates. Does this show that a linear model cannot fit these data? Which constant rates would converge, and what would you change?

Show the answer

No. The coefficients $1$ and $0.01$ on the centered features fit every row exactly, with $J = 0$ , which is $\hat{y} = 2 + x_1 + 0.01x_2$ in the original units. The Hessian is diagonal with entries $1$ and $10000$ , so with $\alpha = 0.1$ the error in the second coordinate is multiplied by $1 - 0.1 \cdot 10000 = -999$ at every step and the run diverges. That is a failure of the step size, not of the model. A constant rate converges only when $0 < \alpha < 2/10000 = 0.0002$ , and then the error in the first coordinate shrinks by the factor $1 - \alpha > 0.9998$ per step, so a tenfold reduction takes more than 11,500 steps. Standardizing the features with the training means and standard deviations makes the Hessian the identity. Every rate with $0 < \alpha < 2$ then converges, and $\alpha = 0.1$ multiplies $J$ by $0.81$ per step. The same saved scaling is then applied to validation and test rows.

Stochastic Gradient Descent and Mini-Batches

Key terms

Full-batch gradient descent

Stochastic gradient

Mini-batch

Stochastic gradient descent

Learning-rate schedule

Convergence

Check yourself

How this lesson connects

Builds on

Used later in