Adam Exploits _-geometry of Loss Landscape via Coordinate-wise Adaptivity

Code

• github.com/mohamad-amin/adam-coordinate-adaptivity
  Officialjax★ 4

Abstract

Adam outperforms SGD when training language models. Yet this advantage is not well-understood theoretically -- previous convergence analysis for Adam and SGD mainly focuses on the number of steps T and is already minimax-optimal in non-convex cases, which are both O(T^-1/4). In this work, we argue that the exploitation of nice _-geometry is the key advantage of Adam over SGD. More specifically, we give a new convergence analysis for Adam under novel assumptions that loss is smooth under _-geometry rather than the more common _2-geometry, which yields a much better empirical smoothness constant for GPT-2 and ResNet models. Our experiments confirm that Adam performs much worse when the favorable _-geometry is changed while SGD provably remains unaffected. We also extend the convergence analysis to blockwise Adam under novel blockwise smoothness assumptions.

Reproductions