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

TL;DR

This work systematically studies the Exponential Moving Average (EMA) of weights as an external, lightweight plug in to SGD training. It demonstrates that EMA reduces parameter noise, induces implicit regularization, and delivers better generalization, robustness to label noise, calibration, and transferability, with notably strong early performance that explains EMA based teacher setups. The authors provide practical guidance on hyperparameter tuning, including cosine annealing and early stopping, and reveal interactions with Batch Normalization that shape the feasible averaging windows. Overall, EMA offers a simple, effective means to improve deep learning models without altering the training trajectory. The study confirms that EMA solutions differ from last-iterate SGD solutions and often yield more robust and transferable representations.

Abstract

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.

Figures (10)

• Figure 1: CIFAR-100 on ResNet-18. Left: EMA vs SGD baseline, and learning rate ($\eta$). EMA is the best among the 5 EMA models at any given epoch, without recomputing BN stats (i.e. the maximum among EMA models plotted on \ref{['fig:ema_bn_vs_not']}). We observe that EMA dominates momentum SGD and has a good performance since early on. EMA peaks at epoch 150, at the optimal$\eta$, and then deteriorates. Right: Breakdown of the 5 EMA models per decay (with and without BN recomputation after every epoch). EMAs with the largest averaging windows fail unless BN stats are recomputed. Sliding window of 5 used for smoothing. All results are the mean of 3 runs.
• Figure 2: CIFAR-100N on ResNet-34. EMA vs SGD baseline, and learning rate $\eta$. EMA dominates SGD throughout training and peaks at epoch 100 ($\eta = 0.4$), greatly outperforming the best SGD model ($+9.65$ pp). Training on data with 40% of label noise, evaluating on clean test set, mean of 3 runs, $\alpha=0.998$.
• Figure 3: Selection of best-performing methods on CIFAR-10N (Worse) and CIFAR-100N, with $40\%$ label noise in train data, using a ResNet-34. Ours is highlighted, all other results are from wei2022learning. Leaderboard available at http://www.yliuu.com/web-cifarN/Leaderboard.html
• Figure 4: CIFAR-100 on ResNet-18, with step decay of the learning rate by a factor of 5 at epochs $[60, 120, 160]$. At each epoch we report the best EMA out of the 5 parallel EMAs kept, and do not recompute BN stats.
• Figure 5: EMA and momentum SGD training dynamics for the different datasets and models used in our work. Learning rate ($\mu$) follows a cosine annealing. Training on the full dataset, results are mean of 3 runs. At each epoch we report the best EMA out of the 5 parallel EMAs kept and do not recompute BN stats.