Optimiseur LAMB

LAMB (Layer-wise Adaptive Moments for Batch training) est un algorithme d'optimisation conçu pour l'entraînement de réseaux de neurones profonds, en particulier dans les environnements de calcul distribué à grande taille de lot. Il a été introduit en 2019 par des chercheurs de Google et l'Université de Toronto pour relever les défis de la mise à l'échelle de l'entraînement sur des milliers d'accélérateurs tout en maintenant la précision du modèle et la vitesse de convergence.

L'algorithme étend l'optimiseur Adam en calculant des taux d'apprentissage adaptatifs par couche plutôt que par paramètre. Cette adaptation par couche permet à LAMB de gérer plus efficacement les échelles de gradients variables entre les différentes couches du réseau, ce qui est particulièrement important dans les architectures profondes comme les Transformers et les réseaux résiduels. En normalisant les mises à jour en fonction de la norme des poids et des gradients de la couche, LAMB assure un entraînement stable et efficace même avec de très grandes tailles de lot.

Contexte et Motivation

L'entraînement de réseaux de neurones à grande échelle nécessite généralement des ressources de calcul massives, souvent réparties sur de nombreux GPU ou TPU. Augmenter la taille du lot est une stratégie courante pour utiliser efficacement ces ressources, mais cela conduit souvent à une dégradation des performances du modèle ou à une convergence plus lente. Les optimiseurs traditionnels comme la descente de gradient stochastique (SGD) et Adam peinent avec les grandes tailles de lot car ils reposent sur des taux d'apprentissage globaux qui ne tiennent pas compte de l'hétérogénéité des échelles de gradients entre les couches.

L'optimiseur LAMB a été développé pour surmonter ces limitations. Sa conception s'inspire de l'algorithme LARS (Layer-wise Adaptive Rate Scaling), précédemment utilisé pour l'entraînement à grande taille de lot de réseaux convolutifs. LAMB généralise ce concept pour fonctionner avec l'estimation adaptative des moments, combinant les avantages des deux approches.

Détails de l'Algorithme

LAMB calcule une mise à jour pour chaque couche en fonction du rapport entre la norme des poids de la couche et la norme de ses gradients. La règle de mise à jour principale pour un tenseur de paramètres à l'étape t est :

1. Calculer les estimations des premier et deuxième moments (moyenne et variance des gradients) comme dans Adam.
2. Calculer la direction de mise à jour comme le gradient corrigé par les moments divisé par la racine carrée du deuxième moment plus un petit epsilon.
3. Mettre à l'échelle cette direction par le rapport entre la norme des poids de la couche et la norme de la direction de mise à jour.
4. Multiplier par un taux d'apprentissage global et appliquer la mise à jour.

Cette mise à l'échelle par couche garantit que les couches avec de grandes normes de poids reçoivent des mises à jour proportionnellement plus importantes, tandis que les couches avec de petites normes sont mises à jour de manière conservatrice. L'algorithme intègre également un ratio de confiance qui peut être plafonné pour éviter des mises à jour extrêmes, similaire aux techniques de limitation de gradient.

Les auteurs ont démontré que LAMB peut entraîner ResNet-50 sur ImageNet avec une taille de lot de 32 768 tout en atteignant la même précision que la référence avec une taille de lot de 256, mais en beaucoup moins d'étapes. Cela le rend particulièrement adapté à l'entraînement distribué sur des centaines ou des milliers d'accélérateurs.

Applications et Impact

LAMB a été largement adopté dans l'entraînement de grands modèles de langage et d'autres modèles de apprentissage profond. Par exemple, il a été utilisé pour entraîner BERT et d'autres modèles basés sur Transformer à grande échelle, réduisant le temps d'entraînement de jours à heures. L'algorithme est particulièrement précieux dans les environnements où les ressources matérielles sont abondantes, comme les plateformes cloud AWS et Microsoft Azure, ainsi que le matériel IA spécialisé comme AWS Trainium et les IPU Graphcore.

De nombreux optimiseurs ultérieurs, tels que les variantes et successeurs de LAMB, se sont appuyés sur ses principes. Il a également influencé la recherche sur les programmes de taux d'apprentissage et les méthodes d'optimisation adaptatives. L'implémentation open-source dans des frameworks comme TensorFlow et PyTorch l'a rendu accessible à la communauté plus large du machine learning.

Comparaison avec d'Autres Optimiseurs

Comparé à Adam, LAMB atteint généralement une convergence plus rapide et de meilleures performances finales lors de l'utilisation de grandes tailles de lot. Le taux d'apprentissage global d'Adam nécessite souvent un réglage minutieux et peut entraîner une instabilité avec de grands lots. L'adaptation par couche de LAMB atténue ces problèmes, permettant une mise à l'échelle plus agressive.

Comparé à LARS, qui est conçu pour la SGD avec momentum, LAMB intègre l'estimation adaptative des moments, ce qui le rend plus robuste aux gradients bruités et aux caractéristiques éparses. Cela fait de LAMB un choix plus polyvalent pour une large gamme d'architectures, y compris les modèles séquence à séquence et les frameworks encodeur-décodeur.

Limitations et Considérations

Malgré ses avantages, LAMB n'est pas sans limitations. L'algorithme introduit des hyperparamètres supplémentaires, tels que le seuil de plafonnement du ratio de confiance et le terme epsilon, qui peuvent nécessiter un réglage pour des tâches spécifiques. Il suppose également que la mise à l'échelle par couche est bénéfique, ce qui peut ne pas toujours être le cas pour des architectures avec des couches fortement corrélées ou lors de l'utilisation de certains schémas de initialisation des poids.

De plus, bien que LAMB excelle dans les contextes de grandes tailles de lot, ses avantages diminuent pour les petites tailles de lot où des optimiseurs plus simples comme Adam peuvent suffire. Les chercheurs ont également noté que les performances de l'algorithme peuvent être sensibles au choix du taux d'apprentissage global, et il peut nécessiter un échauffement du taux d'apprentissage pour obtenir des résultats optimaux.

Conclusion

LAMB représente une avancée significative dans l'optimisation pour l'entraînement à grande échelle en apprentissage profond. En combinant l'adaptation par couche avec les moments adaptatifs, il permet un entraînement efficace et stable avec des tailles de lot massives, ce qui en fait une technique fondamentale dans l'ère de l'IA générative et de la recherche en intelligence artificielle. Son influence s'étend au-delà de son application originale, façonnant le développement d'optimiseurs ultérieurs et de méthodologies d'entraînement.