Riemannian Gradient Descent for Low-Rank Architectures

Résumé Technique : Descente de Gradient Riemannienne pour les Architectures à Faible Rang

Énoncé du Problème
Dans l'apprentissage profond moderne, les paramètres du modèle sont souvent des matrices contraintes par un faible rang pour réduire le coût computationnel et le nombre de paramètres (par exemple, dans les mécanismes d'attention multi-têtes où $W = AB^T$). Les mises à jour standard de l'entraînement utilisent les gradients euclidiens pour les facteurs $A$ et $B$ de manière indépendante, ignorant la structure géométrique de l'ensemble des matrices de rang $r$. Cette publication examine si le fait de traiter la matrice de rang contraint $W$ comme un point sur une sous-variété lisse et d'appliquer une optimisation riemannienne peut produire une meilleure convergence ou de meilleures solutions. Le défi central est que la factorisation $W=AB^T$ n'est pas unique (invariante sous $A \to AS, B \to BS^{-T}$) et que l'ensemble des matrices de rang $r$ est une variété, et non un espace vectoriel. Les auteurs visent à restreindre les modifications uniquement à l'optimiseur, en évitant de modifier l'architecture du modèle, tout en maintenant une complexité par itération comparable aux méthodes standards ($O(mr^2 + nr^2)$).

Méthodologie
Les auteurs proposent un cadre de descente de gradient riemannien adapté aux paramètres factorisés en rang. L'itération générale comprend :

1. La récupération du gradient euclidien ambiant $G_W$ à partir des gradients des facteurs $G_A$ et $G_B$.
2. La projection de ce gradient sur l'espace tangent de la variété spécifique pour obtenir le gradient riemannien $G_W^R$.
3. La mise à jour de l'élan (momentum) et la normalisation de la taille du pas.
4. Le déplacement le long de la géodésique (en utilisant une rétraction) vers l'itération suivante.
5. Le transport parallèle du vecteur d'élan vers le nouvel espace tangent.

La paperelle instancie ce modèle à travers dix conceptions algorithmiques spécifiques catégorisées par géométrie et contraintes :

1. Géométries pour Matrices à Rang Fixe ($R_{r}^{m \times n}$) :

   • Géométrie Imbriquée : Utilise la métrique de Frobenius ambiante. La rétraction est effectuée via la SVD tronquée de rang $r$ de $W + \Xi$. Le transport parallèle utilise la projection métrique.
   • Géométrie de Quotient : Traite la variété comme un quotient de l'espace produit des facteurs $(A, B)$. La métrique est définie « en haut » sur les facteurs, induisant une métrique sur le quotient. La rétraction est additive sur les facteurs ; le transport implique une projection horizontale.

2. Géométries pour Isométries Partielles de Rang $r$ ($P_{r}^{m \times n}$) :

   • Ici, $W$ est contrainte de telle sorte que ses valeurs singulières non nulles soient toutes égales à 1 (équivalent à ce que $A$ et $B$ soient orthonormaux en colonnes).
   • Géométrie Imbriquée : Restreint la métrique de Frobenius. La rétraction utilise la SVD tronquée avec des contraintes d'orthogonalité spécifiques.
   • Géométrie de Quotient : Utilise le produit des variétés de Stiefel ($St_{m,r} \times St_{n,r}$) comme espace « en haut ». Les rétractions utilisent la décomposition QR ou polaire.
   • Géométrie Canonique : Utilise la métrique canonique sur les facteurs de Stiefel, induisant une métrique produit sur le quotient.

