Convergence of Steepest Descent and Adam under Non-Uniform Smoothness

Recent work has analyzed the convergence of first-order methods under non-uniform smoothness assumptions that better model the loss landscape in machine learning tasks. We generalize this assumption to objectives whose curvature is an affine function of the objective value. This property is satisfied by a broad class of problems, including logistic regression, generalized linear models with a logistic link function, softmax policy gradient in reinforcement learning, and a class of neural networks. Under this assumption and gradient domination conditions, we establish a general convergence rate for the steepest descent method, and deterministic, diagonal variants of RMSProp and Adam. Our results imply that for logistic regression on separable data and the softmax policy gradient objective, sign GD converges linearly and is provably faster than GD. Furthermore, we show that for a class of two-layer neural networks on separable data, RMSProp and Adam can converge at a linear rate with a constant step-size and momentum parameter. Finally, we present a lower bound demonstrating that, under our assumption, RMSProp and Adam are provably faster than AdaGrad, AMSGrad, gradient descent, and heavy-ball momentum.

Gradient descent (GD) is a standard algorithm for training machine learning models. However, in practice, large models are trained with adaptive methods such as RMSProp and Adam. These methods often work better than GD, but it is still not fully understood when they are truly faster, or why. In this paper, we identify a rich class of machine learning problems where RMSProp and Adam can have a real and provable advantage. This family ranges from simple problems, such as logistic regression, to more complex ones, such as neural networks and reinforcement learning objectives. For example, on logistic regression with separable data, we prove that Adam converges faster than classical methods such as GD and heavy-ball momentum. It also outperforms other adaptive methods, such as AdaGrad and AMSGrad, which are less commonly used in modern practice. Overall, our work helps explain why Adam and RMSProp are so effective in machine learning.