Neural networks, deep learning, NLP · insight-0472
Scale the loss and the schedule correctly when you accumulate gradients
- Evidence type
- practitioner heuristic
- Difficulty
- hard
- Roles
- ML engineer
- Industries
- cross-industry
Summary: situation, action, reason, boundary
Situation
You emulate a batch of 256 with 8 micro-batches of 32 because the full batch does not fit in memory.
Action
Divide each micro-batch loss by the accumulation steps (or by total tokens for variable lengths), step the optimizer and scheduler once per accumulated batch, and use a normalization layer that does not depend on batch size.
Reason
Accumulation reproduces the gradient of a larger batch only if the losses are averaged consistently; batch statistics are computed per forward pass, so they are not accumulated.
Boundary
Memory-limited training of CNNs and Transformers.
Symptom or risk
Summing unscaled micro-batch losses makes the gradient 8 times too large (an 8x larger effective learning rate under plain SGD; Adam-style optimizers cancel most of this, but gradient clipping thresholds see the inflated norm), stepping the scheduler every micro-batch shortens warmup and decay, and BatchNorm still sees only 32 examples.
Exceptions
Models without batch-dependent layers avoid the normalization caveat.
Trade-offs
Accumulation slows wall-clock time per update compared with a true large batch.
Pitfalls
Logging the micro-batch loss as if it were the full-batch loss.
Lessons and concepts
Claims and their sources
PyTorch BatchNorm normalizes with statistics of the current batch during training and keeps running estimates for evaluation.