μLO : Méta-généralisation efficace en calcul des optimiseurs appris

μLO: Compute-Efficient Meta-Generalization of Learned Optimizers

May 31, 2024

Auteurs: Benjamin Thérien, Charles-Étienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, Eugene Belilovsky

cs.AI

Résumé

Les optimiseurs appris (LOs) peuvent réduire considérablement le temps d'entraînement en temps réel des réseaux de neurones, diminuant ainsi substantiellement les coûts de formation. Cependant, ils souffrent souvent d'une mauvaise méta-généralisation, en particulier lors de l'entraînement de réseaux plus grands que ceux rencontrés pendant la méta-formation. Pour remédier à cela, nous utilisons la Paramétrisation de Mise à Jour Maximale (muP) récemment proposée, qui permet une généralisation à zéro coup des hyperparamètres de l'optimiseur des modèles plus petits vers les modèles plus grands. Nous étendons la théorie muP aux optimiseurs appris, en traitant le problème de méta-formation comme la recherche de l'optimiseur appris sous muP. Notre évaluation montre que les LOs méta-formés avec muP améliorent considérablement la méta-généralisation par rapport aux LOs entraînés sous paramétrisation standard (SP). Notamment, lorsqu'ils sont appliqués à des modèles de grande largeur, notre meilleur muLO, entraîné pendant 103 heures GPU, égale ou dépasse les performances de VeLO, le plus grand optimiseur appris publiquement disponible, méta-formé avec 4000 mois TPU de calcul. De plus, les muLOs démontrent une meilleure généralisation que leurs homologues SP pour des réseaux plus profonds et pour des horizons d'entraînement beaucoup plus longs (25 fois plus longs) que ceux rencontrés pendant la méta-formation.

English

Learned optimizers (LOs) can significantly reduce the wall-clock training time of neural networks, substantially reducing training costs. However, they often suffer from poor meta-generalization, especially when training networks larger than those seen during meta-training. To address this, we use the recently proposed Maximal Update Parametrization (muP), which allows zero-shot generalization of optimizer hyperparameters from smaller to larger models. We extend muP theory to learned optimizers, treating the meta-training problem as finding the learned optimizer under muP. Our evaluation shows that LOs meta-trained with muP substantially improve meta-generalization as compared to LOs trained under standard parametrization (SP). Notably, when applied to large-width models, our best muLO, trained for 103 GPU-hours, matches or exceeds the performance of VeLO, the largest publicly available learned optimizer, meta-trained with 4000 TPU-months of compute. Moreover, muLOs demonstrate better generalization than their SP counterparts to deeper networks and to much longer training horizons (25 times longer) than those seen during meta-training.