🤖 apprentissage automatique

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

Cet article présente MKOR, un optimiseur du second ordre doté de momentum qui utilise des mises à jour de rang 1 et une inversion basée sur la formule de Sherman-Morrison pour réduire la complexité de calcul et de communication de cubique à quadratique et linéaire respectivement, accélérant ainsi considérablement l'entraînement et la convergence des grands modèles de langage par rapport aux méthodes de premier et second ordre de pointe.

Auteurs originaux : Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Publié 2024-01-30✓ Author reviewed ⓘ

📖 7 min de lecture🧠 Analyse approfondie

Auteurs originaux : Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Dans le vaste paysage de l'intelligence artificielle moderne, les systèmes les plus puissants sont construits sur des réseaux de neurones profonds, des structures mathématiques complexes qui apprennent à reconnaître des motifs, à traduire des langues et à générer du texte. Pour enseigner à ces systèmes, les chercheurs doivent résoudre un problème d'optimisation massif : trouver l'ensemble parfait de nombres internes, ou poids, qui permettent au réseau de faire des prédictions précises. Pendant des décennies, l'approche standard a consisté à pousser ces poids dans la direction qui réduit le plus rapidement l'erreur, une méthode semblable à un randonneur qui cherche son chemin en descendant une montagne embrumée en vérifiant la pente sous ses pieds. Cette méthode du premier ordre est fiable et largement utilisée, mais elle peut être lente, nécessitant souvent des milliers d'étapes pour atteindre le fond de la vallée.

Une stratégie plus ambitieuse, connue sous le nom d'optimisation du second ordre, tente de donner au randonneur une carte de tout le terrain, y compris la courbure du sol, permettant des descentes beaucoup plus rapides. En comprenant comment la pente change, ces méthodes peuvent théoriquement atteindre la solution optimale en beaucoup moins d'étapes. Cependant, cette approche a historiquement été trop coûteuse en termes de calcul pour être pratique pour les modèles massifs utilisés aujourd'hui. Calculer les informations de courbure nécessaires pour un grand réseau requiert une puissance de traitement et une mémoire immenses, rendant souvent la méthode plus lente que l'approche plus simple qu'elle vise à remplacer. Le défi a été de capturer la vitesse du randonneur lisant la carte sans le fardeau de porter la carte.

Une équipe de chercheurs de l'Université de Toronto, de l'Université Rutgers et du Texas Advanced Computing Center a développé un nouvel outil appelé MKOR qui résout ce dilemme. Ils ont créé un optimiseur doté de momentum qui utilise un raccourci mathématique spécifique pour approximer la forme du terrain sans calculer chaque détail. Au lieu d'essayer de calculer la courbure complète et complexe du paysage de perte, MKOR utilise une technique appelée mises à jour de rang 1. Cette approche simplifie le calcul en se concentrant sur la direction de changement la plus significative, réduisant efficacement le coût de calcul d'une relation cubique à une relation quadratique. En termes pratiques, cela signifie que la méthode passe beaucoup mieux à l'échelle à mesure que les modèles grandissent, évitant les goulots d'étranglement qui empêchaient auparavant l'utilisation des méthodes du second ordre dans les grands modèles de langage.

Les chercheurs ont testé ce nouvel optimiseur sur un grand modèle de langage connu sous le nom de BERT-Large-Uncased, en effectuant des expériences sur 64 processeurs graphiques haute performance. Ils ont comparé MKOR aux meilleures méthodes du premier ordre existantes et aux implémentations du second ordre les plus avancées disponibles. Les résultats ont été frappants. Dans une tâche de réponse aux questions, la méthode standard du premier ordre nécessitait 1 536 étapes pour atteindre un niveau de précision spécifique. La nouvelle méthode MKOR a atteint la même précision, ou une meilleure, en seulement 600 étapes. En mesurant le temps total nécessaire pour terminer l'entraînement, MKOR a terminé la tâche en un peu plus de trois heures, tandis que la méthode standard a pris près de huit heures. Cela représente une accélération de 2,57 fois, réduisant efficacement le temps d'entraînement de plus de la moitié.

