MatryoshkaLoRA: Learning Accurate Hierarchical Low-Rank Representations for LLM Fine-Tuning

Résumé Technique : MATRYOSHKALORA

Énoncé du Problème
À mesure que les modèles d'apprentissage profond évoluent vers des milliards de paramètres, le coût computationnel du fine-tuning demeure un obstacle majeur à leur déploiement. Bien que l'Adaptation à Faible Rang (LoRA) soit devenue la norme pour le fine-tuning efficace en paramètres (PEFT), elle impose une contrainte structurelle : le rang $r$ doit être prédéfini. Trouver l'équilibre optimal entre efficacité et performance nécessite actuellement des recherches exhaustives par grille sur plusieurs exécutions d'entraînement, chacune avec un rang statique différent.

Les solutions existantes d'adaptation de rang, telles que DyLoRA, tentent d'atténuer ce problème en échantillonnant des rangs durant l'entraînement à partir d'une distribution prédéfinie. Cependant, les auteurs soutiennent que ces méthodes produisent souvent des résultats sous-optimaux aux rangs supérieurs. Plus précisément, DyLoRA contraint l'apprentissage à un seul rang échantillonné aléatoirement $k$ par passage avant, ce qui signifie que les rangs $r > k$ ne reçoivent aucun signal de gradient. Ce manque de signaux de gradient cohérents à travers toute la hiérarchie des rangs rend ces méthodes inefficaces en termes de données et empêche l'apprentissage de ce que les auteurs qualifient de « caractéristiques hiérarchiques à faible rang ».

Méthodologie
L'article propose MATRYOSHKALORA, un cadre d'entraînement général inspiré de l'apprentissage de représentations Matryoshka. L'innovation centrale réside dans l'insertion d'une matrice diagonale fixe et soigneusement conçue $P$ entre les adaptateurs LoRA existants $A$ et $B$.

• Mécanisme : Au lieu d'utiliser uniquement une seule tranche d'adaptateurs (comme dans DyLoRA) ou les adaptateurs complets (comme dans LoRA standard), MATRYOSHKALORA utilise toutes les tranches possibles $A_r$ et $B_r$ (où $r \in \{1, 2, 4, ..., R\}$) au sein d'un seul passage avant. Le passage avant est formulé comme suit :
  $$Y = x \left( W_0 + \sum_{r \in S} s_r A_r B_r \right)$$
  Pour implémenter cela efficacement sans stocker de masques booléens ni boucler sur les rangs, les auteurs dérivent une formulation simplifiée utilisant une matrice diagonale $P$ (représentée comme un vecteur) :
  $$Y = x (W_0 + (A * P)B)$$
  Ici, $*$ désigne la multiplication élément par élément des lignes de $A$ par le vecteur $P$. Le vecteur $P$ est précalculé en fonction de l'ensemble des rangs d'entraînement $S$ et des facteurs d'échelle $s_r$, garantissant que les représentations de dimension inférieure sont contenues comme préfixes dans celles de dimension supérieure.

• Flux de Gradient : Cette modification assure que tous les sous-rangs intègrent efficacement les informations de gradient disponibles. Contrairement à DyLoRA, où seul le rang échantillonné reçoit des mises à jour, MATRYOSHKALORA met à jour toute la structure de l'adaptateur en utilisant les gradients calculés à partir de la somme de toutes les contributions de rang.

• Généralisation : Le cadre est conçu pour être général. En manipulant le vecteur $P$, la méthode peut retrouver LoRA standard (où $P$ est un vecteur constant) ou DyLoRA (où $P$ est un vecteur creux correspondant à un rang échantillonné), unifiant ainsi ces approches sous une paramétrisation partagée.

• Métrique d'Évaluation : Les auteurs introduisent la Surface Sous la Courbe de Précision par Rang (AURAC), une métrique qui évalue la performance des adaptateurs sur un ensemble de rangs. L'AURAC utilise la règle des trapèzes pour résumer le compromis rang-performance, pondérant davantage les rangs plus élevés pour refléter l'attente que les rangs plus grands devraient généralement atteindre des performances supérieures.

Contributions Clés

