Batch Size Effects

You've trained models that either crawl or blow up your GPU memory, and you've heard that batch size is the culprit — but no one explained why changing it from 16 to 64 flips your training curve from jagged to smooth, or why a tiny batch makes your loss bounce like a pinball. This lesson cracks that black box. By the end, you'll experiment with batch size effects like a scientist: you'll know what to measure, how to keep results comparable, and which batch size to pick for your dataset, model, and hardware.

The problem this lesson solves

Batch size is the single most impactful training hyperparameter after learning rate, yet most tutorials just say "use 32" without showing why. When you change batch size, you're not just changing a number — you're changing the statistics of your gradients, the speed of convergence, the stability of training, and the memory footprint of your model. Get it wrong and you get:

• Loss explosion because gradients are too noisy for a large learning rate.
• Slow, stuck training because the optimizer takes tiny, correlated steps.
• Out-of-memory (OOM) crashes because your batch doesn't fit on the GPU.

If you're building an applied AI product — a fine-tuned LLM, a computer vision API, or a recommendation engine — batch size is a lever you pull daily. This lesson gives you a repeatable methodology to find the right batch size for your problem, instead of copying someone else's config blindly.

Core concept / mental model

Think of training as finding the lowest point in a foggy valley. The loss landscape is the terrain, your weights are your position, and the gradient tells you which way is downhill. But you only have a flashlight — you can't see the whole valley at once.

• Batch size = how many data points you use to estimate the downhill direction before taking a step.
• Mini-batch gradient descent: sample a subset of data, compute the average gradient, update weights once.

define: Batch size is the number of training examples processed before the model's weights are updated. It's distinct from epoch (one full pass over the dataset) and iteration (one update step).

Mental picture: Smaller batch -> dim, shaky flashlight (noisy gradient estimate). Larger batch -> bright, steady flashlight (accurate estimate) but you move slower because you wait longer to see.

The bias-variance tradeoff in gradient estimation drives everything: - Small batch (e.g., 8–32): high-variance gradient estimate, acts as implicit regularization, can escape sharp minima, needs lower learning rate to avoid divergence. - Large batch (e.g., 128–1024): low-variance gradient, stable but prone to sharp minima (worse generalization), needs higher learning rate (e.g., linear scaling rule) to match convergence speed.

How it works step by step

1. The forward/backward cycle

For each mini-batch, the model: 1. Forwards: computes predictions for all examples in the batch. 2. Computes loss: averages (or sums) the individual losses across the batch. 3. Backwards: computes gradients of the average loss w.r.t. each weight. 4. Updates: applies optimizer step (e.g., SGD, Adam) using that averaged gradient.

2. Effect on gradient noise

The gradient noise scale increases as batch size decreases. Mathematically, the noise standard deviation is roughly:

std ≈ 1 / sqrt(batch_size)

So a batch of 8 has ~3.5× the noise of a batch of 100. This noise can be helpful early in training (escapes bad local minima) but harmful late (bounces around the optimum).

3. The linear scaling rule (for SGD)

If you double the batch size, you can often double the learning rate to maintain the same effective update step. This works only when gradients are unbiased and your loss landscape is smooth enough. For Adam/AdamW, the rule is less precise — many practitioners keep LR constant and only increase batch size, then tune LR separately.

4. Memory implications

Activations, gradients, and optimizer states scale linearly with batch size. A batch of 128 uses 4× the memory of a batch of 32 for the same model. If you OOM, you either shrink the batch, use gradient accumulation, or reduce model size.

5. Throughput and parallelism

GPUs are optimized for big, batched matrix multiplications. Tiny batches (e.g., 1–4) underutilize the hardware, while very large batches may hit a diminishing returns point where throughput per second plateaus. Use profiling to find the sweet spot.

6. The golden rule of experimentation

Change only one variable at a time. If you alter batch size and learning rate and data augmentation, you can't attribute results to any single factor. Keep everything else fixed.

Hands-on walkthrough

Let's run a controlled experiment on a toy dataset to see batch size effects. We'll train a small neural network on synthetic 2D data and compare loss curves.

Setup: synthetic dataset and model

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt

# Fix seed for reproducibility
torch.manual_seed(42)

# Generate synthetic data: two interleaving half-circles
from sklearn.datasets import make_moons
X, y = make_moons(n_samples=2000, noise=0.2, random_state=42)
X_t = torch.tensor(X, dtype=torch.float32)
y_t = torch.tensor(y, dtype=torch.int64)

dataset = TensorDataset(X_t, y_t)

