Understanding Adam Requires Better Rotation Dependent Assumptions

Tianyue H. Zhang, Lucas Maes, Alexia Jolicoeur-Martineau, Ioannis Mitliagkas, Damien Scieur, Simon Lacoste-Julien, Charles Guille-Escuret

Keywords: Optimization, Adam, Deep Learning, Rotation Dependency

TL;DR: Adam's success in training large language models depends on its sensitivity to parameter space rotations, revealing the need for basis-dependent theoretical frameworks to explain its advantages.

Abstract: Despite its widespread adoption, Adam's advantage over Stochastic Gradient Descent (SGD) lacks a comprehensive theoretical explanation. This paper investigates Adam's sensitivity to rotations of the parameter space. We demonstrate that Adam's performance in training transformers degrades under random rotations of the parameter space, indicating a crucial sensitivity to the choice of basis. This reveals that conventional rotation-invariant assumptions are insufficient to capture Adam's advantages theoretically. To better understand the rotation-dependent properties that benefit Adam, we also identify structured rotations that preserve its empirical performance. We then examine the rotation-dependent assumptions in the literature, evaluating their adequacy in explaining Adam's behaviour across various rotation types. This work highlights the need for new, rotation-dependent theoretical frameworks to understand Adam's empirical success in modern machine learning fully.

Submission Number: 85