Neural networks, deep learning, NLP · insight-1417
Average the weights of late checkpoints, then recompute BatchNorm statistics
- Evidence type
- empirical finding
- Difficulty
- medium
- Roles
- ML engineer, Data scientist
- Industries
- cross-industry
Summary: situation, action, reason, boundary
Situation
Validation loss of a small network fluctuates by a few percent between neighbouring epochs and you must pick one checkpoint to ship.
Action
Average the weights of the last N epochs (stochastic weight averaging or an exponential moving average), run one pass over training data to recompute BatchNorm statistics, and evaluate the averaged model once.
Reason
Averaging points along an SGD trajectory was reported to generalise better and land in flatter regions at almost no extra cost; averaged weights need statistics recomputed, which PyTorch provides as a utility.
Boundary
Networks trained with a constant or cyclical learning rate near the end of training.
Symptom or risk
The chosen checkpoint is partly luck, and model selection on a small validation set overfits to its noise.
Exceptions
Models whose checkpoints are far apart in weight space, for example after a large restart.
Trade-offs
One extra pass over training data and an extra copy of the weights in memory.
Pitfalls
Evaluating the average with the statistics of the last raw checkpoint.
Lessons and concepts
Claims and their sources
Izmailov et al. report that averaging several points along an SGD trajectory with a cyclical or constant learning rate generalises better than conventional training, at almost no extra cost.
The PyTorch swa_utils module provides AveragedModel, SWALR and update_bn, which recomputes BatchNorm activation statistics with one pass over a data loader.