μLO: Rekenkundig efficiënte meta-generalizatie van geleerde optimalisatoren

μ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

Samenvatting

Geoptimaliseerde optimalisatoren (LOs) kunnen de werkelijke trainingstijd van neurale netwerken aanzienlijk verkorten, waardoor de trainingskosten aanmerkelijk worden verlaagd. Ze lijden echter vaak aan slechte meta-generalizatie, vooral bij het trainen van netwerken die groter zijn dan die tijdens de meta-training zijn gezien. Om dit aan te pakken, gebruiken we de recent voorgestelde Maximal Update Parametrization (muP), die zero-shot generalizatie van optimalisatorhyperparameters van kleinere naar grotere modellen mogelijk maakt. We breiden de muP-theorie uit naar geleerde optimalisatoren door het meta-trainingsprobleem te behandelen als het vinden van de geleerde optimalisator onder muP. Onze evaluatie toont aan dat LOs die met muP zijn getraind, de meta-generalizatie aanzienlijk verbeteren in vergelijking met LOs die onder standaard parametrizatie (SP) zijn getraind. Opmerkelijk is dat, wanneer toegepast op modellen met grote breedte, onze beste muLO, getraind voor 103 GPU-uren, de prestaties evenaart of overtreft van VeLO, de grootste openbaar beschikbare geleerde optimalisator, die is getraind met 4000 TPU-maanden aan rekenkracht. Bovendien tonen muLOs betere generalizatie dan hun SP-tegenhangers naar diepere netwerken en naar veel langere trainingshorizons (25 keer langer) dan die tijdens de meta-training zijn gezien.

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.