Résumé Technique : LoRDO – Optimisation à Faible Rang Distribuée avec Communication Peu Fréquente

1. Énoncé du Problème

L'entraînement distribué de modèles de fondation via le mode Parallèle de Données Distribuées (DDP) est de plus en plus contraint par la bande passante d'interconnexion et la surcharge de mémoire requise pour stocker les états de l'optimiseur (par exemple, les moments d'Adam). Bien que les stratégies employant une communication peu fréquente (mises à jour locales) réduisent la fréquence de synchronisation, elles restent limitées par les coûts de mémoire et de communication des états d'optimiseur à rang complet.

Les optimiseurs adaptatifs à faible rang (par exemple, GaLore, LDAdam) offrent une voie pour atténuer ces contraintes en projetant les gradients dans un sous-espace de faible rang. Cependant, généraliser ces méthodes au régime de synchronisation peu fréquente présente un défi d'optimisation critique :

1. Projections Locales : Calculer des projections de faible rang localement sur chaque travailleur (worker) en utilisant des gradients stochastiques introduit un bruit significatif dû aux tailles de lots effectives par travailleur trop réduites. Cela conduit à des sous-espaces de projection sous-optimaux et à une dégradation des performances.
2. Projections Globales : Utiliser une matrice de projection globale dérivée de pseudo-gradients agrégés (comme dans le DDP) fournit une base plus stable. Cependant, dans le régime de mise à jour locale, cette approche restreint de manière permanente la trajectoire d'optimisation à un sous-espace de faible rang fixe. Sans signaux à rang complet, l'optimisation stagne car la base de projection ne parvient pas à se « rafraîchir » ou à explorer de nouvelles directions en dehors du sous-espace initial.

La question centrale abordée est la suivante : Comment concevoir des optimiseurs à faible rang de haute performance pour un entraînement économique en communication, tout en évitant la stagnation du sous-espace et en maintenant une faible surcharge de communication et de mémoire ?

2. Méthodologie : LoRDO

Les auteurs proposent LoRDO (Low-Rank Distributed Optimization), un cadre qui unifie l'optimisation à faible rang avec la synchronisation peu fréquente. La méthodologie répond à la tension entre stabilité de projection et exploration de sous-espace par trois mécanismes clés :

A. Matrices de Projection Globales

