Understanding SGD with Exponential Moving Average: A Case Study in Linear Regression
This provides theoretical insights into EMA's effectiveness, which is incremental but addresses a known bottleneck in understanding popular deep learning training techniques.
The paper tackles the lack of theoretical understanding of exponential moving average (EMA) in SGD by establishing risk bounds for high-dimensional linear regression, showing that EMA reduces variance error compared to non-averaged SGD and that its bias error decays exponentially across eigen-subspaces.
Exponential moving average (EMA) has recently gained significant popularity in training modern deep learning models, especially diffusion-based generative models. However, there have been few theoretical results explaining the effectiveness of EMA. In this paper, to better understand EMA, we establish the risk bound of online SGD with EMA for high-dimensional linear regression, one of the simplest overparameterized learning tasks that shares similarities with neural networks. Our results indicate that (i) the variance error of SGD with EMA is always smaller than that of SGD without averaging, and (ii) unlike SGD with iterate averaging from the beginning, the bias error of SGD with EMA decays exponentially in every eigen-subspace of the data covariance matrix. Additionally, we develop proof techniques applicable to the analysis of a broad class of averaging schemes.