Batch size ↔ learning rate

Your image classifier trains acceptably with a batch of 64 and a learning rate of 0.001. You move to a larger GPU, raise the batch to 256, and keep everything else unchanged.

Training now crawls.

The dataset still contains the same number of images, but each epoch has one quarter as many parameter updates. The optimizer sees cleaner gradients, yet it takes only one step where it used to take four. You try raising the learning rate by four. The loss becomes nan after 30 steps.

Both results are predictable. Batch size controls how noisy a gradient estimate is. Learning rate controls how far the optimizer moves using that estimate. Change the first and the useful range of the second moves with it.

This is why batch size and learning rate are a pair, not two independent knobs.

Bigger batches average away noise

A training example produces a gradient: a direction saying how each weight should change to reduce the loss. A batch gradient averages the gradients from its examples:

g_B = (g₁ + g₂ + ... + g_B) / B

Under the rough assumption that examples are sampled independently, the standard deviation of sampling noise falls approximately as 1/√B. Multiplying the batch by four makes the noise roughly half as large, not four times smaller.

Examples are not always independent. Correlated sequences, repeated samples, class imbalance, and similar photographs reduce the information in a batch. The square-root relationship is a useful model, not a guarantee.

An optimizer update looks roughly like:

new weights = old weights − learning rate × batch gradient

A smoother gradient can support a larger step, but small-batch noise can also help the optimizer escape narrow regions and regularize training. A large batch may reach low training loss while validation accuracy gets worse if the learning rate and schedule are not retuned.

The linear scaling rule

For SGD and SGD-like training, the usual starting point is the linear scaling rule:

Multiply the learning rate by the same factor as the batch size.

If the batch grows from 64 to 256, a learning rate of 0.001 becomes approximately 0.004.

This compensates mechanically for fewer updates. For a fixed number of examples, batch 64 makes four updates where batch 256 makes one. If the large-batch update is four times bigger, the first-order movement is similar:

• Four small updates: 4 × 0.001 × gradient = 0.004 × gradient
• One large update: 1 × 0.004 × gradient = 0.004 × gradient

The runs are not identical. Small-batch gradients are evaluated after each parameter change, while the large-batch gradient is evaluated at the old weights and applied once.

Imagine 49,152 training examples:

• Batch 64 gives 768 updates per epoch.
• Batch 256 gives 192 updates per epoch.
• With a base learning rate of 0.001, scaling suggests 0.004.
• If the average gradient over one group is [0.8, -0.4], four small updates move the weights approximately [-0.0032, +0.0016].
• One large update with learning rate 0.004 makes the same first-order move.

The arithmetic matches; the trajectories differ because the four small gradients contain more noise and are evaluated at successive parameter states.

Changing batch size also changes anything applied per optimizer step. A schedule that decays over 10,000 steps now covers a different number of examples. AdamW weight decay is applied at optimizer steps, so its cumulative effect per epoch changes too. Record both examples per update and the planned number of updates; “10 epochs” may otherwise describe a different optimization process.

Warmup prevents the first update from being a punch

A scaled learning rate can be sensible later and dangerous at initialization. Weights, activations, gradient scales, and optimizer moments are still settling, so 0.004 may make the first update oversized. Loss spikes, unstable activations, or immediate divergence can follow.

Learning-rate warmup starts lower and increases gradually to the target. For a target of 0.004 and a 1,000-update linear warmup:

• update 250: learning rate 0.001;
• update 500: learning rate 0.002;
• update 1,000: learning rate 0.004.

The normal schedule then takes over, such as cosine decay or step decay. Warmup is common with large-batch SGD but is not mandatory for every model.

Count warmup in optimizer updates, not data-loader iterations. With eight micro-batches per optimizer update, 1,000 warmup updates consume 8,000 micro-batches. Advancing the scheduler on every micro-batch makes warmup eight times shorter than intended.

Gradient accumulation: a large effective batch on a small GPU

If a GPU holds 32 examples but the target batch is 256, gradient accumulation runs eight forward and backward passes, adds their gradients, and performs one optimizer update:

32 examples × 8 accumulation steps = effective batch 256

