Convergence of Rotation-based Matrix Optimizers: A Unified Analysis of SOAP, Conda, and SPlus

Abstract

In this work, we develop a unified theoretical framework for analyzing the convergence of rotation-based matrix optimizers, which apply orthogonal transformations to map the momentum into a rotated space, perform coordinate-wise or normalized updates there, and then rotate the updates back. Our framework encompasses prominent matrix optimizers, including SOAP, Conda, and truncated SPlus, as well as matrix-parameterized Adam as a special case. As our main result, we establish, for the first time, the convergence rate of SOAP with sharp dimensional dependence, as well as the convergence rate of Adam measured by the nuclear norm. Our framework also covers a more general variant of rotation-based optimizers that allows arbitrary orthogonal rotation matrices, allowing these matrices to depend on the current stochastic gradient at each iteration. Technically, our framework leverages a row-column trace-control argument that converts elementwise bounds into bounds on two diagonal control matrices, thereby improving the dimension dependence of the resulting bound.