Skip to content
AI.info

Neural networks

The Training Loop: From Mini-Batch to Parameter Update

Assemble forward computation, loss reduction, backpropagation, optimizer steps, gradient accumulation, and evaluation mode into one auditable loop.

By the end you can

Visual

One optimization step, in order

A training loop is short enough to memorize and important enough to instrument. Each of the five stages below has a dated, published failure attached to it.

Loss reduction is one. Hugging Face corrected it in its Trainer in October 2024, after the usual gradient-accumulation arithmetic turned out not to reproduce the full-batch loss. The parameter update is another. Reddi and his co-authors exhibited a convex problem on which Adam does not reach the optimum — an ICLR 2018 best paper — and PyTorch still ships the correction as the amsgrad flag on torch.optim.Adam.

That is the reason to name the stages separately rather than think of the loop as one act. Each one can be wrong on its own. Each one can be wrong while every other stage runs, converges and logs.

FigureProcess · 5 steps
  1. 1

    Load a mini-batch

    Fetch inputs, targets, masks, and metadata from the training stream.

  2. 2

    Run the forward pass

    Produce outputs under training-mode behavior.

  3. 3

    Compute and reduce loss

    Apply targets, masks, weights, and the chosen reduction.

  4. 4

    Backpropagate

    Populate gradients for trainable parameters.

  5. 5

    Apply the update

    Let the optimizer modify parameters, then prepare for the next step.

Comparison

Batch, micro-batch, step, and epoch are different units

Logs become misleading when these terms are used interchangeably. The clearest evidence that an epoch is not a unit of progress is that the industry's own training benchmark refuses to count them. MLPerf scores time-to-train against a fixed quality target: “The MLPerf Training benchmark suite measures the time it takes to train one of six machine learning models to a standard quality target” — MLCommons, 10 July 2019. That round, v0.6, raised the image-classification target to 75.9% top-1 for ResNet.

Now watch the epoch count at that frozen target. Graphcore reported reaching the identical 75.9% figure in 38 epochs and 28.3 minutes on an IPU-POD16, down from roughly 65 epochs before its optimisation work (17 January 2022). Same model, same target quality, and the number of passes over the data nearly halved. An epoch counts trips through a dataset. It does not count progress, and it does not count compute.

A mini-batch is the set of examples whose gradients contribute to one optimizer update. A micro-batch is what fits on the device before accumulation. A step is the moment parameters change. Only the last of these is an event in the model's history.

FigureComparison · 3 columns

Mini-batch

Examples whose gradients contribute to one optimizer update.

  • Defines one stochastic gradient estimate
  • May combine several micro-batches
  • Affects normalization and noise
  • Not always equal to device batch size

Micro-batch

Examples processed together before gradient accumulation.

  • Limited by device memory
  • Several can form one effective batch
  • Requires correct loss scaling
  • May alter batch-dependent layers

Epoch

One pass through a defined training dataset or sampler cycle.

  • Depends on dataset definition
  • Can be ambiguous for streams
  • Contains many optimizer steps
  • Not a universal measure of compute

Case

168,160 trained models mapped what a larger batch actually buys

One step is not a fixed amount of progress. How much it buys depends on the batch, and not linearly. Shallue and five co-authors measured the number of steps needed to reach a fixed out-of-sample error across 35 workloads, then released the data: 71,638,836 loss measurements taken over 168,160 individual trained models (Journal of Machine Learning Research, 2019).

One shape appeared in every workload. “Specifically, for each workload (model, training algorithm, and data set), increasing the batch size initially decreases the required number of training steps proportionally, but eventually there are diminishing returns until finally increasing the batch size no longer changes the required number of training steps.”

It is not one team's claim. OpenAI had come at the same ceiling from a different direction in 2018, characterising it as a predictable “largest useful batch size” set by the gradient noise scale. Two unaffiliated groups, two methods, one bend in the curve.

Where those bends fall varies with model, training algorithm and dataset. So a step count quoted without its batch size states nothing that transfers to your run.

Training mode changes the function being executed

Dropout, batch normalization, and some stochastic modules do not merely behave differently in training and evaluation. They compute different functions. In each case the difference is a written-down constant you can look up.

Srivastava and four co-authors fixed the dropout rule in 2014: “If a unit is retained with probability p during training, the outgoing weights of that unit are multiplied by p at test time.” PyTorch's nn.Dropout ships the inverted form of that same arithmetic — “the outputs are scaled by a factor of 1/(1-p) during training”, while “during evaluation the module simply computes an identity function”. Validation left in training mode is therefore not slightly noisier. It scores a randomly thinned network whose surviving activations have been multiplied by 1/(1-p).

