Adaptive Methods through the Lens of SDEs: Theoretical Insights on the Role of Noise

TL;DR

This paper develops novel stochastic differential equations to accurately model adaptive optimizers like SignSGD, RMSprop, and AdamW, revealing their dynamics, robustness, and the complex role of noise in training deep neural networks.

Contribution

It introduces new SDE models for adaptive optimizers, providing a deeper theoretical understanding and empirical validation of their behavior in deep learning.

Findings

SignSGD converges faster and is more robust to heavy-tail noise than SGD.

The role of noise in AdamW and RMSpropW is complex and differs from SignSGD.

The SDE models accurately predict optimizer behavior across various neural network architectures.

Abstract

Despite the vast empirical evidence supporting the efficacy of adaptive optimization methods in deep learning, their theoretical understanding is far from complete. This work introduces novel SDEs for commonly used adaptive optimizers: SignSGD, RMSprop(W), and Adam(W). These SDEs offer a quantitatively accurate description of these optimizers and help illuminate an intricate relationship between adaptivity, gradient noise, and curvature. Our novel analysis of SignSGD highlights a noteworthy and precise contrast to SGD in terms of convergence speed, stationary distribution, and robustness to heavy-tail noise. We extend this analysis to AdamW and RMSpropW, for which we observe that the role of noise is much more complex. Crucially, we support our theoretical analysis with experimental evidence by verifying our insights: this includes numerically integrating our SDEs using Euler-Maruyama discretization on various neural network architectures such as MLPs, CNNs, ResNets, and Transformers. Our SDEs accurately track the behavior of the respective optimizers, especially when compared to previous SDEs derived for Adam and RMSprop. We believe our approach can provide valuable insights into best training practices and novel scaling rules.