Gradient descent batch processing

Three regimes

• Batch size 1 (stochastic). One sample per update. Very noisy gradient, very many updates. And the noise it sometimes helps escape bad regions.
• Full batch. The whole dataset per update. The gradient is exact. But you get one update per epoch and it is slow.
• Mini-batch (32–256). What in effect everyone uses. Enough samples for a usable gradient estimate, enough updates to make progress. And it fits on a GPU.

The batch size controls a three-way trade-off between gradient accuracy, update frequency and hardware utilisation. A batch of 1 gives the noisiest gradient and the most updates per epoch. A full-batch gradient is exact but gives just one update per epoch. Mini-batches of 32 to 256 sit at the sweet spot: the gradient estimate is close enough to the true gradient for stable convergence, updates are frequent enough for fast progress. And the batch is large enough to saturate the GPU's parallel cores.

Gradient noise is not purely harmful. The stochastic noise from small batches helps the optimiser escape sharp local minima and find flatter minima that generalise better. This is why raising batch size often needs raising the learning rate proportionally (linear scaling rule) to keep the same training dynamics. And why very large batch sizes sometimes give nets that generalise worse still though the training loss is less.

The real upper bound on batch size is GPU memory. A batch of pics must fit in VRAM along with the net params, activations and gradients. When the batch does not fit, the usual fixes are to reduce the batch size, use gradient accumulation (run a few less forward passes and sum the gradients before updating), use mixed-precision training (float16 halves the memory per activation). Or distribute across multiple GPUs. Gradient accumulation is the simplest — it simulates a large batch without needing more memory.

How the counts work out

With 1,024 training samples:

• Batch size 32 → 32 updates per epoch
• Batch size 256 → 4 updates per epoch

Gradient noise falls roughly with the square root of batch size. Going from 32 to 128 — a 4× increase — halves the noise, not quarters it. That diminishing give is why batch sizes stop growing well before memory runs out.

Why mini-batch won

Batch gradient descent computes an exact gradient and needs a full pass over the data for a single update. On a million samples that is a million forward passes to move the weights once — exact and unusable.

Stochastic gradient descent updates after each sample. So it makes a million updates per epoch. But each gradient is a single-sample estimate and the path jitters badly. It also wastes the hardware: one sample cannot occupy a GPU.

Mini-batches of 32 to 256 sit at the point where both problems disappear. The gradient estimate is close enough to the true one for stable progress, updates are frequent. And the batch is large enough to fill the device. This is not a compromise nobody is happy with — it is better than either extreme on both axes that matter.

Batch size and learning rate move together

These two cannot be tuned independently. The gradient of a batch is an average. So a more batch gives a less-variance estimate — and a more reliable estimate can support a more step.

The linear scaling rule is the standard heuristic: multiply the batch size by k and multiply the learning rate by k. Going from batch 32 at learning rate 0.01 to batch 256 suggests roughly 0.08. Very large batches often need a warmup period as well, ramping the rate up over the first few epochs, since the scaled rate is unstable while the weights are still random.

The corollary is that checking two runs at other batch sizes with the same learning rate checks nothing useful.

Shuffling is not optional

Mini-batches must be drawn from a shuffled dataset, reshuffled each epoch. If the data arrives sorted by class — all the cats then all the dogs — then each batch holds one class, each gradient points toward predicting that class. And the net oscillates instead of converging.

Reshuffling each epoch also means the net never sees the same batch composition twice. It adds a small amount of useful noise. The exception is time series. There shuffling destroys the temporal ordering the net is supposed to learn from. There the batches must be contiguous windows.

Three sizes of step

The gradient is an average over training cases. How many you average over before updating sets three variants of the same algorithm.

Full-batch descent takes the true gradient and one step per pass through the data. On a million rows that is a million forward and backward passes for a single weight update, which is unusable.

Stochastic descent updates after each case. Steps are cheap and the gradient is a very noisy estimate. Worse, a single 784-item vector uses a fraction of a percent of a GPU's capacity.

