Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization

Résumé technique : Redémarrage du momentum externe dans l'optimisation en deux phases de haute dimension

Énoncé du problème
L'entraînement moderne des grands modèles de langage (LM) est de plus en plus limité par les coûts de communication dans les environnements distribués. Pour atténuer ce problème, des méthodes d'optimisation en deux phases comme DiLoCo gagnent en popularité, où les workers effectuent plusieurs mises à jour locales (phase interne) avant d'agréger les progrès via un optimiseur externe (phase externe). Bien que efficaces, ces méthodes introduisent une nouvelle couche de sensibilité aux hyperparamètres. Des travaux théoriques récents suggèrent que l'optimiseur externe opère sur un « spectre effectif » induit par la boucle interne. Plus précisément, le momentum externe de type « heavy-ball » (HB) peut contraindre les directions où la boucle interne a déjà réalisé des progrès significatifs, entraînant une convergence lente ou une instabilité, en particulier lorsque les périodes de communication sont longues ou que les valeurs de momentum externe sont élevées. Le problème central abordé est la fragilité des optimiseurs en deux phases face aux choix de momentum externe et de taux d'apprentissage, ce qui complique le réglage dans les scénarios de préentraînement.

Méthodologie
Les auteurs proposent un redémarrage périodique du tampon de momentum externe comme mécanisme complémentaire pour contrôler la « mémoire externe ». La méthodologie combine une analyse théorique et une validation empirique :

1. Cadre théorique :

   • Les auteurs adoptent un modèle déterministe, mode par mode, basé sur l'approximation linéarisée de la perte quadratique sous le noyau tangentiel neuronal (NTK) empirique.
   • Ils analysent la dynamique des coordonnées propres résiduelles à travers les tours de communication. La boucle interne produit un pseudo-gradient $g_t = \sigma x_t$, où $\sigma$ représente le progrès effectif.
   • Pour le momentum Heavy-Ball (HB), la matrice de transition $T_{HB}$ a un déterminant de $\beta_{out}$, conduisant à un taux d'amortissement indépendant du progrès $\sigma$. Cela implique que même les modes à fort progrès sont soumis à la même enveloppe lente.
   • Les auteurs dérivent une contraction de redémarrage mode par mode. En réinitialisant le tampon de momentum externe ($m \leftarrow 0$) tous les $K$ tours, le système exploite l'annulation de phase. Dans le régime complexe (typique d'un momentum élevé), la coordonnée résiduelle oscille au sein d'une enveloppe. Une période de redémarrage $K$ bien choisie peut aligner la réinitialisation avec une phase où le résidu projeté est minimisé, éliminant efficacement le momentum obsolète tout en préservant les progrès de la boucle interne.
   • Ils dérivent une relation de récurrence pour le facteur de contraction $\chi_K(\sigma)$ et montrent que le redémarrage améliore les performances par rapport au HB sans redémarrage lorsque le terme de projection dépendant de la phase a une magnitude inférieure à un.

2. Configuration expérimentale :

   • La méthode est évaluée sur le préentraînement d'un modèle LLaMA de 150 millions de paramètres utilisant une configuration DiLoCo à 2 réplicas sur des GPU H200.
   • Les expériences couvrent à la fois les optimiseurs externes Heavy-Ball et Nesterov (NAG) sur diverses périodes de communication ($S$) et grilles d'hyperparamètres.
   • L'étude compare le DiLoCo standard au DiLoCo avec redémarrages périodiques, testant des stratégies de redémarrage global, par mode et par blocs.

Contributions clés

• Identification du mécanisme : L'article identifie que le redémarrage périodique dans l'optimisation en deux phases agit en exploitant l'annulation de phase dans la dynamique du momentum externe, distinctement de la manière dont le momentum Nesterov modifie l'enveloppe spectrale.
• Dérivation théorique : Une analyse sous forme fermée du facteur de contraction de redémarrage est fournie, démontrant que les réinitialisations peuvent réduire le taux d'amortissement effectif en éliminant les états de momentum obsolètes causant des oscillations ou un désalignement.
• Amélioration de la robustesse : La contribution principale est la démonstration que le redémarrage périodique élargit considérablement la région stable des hyperparamètres externes (taux d'apprentissage $\nu$ et momentum $\beta_{out}$).
• Référence pratique : Les auteurs établissent le redémarrage périodique comme une base simple et robuste réduisant le besoin d'un réglage délicat des hyperparamètres, une propriété critique pour le préentraînement à grande échelle.

Résultats

• Stabilité et robustesse : Les résultats empiriques sur le modèle LLaMA de 150 millions de paramètres montrent que sans redémarrages, des valeurs élevées de momentum externe ($\beta_{out}$) combinées à de longues périodes de communication entraînent une dégradation brutale des performances ou une divergence. Les redémarrages périodiques éliminent largement ces zones de défaillance, permettant un entraînement stable sur une plage plus large de $\beta_{out}$ et $\nu$.
• Sensibilité aux hyperparamètres : L'étude démontre que la période de redémarrage $K$ est un hyperparamètre plus « indulgent » que le coefficient de momentum externe. Alors que les exécutions sans redémarrage nécessitent de retuner $\beta_{out}$ lorsque la période de communication $S$ change, les exécutions avec redémarrage maintiennent des performances stables sur différentes valeurs de $S$ avec une $K$ fixe.
• Performance : En termes de performance maximale, les redémarrages périodiques préservent la perplexité de validation optimale atteignable tout en réduisant considérablement la variance des performances sur la grille d'hyperparamètres.
• Redémarrages doux vs durs : Les expériences avec des « redémarrages doux » (préservant partiellement le momentum) ont montré des améliorations marginales par rapport aux réinitialisations dures mais ont introduit des hyperparamètres supplémentaires ($\alpha, \beta$) sensibles aux configurations d'entraînement, conduisant les auteurs à privilégier l'approche de redémarrage dur plus simple.

Signification et affirmations
L'article affirme que le redémarrage périodique du momentum externe est un mécanisme simple mais efficace pour stabiliser l'optimisation en deux phases de haute dimension. Sa signification réside dans la résolution du caractère « délicat » du réglage des optimiseurs en préentraînement, où les méthodes en deux phases introduisent une couche supplémentaire de complexité. En limitant la mémoire externe inutile, la méthode rend l'optimiseur externe moins sensible au choix du momentum et du taux d'apprentissage.

Les auteurs présentent modestement leur travail comme fournissant une base robuste plutôt qu'une solution universelle. Ils notent que, bien que leur analyse se concentre sur des réinitialisations à période fixe, les résultats empiriques suggèrent que la méthode est assez insensible à la période précise choisie. Ils concluent en posant une question plus large pour les travaux futurs : la boucle externe peut-elle décider adaptivement quand rafraîchir sa mémoire, suggérant des directions potentielles pour des règles de redémarrage adaptatives, des périodes par couche et des études à plus grande échelle.