L’optimiseur LAMB (Layer-wise Adaptive Moments for Batch training) est un algorithme d’optimisation stochastique pour l’entraînement de réseaux de neurones profonds qui étend l’optimiseur Adam avec une étape de normalisation par couche. Introduit en 2019 par des chercheurs de Google (notamment Yang You, Jing Li, Jonathan Hseu et al.), LAMB a été conçu pour permettre l’utilisation efficace de très grandes tailles de mini-lots (par exemple, 32 768 ou plus) sans dégrader la précision du modèle ni nécessiter un réglage exhaustif des hyperparamètres. Il y parvient en ajustant l’ampleur de la mise à jour pour chaque couche en fonction du rapport entre la norme des poids de la couche et sa norme de mise à jour, découplant ainsi efficacement le taux d’apprentissage de l’échelle des paramètres de chaque couche.
LAMB a été particulièrement influent dans l’entraînement de modèles basés sur Transformer (architecture), y compris les premiers grands modèles de langage (LLM) et des architectures de vision comme ResNet. Son adoption dans les cadres d’entraînement distribué, tels que TensorFlow (via tf.keras.optimizers.LAMB) et PyTorch (via l’implémentation de LAMB dans des bibliothèques comme NVIDIA, Hugging Face et FairScale), en a fait un outil standard pour faire passer à l’échelle les sessions d’entraînement. En permettant des lots plus grands, LAMB réduit le temps réel nécessaire pour entraîner des modèles de pointe, ce qui est crucial pour des organisations comme OpenAI, Anthropic et Google DeepMind qui dépendent de grappes de calcul massives.
Motivation et contexte
Entraîner des réseaux de neurones profonds avec de grands mini-lots réduit le nombre de mises à jour de paramètres par époque et les étapes de calcul, mais une mise à l’échelle naïve de la taille du lot entraîne souvent une mauvaise généralisation et une convergence instable. C’est ce qu’on appelle le « problème de l’entraînement par grands lots ». Les optimiseurs standard comme la descente de gradient stochastique (SGD) avec momentum ou l’Adam nécessitent un ajustement minutieux du taux d’apprentissage lorsque la taille du lot augmente, et même alors, la précision se dégrade souvent. L’optimiseur LAMB a été développé pour remédier à cela en rendant l’optimiseur plus robuste aux changements de taille de lot.
L’idée clé derrière LAMB est que différentes couches d’un réseau profond présentent des échelles très différentes de gradients et de normes de poids. Par exemple, les premières couches convolutives ont de petits poids, tandis que les dernières couches entièrement connectées en ont de grands. Un taux d’apprentissage global unique dans Adam peut conduire soit à des mises à jour trop grandes pour certaines couches (provoquant une divergence), soit à des mises à jour trop petites pour d’autres (ralentissant la convergence). LAMB introduit un taux adaptatif par couche qui normalise la mise à jour en fonction du rapport entre les normes des poids et des gradients de la couche, garantissant que chaque couche se déplace à un rythme stable par rapport à son ampleur.
Description de l’algorithme
LAMB peut être considéré comme une variante d’Adam avec une étape de normalisation supplémentaire. Soit \(\theta_t\) les paramètres à l’itération \(t\), et \(g_t\) le gradient de la perte par rapport à \(\theta_t\). LAMB maintient les premier et second moments (\(m_t\) et \(v_t\)) des gradients, comme Adam, avec des taux de décroissance exponentielle \(\beta_1\) et \(\beta_2\) (typiquement 0,9 et 0,999). Après correction du biais, il calcule une mise à jour d’Adam \(\frac{m_t}{\sqrt{v_t} + \epsilon}\).
La différence cruciale est le rapport de confiance \(\phi\) : pour chaque couche \(i\), \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), où \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) est la mise à jour d’Adam (sans taux d’apprentissage) pour cette couche, et \(||\cdot||\) désigne la norme L2. La mise à jour finale pour la couche \(i\) est \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), où \(\eta\) est le taux d’apprentissage global. Ce rapport de confiance met à l’échelle la mise à jour proportionnellement à la norme des poids de la couche, de sorte que les petites couches reçoivent des mises à jour petites (mais non négligeables), et les grandes couches reçoivent des mises à jour plus grandes, mais stables.
En pratique, une petite constante (par exemple, 1e-6) est ajoutée au dénominateur pour éviter une division par zéro. L’algorithme inclut également une décroissance de poids optionnelle (régularisation L2) fusionnée dans la mise à jour, suivant l’approche de décroissance de poids découplée utilisée dans AdamW. Lorsque le rapport de confiance est défini sur 1 pour toutes les couches, LAMB se réduit à l’Adam standard (avec correction du biais).
Hyperparamètres et réglage
LAMB hérite de la plupart des hyperparamètres d’Adam : \(\beta_1\) (momentum), \(\beta_2\) (décroissance de variance), \(\epsilon\) (stabilité numérique) et le taux de décroissance de poids. Le principal nouvel hyperparamètre est le taux d’apprentissage global \(\eta\), souvent défini dans la plage 0,01-0,1 pour l’entraînement par grands lots, nettement plus élevé que ce qui est typique pour Adam (par exemple, 1e-3). Les auteurs ont constaté que pour de très grands lots (par exemple, 32 768 pour BERT), un taux d’apprentissage de 0,01 avec un échauffement linéaire sur les premiers 10 % des étapes fonctionne bien, et ils ont également recommandé d’utiliser une décroissance cosinusoïdale schéma de taux d’apprentissage sur les étapes restantes.
De plus, le choix de \(\beta_2\) peut affecter la stabilité ; pour les modèles avec des gradients creux, un \(\beta_2\) plus élevé (par exemple, 0,99) peut être utilisé. Les auteurs ont également suggéré que la taille du lot peut être mise à l’échelle proportionnellement au taux d’apprentissage (règle de mise à l’échelle linéaire), une ligne directrice qui fonctionne bien avec LAMB. Par exemple, si la taille du lot est doublée, le taux d’apprentissage peut également être doublé sans perte de précision.
Performances et repères
Dans l’article original, LAMB a été évalué sur deux tâches principales : l’entraînement de ResNet-50 sur ImageNet (classification d’images) et BERT (un modèle de langage basé sur transformer) pour le masquage de langage. En utilisant LAMB, les auteurs ont atteint une précision top-1 sur ImageNet de 76,0 % en seulement 2 048 itérations avec une taille de lot de 32 768, égalant la précision de pointe obtenue avec des lots plus petits (par exemple, 256) en beaucoup moins d’époques. Pour BERT, ils ont entraîné le modèle à la même précision (par exemple, score F1 de 1,0 sur SQUAD) en environ 3,5 minutes en utilisant 1 024 TPU, soit une accélération de 10 fois par rapport aux méthodes précédentes.
Par la suite, LAMB est devenu l’optimiseur par défaut pour l’entraînement des modèles basés sur BERT dans les flux de travail internes de Google. L’article a rapporté que LAMB surpassait à la fois Adam et SGD avec momentum lors de la mise à l’échelle des tailles de lot de 1 024 à 65 536. Les auteurs ont également montré que LAMB fonctionne bien avec le Gradient Clipping (utilisé pour empêcher les gradients explosifs) et est compatible avec l’entraînement en précision mixte, comme utilisé sur du matériel moderne tel que GPU NVIDIA et TPU Google.
Applications dans l’entraînement à grande échelle
La principale application de LAMB est l’entraînement distribué où la taille du lot est trop grande pour tenir dans la mémoire d’un seul appareil. Dans de telles configurations, les gradients sont moyennés sur plusieurs GPU ou TPU en utilisant le parallélisme de données. Par exemple, OpenAI et Google DeepMind utilisent des optimiseurs analogues à LAMB lors de l’entraînement de grands modèles Transformer (architecture) avec des longueurs de séquence de plusieurs milliers. Bien que de nouveaux optimiseurs comme LAMB (et son successeur LAMB2) aient été proposés, LAMB reste un choix fiable dans de nombreux efforts open-source, y compris l’entraînement de transformers de vision et de LLM par des groupes de recherche et des entreprises comme AI21 Labs et SambaNova.
Dans le contexte du Machine learning sur Amazon Web Services (avec du matériel AWS Trainium), LAMB est pris en charge dans des noyaux personnalisés pour l’efficacité. De même, Intel et AMD ont évalué LAMB sur leurs accélérateurs. La capacité de l’optimiseur à gérer des tailles de lot extrêmes le rend précieux pour le pré-entraînement de modèles sur des ensembles de données massifs, où le coût d’une seule époque est élevé et où la réduction du nombre d’époques est primordiale.
Relation avec d’autres optimiseurs
LAMB fait partie d’une famille d’optimiseurs adaptatifs qui incluent les variantes de SGD, Adam et ses successeurs comme AdamW (décroissance de poids découplée) et LARS (Layer-wise Adaptive Rate Scaling). LARS, introduit par You et al. en 2017 pour l’entraînement par grands lots de CNN, utilise un rapport de confiance par couche similaire mais ne maintient pas de seconds moments ; il repose sur les premiers moments (momentum) et les normes de gradient. LAMB combine les avantages de LARS (mise à l’échelle par couche) avec les taux d’apprentissage adaptatifs par paramètre d’Adam, le rendant plus robuste pour les modèles avec des gradients creux (comme les transformers).
Un autre optimiseur étroitement lié est NVLAMB (de Nvidia), qui intègre une réduction de variance. Cependant, LAMB reste plus simple et largement utilisé. Pour les tâches Sequence-to-Sequence (Seq2Seq) avec Beam Search, LAMB n’affecte pas directement l’inférence, mais il aide à la convergence de l’entraînement, ce qui améliore indirectement le décodage de séquences.
Extensions et variantes
Depuis son introduction, plusieurs variantes ont été proposées. LAMB2 (également de Google) ajoute un facteur de normalisation basé sur la variance du gradient, améliorant la stabilité pour certains problèmes. LARS avec correction du biais et d’autres modifications sont également courants. En pratique, de nombreux cadres implémentent LAMB avec une correction du biais optionnelle pour les moments, ce qui est bénéfique pendant les premières étapes. Certaines implémentations, comme dans PyTorch torch.optim.Lamb (dans le paquet torch_optimizer), permettent d’ajuster le paramètre de rapport de confiance ou d’utiliser un taux d’apprentissage personnalisé par couche.
Malgré l’émergence de nouveaux optimiseurs comme AdamW avec différentes stratégies de mise à l’échelle (par exemple, les schémas 1cycle), LAMB reste une référence solide pour l’entraînement par grands lots. La communauté de recherche a exploré la combinaison de LAMB avec la Data Augmentation et le Gradient Clipping pour améliorer encore la généralisation.
Considérations pratiques et limites
Bien que LAMB excelle dans les contextes de grands lots, il n’est pas toujours le meilleur choix pour les petites tailles de lot (par exemple, en dessous de 1 024). Dans ces régimes, l’Adam standard ou la SGD avec momentum peuvent être plus simples et tout aussi efficaces. LAMB ajoute un surcoût de calcul pour calculer les normes par couche, ce qui est négligeable sur du matériel moderne mais peut être non trivial pour les modèles avec de nombreuses petites couches (par exemple, les architectures U-Net).
Une autre limitation est que le rapport de confiance de LAMB peut occasionnellement conduire à un entraînement instable si certaines couches ont des normes de poids très petites (proches de zéro). Cela est généralement atténué en ajoutant un terme epsilon au dénominateur et en utilisant une décroissance de poids, qui empêche les poids de dériver vers zéro. De plus, LAMB nécessite un réglage minutieux du taux d’apprentissage et des étapes d’échauffement ; un schéma inapproprié peut conduire à une divergence.
L’utilisation de la mémoire est similaire à celle d’Adam (deux vecteurs de moments par paramètre), donc il n’est pas plus gourmand en mémoire. Pour de très grands modèles, le Model Pruning ou le Gradient Clipping peuvent être utilisés en parallèle avec LAMB, mais ce sont des techniques orthogonales.
Conclusion
LAMB est devenu une pierre angulaire dans la boîte à outils des algorithmes d’optimisation pour l’apprentissage profond à grande échelle. En permettant un entraînement efficace avec des mini-lots massifs, il a accéléré le développement de nombreux modèles de référence et réduit le coût de l’expérimentation. Son principe d’adaptation par couche a influencé les conceptions d’optimiseurs ultérieures et reste une solution pratique et bien comprise pour les praticiens confrontés aux défis de l’entraînement distribué. Alors que l’Artificial intelligence continue de croître, des optimiseurs comme LAMB évolueront probablement, mais ses idées centrales de confiance par couche et de moments adaptatifs sont là pour rester.