Fast Global Convergence via Preconditioning

View a PDF of the paper titled Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning, by Gautam Goel and 2 other authors

View PDF
HTML (experimental)

Abstract:We study the training dynamics of gradient descent in a multilayer softmax self-attention network trained to perform linear regression and show that in the infinite-data limit this regression problem is equivalent to a matrix factorization problem. This connection motivates the design of Hera, a new first-order optimization algorithm for training self-attention networks. Key features of Hera include a reparameterization of the model arising from the matrix factorization geometry and a rescaling of the parameters which ensures that the empirical loss gradients remain close to the auxiliary matrix factorization gradients throughout training. Hera updates in the reparameterized model are equivalent to a certain preconditioned update in the original parameterization. We also propose a variant of Hera called Hera-PSD; our main theoretical result is that under small Gaussian initialization the excess risk of the model trained by Hera-PSD decreases exponentially fast in the number of gradient descent iterations, up to an additional error term that decreases as $1/n$, up to logarithmic factors, in the size of the training set. A particularly appealing aspect of our excess risk bound is that it does not degrade as the number of self-attention layers increases. In a stylized time-series prediction experiment in a nonlinear dynamical system, we find that a model trained by Hera achieves $98.1\%$ lower mean-squared test error than a model trained by AdamW and $97.7\%$ lower than a model trained by Muon at the same iteration budget.