Batch normalization swaps one set of statistics for another. Ioffe and Szegedy's 2015 paper normalises with mini-batch statistics during training and substitutes fixed population statistics at inference. PyTorch implements exactly that: “During training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation”, with a default momentum of 0.1. Validate in training mode and every example is normalised by whichever other examples happened to share its batch. That is the mechanism behind the phrase “batch-statistic leakage”, and it is a specific substitution you can check in your own code. Why the layer is worth getting right is the paper's own headline claim: “Applied to a state-of-the-art image classification model, Batch Normalization achieves the same accuracy with 14 times fewer training steps, and beats the original model by a significant margin.” The same paper reports 4.9% top-5 validation error on ImageNet.

Switching mode does not disable gradient tracking. The framework maintainers say so in their own reference documentation, rather than leaving it to folklore. PyTorch: “Evaluation mode is not a mechanism to locally disable gradient computation. It is included here anyway because it is sometimes confused to be such a mechanism.” The same page calls module.eval() “completely orthogonal to no-grad mode and inference mode”. TensorFlow's custom-training-loop guide sets the two controls independently, calling the model as model(x_batch_train, training=True) inside a separate tf.GradientTape scope, and as model(x_batch_val, training=False) at validation. Two toolchains, one distinction.

“Eval mode” and “no gradient” answer different questions.

Example

Metrics need correct denominators

Averaging batch metrics can mislead. When batches differ in size or valid-token count, the epoch metric is biased. It is the same arithmetic error Hugging Face had to correct one level down, inside the loss itself, in October 2024.

  • Accumulate the sum of losses and the number of contributing examples or tokens.
  • For accuracy, accumulate correct predictions and evaluated predictions separately.
  • Ignore padding and missing labels using the same validity mask as the task definition.
  • Report per-class or per-slice counts when aggregate performance can hide rare failures.
  • Keep training metrics separate from validation metrics and from the objective used for updates.

Key idea

Gradient accumulation changes timing, not always behavior equivalently

Summing gradients over micro-batches can approximate a larger batch when loss scaling is correct. Exact equivalence can still fail, with dropout randomness, batch normalization, data augmentation, or optimizer behavior tied to step count.

Document whether schedules and regularization advance per micro-batch or per optimizer step.

This is not a hypothetical failure, and the chain of events is dated and traceable. On 15 October 2024 Unsloth reported that the usual loop — average each micro-batch's cross-entropy over its own token count, then combine those averages — does not reproduce the full-batch loss whenever the micro-batches hold different numbers of non-padding tokens. Hugging Face confirmed the next day, 16 October 2024, that its Trainer did exactly this, and stated the correct reduction: “To be precise, for gradient accumulation across token-level tasks like causal LM training, the correct loss should be computed by the total loss across all batches in a gradient accumulation step divided by the total number of all non padding tokens in those batches. This is not the same as the average of the per-batch loss values.” ArthurZucker's PR #34191 to huggingface/transformers was opened on 16 October and merged on 17 October 2024. Benjamin Marie, who had reported the discrepancy first, wrote it up independently in The Weekly Kaitchup #63 on 18 October 2024, noting that nearly all language models were affected and that the same accumulation happens across devices in multi-GPU training. Report to merged patch took a day. Before that, the loop had been running, converging and logging a loss.

Loss scaling is the second route by which the scalar you backpropagate differs from the one you meant. Micikevicius and his co-authors proposed “scaling the loss appropriately to handle the loss of information with half-precision gradients” in 2017. PyTorch's automatic-mixed-precision documentation says what happens without it: “Gradient values with small magnitudes may not be representable in float16. These values will flush to zero ("underflow"), so the update for the corresponding parameters will be lost.” Those parameters simply do not move, unless the loss is scaled before backward and the gradients unscaled before the optimizer step. Nothing in the loop announces the omission.

An “effective batch size” is a useful approximation, not a universal identity.

Analogy

A laboratory cycle with a signed handoff

A laboratory receives samples, runs an assay, compares results with references, records deviations, and authorizes one equipment adjustment. Every stage has a named input and a named output. A deviation recorded at one stage does not silently propagate as a valid reading at the next.

Laboratory calibration cannot represent a high-dimensional update assembled from many parameter derivatives. Randomized modules also make repeated runs differ: dropout retains each unit with probability p during training and computes an identity function at evaluation, so the same batch is not the same measurement twice.

The loop becomes reliable when every handoff can be inspected and reproduced.

Example

Loop bugs that mimic model problems