1. Cadre : Un cadre général inspiré de Matryoshka pour l'entraînement d'adaptateurs LoRA où les préfixes de rang inférieur restent précis et utilisables indépendamment au moment de l'inférence.
2. Simplification : Une dérivation montrant que l'apprentissage de représentations hiérarchiques à faible rang conduit naturellement à un pondération diagonale simple entre les adaptateurs LoRA standard, maintenant l'implémentation proche de LoRA standard.
3. MATRYOSHKALORA : Une instance concrète de ce cadre qui apprend des préfixes imbriqués précis à faible rang via la matrice diagonale $P$, permettant l'évaluation ou le déploiement à plusieurs rangs à partir d'un seul point de contrôle sans réentraînement.
4. Métrique : La proposition de l'AURAC pour évaluer de manière cohérente la performance des adaptateurs hiérarchiques à faible rang.

Résultats Expérimentaux
Les auteurs ont évalué MATRYOSHKALORA sur les modèles Llama-3.2-1B-Instruct et Llama-3.1-8B-Instruct en utilisant les ensembles de données GSM-8k, OpenPlatypus, ARC-C et HellaSwag.

• Performance : MATRYOSHKALORA a systématiquement surpassé à la fois LoRA standard et DyLoRA sur l'ensemble du spectre de rangs.
  • Sur GSM-8k avec Llama-3.2-1B, MATRYOSHKALORA a atteint jusqu'à 39 % de précision au rang 16 et 38 % au rang 32, dépassant les meilleurs résultats de LoRA (35 %) et de DyLoRA (35,5 %).
  • Sur GSM-8k avec Llama-3.1-8B (3 coups), la méthode a atteint 78,6 % de précision au rang 128, soit une amélioration de plus de 4 % par rapport aux modèles de base et au modèle pré-entraîné.
  • Sur HellaSwag, MATRYOSHKALORA a montré des gains réguliers avec l'augmentation du rang, culminant à 62,8 % pour le rang 256, soit une augmentation de plus de 3 % par rapport aux bases.
• AURAC : La méthode a obtenu des scores AURAC significativement plus élevés (par exemple, 38,4 % contre ~35 % pour les bases sur GSM-8k 1B), indiquant des gains plus cohérents sur tous les rangs plutôt qu'une performance de pointe à un rang spécifique.
• Efficacité : L'approche ajoute une surcharge minimale, ne nécessitant que la multiplication des colonnes de l'adaptateur par le vecteur $P$, résultant en des coûts de mémoire et d'exécution identiques à LoRA et DyLoRA.
• Ablation de l'Échelle : Des expériences sur les paramètres d'échelle ($s_k \in \{1, 1/k, 1/\sqrt{k}\}$) ont révélé que $s_k=1$ permet un apprentissage efficace avec des taux d'apprentissage plus faibles par rapport aux autres facteurs d'échelle, qui nécessitaient des taux d'apprentissage nettement plus élevés pour atteindre une précision similaire.

Signification et Revendications
L'article revendique que MATRYOSHKALORA apprend avec succès des représentations hiérarchiques à faible rang plus précises que les approches d'adaptation de rang antérieures. La signification de ce travail réside dans trois avantages principaux :

1. Réduction des Coûts : Il élimine le besoin de multiples exécutions d'entraînement pour effectuer une recherche par grille de différents rangs.
2. Déploiement Flexible : Il permet le déploiement de rangs ciblés sur différents appareils en fonction de la puissance de calcul, fournissant des adaptateurs haute performance à chaque rang $r$.
3. Adaptation Dynamique : Il prend en charge la sélection dynamique de rangs sous des charges de cluster variables pour servir les demandes au même rythme avec une baisse de précision minimale.

Les auteurs concluent que leur méthode atteint des compromis précision-performance supérieurs à travers les rangs et fournit un cadre unifié pour LoRA, DyLoRA et les adaptateurs de style Matryoshka. Ils reconnaissent des limites, notant que le processus d'évaluation nécessite de nombreuses évaluations de modèle par rang, augmentant le temps d'exécution pour l'évaluation des performances, et qu'ils n'ont pas encore expérimenté la sensibilité du rang par couche (attribution de rangs différents à différentes couches).