PaperlayerView exact revision on arXiv

Source-linked AI summary

Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits

Daniel Morales-Brotons, Thijs Vogels, Hadrien Hendrikx

arXiv:2411.18704v1cs.LG

TL;DR

Weight averaging is widely used in deep learning, but its standalone properties have not been systematically evaluated. This paper empirically studies EMA dynamics and final solutions, finding broad benefits over last-iterate SGD while noting that the evidence is limited to image-classification benchmarks.

  • Problem

    Weight averaging is often embedded in complex training pipelines or used as a teacher, leaving its standalone properties insufficiently studied.

  • Method

    The paper empirically studies EMA models outside the training loop along SGD trajectories, including training dynamics, hyperparameter tuning, and final-model properties.

  • Results

    EMA models differ from last-iterate SGD solutions and improve generalization, label-noise robustness, prediction consistency, calibration, and transfer learning.

  • Takeaways & Limitations

    EMA is presented as a simple plug-in that can improve deep-learning performance while reducing reliance on late learning-rate decay.

  • Takeaways & Limitations

    The empirical study focuses solely on image-classification benchmarks, so whether EMA properties hold for other tasks remains unexplored.

Abstract

from arXiv · show

Weight averaging of Stochastic Gradient Descent (SGD) iterates is a popular method for training deep learning models. While it is often used as part of complex training pipelines to improve generalization or serve as a `teacher' model, weight averaging lacks proper evaluation on its own. In this work, we present a systematic study of the Exponential Moving Average (EMA) of weights. We first explore the training dynamics of EMA, give guidelines for hyperparameter tuning, and highlight its good early performance, partly explaining its success as a teacher. We also observe that EMA requires less learning rate decay compared to SGD since averaging naturally reduces noise, introducing a form of implicit regularization. Through extensive experiments, we show that EMA solutions differ from last-iterate solutions. EMA models not only generalize better but also exhibit improved i) robustness to noisy labels, ii) prediction consistency, iii) calibration and iv) transfer learning. Therefore, we suggest that an EMA of weights is a simple yet effective plug-in to improve the performance of deep learning models.

1 Introduction

The paper studies EMA of weights outside the training loop as a simple alternative to relying on learning-rate decay for reducing SGD noise. It finds distinct training dynamics and final solutions, with benefits spanning early performance, generalization, robustness, calibration, consistency, and transfer learning.

  • Study scope: The work addresses a gap in isolated evaluation by studying EMA outside the training loop, where it does not alter the underlying SGD trajectory.This setup separates the effects of averaging from changes to optimization.
  • Training dynamics of EMA: EMA reduces parameter noise, allowing higher learning rates and less late-stage decay while providing implicit regularization.The paper contrasts averaging-based noise reduction with learning-rate decay as routes to good solutions.
  • Training dynamics of EMA: EMA performs strongly early in training, helping explain its success as a teacher and potentially reducing compute by avoiding the final low-learning-rate phase.The claimed compute benefit depends on combining early stopping with EMA.
  • Training dynamics of EMA: Recomputing Batch Normalization statistics permits larger EMA windows, which may further improve generalization.Without recomputation, Batch Normalization statistics constrain the usable averaging decay.
  • Properties of the final EMA model: EMA models differ from last-iterate SGD solutions and improve generalization, robustness to label noise, calibration, prediction consistency, and transfer learning.The study examines EMA independently rather than as one component of a complex teacher-training framework.

2 Related Work

Prior work uses weight averaging for generalization and student-teacher learning, but its benefits are often justified informally or entangled with other mechanisms. This paper positions EMA as an under-studied alternative whose online use avoids some practical limitations of SWA.

  • Overview: Weight averaging has been widely used, yet its alleged benefits have lacked systematic, isolated evaluation across domains.The paper reviews existing folklore benefits and tests alternative explanations in later sections.
  • Weight averaging to improve generalization: Prior stochastic-approximation analyses generally assume quadratic or strongly convex objectives, limiting their direct applicability to deep-learning training.The paper notes this theoretical gap before studying EMA empirically in non-convex settings.
  • Weight averaging to improve generalization: SWA averages late SGD checkpoints at a high constant learning rate and is argued to find flatter, better-generalizing solutions.SWA is a uniform checkpoint average rather than an online EMA update.
  • Weight averaging in Student-Teacher methods: EMA teachers generate accurate, consistent predictions for student-teacher methods, while SWA is less suitable online because it requires recomputing Batch Normalization statistics.EMA has been used across semi-supervised, adaptation, continual-learning, and robustness settings.
  • Weight averaging in Student-Teacher methods: EMA is generally preferred for teacher models because it avoids Batch Normalization recomputation for short averaging windows, though longer windows can change this trade-off.The paper specifically investigates how averaging-window length interacts with Batch Normalization statistics.

3 Insights on Weight Averaging during Training

EMA is a lightweight plug-in whose averaged solutions differ from SGD, improving training dynamics and often requiring less learning-rate decay. Its benefits depend on averaging decay and Batch Normalization handling, especially for large averaging windows.

  • 3.2 Implicit Regularization with SGD Noise and Learning Rate Schedule: EMA outperforms SGD throughout training, peaks at epoch 150 under an optimal learning rate, and later deteriorates as the learning rate is decayed further.The best EMA model rises quickly, then improves more slowly before reaching its peak.
  • 3.2 Implicit Regularization with SGD Noise and Learning Rate Schedule: EMA solutions generalize similarly or better than SGD while reaching their best performance in fewer than three-quarters of the cosine-annealing epoch budget.The authors suggest that averaging reduces the need for the final phase of SGD training and extensive learning-rate decay.
  • 3.3 EMA in early training: EMA models perform well early because averaging reduces noise, allowing large learning rates to support fast progress and helping explain EMA teachers’ effectiveness.Student-teacher methods can distill knowledge from this early-performing EMA model.
  • 3.3 EMA in early training: Bootstrapping SGD with EMA parameters provides no benefit because noisy SGD updates quickly deteriorate the model after replacement.The authors therefore identify distillation as a more effective way to exploit EMA’s early performance.
  • 3.4 Batch Norm Statistics and EMA decay: Very slow EMA decay can cause divergence when Batch Normalization statistics are not recomputed, while recomputation restores performance and favors larger averaging windows.Recomputing Batch Normalization statistics also consistently improves generalization.
  • 3.4 Batch Norm Statistics and EMA decay: For online EMA use, faster decay can avoid Batch Normalization failures, whereas final-performance averaging favors slower decay with post-training statistic recomputation.Models without Batch Normalization naturally avoid this specific problem.

4 Results

Across image-classification experiments, EMA consistently improves generalization over momentum-SGD baselines and also improves robustness to label noise, prediction consistency, calibration, and transferability. These benefits arise alongside low-overhead evaluation using held-out validation data and differ from simply decaying the learning rate toward zero.

  • 4.1 Experimental Setup: The study uses multiple image-classification datasets and architectures, three independent runs, and 80/20 hold-out splits for hyperparameter selection and EMA early stopping.The final models are retrained on the full training data using selected hyperparameters, avoiding direct test-set tuning of early stopping.
  • 4.2 Generalization: EMA consistently outperforms the momentum-SGD baseline in test accuracy and loss, while performing comparably to SWA without a clearly superior averaging method.The comparison uses EMA early-stopped by either best validation accuracy or lowest validation loss, with batch-normalization statistics recomputed once.
  • 4.3 Label Noise: 65.15% accuracy at epoch 100 with learning rate η around 0.4 under 40% label noise, before later memorization reduces clean-test generalization.The EMA model peaks while averaging at a relatively high learning rate, illustrating the interaction between averaging and implicit regularization.
  • 4.4 Prediction consistency: EMA substantially improves prediction consistency across independent runs, reducing classification churn and Jensen–Shannon divergence, and outperforming a specialized method with twice the training cost.The experiments train three models with different seeds and evaluate both pairwise churn and JS divergence.
  • 4.5 Transfer Learning: 57.78% versus 52.77% linear-evaluation accuracy on CIFAR-100 is obtained by EMA versus SGD features after TinyImageNet pretraining.Across transfer tasks, EMA representations are reported as more linearly separable and more transferable than baseline representations.
  • 4.6 Calibration: EMA considerably reduces calibration error across tested models and datasets, while combining EMA with temperature scaling generally gives the best calibration.The paper treats the two improvements as complementary and hypothesizes that temporal weight averaging improves uncertainty estimation.

5 Conclusion

The paper systematically studies EMA weight averaging, covering its training dynamics and final-model properties. Across image-classification experiments, EMA improves generalization and several robustness and representation qualities, but the evidence is limited to image classification.

  • Scope and contribution: The study addresses a literature gap by examining EMA itself rather than treating it only as part of complex teacher-model frameworks.The authors organize the work around training dynamics and final EMA-model properties.
  • Training dynamics: EMA reduces SGD noise, supports higher learning rates, and enables implicit regularization that favors more general representations.The paper also presents one-shot tuning and trades learning-rate annealing for averaging to reduce training time.
  • Training dynamics: EMA shows strong early performance, helping explain its success as a teacher, but overly large averaging windows require recomputing Batch Norm statistics.This connects early EMA quality with practical constraints in teacher use.
  • Final EMA models: EMA models differ from last-iterate SGD solutions and improve generalization, robustness to label noise, prediction consistency, transferability, and calibration.The study reports performance comparable to SWA for generalization and benefits from a less complex algorithm for label-noise robustness.
  • Limitations: The empirical study evaluates image-classification benchmarks only, so whether these EMA properties extend to other tasks remains unresolved.The authors identify broader task coverage as future work.
  • Practical implication: The authors propose EMA as a simple plug-in with practical value for improving deep-learning performance and providing training-dynamics guidance.The conclusion frames the contribution as both empirical understanding and practitioner guidance.

A Additional examples of EMA training dynamics

Additional experiments show that EMA’s training dynamics depend on the learning-rate schedule and remain consistent across datasets and architectures. Under step decay, EMA matches rather than exceeds SGD, while cosine-annealed runs show an earlier EMA peak at higher learning rates.

  • A.1 EMA dynamics with a step decay: Under step decay on CIFAR-100 with ResNet-18, EMA does not outperform SGD and only matches the baseline.The learning rate is reduced by a factor of 5 at epochs 60, 120, and 160.
  • A.1 EMA dynamics with a step decay: The step-decay experiment reports the best of five parallel EMAs at each epoch without recomputing Batch Norm statistics.This evaluation setup accompanies the CIFAR-100 ResNet-18 comparison.
  • A.2 Other datasets and architectures: Across datasets and architectures with cosine annealing, EMA peaks well before training ends when averaged at a higher learning rate.Excessive learning-rate reduction weakens implicit regularization and degrades EMA performance.

B Additional results

The appendix provides additional EMA results, variability information, and full experiment coverage across architectures, datasets, and label-noise settings. It compares SGD and EMA variants with and without Batch Norm recomputation across 80% and full-data training.

  • Data splits and tables: The supplementary tables distinguish 80% training with 20% hold-out evaluation from full-data training and report three-run results for each setting.Table 6 summarizes 80% training and complements the subsequent full-data results in Table 1.
  • Evaluation setup: The appendix reports five model variants: momentum SGD, EMA stopped at best accuracy, EMA stopped at lowest loss, and corresponding versions with recomputed Batch Norm statistics.Results include means and standard deviations over three independent runs, plus selected epochs, decays, and learning rates.
  • Label-noise experiments: Noisy-label supplementary evaluations include CIFAR-100-N and CIFAR-10-N with 40% noisy labels, under both hold-out and full-data settings.The CIFAR-10-N setting is identified as the Worse variant.

C Bootstrapping on EMA

The bootstrapping experiment tests whether EMA’s strong early performance can accelerate training by periodically replacing SGD parameters with EMA parameters. Instead, bootstrapping does not help and decreases performance.

  • C Bootstrapping on EMA: Bootstrapping SGD with EMA weights once per epoch decreases performance compared with normal EMA use outside the training loop.The experiment evaluates CIFAR-100 on an 80%/20% train–hold-out split with T = 16 and α = 0.992.
  • C Bootstrapping on EMA: The experiment compares regular SGD and its EMA with bootstrapped SGD and its EMA.Normal EMA remains a slow-moving average of the SGD sequence rather than a training-loop update.

D Learning rate tuning for SGD vs EMA

EMA and SGD share the same best initial learning rate despite differing dynamics and final solutions. EMA also memorizes noisy labels less than SGD relative to clean-label accuracy, with its best performance appearing as noise memorization begins increasing.

  • The best initial learning rate η was the same for EMA and SGD despite differences in their training dynamics and final solutions.The authors hypothesize that initial learning-rate tuning mainly affects early progress and generalization.
  • EMA’s lower noisy-label memorization supports an implicit-regularization effect that improves its behavior under label noise.
  • At 90% clean accuracy, EMA memorized 28.8% of noisy labels versus 68.2% for SGD, while EMA’s best performance occurred around epoch 100.The comparison uses epoch 99 for EMA and epoch 142 for SGD.

E.2 Continued training at constant Learning Rate

The ablation tests whether overfitting during continued training is caused by learning-rate decay rather than training duration alone. Its setup compares cosine-decay training with constant learning rate after a stopping epoch.

  • The experiment is designed to distinguish overfitting caused by learning-rate decay from overfitting caused simply by continued training.
  • The ablation compares test accuracy with cosine learning-rate decay against constant learning rate after the stopping epoch.Experiments use ResNet-18 on CIFAR-100 with 40% label noise.
  • The plotted curves use a sliding window of 5 for smoothing.

F Detailed experimental setup

The experiments tune hyperparameters on an 80/20 train-validation split, retrain on all training data, and evaluate on an untouched test set. EMA-specific choices include searching decay rates, warming up decay, and sampling every 16 steps.

  • The protocol tunes hyperparameters on an 80/20 train-validation split, retrains on 100% of the training data, and reports test performance without using the test set for tuning.
  • For EMA, the search distinguishes early stopping by best accuracy from early stopping by lowest loss.
  • The EMA search uses τ ∈[0.968, 0.984, 0.992, 0.996, 0.998], warms up decay during initial steps, and samples every T = 16 steps.
  • Sampling every T > 1 reduces overhead but changes the effective decay, so equivalent decay rates must be adjusted across sampling periods.

G Sensitivity analysis to EMA decay rate α

Slower EMA decay produces later peak performance and higher accuracy when Batch Norm statistics are recomputed. Without Batch Norm recomputation, large averaging windows can fail and faster decay becomes preferable.

  • Slower EMA decay reaches peak performance later and achieves higher accuracy when Batch Norm statistics are recomputed after every epoch.The analysis covers the decay range τ ∈[0.968, 0.984, 0.992, 0.996, 0.998].
  • Without Batch Norm recomputation, the best decay is faster, indicating that weights tolerate larger averaging windows better than Batch Norm statistics.
  • The decay sensitivity analysis evaluates five EMA models per decay with and without Batch Norm recomputation, averaging results over three runs.
Loading 2411.18704v1…