Mini-batches and data loading

Introduction

Estimate useful training directions from small groups of examples while controlling memory and gradient noise.

Learning goal: Explain how batch size, sampling, and update count change optimization.
Before you start: Tensor shapes, gradients, mean loss, and optimizer updates.

Lesson plan

Compare batch and full gradients
Track batch tensor shapes
Audit shuffling and final batches

The problem

The full training gradient averages contributions from every example. It is stable, but one update may be slow or too large for memory. Updating from one example uses little memory, but its direction can be noisy.

A mini-batch is a middle choice. For input features shaped (N, D), a batch has shape (B, D), where N is dataset size, D is feature count, and B is batch size.

Work through a small example

Let one scalar parameter be w=0. For examples x=[1,2,3,4] and targets y=[2,4,6,8], use mean squared error. The two-example batch (1,2) gives gradient -10; batch (3,4) gives -50. Their average is the full-data gradient -30.

Either mini-batch is a noisy estimate, but it is unbiased under uniform sampling: across random batches, the expected estimate equals the full gradient.

The general rule

For per-example losses l_i(w), mini-batch loss is L_B=(1/B) sum l_i. Its gradient is the same average of per-example gradients. Larger B usually lowers variance but costs more memory and produces fewer updates per epoch.

Noise is not purely harmful. It can help optimization move through flat or narrow regions, but excessive noise makes progress unstable. Learning rate and batch size interact: changing one may require retuning the other.

Implement and inspect

import torch

x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = 2 * x
w = torch.tensor([[0.0]], requires_grad=True)

for ids in (torch.tensor([0, 1]), torch.tensor([2, 3])):
    loss = ((x[ids] @ w - y[ids]) ** 2).mean()
    (gradient,) = torch.autograd.grad(loss, w)
    print(loss.detach().item(), gradient.detach().item())

The expected lines are 10.0 -10.0 and 50.0 -50.0: each contains that batch's mean loss and scalar gradient.

x[ids] has shape (2,1), w has shape (1,1), and the loss is scalar. autograd.grad returns a gradient with the same shape as w.

Engineering checks

Shuffle training examples each epoch, but never shuffle validation labels away from inputs.
Decide whether to keep the final smaller batch. Dropping it changes which examples contribute.
Clear accumulated gradients before every ordinary update.
Report both epochs and optimizer steps; they are not interchangeable.

Go deeper

Examples inside a batch may be correlated. Adjacent text windows or frames from one video can make the gradient less informative than B independent examples. Sampling strategy therefore matters as much as batch size.

Batch-dependent layers add another constraint. Batch normalization estimates moments from the current batch, so very small batches produce unstable statistics. Gradient accumulation can imitate a larger update batch for memory, but it does not reproduce batch-normalization statistics from one large forward pass.

Practice and recap

Question: A dataset has 130 examples and batch size 32. How many updates occur if the final batch is kept?

Worked answer

Four full batches cover 128 examples. One final batch contains two, so the epoch has five updates. With mean loss, that last update gives two examples the same update-level influence as a full batch; consider this when batches are extremely uneven.

Mini-batches trade exact gradients for practical memory and frequent updates. Inspect shapes, sampling, final-batch behavior, and gradient scale together.