Optimisation par lots importants pour l'apprentissage profond : formation de BERT en 76 minutes

L'entraînement de grands réseaux de neurones profonds sur des ensembles de données massifs est un défi informatique majeur. Ces dernières années, l'intérêt pour l'utilisation de méthodes d'optimisation stochastique à grands échantillons pour relever ce défi a considérablement augmenté. L'algorithme le plus notable dans cette lignée de recherche est LARS, qui, en utilisant des taux d'apprentissage adaptatifs par couche, entraîne ResNet sur ImageNet en quelques minutes. Cependant, LARS se montre inefficace pour les modèles d'attention comme BERT, indiquant que ses gains de performance ne sont pas constants entre les tâches. Dans cet article, nous étudions tout d'abord une stratégie d'adaptation par couche rigoureusement fondée pour accélérer l'entraînement des réseaux de neurones profonds à l'aide de grands mini-échantillons. En utilisant cette stratégie, nous développons une nouvelle technique d'optimisation à grands échantillons adaptative par couche appelée LAMB ; nous fournissons ensuite une analyse de la convergence de LAMB ainsi que celle de LARS, montrant la convergence vers un point stationnaire dans des configurations générales non convexes. Nos résultats empiriques démontrent la supériorité de LAMB sur diverses tâches telles que l'entraînement de BERT et ResNet-50 avec très peu d'ajustement des hyperparamètres. En particulier, pour l'entraînement de BERT, notre optimiseur permet l'utilisation de très grands échantillons de 32868 sans aucune dégradation des performances. En augmentant la taille du batch jusqu'à la limite mémoire d'un TPUv3 Pod, le temps d'entraînement de BERT peut être réduit de 3 jours à seulement 76 minutes (Tableau 1). La mise en œuvre de LAMB est disponible à l'adresse suivante : https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py