What Is Mini-Batch Gradient Descent?
- 13 hours ago
- 30 min read
Every neural network that has ever learned anything -- from a spam filter to a large language model -- got there by repeating one humble move millions of times: look at a small batch of examples, measure how wrong the guess was, and nudge the numbers a little closer to right. That repeated nudge is mini-batch gradient descent, and understanding it is the fastest way to understand how machine learning models actually train.
TL;DR
Mini-batch gradient descent trains a model by splitting the training data into small groups (mini-batches), computing the loss and gradient on one group, updating the model parameters, and repeating.
It sits between full-dataset batch gradient descent (accurate but slow to update) and one-example stochastic gradient descent (fast but noisy).
Batch size is a tunable choice, not a fixed law -- 16, 32, 64, and 128 are common starting points, not universal answers.
Batch size and learning rate are linked: change one and you often need to revisit the other.
PyTorch's `DataLoader` and NumPy slicing both implement the same core loop: shuffle, split, forward pass, compute loss, backpropagate, update.
Gradient accumulation and distributed training let teams simulate a larger effective batch size than a single accelerator's memory would normally allow.
What Is Mini-Batch Gradient Descent?
Mini-batch gradient descent is an optimization method that trains a model on small groups of training examples, called mini-batches, instead of the whole dataset or a single example at a time. The model's parameters are updated after each mini-batch, which balances gradient accuracy, memory use, and training speed.
Table of Contents
Mini-Batch Gradient Descent in Simple Terms
Mini-batch gradient descent is a way of training a machine learning model by showing it small groups of examples, one group at a time, instead of the entire dataset or a single example at once. After each small group -- called a mini-batch -- the model checks how far its predictions were from the correct answers, calculates a gradient that points toward better parameter values, and takes a small step in that direction. It then moves on to the next mini-batch and repeats the process.
The word "batch" describes *how the data is grouped* for a single update, not a type of model or a specific algorithm. A convolutional network, a transformer, and a simple linear regression model can all be trained with mini-batches. Mini-batching is an answer to one narrow question: how much data should the optimizer look at before it updates the model's weights?
Note: In modern deep-learning frameworks, an optimizer literally named "SGD" (as in PyTorch's `torch.optim.SGD` or scikit-learn's `SGDClassifier`) almost always operates on mini-batches in practice, not on one example at a time. Batch size and optimizer choice (SGD, Adam, RMSprop) are two separate decisions that happen to be made together.
Mini-batch training is the default for most neural network training today because it strikes a workable balance: gradients from a mini-batch are far more stable than gradients from a single example, and mini-batches are small enough to fit in GPU memory and process with fast, vectorized math.
The Foundations: Loss, Gradients, and Parameter Updates
Before mini-batching makes sense, it helps to be clear about what training is actually doing. A model -- whether it is a single line, a decision tree ensemble, or a deep neural network -- has model parameters (also called weights) that determine what it predicts for a given input. Training means searching for parameter values that make the model's predictions match reality as closely as possible.
To measure "as closely as possible," every training example gets a loss value: a number that is small when the prediction is close to the true target and large when it is far off. Averaging the loss across the training set gives the objective function, sometimes written J(theta), which the training process tries to minimize.
A gradient is a vector of partial derivatives that says, for each parameter, which direction increases the loss and by roughly how much. Moving parameters in the opposite direction of the gradient, scaled by a small step size called the learning rate, is what "descent" means in gradient descent.
Training repeats this update many times because the loss surface is complicated -- one step rarely reaches the best parameters. Each pass over the full dataset is called an epoch, and most models need several to many epochs before performance stabilizes. The remaining design question is how the gradient at each step gets calculated: from every training example, from one example, or from a small group. That choice is what separates batch, stochastic, and mini-batch gradient descent.
Batch, Mini-Batch, Epoch, and Iteration
These four terms get mixed up constantly, so it helps to define them side by side using a concrete example: a training set of 1,000 labeled examples.
Term | Definition | In this example (1,000 examples, batch size 100) |
Sample / example | One row of training data: one input and its target. | One of the 1,000 rows. |
Dataset size (N) | Total number of training examples available. | 1,000. |
Batch size (B) | Number of examples used to compute one gradient and one parameter update. | 100. |
Iteration / step | One forward pass, loss calculation, gradient computation, and parameter update on one mini-batch. | 1 iteration processes 100 examples. |
Epoch | One complete pass through the entire training dataset. | 1 epoch = 1,000 / 100 = 10 iterations. |
The general relationship is: steps per epoch = ceiling(dataset size / batch size). "Ceiling" matters because the dataset rarely divides evenly by the batch size. If this same dataset had 1,050 examples with a batch size of 100, there would be 10 full mini-batches of 100 and one final mini-batch of only 50, giving 11 steps per epoch in total -- unless the final incomplete batch is discarded with a setting such as `drop_last=True`, in which case the epoch would have only 10 steps and would quietly skip those last 50 examples every epoch.
"Mini-batch" and "batch" are often used interchangeably in casual conversation, but in the strict sense used by batch gradient descent, "batch" means the entire training set, while "mini-batch" means a small subset of it. This guide uses "mini-batch" for the small-subset case throughout, and calls out full-dataset batches explicitly.
How Mini-Batch Gradient Descent Works Step by Step
A full mini-batch training cycle follows the same ten steps regardless of framework or model architecture:
Initialize the model's parameters, usually with small random values or a principled scheme such as Xavier or He initialization.
Shuffle the training examples at the start of the epoch, unless the data has meaningful order (such as time series) that must be preserved.
Split the shuffled data into mini-batches of the chosen batch size, handling the final partial batch according to policy.
Forward pass: feed one mini-batch through the model to produce predictions.
Compute the mini-batch loss, typically the average loss across the examples in that mini-batch.
Backpropagate: compute the gradient of the mini-batch loss with respect to every parameter, using automatic differentiation.
Update the parameters by moving them a small step in the direction that reduces the loss, scaled by the learning rate.
Move to the next mini-batch and repeat steps 4 through 7 until every mini-batch in the epoch has been used.
Complete the epoch once every example has contributed to exactly one gradient update (or has been dropped, if `drop_last=True`).
Repeat for further epochs until a stopping condition is reached, such as a fixed epoch count, a validation metric that stops improving, or a training-time budget.
Shuffling matters because, without it, the model sees examples in the same fixed order every epoch. If the data happens to be sorted -- by class label, by date, or by any other pattern -- the model can pick up spurious patterns tied to that order rather than the underlying signal, and mini-batches become poor, biased samples of the full dataset. The main exception is sequential data, such as time series or certain online machine learning settings, where shuffling would destroy the temporal structure the model needs to learn; specialized samplers or windowing schemes are used instead.
The Mathematics of Mini-Batch Gradient Descent
Consider a training set of N examples and a model with parameters theta. Each example i has an individual loss, written l_i(theta), that measures how wrong the model's prediction is on that example. The full-dataset objective function averages loss across every example:
J(theta) = (1 / N) * sum_{i=1}^{N} l_i(theta)
Here, N is the total number of training examples, theta represents every trainable parameter in the model (collected into one vector for convenience), J(theta) is the overall objective, and l_i(theta) is the loss contributed by a single example i.
The gradient of this objective with respect to theta -- the direction that increases J the fastest -- is the average of each example's individual gradient:
grad J(theta) = (1 / N) * sum_{i=1}^{N} grad l_i(theta)
Computing this exact gradient requires passing the entire dataset through the model once per update, which is accurate but potentially very slow for large N. Mini-batch gradient descent instead estimates this quantity using a small, randomly drawn subset. For a mini-batch B_t of size B drawn at step t, the mini-batch gradient estimate is:
g_t = (1 / B) * sum_{i in B_t} grad l_i(theta)
The parameter update then moves theta a small step in the opposite direction of this estimate, scaled by a learning rate eta_t:
theta_{t+1} = theta_t - eta_t * g_t
Every symbol here has a specific role: t indexes the optimization step (not the epoch); B is the mini-batch size; B_t is the specific set of example indices sampled at step t; g_t is the gradient estimate computed from that mini-batch; and eta_t is the learning rate at step t, which may be constant or may follow a schedule that changes over training.
Why the Mini-Batch Gradient Is an Estimate
If mini-batches are drawn uniformly at random from the training set, g_t is an unbiased estimator of the true gradient grad J(theta) -- in expectation, averaging over many random mini-batches would reproduce the full-dataset gradient [1]. In any single step, g_t differs from the true gradient by some amount of noise, and that noise shrinks as batch size grows because averaging over more examples reduces the estimate's variance [2].
This is the central trade-off of mini-batch training: a larger B gives a gradient estimate closer, on average, to the true full-dataset gradient, at the cost of more computation per step and fewer updates per epoch. A smaller B gives a noisier estimate but allows many more updates on the same data. Reduced variance does not automatically translate into faster wall-clock training or better test performance -- the outcome depends on model architecture, dataset, optimizer, hardware, learning-rate schedule, and training budget, not on batch size alone [1][2].
Mean vs. Sum Reduction
Most training libraries, including PyTorch's built-in loss functions, average the loss across a mini-batch by default rather than summing it. This matters: if a mini-batch loss is summed instead of averaged, gradient magnitude scales up directly with batch size, changing how large a step the optimizer takes for a given learning rate. Switching between mean and sum reduction without adjusting the learning rate is a common source of training instability -- the same scaling logic that makes gradient accumulation, covered later, divide accumulated gradients by the number of accumulation steps.
A Worked Numerical Example
A small, fully manual example makes the update rule concrete. Suppose the model is the simplest possible regression: predicting y from x using a single parameter theta, with prediction y_hat = theta * x (no intercept, to keep the arithmetic short).
The mini-batch contains two examples:
Example | x | True target (y) |
1 | 2 | 6 |
2 | 4 | 9 |
Start with an initial parameter value of theta = 1, and use the mean squared error loss for each example, l_i = (y_hat_i - y_i)^2.
Step 1: Forward Pass
With theta = 1: y_hat_1 = 1 * 2 = 2, and y_hat_2 = 1 * 4 = 4.
Step 2: Individual Errors and Losses
Error_1 = y_hat_1 - y_1 = 2 - 6 = -4, so l_1 = (-4)^2 = 16. Error_2 = y_hat_2 - y_2 = 4 - 9 = -5, so l_2 = (-5)^2 = 25.
Step 3: Mini-Batch Loss (Averaged)
Mini-batch loss = (16 + 25) / 2 = 20.5.
Step 4: Gradient
For mean squared error with y_hat = theta * x, the derivative of l_i with respect to theta is 2 * x_i * (y_hat_i - y_i). Averaging across the mini-batch:
grad = (1/2) * [2*2*(-4) + 2*4*(-5)]
= (1/2) * [-16 + -40]
= (1/2) * (-56)
= -28
Step 5: Parameter Update
Using a learning rate of eta = 0.01: theta_new = theta_old - eta * grad = 1 - 0.01 * (-28) = 1 + 0.28 = 1.28.
After one mini-batch update, theta moved from 1.0 to 1.28 -- a step in the right direction, since the true relationship in this toy example is closer to y = 2.2x. A different mini-batch from the same dataset would produce a different gradient estimate and a slightly different update -- the gradient noise described earlier. A larger mini-batch would average out more of that noise, typically getting closer to the gradient the entire dataset would produce, at the cost of touching more data per step. Repeating this update across many mini-batches and epochs gradually pulls theta toward the value that minimizes loss across the whole dataset.
Batch vs. Stochastic vs. Mini-Batch Gradient Descent
All three approaches use the same update rule; they differ only in how many examples contribute to each gradient calculation.
Property | Batch gradient descent | Stochastic gradient descent (1 example) | Mini-batch gradient descent |
Examples per update | All N examples | 1 example | B examples (commonly tens to thousands) |
Updates per epoch | 1 | N | ceiling(N / B) |
Gradient noise | None (exact gradient) | Very high | Moderate, decreasing as B grows |
Memory requirement | High -- full dataset per step | Very low | Moderate, tunable via B |
Hardware utilization | Can be efficient per step, but rare updates | Poor -- little to vectorize | Strong -- suited to vectorized, parallel hardware |
Typical use case | Small datasets, convex problems, some classical ML | Strict online/streaming learning, very constrained memory | Neural network training on medium-to-large datasets |
Main advantage | Exact, stable gradient direction | Cheapest possible single step; can adapt online | Balances gradient quality, update frequency, and throughput |
Main limitation | Expensive and slow to update on large datasets | Extremely noisy updates; poor hardware efficiency | Requires tuning batch size and learning rate together |
None of the three is universally superior. Full-batch gradient descent remains a reasonable choice for small datasets or certain convex optimization problems where an exact gradient is cheap to compute. Strict single-example SGD is still relevant for genuine online learning settings where data arrives one record at a time. Mini-batch training became the default for most neural network work because it fits the realities of modern accelerator hardware and large datasets better than either extreme [3].
Why Mini-Batches Work So Well in Practice
Mini-batching earns its default status by balancing several competing concerns at once rather than optimizing any single one of them in isolation.
Gradient quality: averaging over B examples reduces variance compared to single-example gradients, giving a more reliable direction to step in.
Update frequency: unlike full-batch training, mini-batches allow many updates per epoch, so the model starts improving long before it has seen the entire dataset once.
Vectorized computation: modern numerical libraries (NumPy, PyTorch, TensorFlow) are built around matrix and tensor operations that process a batch of examples in parallel far faster than a loop over individual examples.
GPU and accelerator utilization: GPUs contain thousands of cores designed for parallel arithmetic; a mini-batch gives them enough independent work per step to stay busy, while a single example often leaves most of the hardware idle.
Memory usage: a mini-batch's activations and gradients must fit in accelerator memory; mini-batching keeps this bounded and tunable, unlike full-batch training, whose memory footprint grows with the whole dataset.
Data-loading overlap: while the accelerator processes one mini-batch, the next one can be loaded and preprocessed in parallel (this is exactly what PyTorch's `DataLoader` with multiple workers does).
Training stability: mini-batches large enough to average out extreme outliers, but small enough to update often, tend to produce smoother, more stable loss curves than single-example updates.
Hardware throughput does not scale indefinitely with batch size. Below a certain size, an accelerator sits partly idle because there isn't enough parallel work to fill its cores; above a certain size, throughput per example levels off because hardware is already saturated, and further increases mostly consume memory rather than buying speed [4][5].
It helps to separate several measurements that are easy to conflate: time per step, steps per epoch, time per epoch, total training time, and time required to reach a target validation score. A larger batch size can lower time per step and raise throughput while still increasing total time-to-target if it needs more epochs or careful tuning. Optimizing throughput alone can produce a batch-size choice that looks good on a benchmark but performs worse for the actual project goal [5].
How Batch Size Affects Training
Smaller Mini-Batches
Potential advantages: lower memory requirements; more parameter updates per epoch; and gradient noise that some practitioners find helps escape sharp, narrow regions of the loss landscape early in training.
Potential disadvantages: lower hardware utilization; greater variability between updates; a higher risk of unstable training if the learning rate isn't reduced to match; and more per-step overhead accumulated across many more steps.
Larger Mini-Batches
Potential advantages: more stable, lower-variance gradient estimates; stronger vectorization and accelerator throughput up to hardware saturation; fewer updates needed per epoch; and practical advantages in distributed training, where each worker processes a sizable local batch before synchronizing.
Potential disadvantages: higher memory usage; diminishing throughput returns once hardware is saturated; fewer updates per epoch, which can slow early progress; frequent need to retune the learning rate or add a warm-up phase; possible generalization challenges reported in some settings; and added communication overhead in distributed training [6][7].
Warning: Research on large-batch training is genuinely nuanced. Keskar et al. (2016) reported that very large batches, trained with the same hyperparameters as small-batch runs, tended to converge to sharper minima that generalized somewhat worse on held-out data [6]. Goyal et al. (2017) later showed that with a properly scaled learning rate and a warm-up schedule, ResNet-50 could be trained with batch sizes up to 8,192 on ImageNet with no loss of accuracy compared to a small-batch baseline [7]. Neither paper supports a blanket claim that large batches always generalize worse or that small batches always generalize better -- the outcome depends heavily on how the learning rate and schedule are adjusted alongside batch size.
How to Choose a Batch Size
There is no single best batch size that applies to every project. Values such as 16, 32, 64, and 128 are common starting points because they are convenient defaults that fit comfortably in memory for many models, not because they are mathematically optimal. Powers of two can be a convenient choice on some hardware due to memory alignment, but they are not a mathematical requirement.
A practical, repeatable process for choosing a batch size:
Identify the hardware memory limit. Determine the largest batch size that fits in available GPU, TPU, or CPU memory for the current model, including activations, gradients, and optimizer state.
Select a reasonable baseline batch size, typically informed by similar published work on comparable model sizes and datasets, or a mid-range value such as 32 or 64 as a starting point.
Establish a fixed evaluation protocol -- the same validation set, the same number of epochs or training budget, and the same metric -- so different batch sizes can be compared fairly.
Tune the learning rate jointly with batch size rather than in isolation; a batch-size change without a learning-rate adjustment is one of the most common sources of misleading experiments.
Monitor training loss, validation metrics, gradient stability, memory use, throughput, and time to reach the target metric, not just examples processed per second.
Test a small logarithmic range of batch sizes, such as 16, 32, 64, 128, and 256, rather than an exhaustive sweep, to see how performance and stability change.
Keep the training budget and comparison conditions explicit in any write-up or report, since batch-size comparisons are only meaningful when the rest of the setup is held constant.
Choose based on the project's actual objective -- fastest wall-clock time to a target metric, lowest memory footprint, or best final validation score -- since these can point to different batch sizes.
The right starting strategy also depends on the kind of data and model involved:
Situation | Sensible starting strategy |
Small tabular datasets | Small-to-moderate batch sizes (often 16-128); full-batch training is sometimes feasible and worth trying. |
Image models on GPUs | Batch sizes typically in the tens to low hundreds, bounded by image resolution and GPU memory. |
Language models / long sequences | Batch size is often expressed in tokens rather than sequences, since sequence length varies; gradient accumulation is common. |
Highly imbalanced datasets | Consider stratified or weighted sampling alongside batch-size choice so rare classes appear in most mini-batches. |
Very small datasets (hundreds of examples) | Smaller batches or full-batch training; large batches may leave too few updates per epoch to learn effectively. |
Distributed training across many devices | Set a reasonable per-device batch size first, then scale the global batch size with device count, adjusting the learning rate accordingly. |
Online or streaming learning | Small batches or single-example updates that match how data actually arrives in production. |
Memory-constrained edge devices | Small batches, sometimes combined with gradient accumulation to simulate a larger effective batch size. |
Mini-Batch Gradient Descent Pseudocode
The following framework-independent pseudocode captures the general algorithm described above:
initialize parameters theta
for epoch in 1 .. num_epochs:
shuffle(training_data) # skip for strictly sequential data
batches = split_into_minibatches(training_data, batch_size)
for batch in batches:
predictions = forward_pass(batch.inputs, theta)
loss = compute_loss(predictions, batch.targets) # mean over the batch
gradients = backward_pass(loss, theta)
theta = theta - learning_rate * gradients
if stopping_condition_met(theta, validation_data):
break
return theta
NumPy Implementation
This compact NumPy example trains a one-feature linear regression model, y_hat = w*x + b, using mini-batch gradient descent with mean squared error loss.
import numpy as np
rng = np.random.default_rng(seed=42)
# --- Synthetic dataset: y = 3x + 2 plus small noise ---
N = 500
X = rng.uniform(-5, 5, size=(N, 1))
noise = rng.normal(0, 0.5, size=(N, 1))
y = 3 * X + 2 + noise
# --- Initialize parameters ---
w = np.zeros((1, 1))
b = np.zeros((1, 1))
learning_rate = 0.01
batch_size = 32
epochs = 20
n_batches = int(np.ceil(N / batch_size))
for epoch in range(epochs):
# Shuffle indices at the start of every epoch
perm = rng.permutation(N)
X_shuffled, y_shuffled = X[perm], y[perm]
epoch_loss = 0.0
for i in range(n_batches):
start = i * batch_size
end = min(start + batch_size, N) # handles the final partial batch
X_batch = X_shuffled[start:end]
y_batch = y_shuffled[start:end]
m = X_batch.shape[0] # actual batch size (may be < batch_size)
# Forward pass
y_pred = X_batch @ w + b
# Mean squared error loss (mean over the batch)
error = y_pred - y_batch
loss = np.mean(error ** 2)
epoch_loss += loss * m
# Analytical gradients (mean over the batch)
grad_w = (2.0 / m) * (X_batch.T @ error)
grad_b = (2.0 / m) * np.sum(error)
# Parameter update
w -= learning_rate * grad_w
b -= learning_rate * grad_b
if epoch % 5 == 0 or epoch == epochs - 1:
print(f"epoch {epoch:2d} avg_loss={epoch_loss / N:.4f} w={w.item():.3f} b={b.item():.3f}")
A few details are worth checking. `X` has shape (N, 1) and `w` has shape (1, 1), so `X_batch @ w` broadcasts to shape (m, 1), matching `y_batch`. `perm = rng.permutation(N)` re-shuffles indices every epoch. The slice `X_shuffled[start:end]` naturally produces a shorter final batch when N doesn't divide evenly by `batch_size`, and `m = X_batch.shape[0]` uses the *actual* batch size for averaging, so the final partial batch is handled correctly. The gradient formulas follow from differentiating mean squared error with respect to `w` and `b`.
PyTorch Implementation
PyTorch's `DataLoader` handles shuffling and mini-batch creation automatically. This example trains the same style of model using a standard PyTorch training loop.
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
torch.manual_seed(42)
# --- Synthetic dataset: y = 3x + 2 plus small noise ---
N = 500
X = torch.empty(N, 1).uniform_(-5, 5)
y = 3 * X + 2 + torch.randn(N, 1) * 0.5
dataset = TensorDataset(X, y)
loader = DataLoader(
dataset,
batch_size=32,
shuffle=True, # reshuffles the data at the start of every epoch
drop_last=False, # keeps the final, smaller batch by default
)
model = nn.Linear(in_features=1, out_features=1)
loss_fn = nn.MSELoss() # averages loss over the batch by default
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
epochs = 20
for epoch in range(epochs):
running_loss = 0.0
for X_batch, y_batch in loader: # DataLoader yields one mini-batch per iteration
optimizer.zero_grad(set_to_none=True) # clear gradients from the previous step
predictions = model(X_batch) # forward pass
loss = loss_fn(predictions, y_batch) # mean squared error over this mini-batch
loss.backward() # backpropagation: computes gradients
optimizer.step() # parameter update for this mini-batch
running_loss += loss.item() * X_batch.size(0)
if epoch % 5 == 0 or epoch == epochs - 1:
avg_loss = running_loss / N
w, b = model.weight.item(), model.bias.item()
print(f"epoch {epoch:2d} avg_loss={avg_loss:.4f} w={w:.3f} b={b:.3f}")
Each iteration of `for X_batch, y_batch in loader` corresponds to exactly one mini-batch and one optimizer step [8]. Gradients must be cleared with `optimizer.zero_grad()` before every backward pass because PyTorch accumulates gradients into each parameter's `.grad` attribute by default; without clearing them, gradients from previous mini-batches would add to the current one instead of being replaced. `loss.backward()` runs backpropagation, populating `.grad` [8]. `optimizer.step()` then applies the update rule using those gradients, once per mini-batch, not once per epoch. By default, `DataLoader` keeps the final, smaller mini-batch of each epoch; `drop_last=True` discards it instead, which matters most for architectures such as batch normalization-heavy networks that require every batch to have an identical size.
Gradient Accumulation and Effective Batch Size
Gradient accumulation lets a small, memory-friendly "micro-batch" simulate the behavior of a larger batch that would not otherwise fit in accelerator memory. Instead of updating parameters after every micro-batch, the optimizer accumulates (sums, then averages) gradients across several micro-batches and only then takes one parameter-update step.
effective batch size = micro-batch size * accumulation steps
In distributed training, this extends further by multiplying in the number of devices working in parallel:
global effective batch size = per-device batch size * number of devices * accumulation steps
Because most loss functions average across a batch, accumulated gradients typically need to be scaled -- usually divided by the number of accumulation steps -- so the combined update behaves like one larger averaged mini-batch rather than one summed one, which otherwise would produce a much larger effective learning rate than intended.
Gradient accumulation approximates, but does not always perfectly reproduce, a genuinely larger physical batch, because other pipeline components can behave differently when data is processed in smaller pieces:
Batch-dependent layers, such as batch normalization, compute statistics from whatever is in the current micro-batch, not the full accumulated effective batch, unless specifically re-implemented to account for this.
Dropout randomness is applied independently to each micro-batch's forward pass.
Optimizer-step timing changes -- some optimizer states update once per micro-batch step and some only once per full accumulation cycle, depending on implementation.
Gradient clipping, if used, should generally be applied after accumulation across all micro-batches, not after each individual one.
Regularization implementation details, such as weight decay applied per step versus per effective batch, can shift behavior slightly.
Numerical precision differences can accumulate slightly differently when gradients are summed across more, smaller steps versus fewer, larger ones.
Distributed Training and Other Advanced Considerations
Distributed Training
When training is spread across multiple devices, each worker typically processes its own per-device batch, computes local gradients, and synchronizes them -- usually by averaging -- across all devices before the shared parameters update. The global batch size is the sum of all per-device batches, multiplied further by any accumulation steps. Adding devices raises the global batch size even if the per-device batch stays fixed, which is why distributed training often needs its own learning-rate retuning, as demonstrated by Goyal et al.'s linear-scaling and warm-up approach for ResNet-50 across 256 GPUs [7]. Communication overhead also grows with device count and can offset some of the raw compute speedup if not managed carefully.
Batch Normalization at Small Batch Sizes
Batch normalization layers estimate the mean and variance of activations from the current mini-batch. When batch size is very small, these statistics become noisy and unreliable, which can destabilize training. In batch-size-constrained settings, alternative normalization strategies (such as group or layer normalization) or techniques that accumulate statistics across steps are sometimes used instead of standard batch normalization.
Variable-Length Data
Text, audio, and other sequential data often have variable length, which complicates the idea of a fixed batch size. Common solutions include padding shorter sequences to match the longest one in a batch, bucketing similar-length sequences together to minimize wasted padding, and token-based batching, where batch size is measured in total tokens rather than number of sequences -- since a batch of 32 very long sequences uses far more memory than a batch of 32 very short ones.
Imbalanced Data
Randomly drawn mini-batches from a class-imbalanced dataset can systematically underrepresent rare classes, especially with small batch sizes. Stratified sampling, weighted sampling, or explicitly balanced batches can correct this, though each approach changes the effective training distribution the model sees and should be applied deliberately rather than as a default.
Reproducibility
Exactly reproducing a training run depends on more than one random seed. Data shuffling order, multiprocessing worker scheduling, and non-deterministic accelerator operations (some GPU kernels aren't bit-for-bit deterministic by default) can all introduce run-to-run variation even with a fixed seed. Recording software versions, hardware, and configuration is generally more reliable than relying on a seed alone.
Numerical Precision and Memory
Mixed-precision training uses lower-precision number formats (such as 16-bit floating point) for parts of the computation to reduce memory use and increase throughput, which can indirectly allow a larger batch size on the same hardware. Out-of-memory errors during training are most directly addressed by reducing batch size (or using gradient accumulation to compensate), since activation memory during the forward pass scales with batch size; this is a different lever from reducing model size, which affects parameter and optimizer-state memory instead.
Common Mistakes and How to Avoid Them
Mistake | Why it causes problems | Safer practice |
Confusing a batch with an epoch | Leads to miscounted training steps and misread logs. | Keep the definitions from the batch/epoch/iteration table explicit in code and documentation. |
Treating mini-batching as a separate model | Mini-batching is a training procedure, not an architecture; confusing the two muddles debugging. | Remember mini-batching only changes how gradients are computed and how often parameters update. |
Forgetting to shuffle suitable training data | The model can learn spurious order-based patterns. | Shuffle at the start of every epoch unless the data is genuinely sequential. |
Shuffling time-series or sequential data incorrectly | Destroys temporal structure the model needs to learn. | Use sequential or windowed sampling for time-dependent data instead of random shuffling. |
Treating one conventional batch size as universally optimal | Batch size interacts with model, data, and hardware. | Test a small logarithmic range of batch sizes under a fixed evaluation protocol. |
Changing batch size without reconsidering the learning rate | The two are linked; changing one without the other can destabilize or slow training. | Retune the learning rate whenever batch size changes meaningfully. |
Summing loss when a mean-based implementation is assumed elsewhere | Silently changes effective gradient scale and learning-rate behavior. | Be explicit about mean vs. sum reduction and keep it consistent across the codebase. |
Forgetting to clear gradients each step | Frameworks like PyTorch accumulate gradients by default, corrupting later updates. | Call `optimizer.zero_grad()` before every `loss.backward()`. |
Mishandling the final incomplete mini-batch | Can silently skip data or use the wrong batch size in an average. | Decide deliberately whether to keep (`drop_last=False`) or discard (`drop_last=True`) the final batch. |
Comparing experiments trained under different budgets | Makes batch-size or hyperparameter comparisons meaningless. | Hold epoch count, compute budget, and evaluation protocol constant across comparisons. |
Measuring only examples processed per second | Ignores whether more epochs or worse stability offset the raw speed gain. | Track time (or steps) needed to reach a target validation metric, not just throughput. |
Assuming gradient accumulation always equals one larger physical batch | Batch-dependent layers and randomness can behave differently at the micro-batch level. | Check batch normalization, dropout, and gradient-clipping placement when using accumulation. |
Ignoring class imbalance or sequence-length variation | Random mini-batches can underrepresent rare classes or waste memory on padding. | Use stratified or weighted sampling, and bucketing or token-based batching where relevant. |
Using validation or test data to tune training updates | Leaks evaluation information into the training process, inflating reported performance. | Keep validation and test data strictly separate from anything that influences parameter updates. |
When to Use Mini-Batch Gradient Descent
Mini-batch training is a strong default for:
Training neural networks of essentially any architecture.
Medium-to-large datasets that don't comfortably fit, or don't need to fit, entirely in memory at once.
Training on GPUs, TPUs, or other accelerators built for parallel, vectorized computation.
Models trained with automatic differentiation and backpropagation.
Workloads that need to balance memory constraints against training throughput.
Other approaches remain reasonable in specific contexts:
Very small datasets, where computing the exact gradient over the whole dataset is inexpensive and full-batch training is practical.
Strict online-learning settings, where data genuinely arrives one observation at a time and the model must update immediately.
Certain convex optimization problems, which sometimes have specialized solvers that are more efficient than general-purpose mini-batch gradient descent.
Problems where the exact gradient (or the full objective) is already cheap to calculate, removing much of the practical motivation for approximating it with a subset.
None of these alternatives are obsolete; they simply fit narrower situations than mini-batch training was designed to handle.
Practical Batch-Size Checklist
Confirmed the largest batch size that fits in available accelerator memory for this model and input size.
Chose a reasonable baseline batch size (informed by comparable published work or a mid-range default).
Fixed the evaluation protocol: validation set, training budget, and target metric.
Retuned the learning rate whenever batch size changed, rather than leaving it fixed by default.
Tested a small logarithmic range of batch sizes rather than relying on a single guess.
Tracked time-to-target-metric alongside raw throughput, not throughput alone.
Checked whether the dataset needs stratified, weighted, or balanced sampling given its class distribution.
Verified how the final incomplete batch of each epoch is handled (`drop_last` behavior).
Recorded batch size, learning rate, hardware, and software versions for reproducibility.
Reconsidered mixed precision or gradient accumulation if memory, not compute, is the binding constraint.
Frequently Asked Questions
What is mini-batch gradient descent?
Mini-batch gradient descent is an optimization method that trains a model by splitting training data into small groups called mini-batches. It computes the loss and gradient on one mini-batch, updates the model's parameters, then moves to the next mini-batch, repeating across the dataset and multiple epochs. It balances the exact-but-expensive gradient of full-batch training against the fast-but-noisy gradient of single-example training.
How is mini-batch gradient descent different from batch gradient descent?
Batch gradient descent computes the gradient using the entire training set for every parameter update, producing an exact gradient but only one update per epoch. Mini-batch gradient descent uses a small subset of the data per update, producing an estimate of the gradient but allowing many updates per epoch. Batch gradient descent needs more memory and is slower per epoch on large datasets, while mini-batch training scales far better to large datasets and modern accelerator hardware.
How is mini-batch gradient descent different from stochastic gradient descent?
Classical stochastic gradient descent (SGD) uses one training example per parameter update, producing very noisy gradients but extremely cheap individual steps. Mini-batch gradient descent uses a small group of examples (commonly tens to thousands) per update, reducing gradient noise compared to single-example SGD while still allowing frequent updates. In practice, the optimizer libraries called "SGD" in frameworks like PyTorch typically operate on mini-batches rather than single examples.
Why are mini-batches used in deep learning?
Mini-batches balance gradient accuracy, update frequency, memory use, and hardware efficiency. Modern GPUs perform best with vectorized, parallel computation, and mini-batches let hardware process many examples simultaneously rather than one at a time. This typically produces faster training in wall-clock time than either full-batch training or single-example SGD for most deep learning workloads.
What is the best mini-batch size?
There is no single best mini-batch size. It depends on the model architecture, dataset size, hardware memory, and training goals. Common starting points such as 32, 64, or 128 are convenient defaults, not universal answers. The recommended approach is to test a small logarithmic range of batch sizes under a fixed evaluation protocol and choose based on validation performance, training stability, and available memory.
How does batch size affect the learning rate?
Batch size and learning rate are closely linked. Larger batches produce lower-variance gradient estimates, which often allows a larger learning rate without destabilizing training; smaller batches produce noisier gradients that may require a smaller learning rate for stability. A commonly cited heuristic is the linear scaling rule, though it requires careful validation and often a warm-up period rather than blind application.
What is the difference between a batch, epoch, and iteration?
A batch (or mini-batch) is the subset of examples used for one gradient computation and parameter update. An iteration is one such update. An epoch is one complete pass through the entire training dataset, made up of multiple iterations. For a dataset of 1,000 examples and a batch size of 100, one epoch contains 10 iterations.
Should training data be shuffled before creating mini-batches?
Yes, in most cases. Shuffling at the start of every epoch prevents the model from learning patterns tied to the original data ordering and helps each mini-batch be a more representative sample of the full dataset. The main exception is sequential data such as time series, where shuffling can destroy meaningful temporal structure; sequential or windowed sampling is used instead.
What happens to the final incomplete batch?
If the dataset size is not evenly divisible by the batch size, the last mini-batch of each epoch will contain fewer examples than the rest. Frameworks like PyTorch handle this automatically by default, using the smaller final batch. Setting drop_last=True discards it instead, which is sometimes necessary when a model architecture requires every batch to have an identical, fixed size.
Is Adam a form of mini-batch gradient descent?
Adam is an optimizer that uses gradients to update parameters, typically incorporating adaptive per-parameter learning rates and momentum-like terms. It is commonly used with mini-batches, but mini-batching and optimizer choice are separate decisions. Adam mini-batch training, SGD mini-batch training, and RMSprop mini-batch training are all valid combinations of how data is grouped and how the gradient is used to update parameters.
What is effective batch size?
Effective batch size is the total number of examples whose gradients contribute to a single parameter update, even if they were not all processed simultaneously. It commonly arises with gradient accumulation, where effective batch size equals micro-batch size multiplied by the number of accumulation steps, and in distributed training, where it also multiplies by the number of devices involved.
How does gradient accumulation work?
Gradient accumulation computes gradients on several smaller micro-batches without updating parameters after each one, summing and typically averaging the gradients across those micro-batches before a single parameter update. This simulates a larger effective batch size than would otherwise fit in memory, though it does not perfectly replicate every property of one physically larger batch due to factors like batch normalization statistics and dropout randomness.
Do larger batches always train faster?
Not necessarily. Larger batches typically use hardware more efficiently per step and need fewer updates per epoch, but they can require more careful learning-rate tuning and sometimes need more total epochs to reach the same validation performance as a well-tuned smaller batch size. The right comparison metric is usually time to reach a target validation score, not raw examples processed per second.
Can batch size affect generalization?
Some research reports that very large batch sizes, trained with the same hyperparameters as smaller-batch runs, can generalize slightly worse on held-out data -- the generalization gap. Other work shows this gap can be reduced or eliminated with a properly scaled learning rate and a warm-up schedule. The effect is context-dependent, not a fixed rule for every model and dataset.
How can I tell whether my batch size is too large or too small?
Signs a batch size may be too small include highly unstable or erratic training loss and poor hardware utilization. Signs it may be too large include out-of-memory errors, very few updates per epoch leading to slow early progress, or a noticeable gap between training and validation set performance compared to smaller-batch runs. Testing a range of batch sizes under a consistent evaluation protocol is the most reliable way to diagnose this for a specific project.
Key Takeaways
Mini-batch gradient descent trains a model on small groups of examples at a time, updating parameters after each group rather than after the whole dataset or a single example.
It sits between full-dataset batch gradient descent and single-example stochastic gradient descent, trading gradient accuracy for update frequency and hardware efficiency.
Batch, epoch, and iteration are distinct: a batch is a subset of data, an iteration is one update, and an epoch is one full pass through the dataset.
There is no universal best batch size -- common values like 32, 64, and 128 are starting points, not laws.
Batch size and learning rate should be tuned together; changing one without reconsidering the other is a frequent source of unstable training.
PyTorch's DataLoader and a manual NumPy loop implement the same core cycle: shuffle, split into mini-batches, forward pass, compute loss, backpropagate, update.
Gradient accumulation and distributed training both extend the idea of effective batch size beyond what a single device's memory allows.
Research on large-batch generalization is nuanced and depends heavily on learning-rate scaling and warm-up, not on batch size alone.
Actionable Next Steps
Implement the NumPy or PyTorch training loop from this guide on a small dataset you already have, and confirm the loss decreases over several epochs.
Test at least three batch sizes (for example 16, 64, and 256) on the same dataset and model, keeping the number of epochs and evaluation protocol fixed.
Plot training loss and validation metric per epoch for each batch size to see differences in stability and convergence speed.
Adjust the learning rate for each batch size you test rather than reusing one fixed value across all of them.
Record batch size, learning rate, epoch count, hardware, and framework version for every run so results are reproducible later.
Watch for out-of-memory errors or very slow per-step progress, and use them to diagnose whether batch size is too large or too small.
If memory is the binding constraint, try gradient accumulation or mixed-precision training before shrinking the model itself.
Once comfortable with the basics, read the Goodfellow, Bengio, and Courville optimization chapter and the Goyal et al. large-batch training paper linked in the references for deeper theoretical and empirical grounding.
Glossary
Backpropagation -- The algorithm that computes gradients of the loss with respect to every network parameter by propagating error signals backward through the layers.
Batch -- Strictly, the entire training dataset used for one update in batch gradient descent; often used loosely to mean a mini-batch.
Batch gradient descent -- An optimization approach that computes the gradient using the entire training dataset for every update.
Batch normalization -- A technique that normalizes layer activations using mini-batch statistics, which can become unstable at very small batch sizes.
Batch size -- The number of training examples used to compute one gradient estimate and one parameter update.
Convergence -- The point where further updates produce little further improvement in loss or a chosen evaluation metric.
Dataset -- The full collection of training examples available for a task.
Effective batch size -- The total examples contributing to a single update, accounting for gradient accumulation and, in distributed settings, device count.
Epoch -- One complete pass through the entire training dataset.
Forward pass -- The step in which input data passes through a model to produce predictions.
Gradient -- A vector of partial derivatives indicating the direction and rate of steepest increase, used to determine parameter updates.
Gradient accumulation -- Summing gradients across several smaller micro-batches before one parameter update, simulating a larger effective batch size.
Gradient descent -- An iterative optimization method that updates parameters in the direction reducing a loss function, scaled by a learning rate.
Iteration -- One cycle of forward pass, loss computation, gradient computation, and parameter update on one mini-batch; also called a step.
Learning rate -- A scalar controlling how large a step is taken during each parameter update.
Loss function -- A function measuring how far a model's prediction is from the true target.
Mini-batch -- A small, randomly sampled subset of the training data used for one gradient computation and update.
Mini-batch gradient descent -- An optimization procedure that computes gradients and updates parameters using small mini-batches rather than the whole dataset or a single example.
Model parameter -- A trainable value inside a model, such as a weight or bias, adjusted during training.
Optimizer -- An algorithm, such as SGD or Adam, that uses computed gradients to decide how to update parameters.
Parameter update -- The step where model parameters are adjusted based on a gradient and a learning rate.
Sample -- A single training example: an input and, in supervised learning, its target.
Stochastic gradient descent -- Classically, an approach that computes the gradient from one example per update; in modern usage, often refers to optimizers running on mini-batches.
Training loop -- The repeating sequence -- shuffle, batch, forward pass, loss, backpropagation, update -- that trains a model over epochs.
Validation set -- Held-out data, separate from the training set, used to monitor performance without influencing parameter updates.
Variance -- How much a quantity, such as a gradient estimate, fluctuates around its average across different samples.
Sources & References
[1] Bottou, L., Curtis, F. E., and Nocedal, J. "Optimization Methods for Large-Scale Machine Learning." SIAM Review, vol. 60, no. 2, 2018, pp. 223-311. https://arxiv.org/abs/1606.04838
[2] Goodfellow, I., Bengio, Y., and Courville, A. "Deep Learning." Chapter 8: Optimization for Training Deep Models. MIT Press, 2016. https://www.deeplearningbook.org/contents/optimization.html
[3] scikit-learn developers. "Stochastic Gradient Descent." scikit-learn documentation, n.d. https://scikit-learn.org/stable/modules/sgd.html
[4] PyTorch Contributors. "torch.utils.data." PyTorch documentation, n.d. https://docs.pytorch.org/docs/stable/data.html
[5] Masters, D., and Luschi, C. "Revisiting Small Batch Training for Deep Neural Networks." arXiv preprint arXiv:1804.07612, 2018. https://arxiv.org/abs/1804.07612
[6] Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. "On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima." arXiv preprint arXiv:1609.04836, 2016. https://arxiv.org/abs/1609.04836
[7] Goyal, P., Dollar, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour." arXiv preprint arXiv:1706.02677, 2017. https://arxiv.org/abs/1706.02677
[8] PyTorch Contributors. "torch.Tensor.backward." PyTorch documentation, n.d. https://docs.pytorch.org/docs/stable/generated/torch.Tensor.backward.html
[9] Smith, S. L., Kindermans, P.-J., Ying, C., and Le, Q. V. "Don't Decay the Learning Rate, Increase the Batch Size." arXiv preprint arXiv:1711.00489, 2017. https://arxiv.org/abs/1711.00489
[10] Google. "Article (Schema.org Article, NewsArticle, BlogPosting)." Google Search Central documentation, n.d. https://developers.google.com/search/docs/appearance/structured-data/article