Mini-batch is the compromise everyone uses. And it wins on all three counts at once: a usable gradient estimate, a matrix multiplication large enough to saturate the hardware. And enough residual noise to escape saddle points.

Why the noise is useful

It is tempting to treat gradient noise as a defect to be minimised. It is not.

It escapes saddle points. In high dimensions almost each point with zero gradient is a saddle rather than a min. An exact gradient at a saddle is zero and full-batch descent stalls. A noisy one is not. And the net moves on.

It biases towards flat minima. Noise makes it hard to settle into a narrow crevice in the loss surface. And flat minima generalise better since a small shift in the data costs less. This is the main reason very large batches sometimes generalise slightly worse.

It acts as a regulariser. Part of why deep networks generalise better than their param counts suggest is this implicit regularisation from stochastic updates.

So the goal is not the cleanest possible gradient. It is a gradient good enough to point downhill, computed on a batch large enough to use the hardware.

Shuffling, and why it is not optional

Mini-batch descent assumes each batch is a representative sample. If the data arrives sorted — by class, by date, by source — then it is not.

Unshuffled sorted data gives batches holding a single class. The net spends one batch learning "all is class A", the next unlearning it in favour of class B. And the gradients pull in contradictory ways. Batch normalisation makes it worse, since its statistics are computed per batch and turn wildly unrepresentative.

Reshuffling each epoch matters too: it means the net never sees the same partition of cases twice. It adds useful variation.

loader = DataLoader(dataset, batch_size=64, shuffle=True, drop_last=True)

drop_last=True discards a final partial batch. With batch normalisation that matters — a last batch of one or two cases gives meaningless statistics and a visibly worse update.

The exception is time series, where shuffling destroys the temporal ordering the net needs. There, batches are contiguous windows and the split is chronological.

Try this above

1. Drag Batch Size to its min. The descent path visibly jitters — each step is reacting to one or two samples.
2. Raise it toward the max. The path straightens out. But notice it also takes fewer, more steps to cover the same ground.
3. Now hold batch size high and less Learning Rate. Progress stalls — a big smooth gradient with a tiny step size wastes the smoothness.

What often goes wrong

Raising batch size without touching the learning rate. You get fewer updates per epoch, each no more than before. So training slows down and people conclude the more batch "trains worse". The usual remedy is the linear scaling rule: double the batch, double the learning rate, within reason.Checking runs by epoch count. An epoch at batch size 512 holds a fraction of the updates an epoch at batch size 32 does. Check by count of updates, or by wall-clock time, not by epochs.

In one line

Batch size trades gradient noise against how many times per epoch you get to move.

Choosing a batch size in practice

Start with the largest that fits comfortably in memory, up to about 256, then scale the learning rate to match.

The linear scaling rule is the standard guidance: double the batch size, double the learning rate. The reasoning is that a batch twice as large gives a gradient with roughly half the noise. So a step twice as long is equally safe. It holds well up to a few thousand cases per batch, with a warm-up period to survive the first large steps.

Beyond that, the noise that was giving implicit regularisation has largely gone. And specialised optimisers (LARS, LAMB) are needed to keep large-batch training strong. This is why the papers that trained ImageNet in minutes needed more than just more GPUs.

Constraints that override the "as large as fits" rule:

• Batch normalisation needs at least about 8 cases per batch for usable statistics. Below that, use GroupNorm.
• Very large inputs — high-resolution pics, long sequences — may force a batch of 1 or 2. Gradient accumulation recovers the effective batch size.
• Small datasets give few updates per epoch at a large batch size. So a less batch and more updates often trains better.

Gradient accumulation

When the batch you want does not fit, accumulate gradients over a few small batches and update once:

accum = 8                                    # effective batch = 8 x batch_size
opt.zero_grad()
for i, batch in enumerate(loader):
    loss = criterion(model(batch.x), batch.y) / accum   # keep the average right
    loss.backward()                                      # accumulates into .grad
    if (i + 1) % accum == 0:
        opt.step()
        opt.zero_grad()

Dividing the loss by accum is key — without it the accumulated gradient is eight times too large and the effective learning rate is eight times what you set.

This is the same accumulation behaviour that makes a forgotten zero_grad() a bug. Here it is the feature.

The cost is wall-clock time: eight small batches take about as long as one large one would have. You get the large-batch statistics, not the large-batch speed.

Full batch, one row, and all between

The same problem solved by batch, stochastic and mini-batch gradient descent, with the noise each one injects measured right -- since that noise is the whole gap.

import numpy as np

rng = np.random.default_rng(0)
N, D = 2048, 8
X = rng.normal(size=(N, D))
w_true = rng.normal(size=D)
y = X @ w_true + rng.normal(0, 0.5, N)

def full_grad(w):
    return 2 * X.T @ (X @ w - y) / N

def loss(w):
    return ((X @ w - y) ** 2).mean()

print("%d rows. one epoch means seeing all of them once." % N)
print()
print("%-24s %14s %18s %16s"
      % ("method", "batch size", "updates per epoch", "rows per update"))