3. Partage de Poids (Variétés de Grilles) :

   • Étend les points ci-dessus à des architectures où les facteurs sont partagés entre les blocs de lignes et de colonnes (par exemple, le Group-Query Attention).
   • Pour le Rang Fixe, les auteurs montrent qu'empiler les facteurs dans des matrices plus larges permet d'utiliser des algorithmes simples (singleton), à condition que la préservation du rang soit vérifiée bloc par bloc.
   • Pour les Isométries Partielles, l'empilement simple échoue en raison des contraintes d'orthogonalité. Les auteurs développent des algorithmes plus complexes, incluant une approche « hybride » pour le cas imbriqué où ils remontent vers la géométrie de quotient pour effectuer les rétractions et le transport, puis redescendent par projection.

Contributions Clés

• Exploration Systématique : Le papier fournit une enquête exhaustive et une implémentation de dix variantes distinctes d'optimisation riemannienne pour les paramètres matriciels à faible rang, couvrant à la fois les contraintes de rang fixe et d'isométrie partielle, ainsi que les scénarios de partage de poids.
• Détails d'Implémentation : Il détaille les rétractions spécifiques (SVD, QR, Polaire) et les transports vectoriels requis pour chaque géométrie, traitant de la non-unicité des factorisations et de la préservation des contraintes de la variété (rang et orthogonalité).
• Évaluation Empirique : Les méthodes sont appliquées aux matrices Q, K, V, O dans des modèles de langage à petite échelle (6 couches, dimension d'embedding de 512) en utilisant le jeu de données FineWeb.

Résultats Expérimentaux

• Performance : Après ajustement des taux d'apprentissage, les méthodes riemanniennes n'ont pas surpassé de manière concluante la référence AdamW. Dans les configurations Multi-Head Attention (MHA) et Group-Query Attention (GQA), les différences de perte étaient souvent à l'intérieur de la variance des graines aléatoires (random seeds).
• Hyperparamètres : La normalisation de la métrique a apporté de légères améliorations (0,01–0,05 nats). Un élan important ($\nu=7/8$) a généralement nui à la performance ou n'a apporté que des gains négligeables par rapport à la référence.
• Stabilité : Certaines exécutions ont présenté des pics de perte (par exemple, autour de l'étape 6660) qui se sont rétablis, souvent corrélés à des données structurées dans l'ensemble d'entraînement.
• Complexité : Les auteurs reconnaissent que bien que les méthodes soient mathématiquement rigoureuses, leur complexité computationnelle par itération est supérieure à celle d'AdamW, ce qui les rend moins attractives pour les expériences à petite échelle menées.

Signification et Revendications
Le papier maintient une position modeste concernant sa nouveauté et son impact immédiat :

• Redécouverte : Les auteurs déclarent explicitement que toute nouveauté apparente est « presque certainement une redécouverte », s'appuyant fortement sur les manuels établis (Absil et al., Boumal) et la littérature antérieure.
• Motivation Heuristique : Le travail est motivé par l'heuristique selon laquelle une meilleure prise en compte des relations entre paramètres (via la géométrie de la variété) pourrait conduire à une convergence plus rapide ou à de meilleures solutions, bien que cela n'ait pas été prouvé de manière définitive dans les expériences.
• Potentiel Futur : Les auteurs soutiennent que l'absence de surperformance dans les modèles à petite échelle n'exclut pas de mérite à plus grande échelle. Ils suggèrent que l'approche est mathématiquement solide et que le « compromis coût-précision » pourrait évoluer favorablement avec des modèles plus grands ou des architectures différentes (par exemple, les modèles basés sur les tenseurs comme MLA, MoLAE).
• Limites : L'étude admet des limites significatives, notamment la petite taille du modèle, l'absence d'amortissement (annealing) du taux d'apprentissage et l'utilisation d'une seule référence (AdamW) sans ajustement de ses hyperparamètres.

En conclusion, ce papier sert de démonstration technique rigoureuse de l'application de l'optimisation riemannienne aux composants à faible rang du deep learning, fournissant une bibliothèque d'algorithmes et une base de référence pour les recherches futures, tout en reconnaissant que les résultats empiriques actuels ne justifient pas encore de remplacer les optimiseurs standards comme AdamW.