On the Convergence Rate of LoRA Gradient Descent

Résumé technique : Sur le taux de convergence de la descente de gradient LoRA

Énoncé du problème
L'adaptation à faible rang (LoRA) est devenue une méthode standard pour le fine-tuning des grands modèles de langage en raison de son efficacité, mettant à jour uniquement les matrices d'adaptateur de faible rang $A$ et $B$ de sorte que les poids finaux soient $W_0 + BA$. Cependant, la compréhension théorique de sa convergence est limitée. Bien que la fonction de perte originale $\ell(W)$ soit généralement supposée être lisse au sens de Lipschitz, la fonction de perte reparamétrée $L(BA)$ n'est pas lisse au sens de Lipschitz par rapport aux matrices d'adaptateur $A$ et $B$. Ce manque de lissité empêche l'application directe des techniques classiques d'analyse de convergence. Les travaux théoriques existants se divisent en trois catégories : (1) des analyses asymptotiques dans des régimes infinis (par exemple, le noyau tangent neuronal ou une largeur infinie), qui ne fournissent pas de taux non asymptotiques pour des modèles finis ; (2) des analyses d'algorithmes modifiés (par exemple, GaLore, LoRA-One) qui ne traitent pas strictement la formulation LoRA originale ; et (3) des analyses de l'algorithme LoRA original qui reposent sur des hypothèses fortes et artificielles, telles que des matrices d'adaptateur uniformément bornées, pour imposer une lissité de Lipschitz. Par conséquent, le taux de convergence non asymptotique de l'algorithme de descente de gradient LoRA original, sans hypothèses de bornitude, reste une question ouverte.

Méthodologie
Les auteurs fournissent la première analyse de convergence non asymptotique de l'algorithme de descente de gradient LoRA original sous des hypothèses minimales : la fonction de perte originale est lisse au sens de Lipschitz et minorée. La preuve procède en trois étapes clés :

1. Reformulation via des matrices empilées : Le problème est reformulé en empilant les matrices d'adaptateur $B \in \mathbb{R}^{m \times r}$ et $A^T \in \mathbb{R}^{r \times n}$ en une seule variable $V \in \mathbb{R}^{(m+n) \times r}$. La mise à jour originale $BA$ est récupérée comme un bloc spécifique du produit extérieur $VV^T$. Cela permet de considérer l'optimisation comme une descente de gradient sur une nouvelle fonction $J(V) = L(E_1 V V^T E_2)$, où $E_1$ et $E_2$ sont des matrices extractrices.
2. Lemme de descente modifié : Puisque $J(V)$ n'est pas lisse au sens de Lipschitz (en raison du facteur multiplicatif $V$ dans le gradient), les auteurs dérivent un lemme de descente modifié. Ce lemme borne la diminution de la fonction $J(V_{t+1}) - J(V_t)$ en utilisant des termes d'ordre supérieur impliquant les normes des paramètres $\|V_t\|$ et du gradient.
3. Contrôle du taux d'apprentissage : Pour assurer une descente à chaque étape, le taux d'apprentissage $\eta_t$ est contrôlé dynamiquement en fonction de la norme actuelle des paramètres et de la norme du gradient. Plus précisément, $\eta_t$ est défini inversement proportionnel à $\|V_t\|^2 + \|\nabla L(E_1 V_t V_t^T E_2)\|$. Les auteurs prouvent que même si la norme des paramètres $\|V_t\|$ croît de manière non bornée (ce qui force $\eta_t$ à diminuer), la somme des taux d'apprentissage $\sum \eta_t$ diverge suffisamment lentement (comme une série harmonique) pour garantir la convergence, bien qu'à un rythme plus lent que la descente de gradient standard.

Contributions principales

• Taux de convergence non asymptotique : L'article démontre que la norme minimale du gradient sous la descente de gradient LoRA converge vers zéro à un taux de $O(1/\log T)$, où $T$ est le nombre d'itérations. Ce résultat tient sans supposer de matrices d'adaptateur bornées ni de lissité de Lipschitz de la fonction de perte reparamétrée.
• Récupération des taux standards : Les auteurs montrent que si les normes des adaptateurs sont bornées (une condition souvent supposée dans les travaux antérieurs), le taux de convergence s'améliore jusqu'au taux standard $O(1/T)$, correspondant au taux pour les fonctions lisses au sens de Lipschitz.
• Dépendance à la position : L'analyse révèle une « dépendance à la position » dans la dynamique LoRA : la convergence est ralentie si les itérés s'éloignent de l'origine (augmentant $\|V_t\|$) et accélérée s'ils se rapprochent de l'origine. Cela est attribué à la restructuration géométrique du paysage de perte par la paramétrisation LoRA, qui crée un point stationnaire à l'origine indépendamment de la structure de la perte originale.
• Schémas de taux d'apprentissage pratiques : Motivés par l'exigence théorique de contrôler $\eta_t$ en fonction des normes des paramètres et du gradient, les auteurs proposent et valident des calendriers de taux d'apprentissage adaptatifs pratiques ($\eta_{adapt}$, $\eta_{adapt2}$ et $\eta_{norm}$) qui s'ajustent en fonction de la norme des adaptateurs empilés ou de la valeur de la perte.

Résultats

• Théoriques : Le théorème 3.5 établit le taux de convergence $O(1/\log T)$ pour le cas général et $O(1/T)$ pour le cas borné. La preuve démontre que la relation récursive entre le taux d'apprentissage, la norme des paramètres et la norme du gradient introduit un ralentissement de la convergence par rapport à la descente de gradient standard.
• Empiriques : Des expériences ont été menées sur une régression logistique avec des embeddings CIFAR-10 et un ResNet-18 entraîné sur CIFAR-10.
  • Les taux d'apprentissage adaptatifs et normalisés dérivés de la théorie ont démontré une convergence plus rapide et plus stable par rapport aux taux d'apprentissage constants.
  • Dans les expériences de régression logistique, la norme des paramètres $\|V_t\|$ a été observée croître de manière non bornée sur 1000 époques, confirmant le scénario théorique où la convergence ralentit jusqu'à $O(1/\log T)$.
  • Dans les expériences ResNet-18, $\|V_t\|$ a convergé vers un ensemble borné, résultant en une convergence plus rapide.
  • Dans le fine-tuning de grands modèles de langage (LLM) (TinyLlama-1.1B), les taux adaptatifs ont montré des avantages significatifs lorsque les paramètres étaient initialisés avec de petites normes, mais se sont comportés de manière similaire aux taux constants lorsqu'ils étaient initialisés avec de grandes normes (où $\|V_t\|$ est grand et croît lentement par rapport à sa taille).

Importance
L'article prétend combler un écart significatif dans la compréhension théorique de LoRA en fournissant la première analyse de convergence non asymptotique pour l'algorithme original sans hypothèses de bornitude artificielles. Il clarifie que le taux de convergence de LoRA ne dépend pas du rang choisi $r$ mais est régi par l'interaction entre le taux d'apprentissage et la norme des itérés. Le travail souligne que la reparamétrisation de LoRA altère fondamentalement la géométrie de l'optimisation, pouvant mener à des solutions différentes par rapport au fine-tuning à rang complet. De plus, les stratégies de taux d'apprentissage proposées offrent une méthode pratique pour stabiliser et accélérer l'entraînement LoRA en tenant compte de la structure unique induite par la reparamétrisation à faible rang.