for name, bs in (("batch (full)", N), ("mini-batch", 64),
                 ("stochastic (SGD)", 1)):
    print("%-24s %14d %18d %16d" % (name, bs, N // bs, bs))
print()

Building on the same setup, the next block shows the left cases.

import numpy as np

print("how noisy is each one's gradient? cosine to the true full-batch")
print("gradient, and how far off in magnitude:")
w0 = np.zeros(D)
true_g = full_grad(w0)
print("%14s %20s %20s" % ("batch size", "cosine to true", "relative error"))
for bs in (1, 8, 64, 512, 2048):
    cos, err = [], []
    for _ in range(200):
        idx = rng.integers(0, N, bs)
        g = 2 * X[idx].T @ (X[idx] @ w0 - y[idx]) / bs
        cos.append(g @ true_g / (np.linalg.norm(g) * np.linalg.norm(true_g)))
        err.append(np.linalg.norm(g - true_g) / np.linalg.norm(true_g))
    print("%14d %20.4f %20.4f" % (bs, np.mean(cos), np.mean(err)))
print("  the error falls as 1/sqrt(batch size). quadrupling the batch halves")
print("  it, so the returns diminish immediately and never stop diminishing.")
print()

def run(bs, epochs, lr, seed=0):
    r = np.random.default_rng(seed)
    w = np.zeros(D)
    trace = []
    for e in range(epochs):
        order = r.permutation(N)          # reshuffle every epoch
        for i in range(0, N, bs):
            idx = order[i:i + bs]
            w -= lr * 2 * X[idx].T @ (X[idx] @ w - y[idx]) / len(idx)
        trace.append(loss(w))
    return w, trace

The next block covers the left steps.

import numpy as np

print("now run all three for 30 epochs, each at a learning rate that suits it:")
print("%-22s %10s %14s %16s %16s"
      % ("method", "updates", "final loss", "loss at epoch 5", "wobble"))
for name, bs, lr in (("batch (full)", N, 0.15), ("mini-batch 64", 64, 0.05),
                     ("stochastic", 1, 0.01)):
    w, trace = run(bs, 30, lr)
    wobble = np.std(trace[-10:])
    print("%-22s %10d %14.6f %16.6f %16.2e"
          % (name, 30 * (N // bs), trace[-1], trace[4], wobble))
print()
print("full batch does %d updates in 30 epochs; SGD does %d. that is the"
      % (30, 30 * N))
print("trade: full batch takes a perfect step rarely, SGD takes a noisy step")
print("constantly, and mini-batch is the only one that uses the hardware.")
print()

print("look at the wobble column. the noisy methods never fully settle:")
w, trace = run(1, 30, 0.01)
print("   SGD loss over the last 6 epochs:", np.round(trace[-6:], 6))
w, trace = run(N, 30, 0.15)
print("   full batch, same epochs:        ", np.round(trace[-6:], 6))
print("   full batch converges to a point. SGD orbits one, at a radius set")
print("   by its learning rate -- which is why a decay schedule matters far")
print("   more for SGD than for full-batch descent.")
print()

print("and the noise is not purely a cost. on a bumpy surface it is the only")
print("thing that can move you off a bad minimum -- provided it is ANNEALED,")
print("so the run can still settle at the end:")

def bumpy(x):
    return 0.05 * (x - 3.0) ** 2 + np.sin(3 * x)

def bumpy_grad(x):
    return 0.1 * (x - 3.0) + 3 * np.cos(3 * x)

The next block continues with the left demonstrations.

grid = np.linspace(-8, 12, 40001)
best = bumpy(grid).min()
print("   the global minimum of this surface is %.4f" % best)
print()
print("%18s %18s %20s" % ("noise sd", "average final loss", "runs within 0.05"))
hits = {}
for noise in (0.0, 5.0, 20.0, 60.0):
    finals = []
    for seed in range(60):
        r = np.random.default_rng(seed)
        x = r.uniform(-6, 10)              # a different starting basin each run
        for t in range(600):
            anneal = noise * (1 - t / 600)  # noise decays to zero
            x -= 0.05 * (bumpy_grad(x) + r.normal(0, anneal))
        finals.append(bumpy(x))
    finals = np.array(finals)
    hits[noise] = (finals.mean(), (finals < best + 0.05).mean())
    print("%18.1f %18.4f %19.0f%%"
          % (noise, finals.mean(), 100 * hits[noise][1]))
print()
bestn = max(hits, key=lambda k: hits[k][1])
print("   with no noise, each run slides into whichever dip it started next")
print("   to and stops there -- %.0f%% of them find the global minimum."
      % (100 * hits[0.0][1]))
print("   at sd %.0f that rises to %.0f%%: the kicks are large enough to hop"
      % (bestn, 100 * hits[bestn][1]))
print("   a barrier, and the annealing lets the run settle once it has.")
print("   at sd 60 it collapses to %.0f%% and the average loss is %.1f --"
      % (100 * hits[60.0][1], hits[60.0][0]))
print("   still being thrown around when the budget runs out.")
print()
print("   the gain is real but modest, and that is the honest shape of it.")
print("   mini-batch noise is not a magic escape hatch; it is a small,")
print("   free nudge that costs nothing because you were going to have it")
print("   anyway.")
print()
print("the practical answer is almost always mini-batch: 32 to 512, the")
print("largest that fits, shuffled every epoch, with a decaying learning")
print("rate to remove the noise once it has done its job.")

Questions people ask

How many updates does one epoch give? dataset size / batch size. With 50,000 cases and a batch of 64, that is 782 — not 1.

Why did accuracy drop when I raised the batch size? Fewer updates per epoch and an unscaled learning rate. Scale the rate up and add warm-up.

Is a power of two needed? No, though it aligns well with hardware and is marginally faster.

Should I use full-batch descent on a small dataset? Rarely — you lose the noise that helps generalisation. And get very few updates. Still on 1,000 rows, mini-batches often train better.

Does batch size affect the final accuracy? Somewhat. Very large batches tend to generalise slightly worse. Very small ones are slow and noisy.

What if my dataset does not divide evenly? The last batch is less. Drop it if batch normalisation is in use.

Recap in one screen

• Full batch is exact and unusable. Single-case updates are noisy and waste hardware. Mini-batches win on both.
• The residual noise is useful — it escapes saddles and favours flat minima.
• Shuffle each epoch, and drop the final partial batch when with batch normalisation.
• Scale the learning rate linearly with batch size, with warm-up.
• Gradient accumulation buys a large effective batch on small hardware, at the cost of time.