Au-delà de la vitesse brute, l'étude a abordé un problème critique de stabilité qui affecte souvent les méthodes du second ordre. Parce que ces méthodes ajustent le processus d'apprentissage sur la base de facteurs mathématiques complexes, elles sont sujettes aux « gradients explosifs », un phénomène où les ajustements deviennent si grands que le processus d'entraînement plante ou diverge. Les solutions précédentes nécessitaient souvent de ralentir considérablement le taux d'apprentissage pour maintenir la sécurité, ce qui annulait les avantages de vitesse. MKOR inclut un stabilisateur intégré qui surveille la taille de ces ajustements. Si les valeurs deviennent trop grandes, le système les redimensionne automatiquement, empêchant l'entraînement d'exploser tout en maintenant le rythme agressif de l'approche du second ordre. Cela permet à l'optimiseur de se déplacer rapidement sans risque d'instabilité.

Les chercheurs ont également observé que l'avantage des méthodes du second ordre est plus prononcé lors des premières étapes de l'entraînement, lorsque le modèle apprend les contours larges des données. À mesure que le modèle se rapproche de la solution optimale, les bénéfices des calculs de courbure complexes diminuent. Pour remédier à cela, ils ont introduit une version hybride de l'optimiseur, appelée MKOR-H. Cette version commence par la puissante méthode du second ordre pour accélérer la phase d'apprentissage initiale, mais bascule automatiquement vers une méthode plus simple du premier ordre une fois que la perte d'entraînement cesse de diminuer rapidement. Cette approche hybride a permis au système de converger en encore moins d'étapes, atteignant la même précision finale que la méthode standard en moins de la moitié du temps.

L'étude a également examiné pourquoi ce raccourci mathématique spécifique fonctionne si bien. Les chercheurs ont analysé les structures de données au sein du réseau de neurones et ont découvert que les matrices de covariance, qui décrivent comment les différentes parties du réseau interagissent, sont souvent de faible rang. Cela signifie que malgré la haute dimensionnalité des données, l'information est concentrée dans quelques directions clés. En approximant ces matrices à l'aide de mises à jour de rang 1, MKOR capture l'information essentielle nécessaire à l'optimisation sans la surcharge des calculs matriciels complets. L'équipe a vérifié, par des tests approfondis, que cette approximation introduit très peu d'erreurs, confirmant que la simplification ne se fait pas au détriment de la précision.

En plus de la vitesse et de la stabilité, la nouvelle méthode réduit considérablement la quantité de données qui doivent être transférées entre différents processeurs pendant l'entraînement. Dans l'informatique distribuée, où des milliers de processeurs travaillent ensemble, la communication peut devenir un goulot d'étranglement majeur. MKOR réduit cette complexité de communication d'une relation quadratique à une relation linéaire, ce qui signifie qu'à mesure que davantage de processeurs sont ajoutés, le système passe à l'échelle efficacement sans être ralenti par les délais de transfert de données. Cette efficacité permet à la méthode de gérer les tailles de lots massives requises par les modèles de langage modernes, ce qui était une limitation pour les approches précédentes du second ordre.

Les implications de ce travail s'étendent au-delà de l'entraînement d'un modèle spécifique. En rendant l'optimisation du second ordre pratique pour les systèmes à grande échelle, MKOR offre une voie vers un entraînement plus efficace de l'intelligence artificielle. Les chercheurs ont noté que leur méthode pourrait aider à économiser des quantités significatives d'énergie et de temps pour les praticiens de l'apprentissage automatique. Bien que la technologie elle-même soit un outil neutre, son application dans des domaines allant du diagnostic médical à la modélisation climatique pourrait accélérer la découverte scientifique. Cependant, les auteurs ont également reconnu que les algorithmes d'optimisation efficaces peuvent être appliqués à n'importe quel modèle, y compris ceux ayant des impacts sociétaux négatifs, soulignant que l'utilisation responsable d'une telle technologie dépend des choix des personnes qui la manipulent.

En fin de compte, l'article présente une solution concrète à un problème de longue date en apprentissage profond. En combinant le momentum, les approximations de rang 1 et un mécanisme de commutation hybride, les chercheurs ont créé un optimiseur qui est plus rapide, plus stable et plus évolutif que ses prédécesseurs. Les expériences démontrent qu'il est possible d'exploiter la puissance des méthodes du second ordre sans les coûts de calcul prohibitifs qui les ont maintenues en retrait. Alors que le domaine de l'intelligence artificielle continue de progresser vers des modèles plus grands et plus complexes, des outils comme MKOR offrent un moyen d'entraîner ces systèmes plus efficacement, transformant ce qui était autrefois un avantage théorique en une réalité pratique.

Résumé Technique : MKOR – Optimiseur basé sur les facteurs de Kronecker avec activation de momentum utilisant des mises à jour de rang 1

1. Énoncé du Problème

