Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits
This work provides a simple plug-in method to enhance deep learning model performance across multiple metrics, though it is incremental as it builds on existing weight averaging techniques.
The authors systematically studied the Exponential Moving Average (EMA) of weights in deep learning, finding that it improves generalization, robustness to noisy labels, prediction consistency, calibration, and transfer learning compared to standard SGD solutions.
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.