Each optimizer step changes the model based on only a slice of the training data. This article is about how large that slice should be, what makes it useful, and why the answer depends on the training setup.

One update is a noisy gradient estimate

An optimizer does not know the exact direction that would improve the model across the entire data distribution. An optimizer step uses a batch to build a gradient estimate, an approximation of where the model should move. The batch size therefore determines the signal-to-noise ratio of that estimate.

Using the entire dataset would produce the cleanest available estimate, but computing it before every optimizer step is usually too expensive. Small batches make each step cheaper, but their gradient estimates are noisy and will generally differ from the true population gradient. As the batch grows, variation between individual examples tends to average out, so larger batches usually improve the signal-to-noise ratio of the estimate. Mini-batches balance this improvement against the cost of processing more examples per update.

Sometimes the batch needed for a reliable gradient estimate is too large to fit in memory. Gradient accumulation addresses this constraint by processing several smaller microbatches and combining their gradients before a single optimizer step. This allows more training data to contribute to each update without storing all of its activations in memory at once.

Some gradient noise can be useful because it prevents the optimizer from following the training data too precisely and introduces variation that can help it move out of narrow or poorly generalizing regions of the loss landscape. Too much noise, however, obscures the underlying gradient signal. Successive updates then point in inconsistent directions, so they may cancel one another, cause the loss to fluctuate, and require more steps to make reliable progress.

In a supervised problem such as MNIST, the training set is fixed and relatively small, and each image comes with a direct label. A batch can therefore cover a meaningful part of the data distribution and produce a useful gradient estimate with relatively few examples. Reinforcement learning is generally much more stochastic. Actions are sampled from a changing policy, the environment produces different trajectories, and rewards may be delayed, sparse, or noisy. The space of possible trajectories is also far larger and harder to cover than the label space of a simple classification problem. A single batch may therefore need many samples before recurring signal becomes distinguishable from chance variation. Larger batches expose the update to more trajectories and rewards, making its estimate of the population gradient more reliable. The required batch size depends on the particular RL problem, but RL commonly needs much more data per update than a simpler supervised task.

Image models are often trained for many epochs while random augmentation creates a new view of each image on every pass. The underlying files repeat, but their crops, flips, colors, or other transformations may differ, which makes the effective stream of training examples larger and more varied than the stored dataset. Seeing an image in two epochs therefore does not necessarily mean seeing exactly the same input twice. The meaning of an epoch depends on how much the training data changes between passes.

Batch composition matters

A batch should represent the training distribution well enough for its gradient estimate to be useful. If the examples inside one update all come from the same class, domain, task, or pattern, the optimizer step will be biased toward that slice of the data rather than toward the distribution as a whole.

Shuffling helps because it breaks long runs of similar examples and gives each update a better chance of seeing a mixed slice of the dataset. It also reduces correlation between neighboring updates. Instead of moving through one narrow region of the distribution for many steps in a row, the optimizer receives a more varied sequence of gradient estimates.

Suppose the data is ordered by domain, and the learning rate warmup reaches its normal value only by the end of the third domain. The first two domains may still appear in the training stream, but their updates happen while the learning rate is very small. They therefore contribute much less to the model parameters than later domains that are seen after warmup. In effect, the data order changes the weight of each domain in training, even if every domain contains the same number of samples. The same issue can appear with other learning-rate schedules whenever poor ordering makes some domains receive only small parameter updates.

Batch size controls how much data contributes to an update. Batch composition controls which parts of the distribution that update actually sees.

Learning rate as noise amplifier and exploration knob

A noisy gradient estimate becomes more dangerous as the learning rate increases. The estimate does not only choose a direction. The learning rate decides how far the model moves in that direction. When the step size is large, a noisy batch can push the model far away from the direction suggested by the broader training distribution.

This is why a higher learning rate can behave like an exploration knob. It lets the optimizer move more aggressively and can help it escape flat or narrow regions of the loss landscape, but it also amplifies mistakes in the gradient estimate. If the batch is small and noisy, a high learning rate can turn ordinary sampling variation into unstable updates.

Batch size and learning rate should therefore be understood together. Larger batches usually reduce gradient noise, which can make larger learning rates easier to use. Smaller batches produce noisier estimates, so the same learning rate can become much more aggressive in practice.

Effective batch size

Effective batch size is the amount of useful training data that contributes to one optimizer update. For language models, this is closer to non-padding tokens per update than to the configured number of sequences.

Effective tokens per update are roughly non-padding tokens per batch, multiplied by gradient accumulation steps and the number of devices. Packing improves utilization by replacing padding with real tokens. Padding takes space, but it does not add useful gradient signal.

Critical batch size

Critical batch size is the point where making the batch larger starts giving diminishing returns. Below that point, adding more examples to each update can reduce gradient noise and save optimizer steps. Above that point, the update is already supported by enough data that adding more mostly repeats the same signal.

This matters because larger batches improve time efficiency but can hurt data efficiency. They let the training run use more parallel hardware and finish in fewer optimizer steps, but after the critical batch size, each additional example contributes less new information. You spend more data and compute per step without getting a proportional improvement in the update.