Les méthodes d'optimisation du second ordre, telles que la descente de gradient naturel (NGD) et la courbure approximative factorisée par Kronecker (KFAC), offrent des taux de convergence supérieurs par rapport aux méthodes du premier ordre (ex. : SGD, Adam, LAMB) en utilisant des informations de courbure (approximations de la matrice Hessienne ou de l'information de Fisher). Cependant, leur adoption dans l'entraînement des réseaux de neurones profonds (DNN) à grande échelle, particulièrement les Transformers et les grands modèles de langage (LLM), est entravée par trois goulots d'étranglement principaux :

Complexité Computationnelle : Les méthodes KFAC traditionnelles nécessitent l'inversion de matrices de covariance de taille $d \times d$ (où $d$ est la dimension de la couche), entraînant une complexité de $O(d^3)$ . De même, les méthodes basées sur Sherman-Morrison-Woodbury (SNGD) passent souvent à l'échelle avec la taille du batch $b$ ( $O(b^3)$ ). Dans les Transformers, où les tailles de batch augmentent avec la longueur de la séquence, ces complexités deviennent prohibitives.
Surcharge de Communication : L'entraînement distribué nécessite la synchronisation des informations du second ordre (matrices de covariance et leurs inverses) entre les workers. KFAC nécessite la synchronisation de $O(d^2)$ données, tandis que SNGD nécessite $O(bd + b^2)$ , créant des goulots d'étranglement de communication significatifs qui limitent la scalabilité.
Fréquence de Mise à Jour et Stabilité : En raison des coûts d'inversion élevés, les implémentations de pointe (comme KAISA) mettent à jour les informations du second ordre de manière peu fréquente (ex. : toutes les 100 à 1000 itérations) en utilisant des facteurs « obsolètes » (stale). Cela réduit l'efficacité de l'information de courbure. De plus, les méthodes du second ordre sont sujettes à l'instabilité numérique (gradients explosifs, matrices singulières) et nécessitent souvent des facteurs d'amortissement ou une arithmétique de haute précision, ce qui augmente davantage les coûts de calcul.

2. Méthodologie : L'algorithme MKOR

Les auteurs proposent MKOR (Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates), qui répond aux problèmes susmentionnés via quatre mécanismes clés :

Approximation de Rang 1 et Inversion basée sur SM :
Au lieu de calculer et d'inverser des matrices de covariance complètes, MKOR approxime la covariance des activations et des gradients en utilisant des matrices de rang 1 dérivées de la moyenne des échantillons ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). En appliquant l'identité de Sherman-Morrison à ces mises à jour de rang 1, MKOR met à jour directement l'inverse des facteurs de Kronecker. Cela réduit la complexité computationnelle de l'inversion des facteurs de $O(d^3)$ à $O(d^2)$ .
- Règle de Mise à Jour : Les facteurs inverses sont mis à jour en utilisant le momentum ( $\gamma$ ) et la formule de Sherman-Morrison, permettant des mises à jour fréquentes sans le coût d'une inversion de matrice complète.
Complexité de Communication Linéaire :
Dans des contextes distribués, MKOR ne synchronise que les deux vecteurs d'approximation de rang 1 (taille $2d$ ) entre les workers, plutôt que des matrices de covariance complètes ( $4d^2$ ). Cela réduit la complexité de communication d'un niveau quadratique $O(d^2)$ à un niveau linéaire $O(d)$ . De plus, MKOR utilise la demi-précision (FP16) pour ces vecteurs, divisant par deux les coûts de communication, un exploit difficile pour d'autres méthodes du second ordre en raison de la sensibilité numérique.
Mécanismes de Stabilisation :
Pour prévenir l'instabilité numérique et l'explosion des gradients courante dans les méthodes du second ordre :
- Stabilisateur basé sur la Norme : Surveille la norme infini des facteurs inverses. Si ceux-ci dépassent un seuil, les facteurs sont mélangés avec la matrice identité ( $\hat{L} = \zeta L + (1-\zeta)I$ ) pour prévenir la divergence.
- Redimensionnement du Gradient (Rescaling) : Les gradients préconditionnés sont redimensionnés pour correspondre à la norme des gradients originaux, garantissant que les planificateurs de taux d'apprentissage (learning rate schedulers) restent efficaces.
- Stabilité Théorique : Les auteurs prouvent (Lemme 3.1) que les facteurs restent définis positifs sans nécessiter de facteurs d'amortissement artificiels ( $\mu I$ ), et que la division par un scalaire unique dans la règle de mise à jour est garantie de ne pas être nulle.
MKOR Hybride (MKOR-H) :
Reconnaissant que les avantages du second ordre diminuent souvent à mesure que l'entraînement converge (lorsque le paysage de perte approche l'identité), MKOR-H introduit un mécanisme de commutation. Il surveille le taux de réduction de la perte et revient automatiquement à un optimiseur du premier ordre si les mises à jour du second ordre ne fournissent plus d'accélération, optimisant ainsi le compromis entre coût et vitesse de convergence.

