Long-Time Behavior of Zero-Noise Mean-Field Adam

Abstract

Adaptive methods such as Adam are the most commonly used methods for training neural networks, but their long-time behavior is not as well understood as that of non-adaptive methods. This is particularly true in the mean-field limit, where the dynamics of infinitely many neurons are described as a distribution over parameters which satisfies a self-consistency condition. In this paper, we study a deterministic continuous-time model of Adam in the mean-field setting. Under a condition relating the averaging rates of the momentum and second-moment variables, the size of the gradients, and a regularization parameter, we prove that the distribution of parameters converges to the set of equilibria. This condition extends the range known for finite-dimensional models, and we show that it is sharp. Under an additional Lojasiewicz-type inequality for distributions, we prove that individual parameter trajectories converge and obtain quantitative rates for this convergence. We verify the qualitative convergence assumptions for a regularized two-layer neural network. We also show that adding noise to the momentum prevents the system from approaching these equilibria and rules out many free-energy type Lyapunov functionals. Finally, numerical experiments illustrate the energy estimates, give evidence of persistent oscillations outside the sufficient parameter range, and compare the observed decay with the predicted rates.