Mini-batch SGD

Mini-batch stochastic gradient descent is the standard compromise between the two previous regimes:

• full-batch gradient descent, which uses the entire dataset at each update;
• pure SGD, which uses exactly one training example at each update.

In mini-batch SGD, each update is computed from a small subset of training examples, called a mini-batch.
If the mini-batch size is denoted by $B$, then typical choices are $B=32$, $64$, or $128$, although the best value depends on the model, the hardware, and the optimization regime.

On the Dimension of Batch Size

Those values are not special from a purely mathematical point of view. In practice, they are often favored because they interact well with modern accelerator hardware.

On GPUs, parallel work is executed in small groups of threads. In the NVIDIA terminology these groups are called warps, and a warp typically contains $32$ threads. Using batch sizes that are multiples of $32$ often helps keep the hardware more fully occupied, reducing the amount of partially idle parallel work.

There is also a memory-side reason. Batch sizes such as $32$, $64$, and $128$ often lead to tensor shapes that are friendlier to low-level kernels, memory alignment, and high-throughput matrix operations. In modern training stacks, this can improve throughput and make better use of matrix units or tensor-oriented compute paths.

However, this is a hardware-aware heuristic, not a universal law. The best batch size is still a compromise among:

• optimization behavior;
• available memory;
• hardware throughput.

For that reason, powers of two are common, but values such as $96$, $160$, or other non-powers of two may also be entirely reasonable when they better fit the model or the device.

Using the notation introduced in Neural Networks vocabulary, let

$$\mathcal{D}={\left\{\left({\mathbf{x}}^{(i)},{\mathbf{y}}^{(i)}\right)\right\}}_{i=1}^N,$$

and suppose that the objective is

$$\mathcal{L}(\mathbf{\theta })=\frac{1}{N}\sum _{i=1}^N{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

At iteration $t$, instead of selecting a single example, one selects a mini-batch index set

$${\mathcal{B}}_t\subseteq \{1,\dots ,N\},\left|{\mathcal{B}}_t\right|=B,1<B<N.$$

The set ${\mathcal{B}}_t$ identifies the examples used at that optimizer step.
The corresponding batched inputs and targets are

$${\mathbf{X}}_{{\mathcal{B}}_t}\in {\mathbb{R}}^{B\times d_x},{\mathbf{Y}}_{{\mathcal{B}}_t}\in {\mathbb{R}}^{B\times d_y},$$

which matches the batch notation adopted in Neural Networks vocabulary.

The mini-batch loss is the average loss over the examples contained in the current batch:

$${\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })=\frac{1}{B}\sum _{i\in {\mathcal{B}}_t}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

Its gradient is therefore

$${\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })=\frac{1}{B}\sum _{i\in {\mathcal{B}}_t}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }),$$

which is used as a computationally cheaper approximation to the full gradient

$${\nabla }_{\mathbf{\theta }}\mathcal{L}(\mathbf{\theta }).$$

Why The Mini-Batch Gradient Is A Random Vector

Under random batching, the object

$${\mathcal{B}}_t$$

is random: strictly speaking, it is a random subset of $\{1,\dots ,N\}$, or equivalently a set-valued random variable.

Consequently,

$${\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })=\frac{1}{B}\sum _{i\in {\mathcal{B}}_t}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta })$$

is random as well, because it is obtained by evaluating a vector-valued expression at a random mini-batch.

It’s a random vector because its value has one coordinate for each parameter in $\mathbf{\theta }$. Each coordinate is a scalar random variable, and together those coordinates form a vector.

The update rule becomes

$${\mathbf{\theta }}^{(t+1)}={\mathbf{\theta }}^{(t)}-\eta {\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}\left({\mathbf{\theta }}^{(t)}\right)={\mathbf{\theta }}^{(t)}-\eta \frac{1}{B}\sum _{i\in {\mathcal{B}}_t}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}\left({\mathbf{\theta }}^{(t)}\right).$$

