# Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA

arXiv: https://arxiv.org/abs/2605.07959 · HTML: https://mlsanity.com/paper/2605.07959/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 38 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Infinite attention: NNGP and NTK for deep attention networks](https://mlsanity.com/paper/2006.10540/) (2020)
- [LoRA vs Full Fine-tuning: An Illusion of Equivalence](https://mlsanity.com/paper/2410.21228/) (2024)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers](https://mlsanity.com/paper/2409.10559/) (2024)
- [LoRA Training in the NTK Regime has No Spurious Local Minima](https://mlsanity.com/paper/2402.11867/) (2024)
- [On Learning Rates and Schr"odinger Operators](https://mlsanity.com/paper/2004.06977/) (2020)
- [Generative AI for fast and accurate statistical computation of fluids](https://mlsanity.com/paper/2409.18359/) (2024)
- [Provably learning a multi-head attention layer](https://mlsanity.com/paper/2402.04084/) (2024)
- [LoRA Training Provably Converges to a Low-Rank Global Minimum or It Fails Loudly (But it Probably Won't Fail)](https://mlsanity.com/paper/2502.09376/) (2025)
- [Global Convergence of SGD On Two Layer Neural Nets](https://mlsanity.com/paper/2210.11452/) (2022)
- [Global Convergence in Training Large-Scale Transformers](https://mlsanity.com/paper/2410.23610/) (2024)
- [Global Convergence of SGD For Logistic Loss on Two Layer Neural Nets](https://mlsanity.com/paper/2309.09258/) (2023)
- [Langevin Monte-Carlo Provably Learns Depth Two Neural Nets at Any Size and Data](https://mlsanity.com/paper/2503.10428/) (2025)

## Built on by (most cited first, 1 of 1)

- [Certification Frontiers for Gaussian LoRA: Independent Priors, Posterior Risk, and Prediction-Preserving Balancing](https://mlsanity.com/paper/2609.32271/) (2026)

## Abstract (arXiv)

Transformers have revolutionized machine learning and deploying attention layers in the model is increasingly standard across a myriad of applications. Further, for large models, it is common to implement Low Rank Adaptation (LoRA), whereby a factorized parameterization of them is trained, to achieve a surprisingly beneficial accuracy-size trade-off. In this work, via a unified framework we rigorously establish trainability of such models under stochastic methods. We prove that for a class of mild regularizations, the empirical regression loss on a attention layer and LoRA on a shallow neural net, both induce Poincar'e inequality for the corresponding Gibbs' measure. Crucially, we show that the Poincar'e constant is free of the data dimension for LoRA and for multi-head attention it is free of the head dimensions, Then it follows, via invoking recent results, that a certain SDE, which mimics the SGD, minimizes the corresponding losses. In both the cases, our first-of-its-kind results of trainability on attention and nets, do not use any assumptions on the data or the model size.