Accumulation saves memory, not computation. The key requirement is to average gradients rather than merely sum them. If each micro-batch loss is a mean, divide it by the number of accumulation steps before backward().

accum_steps = 8
total_micro_batches = len(loader)

optimizer.zero_grad(set_to_none=True)

for micro_step, (xb, yb) in enumerate(loader):
    # Start a window and record its actual length.
    if micro_step % accum_steps == 0:
        window_steps = min(
            accum_steps,
            total_micro_batches - micro_step,
        )

    # loss_fn returns the mean loss for this micro-batch.
    loss = loss_fn(model(xb), yb) / window_steps
    loss.backward()  # adds this micro-batch's gradient to .grad

    end_of_window = (micro_step + 1) % accum_steps == 0
    end_of_loader = micro_step + 1 == total_micro_batches

    if end_of_window or end_of_loader:
        # Clip once, after the full effective-batch gradient exists.
        # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

If the loss is not divided by eight, the accumulated gradient is about eight times too large. For plain SGD, that acts like an accidental eightfold learning-rate increase. Adam and AdamW normalize updates with running moments, so the parameter update may not grow eightfold, but the optimizer still receives the wrong gradient scale and develops the wrong state. Divide the loss correctly.

In distributed data parallelism, include all workers:

effective batch = micro-batch × accumulation steps × workers

A micro-batch of 32 on four workers with eight accumulation steps gives 32 × 8 × 4 = 1,024. Gradients are normally averaged across workers during backward communication, so forgetting the worker factor can make the learning rate appear much too small or large.

Accumulation matches a single large batch only when:

• the loss is normalized correctly;
• the optimizer steps once after all micro-batches;
• the scheduler advances once per optimizer update;
• clipping happens after accumulation;
• batch-dependent layers do not require statistics from the entire batch;
• examples are weighted consistently.

BatchNorm is the common exception. Eight passes of 32 examples update its statistics using eight small sets, unlike one pass of 256. Accumulation cannot reconstruct those statistics afterward. A sufficiently large micro-batch or a different normalization strategy may be necessary.

Variable-length sequence data adds a loss-normalization trap. If loss_i is the mean loss over n_i valid tokens and T = Σ_i n_i is the total number of valid tokens in the window, backpropagate:

loss_i × n_i / T

This produces the global mean over tokens. Dividing every micro-batch mean only by the number of micro-batches gives equal weight to batches with different token counts.

Finally, handle an incomplete accumulation window. If there are 10 micro-batches and accum_steps is eight, discard the last two deliberately or flush them using a divisor of two. Dividing them by eight underweights those examples. drop_last=True ensures full DataLoader batches, but does not ensure that the loader length divides evenly by the accumulation count.

What breaks first in practice

An overly aggressive learning rate usually shows up as a loss jump at the end of warmup, a gradient-norm spike, or poor validation despite fast training-loss reduction. Reduce the target rate, lengthen warmup, or scale up more gradually. Clipping can contain an outlier; it cannot fix a consistently wrong rate.

A scheduler running on micro-steps reaches its minimum after too little data. Call it after optimizer.step() and define warmup and decay in a consistent unit: optimizer updates, examples, or tokens. Log the gradient norm immediately before each optimizer step; an unexpected factor near accum_steps often indicates missing loss scaling.

Larger batches also hit a throughput ceiling. Once the GPU, memory bandwidth, or inter-GPU communication is saturated, a larger batch only makes each update wait longer. Measure examples per second and validation quality.

The honest trade-off is that a larger effective batch gives fewer, less noisy decisions. It can improve hardware efficiency and distributed training, but it can also remove useful stochasticity and hurt validation performance. If batch 256 matches batch 64’s training loss but has worse validation accuracy, compare learning rate, warmup, total examples, weight decay, augmentation, and optimizer updates before increasing the batch again.

A production change should therefore identify the largest reliable micro-batch, choose the effective batch separately, calculate worker and accumulation factors, tune the learning rate for the optimizer, warm up the schedule, and log effective batch beside learning-rate and gradient-norm charts.

Quick check

Next

Once one GPU is no longer enough, distributed training (DDP & FSDP) shows how batches and model state are split across devices without losing track of what one optimizer update actually means.