If the compact parameter vector $\mathbf{\theta }$ is unpacked into weights and biases, the same rule can be written component wise as

$$w_k^{(t+1)}=w_k^{(t)}-\eta \frac{1}{B}\sum _{i\in {\mathcal{B}}_t}\frac{\partial {\mathcal{L}}_{{\mathbf{x}}^{(i)}}}{\partial w_k},$$ $$b_{\ell }^{(t+1)}=b_{\ell }^{(t)}-\eta \frac{1}{B}\sum _{i\in {\mathcal{B}}_t}\frac{\partial {\mathcal{L}}_{{\mathbf{x}}^{(i)}}}{\partial b_{\ell }}.$$

These are not different learning rules: they are simply the componentwise form of the same mini-batch update.

Why Mini-Batch SGD Is A Compromise

Mini-batch SGD does not use the whole dataset at once, as in full-batch gradient descent, and it does not use only one training example at a time, as in pure SGD.
Instead, it uses a fixed-size mini-batch containing $B$ training examples.
For this reason, it sits naturally between the two extremes:

• it is much cheaper per update than full-batch gradient descent;
• it is less noisy than pure SGD;
• it maps efficiently onto modern vectorized hardware such as GPUs.

Modern Viewpoint: Noise Can Be Beneficial

In deep learning, the stochasticity of mini-batch gradients is not merely a defect to be reduced as much as possible.
It can also act as a form of implicit regularization.

Because each mini-batch provides only a partial view of the whole dataset, the corresponding gradient typically fluctuates around the full-gradient direction. This noise perturbs the optimization trajectory and can help the iterates avoid becoming too tightly trapped in sharp basins of the loss landscape.

From a modern optimization viewpoint, this is one reason mini-batch training often favors flatter solutions over sharper ones. Flatter minima are frequently associated with better robustness to perturbations of the parameters and, in many practical settings, with better generalization. In this sense, the residual noise of mini-batch SGD is often a useful feature rather than only a source of error.

Batch Size And Learning Rate

The batch size $B$ should not be thought of independently of the learning rate $\eta$.
When $B$ increases, the mini-batch gradient usually becomes less noisy, so the update direction is estimated with lower variance. In that regime, one can often take larger optimization steps without immediately destabilizing training.

This is the intuition behind the linear scaling rule often used in large-batch training: if the batch size is multiplied by a factor $k$, one tries increasing the learning rate by approximately the same factor. In informal form,

$$B\mapsto kB⟹\eta \mapsto k\eta .$$

This should be interpreted as a practical heuristic, not as a universal law. Its usefulness depends on the optimizer, the model, the training phase, and the overall schedule. In particular, large-batch training often also requires additional stabilization, such as a warm-up phase, before the scaled learning rate can be used safely.

Increasing $B$ therefore tends to make each individual update more accurate, but it does not imply that the entire training run will become proportionally faster in wall-clock time. In practice, diminishing returns usually appear: beyond a hardware- and problem-dependent regime, enlarging the batch may keep reducing gradient noise while providing little additional reduction in the number of optimization steps required to reach a target loss.

This practical saturation point is often discussed under the name critical batch size. The expression should not be interpreted as a single universal threshold valid for every model and dataset; rather, it refers to the empirical observation that, past a certain scale, using still larger mini-batches often consumes more compute than it saves in optimization time.

The broader lesson is that increasing $B$ does not merely change computational throughput: it also changes the optimization dynamics. Batch size and learning rate must therefore be tuned jointly rather than treated as unrelated hyperparameters.

One Epoch In Mini-Batch Training

In the common finite-dataset implementation, a single epoch proceeds as follows:

Training pipeline with mini-batches

