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.
Quentin Berthet, Peter Bartlett, Romuald Elie, Yu-Han Wu
原文摘要
We study variational inference in the presence of distributional symmetry, focusing on the interaction between two distinct group actions on a sample space: a symmetry action that leaves a reference distribution invariant, and a transport action that generates a parametric variational family by pushforward. By analyzing the algebraic structure of the total transformation group generated by both actions and the induced adjoint representation of the symmetry group on the transport Lie algebra, we establish that the Riemannian gradient of the Kullback--Leibler divergence at the true transport parameter is constrained to lie in the trivial representation of the symmetry group. Consequently, symmetry forces the gradient to vanish along all non-trivial representations. Combining this representation-theoretic constraint with parameter identifiability yields an orthogonal decomposition of the transport Lie algebra into three interpretable subspaces: an identifiability kernel (along which the objective is locally constant), a bias subspace (where the gradient may be nonzero due to model misspecification), and a symmetry-breaking subspace (where the gradient is forced to vanish by symmetry alone). Finally, we characterize the residual bias directions in terms of distributional mismatch between the base model and the standardized target, and derive an explicit Stein-type moment-matching condition that eliminates this bias, extending the location-scale results of (Margossian and Saul, 2025) to general Lie group transformation families.