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

Jean Kaddour|arXiv (Cornell University)|Sep 29, 2022

Topic Modeling4 citations

TL;DR

This paper proposes Latest Weight Averaging (LAWA), a simple, model- and optimizer-agnostic method that accelerates training convergence by averaging the weights of the k latest model checkpoints collected at the end of each training epoch. LAWA reduces training time by up to 68 GPU hours on ImageNet and 30 GPU hours on WikiText-103 by enabling models to reach target loss and accuracy levels dozens of epochs earlier, without modifying the training loop beyond checkpoint collection and 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.

Motivation & Objective

• To address the high computational cost and long training times of large-scale vision and language models, especially for researchers without access to high-end hardware.
• To improve training convergence speed without modifying model architecture, optimizer, or hyperparameters.
• To explore whether averaging recent model weights during the middle phase of training—when loss landscapes are still evolving—can accelerate convergence.
• To provide a practical, plug-and-play method that enables faster iterative experimentation in deep learning research.

Proposed method

• LAWA collects model weights at the end of each training epoch and maintains a sliding window of the k most recent checkpoints in a queue.
• At each epoch, the model parameters of the LAWA ensemble are updated by computing the uniform average of the k latest checkpoints: θ^LAWA_E = (1/k) * Σ_{i=E-k+1}^{E} θ_i.
• The averaging operation is efficiently implemented using a circular buffer to reduce time complexity, and the averaged model is evaluated only when needed, minimizing memory and computation overhead.
• Batch normalization statistics for the averaged model are recomputed via a full inference pass through the training data to ensure consistent batch statistics.
• The method is agnostic to model architecture, optimizer (e.g., SGD or Adam), and task, making it broadly applicable across vision and NLP benchmarks.
• The hyperparameter k (number of latest checkpoints) is tuned empirically, with k=6 showing strong performance across both ImageNet and WikiText-103 experiments.

Experimental results

Research questions

• RQ1 Can averaging the k latest model checkpoints during the middle phase of training significantly accelerate convergence in terms of loss and accuracy?
• RQ2 Does LAWA outperform standard training with SGD or Adam in terms of time-to-accuracy on large-scale vision and language models?
• RQ3 How sensitive is LAWA’s performance to the choice of k, and does it exhibit optimal performance at a specific k value across different training stages?
• RQ4 Can LAWA be applied effectively to both computer vision (ResNet50 on ImageNet) and NLP (RoBERTa-Base on WikiText-103) without architectural or optimizer modifications?
• RQ5 Does LAWA improve final model generalization performance, or is it purely a convergence acceleration technique?

Key findings

• LAWA reduced the time to reach a given validation loss on ResNet50/ImageNet by up to 68 GPU hours compared to standard SGD training.
• On RoBERTa-Base/WikiText-103, LAWA enabled the model to match Adam’s final validation loss approximately 45 epochs earlier, saving around 30 GPU hours.
• The final generalization performance of the LAWA-averaged model was consistently better than that of the final Adam-optimized model, confirming improved generalization.
• LAWA was most effective when k was moderate (e.g., k=6), and performance degraded when k exceeded 16, indicating that averaging too many past checkpoints can be detrimental.
• The method is robust across different optimizers and tasks, with consistent convergence speedups observed in both vision and language modeling benchmarks.
• Empirical results show that averaging during the early training phase worsens performance, confirming that LAWA is most effective during the middle phase of training when the loss landscape is still evolving.