Exponential weight averaging as damped harmonic motion

TL;DR

This paper establishes a physical analogy between exponential moving averages in deep learning and damped harmonic motion, leading to a new algorithm called BELAY that improves training stability and performance.

Contribution

It introduces a novel physical interpretation of EMA as a damped harmonic system and proposes the BELAY algorithm for enhanced training stability.

Findings

BELAY outperforms standard EMA in stability and accuracy.

The physical analogy provides new insights into EMA dynamics.

Empirical results show improved convergence with BELAY.

Abstract

The exponential moving average (EMA) is a commonly used statistic for providing stable estimates of stochastic quantities in deep learning optimization. Recently, EMA has seen considerable use in generative models, where it is computed with respect to the model weights, and significantly improves the stability of the inference model during and after training. While the practice of weight averaging at the end of training is well-studied and known to improve estimates of local optima, the benefits of EMA over the course of training is less understood. In this paper, we derive an explicit connection between EMA and a damped harmonic system between two particles, where one particle (the EMA weights) is drawn to the other (the model weights) via an idealized zero-length spring. We then leverage this physical analogy to analyze the effectiveness of EMA, and propose an improved training algorithm, which we call BELAY. Finally, we demonstrate theoretically and empirically several advantages enjoyed by BELAY over standard EMA.