1. the indices $\{1,\dots ,N\}$ are shuffled;
2. the shuffled list is split into mini-batches of size $B$;
3. one mini-batch ${\mathcal{B}}_t$ is selected;
4. the corresponding batch ${\mathbf{X}}_{{\mathcal{B}}_t},{\mathbf{Y}}_{{\mathcal{B}}_t}$ is fed through the network;
5. the mean mini-batch loss and the mean mini-batch gradient are computed;
6. the update rule for $\mathbf{\theta }$ is applied;
7. the same steps are repeated for every mini-batch created from that shuffled pass through the dataset.

Once all those mini-batches have been processed, one epoch is completed.

If the final smaller batch is kept, then the number of iterations per epoch is

$$\lceil \frac{N}{B}\rceil ;$$

if the final smaller batch is dropped, it is

$$\lfloor \frac{N}{B}\rfloor .$$

Shuffling is important here for the same reason as in pure SGD: without it, the order of the examples inside successive mini-batches would inherit the storage order of the dataset, which can introduce persistent ordering effects across epochs. If no shuffling and no other random batching mechanism are used, then the sequence of mini-batches is deterministic, and the randomness discussed below is not generated by the batching process itself.

Expectation And Unbiasedness Of The Mini-Batch Gradient

In the idealized analysis, ${\mathcal{B}}_t$ is treated as a uniformly chosen subset of $\{1,\dots ,N\}$ having cardinality $B$.
For each fixed parameter value $\mathbf{\theta }$, define

$${\mathbf{G}}_t(\mathbf{\theta })={\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })=\frac{1}{B}\sum _{i\in {\mathcal{B}}_t}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

Since ${\mathbf{G}}_t(\mathbf{\theta })$ is a random vector, its expectation is taken componentwise. The expectation is with respect to the random choice of the mini-batch ${\mathcal{B}}_t$.

A convenient way to write the batch average is

$${\mathbf{G}}_t(\mathbf{\theta })=\frac{1}{B}\sum _{i=1}^N{1}_{\{i\in {\mathcal{B}}_t\}}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }),$$

where $1_{\{i\in {\mathcal{B}}_t\}}$ is the indicator of the event that example $i$ belongs to the current mini-batch.

Taking expectations and using linearity,

$$\mathbb{E}\left[{\mathbf{G}}_t(\mathbf{\theta })\right]=\frac{1}{B}\sum _{i=1}^N\mathbb{P}(i\in {\mathcal{B}}_t){\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

Under uniform sampling of a batch of size $B$, each example has inclusion probability

$$\mathbb{P}(i\in {\mathcal{B}}_t)=\frac{B}{N}.$$

Therefore

$$\mathbb{E}\left[{\mathbf{G}}_t(\mathbf{\theta })\right]=\frac{1}{B}\sum _{i=1}^N\frac{B}{N}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta })=\frac{1}{N}\sum _{i=1}^N{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

But, by linearity of differentiation,

$${\nabla }_{\mathbf{\theta }}\mathcal{L}(\mathbf{\theta })=\frac{1}{N}\sum _{i=1}^N{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }).$$

Hence

$$\mathbb{E}\left[{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })\right]={\nabla }_{\mathbf{\theta }}\mathcal{L}(\mathbf{\theta }).$$

This is the precise sense in which the mini-batch gradient is an unbiased estimator of the full gradient. It does not mean that the gradient computed from one specific mini-batch is itself equal to the full gradient; the equality holds only at the level of expectation over the random batch selection.

When mini-batches are instead formed by shuffling once per epoch and then partitioning the permuted dataset, the same intuition remains valid, but the exact conditional expectation at a given within-epoch iteration is the average over the examples that remain unused in that epoch rather than over the whole dataset.

Variance Reduction With Batch Size

Unbiasedness answers only one question: is the mini-batch gradient correct on average? The answer is yes. If one could average the gradients produced by many random mini-batches of the same size, that average would match the full gradient.

A different question is how much the gradient fluctuates from one random mini-batch to another. This is the role of variance. In scalar problems one speaks of variance; for vector-valued gradients, the corresponding object is the covariance matrix. To make this precise, let