# A tiny 2-layer MLP
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64),
            nn.ReLU(),
            nn.Linear(64, 2)
        )
    def forward(self, x):
        return self.net(x)

Training function that takes batch size as input

def train_with_batch_size(batch_size, lr=0.1, epochs=30):
    torch.manual_seed(0)  # reset weights each run
    model = SimpleMLP()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()

    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    losses = []
    for epoch in range(epochs):
        epoch_loss = 0.0
        for Xb, yb in dataloader:
            optimizer.zero_grad()
            out = model(Xb)
            loss = loss_fn(out, yb)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        losses.append(epoch_loss / len(dataloader))
    return losses

# Run experiment: compare batch sizes 8, 32, 128
batch_sizes = [8, 32, 128]
all_losses = {}
for bs in batch_sizes:
    print(f"Training with batch size {bs}...")
    all_losses[bs] = train_with_batch_size(bs, lr=0.05)

# Plot the loss curves
plt.figure(figsize=(10,6))
for bs, losses in all_losses.items():
    plt.plot(losses, label=f"batch={bs}")
plt.xlabel("Epoch")
plt.ylabel("Training loss")
plt.title("Effect of Batch Size on Training Loss")
plt.legend()
plt.grid(True)
plt.show()

Expected output: The batch size 8 curve will be noisy and may oscillate; batch size 128 will be smooth but may converge slower or plateau. Batch size 32 sits between. If you see the loss for batch 8 increase with LR=0.1, that's the noise — you'll need to lower LR.

Measuring generalization: validation accuracy

from sklearn.model_selection import train_test_split

# Split data once
Xtr, Xval, ytr, yval = train_test_split(X, y, test_size=0.2, random_state=1)
# (repeat training with the split to compare final validation accuracy)

def evaluate_model(batch_size, lr=0.05, epochs=50):
    # ... same as above but returns val accuracy after training ...
    pass

# You'd see: batch=8 often gets better val accuracy (~0.95) than batch=128 (~0.90) on this problem.

Pro tip: Always track validation accuracy alongside training loss. A smooth training loss with poor val accuracy means your large batch is overfitting to sharp minima.

Compare options / when to choose what

When to choose what: - Limited GPU memory → start small (16–32) and use gradient accumulation if needed. - Very large dataset (millions) → choose 64–256 to keep training time sane; consider distributed training with larger batches per GPU. - Fine-tuning a pretrained model → small batches (8–16) perform better for adaptation, as seen in many LLM fine-tuning recipes. - Reproducibility → use a single batch size and seed; extreme differences can be non-deterministic.

Troubleshooting & edge cases

1. Loss spikes or NaN

• Symptom: With small batch size, loss jumps to NaN or oscillates.
• Cause: High gradient noise + learning rate too high.
• Fix: Lower learning rate by 2–10×, or increase batch size. Use gradient clipping.

2. OOM (out-of-memory) error

# This raises CUDA OOM on a small GPU:
dataloader = DataLoader(dataset, batch_size=1024)  # too big

• Fix: Reduce batch size, use torch.cuda.amp (mixed precision), or enable gradient accumulation to simulate larger batch:

accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # normalize
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

This effectively gives you a batch size of actual_batch * accumulation_steps without the memory spike.

3. Training is too slow with small batches

• Cause: GPU underutilization.
• Fix: Increase batch size until GPU utilization is >80%, or use torch.utils.data.DataLoader(num_workers=...) to speed data loading.

4. Changing batch size changes other hyperparameters

• Effect: If you increase batch size, you may need to increase epochs to see the same number of weight updates (because your dataset has fewer iterations per epoch). Compare total steps, not epochs, for fair comparison.

5. Batch size affects learning rate schedules

• Fix: If you use a learning rate scheduler, adjust its parameters when batch size changes. Many schedules decay with total steps; recalc accordingly.

What you learned & what's next

You now have a scientific method to experiment with batch size effects: you understand the mental model (gradient noise vs. stability), the mechanics (forward/backward/update), the linear scaling rule, and how to compare options on your own hardware. You can troubleshoot OOM, loss spikes, and slow training, and you know to change one variable at a time with reproducible seeds.

You've met both learning objectives: you can explain the core idea behind batch size effects, and you've completed a practical exercise that measures them. Use this knowledge to tune any DL model going forward.

What's next: Now that you can manipulate batch size, the next lesson in this track is Learning rate scheduling, where you'll combine batch size with dynamic learning rates to squeeze out even more performance. You'll build on this experiment's methodology to test LR warmup, cosine decay, and cyclical schedules.

Final tip: Keep a training log with batch size, LR, loss curves, and validation accuracy for every experiment. Your future self will thank you.