Robustness early in training
We study when stochastic gradient descent helps models avoid shortcuts that work during training but fail when the data change. Our finite-time analysis shows that, under the conditions we study, SGD’s implicit regularization is strongest early in training and then weakens. Its strength depends on the learning-rate-to-batch-size ratio.
Small-batch warmup
Based on this analysis, we use small-batch warmup: train with small batches at first, then switch to larger ones. Across several benchmarks and architectures, this improves worst-group accuracy, often recovering much of the robustness of full small-batch training while retaining most of the efficiency of large-batch training.
Small-batch warmup
- Early trainingUse small batches
At a fixed learning rate, smaller batches increase the learning-rate-to-batch-size ratio.
- After warmupSwitch to larger batches
Use larger batches for the rest of training to save computation.
The schedule targets the early period when implicit regularization is strongest in the paper’s analysis.
Related paper
Understanding When SGD Improves Robustness to Spurious Correlations
BibTeX
@unpublished{raoof_spurious,
title = {{Understanding When SGD Improves Robustness to Spurious Correlations}},
author = {Mirzaie, Nahal and Zare Moayedi, Raoof and Jafarinia, Hossein and Rohban, Mohammad Hossein and Soleymani Baghshah, Mahdieh},
year = {2026},
note = {Under review}
}