Several failures appear as “the network will not learn” even though the architecture is fine. Each one below has a mechanism you can name and inspect, not just a warning to heed.

  • Gradients are never cleared, so updates unintentionally include earlier batches.
  • Validation runs in training mode, so dropout scales the surviving outputs by 1/(1-p) instead of computing an identity function, and batch normalization normalises each example by its own batch instead of the running estimates it keeps with a default momentum of 0.1.
  • Targets use a different class index convention from the output head.
  • The mask excludes valid examples or includes padded positions, so the denominator stops being the number of non-padding tokens.
  • A scheduler advances at the wrong frequency relative to optimizer updates — the per-micro-batch versus per-step ambiguity that gradient accumulation introduces.
  • Metrics average batch averages instead of underlying counts, the same arithmetic Hugging Face removed from the Trainer in PR #34191 on 17 October 2024, one level up from the loss.

Position

A falling loss curve is evidence about the loop before it is evidence about the model

The descending curve is the standard proof that something is working — in a paper figure, on a dashboard, in a demo. Be blunt about what it certifies. The loop ran. The number the loop computed went down. Both of those can be true while the quantity being minimised is not the one anyone intended.

October 2024 is the clean case. The usual gradient-accumulation loop averaged each micro-batch's cross-entropy over its own token count, then combined those averages. That does not reproduce the full-batch loss when the micro-batches hold different numbers of non-padding tokens. Hugging Face confirmed it on 16 October 2024 and merged ArthurZucker's PR #34191 the next day, writing that the correct loss is “the total loss across all batches in a gradient accumulation step divided by the total number of all non padding tokens in those batches”, and that “This is not the same as the average of the per-batch loss values.” Benjamin Marie, reporting it first, noted that nearly all language models were affected. Nothing in any curve had said so.

The second case is older and lives in the update rule itself. “We provide an explicit example of a simple convex optimization setting where Adam does not converge to the optimal solution, and describe the precise problems with the previous analysis of Adam algorithm.” That is the abstract of an ICLR 2018 best paper by Reddi and his co-authors, and PyTorch ships the correction as a flag: torch.optim.Adam's amsgrad option is documented as “whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond”. Two independent, dated defects in code everybody was running. Neither was visible in a loss curve. The bugs listed just above are the same failure in miniature — gradients never cleared, validation running in training mode, a mask that keeps padded positions, metrics that average batch averages. Every one of them produces a curve.

The axes are no safer than the line. Horizontally, Shallue and five co-authors trained 168,160 models across 35 workloads and found the same shape in every one: steps to a fixed error fall proportionally with batch size, then return less, then stop falling. The bends land in places that vary by model, training algorithm and dataset, and OpenAI reached the same ceiling from the gradient noise scale. A step count quoted without its batch size transfers nothing. Nor does an epoch count. Graphcore reached MLPerf's fixed 75.9% target in 38 epochs where roughly 65 had been needed before, for the same model at the same quality.

None of this is an argument against reading curves. It is an argument that a curve is evidence about an architecture only after the loop underneath it has been checked by something that is not the curve.

A curve that goes down tells you the loop runs. It does not tell you what the loop minimised.

Steps

A reference implementation checklist

Keep the first working loop boring, explicit, and testable. Set modes deliberately, remembering that mode and gradient recording are separate switches: PyTorch documents module.eval() as “completely orthogonal to no-grad mode and inference mode”. Clear gradients at the chosen accumulation boundary. Validate a batch by asserting shapes, labels, masks, ranges, and finite values.

At the forward-loss-backward step, log each scalar and detect non-finite results immediately. State the reduction explicitly: the total loss over the accumulation window divided by the total number of non-padding tokens in it, not the average of per-batch averages. If the run is in half precision, scale the loss before backward and unscale before the step. Gradient values too small for float16 “will flush to zero”, and those parameters are then not updated at all.

At the update, apply clipping or optimizer logic and store step-level diagnostics. Record which optimizer variant you are running, since the difference between Adam and its AMSGrad correction is a documented convergence result rather than a preference. And record the batch size beside every step count you report. Without it the step count says nothing that transfers.

FigureProcess · 5 steps

  1. 1. Set modes

    Put the model and data transforms in their intended state.

  2. 2. Clear gradients

    Reset buffers at the chosen accumulation boundary.

  3. 3. Validate a batch

    Assert shapes, labels, masks, ranges, and finite values.

  4. 4. Forward, loss, backward

    Log each scalar and detect non-finite results immediately.

  5. 5. Update and record

    Apply clipping or optimizer logic, then store step-level diagnostics.

Key takeaways