Skip to content
datarekha

How are batch size and learning rate related, and what is learning-rate warmup?

The short answer

A larger batch produces a lower-variance gradient estimate, so it can often support a proportionally larger learning rate, but linear scaling is only a starting heuristic. Warmup ramps the learning rate from a small value to its target over early optimizer steps, reducing the risk that unstable initial updates derail training.

How to think about it

Batch size is the number of training examples used for one parameter update. A larger batch gives a less noisy gradient estimate, so it often permits a larger learning rate; warmup gradually increases that rate because the target rate can be unstable during the first few updates.

A model learns by changing its weights in the direction that reduces the loss. The gradient is that direction, with its magnitude; the learning rate controls how large a step to take.

A simplified update is:

new weights = old weights - learning rate × gradient

With a mini-batch of B examples, the training code usually averages their individual gradients:

g_B = (1 / B) × sum of the individual gradients

Small batches produce a cheap but noisy estimate. One batch may contain unusually easy images, unusual customers, or a few outliers, so its gradient can point partly in the wrong direction. Increasing B averages out more of that randomness. If examples are reasonably independent, the variance of the average falls roughly in proportion to 1 / B, while the typical sampling error falls in proportion to 1 / square root of B.

That creates the practical relationship:

  • A small batch has a noisy direction, so a large step can amplify the noise and make training bounce or diverge.
  • A larger batch has a more reliable direction, so a larger step can move the model farther without reacting as strongly to one unlucky batch.

The common starting heuristic is the linear scaling rule:

new learning rate ≈ old learning rate × (new batch size / old batch size)

The reasoning is not magic. If the batch becomes eight times larger, an epoch contains about eight times fewer optimizer updates. Multiplying the learning rate by eight roughly preserves the total amount of parameter movement over one pass through the data, assuming the average gradient is similar.

But this is a heuristic, not a law of nature. Neural-network updates are nonlinear, gradients change after every update, and very large batches eventually provide diminishing returns. At some point, adding more examples to a batch makes each update more expensive without making the direction meaningfully better.

A concrete example

Suppose an image classifier is trained on 1,280,000 images.

The original configuration is:

  • Batch size: 256
  • Learning rate: 0.1
  • Updates per epoch: 1,280,000 / 256 = 5,000

You move to eight GPUs. Each GPU processes 32 images, and gradients are accumulated for eight forward-backward passes before the optimizer steps. The effective batch size is therefore:

32 × 8 GPUs × 8 accumulation steps = 2,048

That is eight times the original batch. Linear scaling suggests a starting learning rate of:

0.1 × (2,048 / 256) = 0.8

The new run performs only 625 optimizer updates per epoch. If you warm up over five epochs, the warmup lasts 3,125 optimizer steps. A linear warmup reaches 0.8 gradually:

import math

def learning_rate_at_step(step, warmup_steps, peak_lr, total_steps):
    if step < warmup_steps:
        return peak_lr * (step + 1) / warmup_steps

    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return 0.5 * peak_lr * (1 + math.cos(math.pi * progress))

print(learning_rate_at_step(0, 3125, 0.8, 50000))      # 0.000256
print(learning_rate_at_step(999, 3125, 0.8, 50000))    # 0.256
print(learning_rate_at_step(3124, 3125, 0.8, 50000))   # 0.8

After warmup, this example uses cosine decay. The exact decay schedule is a separate choice. Warmup answers the question, “How do I reach the target learning rate safely?” It does not determine what happens for the rest of training.

Count warmup in optimizer steps, not raw mini-batches. With gradient accumulation, eight mini-batches may produce only one optimizer update. Warming up for 3,125 raw mini-batches in the example would reach the target eight times too quickly.

What learning-rate warmup actually fixes

At the beginning of training, the model is in a fragile state. Its weights may be randomly initialized, its activations may have poor scales, and its optimizer state contains little history.

For SGD with momentum, the velocity starts at zero and is built from early gradients. For Adam or AdamW, the first- and second-moment estimates are still being formed. Bias correction helps correct the initialization bias in those estimates, but it does not make a handful of observations statistically stable.

