Gradient Descent and Mini-Batches

Three flavors of gradient descent

Batch (full) gradient descent — compute the gradient on the entire dataset, take one step, repeat. Smooth, but expensive (one gradient per epoch) and impossible for datasets that do not fit in memory.

Stochastic gradient descent (SGD) — compute the gradient on one example at a time, take a step, repeat. Noisy, fast, hard to use the GPU efficiently.

Mini-batch SGD — compute the gradient on a small batch (16, 32, 256, ...), take a step. Noisy enough to escape some local minima, smooth enough to be stable, batched enough to use the GPU. This is what everyone actually does.

Why noise helps generalization

The randomness in mini-batch sampling acts as an implicit regularizer. The optimizer follows a noisy estimate of the true gradient, which keeps it from settling too aggressively into sharp local minima — and sharp minima generalize worse than flat ones. This is why mini-batch SGD often generalizes better than the smoother full-batch update, even when it converges to a higher training loss.

Tip: Larger batches give a less noisy gradient and let you use a slightly larger learning rate (rule of thumb: linear scaling). But after a point, bigger batches stop helping wall-clock time and may hurt generalization. Tune both.

Modern optimizer-aware mini-batches

Modern optimizers (Adam, AdamW) maintain per-parameter running estimates of the gradient mean and variance, which makes the choice of batch size less critical for stable training. You still tune it for hardware (memory, throughput) and for generalization (smaller batches often regularize better), but you no longer have to baby-sit the optimizer through wild gradient noise.

Principle: Pick the largest batch that gives stable training and fits in VRAM. Then revisit it after the model is working — sometimes smaller batches train better at the same wall-clock cost.

Code

Batched SGD step in PyTorch·python

import torch
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset

X = torch.randn(10000, 20)
y = torch.randint(0, 5, (10000,))
loader = DataLoader(TensorDataset(X, y), batch_size=128, shuffle=True)

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
opt   = optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(3):
    running = 0.0
    for xb, yb in loader:
        opt.zero_grad()
        logits = model(xb)
        loss = loss_fn(logits, yb)
        loss.backward()
        opt.step()
        running += loss.item() * xb.size(0)
    print(f"epoch {epoch} avg loss: {running / len(loader.dataset):.4f}")

Three flavors of gradient descent

Why noise helps generalization

Modern optimizer-aware mini-batches

Code

External links

Exercise

Progress

Comments 0