Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging

Abstract

Training vision or language models on large datasets can take days, if not weeks. We show that averaging the weights of the k latest checkpoints, each collected at the end of an epoch, can speed up the training progression in terms of loss and accuracy by dozens of epochs, corresponding to time savings up to 68 and 30 GPU hours when training a ResNet50 on ImageNet and RoBERTa-Base model on WikiText-103, respectively. We also provide the code and model checkpoint trajectory to reproduce the results and facilitate research on reusing historical weights for faster convergence 1 . However, in deep learning, loss functions are highly non-convex [8]. Weight averaging has been mainly used to improve the model's generalization performance at the end of or after training [14,16].Contribution We revisit weight averaging applied to neural networks from a convergence speed perspective. Inspired by Li et al. [22], we focus on the middle stage of training: after the dramatic changes of the local loss landscape during the very first training steps [10,7] but before the optimizer converges. Because the weights still undergo substantial change in that middle phase, averaging all 1 https://github.com/jeankaddour/lawa Preprint. Under review.