Convergent Stochastic Training of Attention and Understanding LoRA

TL;DR

This paper establishes the trainability of attention layers and LoRA in large models using stochastic methods, providing rigorous theoretical guarantees without data or size assumptions.

Contribution

It offers a unified theoretical framework proving trainability of attention and LoRA layers via stochastic differential equations, independent of data or architecture size.

Findings

Proves Poincaré inequality for Gibbs' measure in attention and LoRA models.

Shows SDE mimicking SGD minimizes the loss in these models.

First rigorous results on trainability of attention and LoRA layers.

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\'e 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.