3. Contributions Clés

Réduction de la Complexité : MKOR atteint une complexité computationnelle de $O(d^2)$ et une complexité de communication de $O(d)$ pour les mises à jour du second ordre, atténuant les goulots d'étranglement de $O(d^3)$ et $O(d^2)$ trouvés dans KFAC et $O(b^3)$ dans SNGD.
Mises à Jour à Haute Fréquence : En réduisant les coûts d'inversion, MKOR permet des mises à jour du second ordre à chaque itération (ou très fréquemment), éliminant le besoin de facteurs « obsolètes » et exploitant pleinement l'information de courbure.
Robustesse : L'introduction de la stabilisation basée sur la norme et du redimensionnement du gradient permet à MKOR de fonctionner de manière stable sans facteurs d'amortissement ou arithmétique en double précision, permettant une utilisation efficace du matériel en demi-précision.
Stratégie Hybride : La variante MKOR-H bascule dynamiquement entre les méthodes du second ordre et du premier ordre pour maintenir l'efficacité tout au long du cycle de vie de l'entraînement.

4. Résultats Expérimentaux

Les auteurs ont évalué MKOR sur BERT-Large-Uncased en utilisant 64 GPU A100, en le comparant à l'optimiseur de pointe du premier ordre LAMB et aux implémentations du second ordre KAISA (KFAC) et HyLo (SNGD), ainsi qu'à Eva.

Accélération (Speedup) : MKOR a obtenu une accélération de 2,57× par rapport à LAMB et une accélération de 1,85× par rapport à KAISA sur la tâche de pré-entraînement. Plus précisément, il a réduit le temps d'entraînement de BERT-Large-Uncased de 8 heures à 3 heures.
Convergence : Sur la tâche de réponse aux questions SQuAD v1.1, MKOR-H a convergé en 600 étapes (contre 1 536 pour LAMB) tout en atteignant des scores F1 comparables ou meilleurs (90,64 contre 90,44).
Benchmark GLUE : MKOR a atteint de nouvelles métriques de pointe sur le jeu de données GLUE en 1 500 étapes. Notamment, alors que KAISA n'a pas réussi à converger vers la précision moyenne de référence même après 1 563 étapes, MKOR a convergé avec succès.
Scalabilité : Contrairement à HyLo, qui a divergé ou n'a pas pu tenir dans la mémoire GPU en raison d'une surcharge mémoire élevée ( $O(bd + b^2)$ ), MKOR a réussi l'entraînement avec de grandes tailles de batch (jusqu'à 64k de taille de batch effective via l'accumulation) sans problèmes de mémoire.
Validité de l'Approximation : Les expériences sur ResNet-50 et BERT ont montré que les approximations de rang 1 des matrices de covariance induisent une erreur négligeable, validant l'hypothèse d'approximation centrale.

5. Signification et Revendications

L'article affirme que MKOR comble avec succès le fossé entre les avantages théoriques de l'optimisation du second ordre et les contraintes pratiques de l'entraînement des DNN modernes à grande échelle. En réduant la complexité des mises à jour du second ordre à un niveau comparable aux méthodes du premier ordre (en termes de mise à l'échelle avec la taille du modèle) tout en maintenant les avantages de convergence, MKOR permet :

Un Entraînement Plus Rapide : Réduction significative du temps d'entraînement réel (wall-clock time) pour les LLM.
Une Meilleure Convergence : La capacité de converger vers des métriques de précision plus élevées là où d'autres méthodes du second ordre (comme KAISA) échouent ou stagnent.
Une Scalabilité : La capacité de passer à l'échelle avec de grands batchs et des dimensions de modèles sans les goulots d'étranglement de communication et de mémoire qui limitaient auparavant les méthodes du second ordre.

Les auteurs soulignent que leur travail rend l'optimisation du second ordre viable pour l'ère des « grands modèles », en traitant spécifiquement les problèmes de scalabilité introduits par les tailles de batch dépendantes de la longueur de la séquence dans les architectures Transformer. Ils concluent que MKOR représente un nouveau standard de pointe pour un entraînement efficace et performant des réseaux de neurones profonds.