Saved in:
| Main Authors: | , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2603.01514 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866918365740662784 |
|---|---|
| author | Goel, Gautam Soltanolkotabi, Mahdi Bartlett, Peter |
| author_facet | Goel, Gautam Soltanolkotabi, Mahdi Bartlett, Peter |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2603_01514 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning Goel, Gautam Soltanolkotabi, Mahdi Bartlett, Peter Machine Learning 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. |
| title | Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2603.01514 |