Optimal higher order convergence rates for the MUON optimizer

Abstract

Gradient descent optimization methods are the key tools to train deep neural networks in artificial intelligence systems. In practice, one typically employs suitable accelerated gradient descent optimization methods such as the momentum optimizer, the Adam optimizer, or the MUON optimizer proposed by Jordan et al. in 2024. In 1964 Polyak showed that, in a neighborhood of a local minimizer, the momentum optimizer converges with the rate $(\sqrtχ+1)^{-1}(\sqrtχ-1)$, which is strictly faster than the rate $(χ+1)^{-1}(χ-1)$ of the standard gradient descent optimizer, where $χ\in(1,\infty)$ is the condition number (the ratio of the largest to the smallest eigenvalue) of the Hessian of the objective function at the local minimizer. Recently, the same faster rate has been established for the Adam optimizer. In this work we show that the MUON optimizer achieves the same higher speed of convergence and we establish the optimality of the obtained convergence rates. More specifically, we propose a precise concept for measuring the speed of convergence of gradient descent optimization methods, and we establish optimal convergence rates for the standard gradient descent, the momentum, the RMSprop, the Adam, and the MUON optimizers.