$${\mathbf{g}}^{(i)}(\mathbf{\theta })={\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathbf{x}}^{(i)}}(\mathbf{\theta }),\bar{\mathbf{g}}(\mathbf{\theta })=\frac{1}{N}\sum _{i=1}^N{\mathbf{g}}^{(i)}(\mathbf{\theta }),$$

so that ${\mathbf{g}}^{(i)}(\mathbf{\theta })$ is the gradient contributed by example $i$, while $\bar{\mathbf{g}}(\mathbf{\theta })$ is the full-dataset gradient.

Now define the covariance matrix of the per-example gradients:

$${\Sigma }_{\mathrm{pop}}(\mathbf{\theta })=\frac{1}{N-1}\sum _{i=1}^N({\mathbf{g}}^{(i)}(\mathbf{\theta })-\bar{\mathbf{g}}(\mathbf{\theta }))({\mathbf{g}}^{(i)}(\mathbf{\theta })-\bar{\mathbf{g}}(\mathbf{\theta }){)}^{\top }.$$

This matrix measures how much the example-wise gradients are scattered around their mean. Its diagonal entries are the variances of the individual gradient coordinates, while the off-diagonal entries describe how different coordinates fluctuate together.

If the mini-batch is sampled uniformly without replacement, then the covariance of the mini-batch gradient is

$$\mathrm{Cov}\left[{\mathbf{G}}_t(\mathbf{\theta })\right]=\frac{N-B}{B(N-1)}{\Sigma }_{\mathrm{pop}}(\mathbf{\theta }).$$

The important point is the prefactor in front of ${\Sigma }_{\mathrm{pop}}(\mathbf{\theta })$. As $B$ grows, that prefactor shrinks roughly like $1/B$. So the mini-batch gradient stays centered at the same mean, but it becomes progressively less noisy.

In particular, if one looks at a single gradient coordinate and denotes by ${\sigma }_j^2$ its variance across training examples, then

$$\mathrm{Var}\left[G_{t,j}(\mathbf{\theta })\right]=\frac{N-B}{B(N-1)}{\sigma }_j^2.$$

Hence each coordinate variance decreases on the order of $\mathcal{O}(1/B)$, with the additional finite-population correction factor $\frac{N-B}{N-1}$. When the dataset is large and $B\ll N$, this is often summarized more simply as

$$\mathrm{Cov}\left[{\mathbf{G}}_t(\mathbf{\theta })\right]\approx \frac{1}{B}{\Sigma }_{\mathrm{pop}}(\mathbf{\theta }).$$

Equivalently, doubling the batch size roughly halves the variance of the gradient estimator. This is the mathematical reason why larger mini-batches usually produce a cleaner optimization signal: the estimator remains centered at the full gradient, but its fluctuations become smaller. That reduction in stochasticity is one reason larger learning rates can sometimes be used more safely, although the effect is optimizer- and regime-dependent rather than automatic.

Approximate Equality In Practice

For one realized mini-batch, one often writes informally

$${\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{{\mathcal{B}}_t}(\mathbf{\theta })\approx {\nabla }_{\mathbf{\theta }}\mathcal{L}(\mathbf{\theta }),$$

because the mini-batch average is usually more representative of the full dataset than a single-example gradient. However, this remains an approximation for the specific batch currently used. In general, the approximation tends to improve as the batch size $B$ increases, and it becomes exact when the batch coincides with the whole dataset.

Three Regimes, One Common Formula

The same update pattern covers the three standard regimes:

• full-batch gradient descent: $B=N$;
• pure SGD: $B=1$;
• mini-batch SGD: $1<B<N$.

Convention On Scaling

Some formulations omit the factor $\frac{1}{B}$ and sum gradients rather than averaging them.
This is not a different method in principle; it corresponds to a rescaling of the learning rate $\eta$.
What matters is that the convention be stated explicitly and used consistently.