Au lieu que chaque travailleur détermine sa propre matrice de projection $Q_m$ basée sur des gradients stochastiques locaux (ce qui souffre d'une variance élevée), LoRDO emploie une matrice de projection globale $Q_t$. Cette matrice est calculée à la limite de synchronisation en utilisant le pseudo-gradient agrégé ($\Delta_t$), qui représente le changement total des paramètres du modèle sur $K$ étapes locales.

• Bénéfice : Le pseudo-gradient possède une taille de lot effective de $M \times B$ (où $M$ est le nombre de travailleurs), offrant un signal plus stable avec une variance réduite par rapport aux gradients locaux. Il intègre également l'information de courbure inscrite dans le pseudo-gradient, menant à une estimation de base de meilleure qualité.

B. Momentum Quasi-Hyperbolique à Plein Rang (QHM)

Pour résoudre le problème de la « stagnation du sous-espace » inhérent aux projections globales, LoRDO injecte un terme de momentum quasi-hyperbolique à plein rang dans la mise à jour locale.

• Mécanisme : La règle de mise à jour combine le momentum à faible rang avec un signal de gradient à plein rang mis à l'échelle par le second moment.
  $$\text{Mise à jour} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• Effet : Cette injection empêche le pseudo-gradient agrégé de rester piégé dans un sous-espace fixe de rang-$r$. Elle permet de l'agrégation temporelle des sous-espaces de rang-$r$ provenant de différents travailleurs, permettant ainsi une exploration complète du sous-espace. Cela restaure la capacité d'explorer l'espace de solution complet sans augmenter la surcharge de communication.

C. Retour d'Erreur et Rotation de Momentum

LoRDO incorpore un retour d'erreur (error feedback) local pour tenir compte de l'information perdue lors de la compression à faible rang. De plus, après le calcul d'une nouvelle matrice de projection, les états de l'optimiseur (moments) sont rotés pour garantir qu'ils restent accumulés sur le même sous-espace, une technique adaptée de LDAdam.

3. Contributions Clés

1. Analyse des modes de défaillance de projection : Les auteurs démontrent que les projections locales nuisent aux performances en raison de la variance élevée issue des petites tailles de lots par travailleur, tandis que les projections globales sans signaux à plein rang induisent une stagnation du sous-espace.
2. Restauration de l'exploration complète du sous-espace : LoRDO introduit un optimiseur à faible rang qui injecte un signal de gradient à plein rang (via QHM) dans les mises à jour locales tout en maintenant une projection globale dérivée des pseudo-gradients agrégés. Cela empêche la stagnation sans augmenter la surcharge de communication ou de mémoire.
3. Parité et Efficacité DDP : Le cadre atteint une quasi-parité avec le DDP à faible rang synchrone (écart de perplexité < 1 %) tout en réduisant la mémoire et la communication de l'optimiseur jusqu'à $8\times$–$12\times$. Sous de fortes contraintes de mémoire nécessitant des rangs faibles, LoRDO surpasse même le DDP en termes de perplexité.

4. Résultats Expérimentaux

Les expériences ont été menées sur des transformeurs de type decoder-only mis à l'échelle à 16M, 125M et 720M de paramètres en utilisant le mélange de données SmolLM2.

• Stagnation vs Exploration : Dans les modèles de 16M, les projections globales sans QHM n'ont pas réussi à apprendre (perplexité 46.8) par rapport au DDP (36.6), confirmant la stagnation du sous-espace. L'ajout du QHM à plein rang a restauré les performances, permettant à LoRDO-Global de surpasser les variantes locales.
• Qualité de la Projection : LoRDO-Global a démontré une stabilité et des performances supérieures, particulièrement à bas rang ($r=8, 16$), grâce à la base de projection de meilleure qualité dérivée du pseudo-gradient agrégé. Il a montré une moindre sensibilité aux changements de taille de lot locale par rapport aux méthodes de projection locales.
• Mise à l'échelle vers 720M : À l'échelle 720M, LoRDO-Global ($r=256$) a atteint une quasi-parité avec le DDP à faible rang ($r=256$) avec un écart de perplexité de moins de 1 % (10.41 vs 10.34). Il a également égalé la précision sur les tâches en aval (ARC, HellaSwag, MMLU, PIQA) avec le DDP.
• Gains d'Efficacité :
  • Communication : LoRDO réduit la surcharge de communication d'environ 10$\times$ par rapport au DDP à faible rang et de 25$\times$ par rapport au DDP à plein rang.
  • Mémoire : Il réduit la surcharge de mémoire des états de l'optimiseur de 8$\times$ par rapport aux équivalents à plein rang.
  • Temps d'exécution (Wall-Clock Time) : Dans les environnements à faible bande passante (modélisés à 10 Gbps), LoRDO surpasse significativement le DDP en temps d'exécution grâce à la réduction de la fréquence de synchronisation.

5. Signification et Revendications

L'article affirme que LoRDO résout la tension fondamentale entre la stabilité de la projection et l'exploration du sous-espace dans l'optimisation efficace en communication à faible rang. En utilisant des pseudo-gradients agrégés pour la stabilité de la projection et un momentum quasi-hyperbolique à plein rang pour l'exploration, LoRDO permet :

• Entraînement Scalable : Performances compétitives à grande échelle (jusqu'à 720M de paramètres) avec des coûts de communication et de mémoire considérablement réduits.
• Efficacité des Ressources : Faisabilité de l'entraînement sur du matériel aux ressources limitées (mémoire/bande passante faibles), là où les optimiseurs à plein rang sont intraitables.
• Apprentissage Décentralisé : Ouvrir la voie à un apprentissage efficace dans des environnements décentralisés en étendant l'entraînement des modèles au-delà des centres de données uniques.

Les auteurs notent des limites, notamment l'absence de preuve de convergence formelle et l'accent actuel mis sur les optimiseurs adaptatifs (Adam) plutôt que sur les optimiseurs matriciels (ex: Muon), laissant ces points pour des travaux futurs. Cependant, ils postulent que les tendances observées sur de plus petits modèles se tiendront à des échelles plus grandes.