Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA
Abstract
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é inequality for the corresponding Gibbs' measure. Crucially, we show that the Poincaré 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.
Related
- Understanding axial attention in TSFMs through single location regression
- Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives
- A discrete-continuous-time stochastic framework for attention-deficit/hyperactivity disorder-related symptom modeling and pharmacological intervention via methylphenidate
- Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
- Training with ML-labeled Data via Stochastic First- and Zeroth-Order Optimization
BibTeX
@misc{sun2026convergent,
title = {{Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA}},
author = {Zhengkai Sun and Dibyakanti Kumar and Alejandro F Frangi and Anirbit Mukherjee and Mingfei Sun},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2605.07959},
eprint = {2605.07959},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.07959},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA
AU - Zhengkai Sun
AU - Dibyakanti Kumar
AU - Alejandro F Frangi
AU - Anirbit Mukherjee
AU - Mingfei Sun
PY - 2026
DA - 2026/09/06
PB - arXiv
DO - 10.48550/arxiv.2605.07959
UR - https://arxiv.org/abs/2605.07959
AB - 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é inequality for the corresponding Gibbs' measure. Crucially, we show that the Poincaré 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.
ER -