Convergent Stochastic Training of Attention and Understanding LoRA

2026-05-08

AI Summary

This work addresses the lack of theoretical guarantees for the trainability of attention mechanisms and LoRA under stochastic optimization by proposing a unified analytical framework. By integrating Gibbs measures, the Poincaré inequality, and stochastic differential equations, the study establishes, for the first time, rigorous convergence guarantees for attention layers and shallow neural networks equipped with LoRA under stochastic gradient descent—without requiring assumptions on data distribution or model scale. The key contribution lies in proving that the empirical loss satisfies the Poincaré inequality, thereby providing a theoretical foundation for their stochastic trainability and offering solid justification for the practical efficacy of efficient fine-tuning strategies.

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 any mild regularization, the empirical regression loss on a attention layer and LoRA on a shallow neural net, both induce Poincaré inequality for the corresponding Gibbs' measure. 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 rely on any assumptions on the data or the size of the architecture.

Problem

• Attention
• LoRA
• Stochastic Training
• Trainability
• Transformers

Innovation

• Stochastic Training
• Attention Mechanism
• LoRA
• Poincaré Inequality
• Trainability