Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

• URL: http://arxiv.org/abs/2603.01514v1
• Date: Mon, 02 Mar 2026 06:44:54 GMT
• Title: Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning
• Authors: Gautam Goel, Mahdi Soltanolkotabi, Peter Bartlett,
• Abstract summary: We train dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression.<n>We show that a simple first-order gradient descent can converge to the globally optimal self-attention parameters.
• Score: 17.65459083031186
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.

Related papers

• Towards a Unified Analysis of Neural Networks in Nonparametric Instrumental Variable Regression: Optimization and Generalization [66.08522228989634]
  We establish the first global convergence result of neural networks for two stage least squares (2SLS) approach in nonparametric instrumental variable regression (NPIV)<n>This is achieved by adopting a lifted perspective through mean-field Langevin dynamics (MFLD)
  arXiv  Detail & Related papers  (2025-11-18T17:51:17Z)
• Closed-Form Last Layer Optimization [72.49151473937319]
  Under a squared loss, the optimal solution to the linear last layer weights is known in closed-form.<n>We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer.
  arXiv  Detail & Related papers  (2025-10-06T09:14:39Z)
• Gradient-Normalized Smoothness for Optimization with Approximate Hessians [37.1630298053787]
  We develop algorithms that use approximate second-order information combined with the gradient regularization technique to achieve fast global convergence rates.<n>We demonstrate direct applications of our results in logistic regression problems with Hessians, as well as in non-continuous soft convergence optimization using Fisher and Gauss.
  arXiv  Detail & Related papers  (2025-06-16T17:19:34Z)
• Convergence of Implicit Gradient Descent for Training Two-Layer Physics-Informed Neural Networks [4.554284689395686]
  implicit gradient descent (IGD) outperforms the common gradient descent (GD) algorithm in handling certain multi-scale problems.<n>We show that IGD converges to a globally optimal solution at a linear convergence rate.
  arXiv  Detail & Related papers  (2024-07-03T06:10:41Z)
• Stochastic Zeroth-Order Optimization under Strongly Convexity and Lipschitz Hessian: Minimax Sample Complexity [59.75300530380427]
  We consider the problem of optimizing second-order smooth and strongly convex functions where the algorithm is only accessible to noisy evaluations of the objective function it queries. We provide the first tight characterization for the rate of the minimax simple regret by developing matching upper and lower bounds.
  arXiv  Detail & Related papers  (2024-06-28T02:56:22Z)
• Implicit Bias and Fast Convergence Rates for Self-attention [26.766649949420746]
  We study the fundamental optimization principles of self-attention, the defining mechanism of transformers.<n>We analyze the implicit bias of gradient-baseds in a self-attention layer with a decoder in a linear classification.
  arXiv  Detail & Related papers  (2024-02-08T15:15:09Z)
• Zeroth-Order Hybrid Gradient Descent: Towards A Principled Black-Box Optimization Framework [100.36569795440889]
  This work is on the iteration of zero-th-order (ZO) optimization which does not require first-order information. We show that with a graceful design in coordinate importance sampling, the proposed ZO optimization method is efficient both in terms of complexity as well as as function query cost.
  arXiv  Detail & Related papers  (2020-12-21T17:29:58Z)
• GTAdam: Gradient Tracking with Adaptive Momentum for Distributed Online Optimization [4.103281325880475]
  This paper deals with a network of computing agents aiming to solve an online optimization problem in a distributed fashion, by means of local computation and communication, without any central coordinator. We propose the gradient tracking with adaptive momentum estimation (GTAdam) distributed algorithm, which combines a gradient tracking mechanism with first and second order momentum estimates of the gradient. In these numerical experiments from multi-agent learning, GTAdam outperforms state-of-the-art distributed optimization methods.
  arXiv  Detail & Related papers  (2020-09-03T15:20:21Z)
• Cogradient Descent for Bilinear Optimization [124.45816011848096]
  We introduce a Cogradient Descent algorithm (CoGD) to address the bilinear problem. We solve one variable by considering its coupling relationship with the other, leading to a synchronous gradient descent. Our algorithm is applied to solve problems with one variable under the sparsity constraint.
  arXiv  Detail & Related papers  (2020-06-16T13:41:54Z)
• Optimizing generalization on the train set: a novel gradient-based framework to train parameters and hyperparameters simultaneously [0.0]
  Generalization is a central problem in Machine Learning. We present a novel approach based on a new measure of risk that allows us to develop novel fully automatic procedures for generalization.
  arXiv  Detail & Related papers  (2020-06-11T18:04:36Z)
• Towards Better Understanding of Adaptive Gradient Algorithms in Generative Adversarial Nets [71.05306664267832]
  Adaptive algorithms perform gradient updates using the history of gradients and are ubiquitous in training deep neural networks. In this paper we analyze a variant of OptimisticOA algorithm for nonconcave minmax problems. Our experiments show that adaptive GAN non-adaptive gradient algorithms can be observed empirically.
  arXiv  Detail & Related papers  (2019-12-26T22:10:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.