A high target learning rate can therefore make the first update far too large. The model may jump into a bad region, produce enormous activations, or overflow in mixed-precision arithmetic. Warmup makes those early parameter changes smaller while the optimizer and model settle.

For linear warmup, with zero-based step t and warmup length W:

learning rate at step t = target learning rate × (t + 1) / W

Warmup does not make the gradient itself smaller. It only reduces the size of the weight update. That distinction matters: warmup cannot repair corrupted labels, exploding activations, a broken loss function, or a learning rate whose final value is simply too high.

The nuance that earns the senior signal

The textbook answer is not “increase the batch by four, then increase the learning rate by four.”

First, the optimizer matters. Linear scaling is most natural for SGD-style training. Adam and AdamW normalize gradients using running estimates, so their nominal learning rate does not map as directly to the raw gradient magnitude. Scaling may still work, but it is usually something to validate rather than apply mechanically.

Second, define what “batch size” means. In distributed training, it is often the effective batch across devices and accumulation steps. In language-model training, tokens per optimizer update are frequently more meaningful than sequences per update, because eight sequences of 2,048 tokens contain sixteen times as many tokens as eight sequences of 128 tokens.

Third, decide what budget you are holding constant. With a fixed number of training examples, a larger batch gives fewer optimizer updates. With a fixed number of updates, it consumes more examples and usually more compute. Those are different experiments and can produce different conclusions.

Fourth, large-batch training can change generalization. Gradient noise is not only an annoyance; it can act as a form of regularization. Removing too much noise may lower training loss while hurting validation accuracy. The effect depends on the model, data, optimizer, regularization, and schedule.

Gradient accumulation is also only approximately equivalent to a genuinely large batch. It can match the averaged gradient when parameters and optimizer state remain unchanged until the accumulated update. It differs when batch-dependent layers update statistics per micro-batch, when dropout randomness matters, or when loss weighting changes with sequence length. For example, accumulation does not make BatchNorm see the full effective batch; it still sees each device’s local micro-batch unless synchronized statistics are used.

A failure mode you can recognize

A typical failure log looks like this: the loss falls for 20 steps, jumps from 2.1 to 6.8, the gradient norm becomes enormous, and the loss becomes NaN before the first epoch ends. This usually means the peak learning rate is too high, the warmup is too short, or the effective batch and gradient scaling are not what you think they are.

Log the actual learning rate, gradient norm, loss, parameter-update norm, and mixed-precision overflow events. If the spike occurs exactly when warmup ends, the target rate or the transition is suspect. If it occurs during the first step, inspect data, initialization, loss reduction, and numerical scaling as well. A longer warmup may hide the symptom, but it will not make an excessive final learning rate safe.

Warmup is also not automatically useful. A short fine-tuning run with a conservative learning rate may spend most of its useful training time crawling upward. In that case, no warmup or a very short ramp can be better.

What they’ll ask next

Should I always scale the learning rate linearly?

No. Use linear scaling as the first experiment, then check training stability and validation quality. The rule becomes less reliable for very large batches, adaptive optimizers, unusual loss reductions, and changes in data composition. A small sweep around the proposed value is usually more informative than defending the formula.

Is gradient accumulation the same as increasing batch size?

Only under specific conditions. If gradients are averaged consistently, parameters are updated once per effective batch, and batch-dependent layers or loss weighting do not introduce differences, it is close. It is not equivalent if you step the optimizer after every micro-batch, or if BatchNorm, variable-length sequences, or other batch-dependent operations behave differently.

How do you choose the warmup length?

Choose it in optimizer steps and relate it to the size of the learning-rate jump, the optimizer, and the total run. A large increase in effective batch or target learning rate generally deserves more warmup than a conservative fine-tune. I would start with a short ramp, inspect the first few thousand steps, and compare a shorter and longer version rather than treating a fixed percentage as universal.

Say this in the interview: “Larger batches reduce gradient noise, so they can often support a larger learning rate, with linear scaling as a starting heuristic; warmup then ramps to that rate because the first optimizer steps and optimizer statistics are still unstable.”

Learn it properly Batch size ↔ learning rate

Keep practising

All Deep Learning questions

Explore further