μLO: Meta-Generalizzazione Efficiente dal Punto di Vista Computazionale degli Ottimizzatori Appresi

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

May 31, 2024

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

cs.AI

Abstract

Gli ottimizzatori appresi (LOs) possono ridurre significativamente il tempo di addestramento in termini di clock reale delle reti neurali, diminuendo sostanzialmente i costi di training. Tuttavia, spesso soffrono di una scarsa meta-generalizzazione, specialmente quando addestrano reti più grandi di quelle viste durante il meta-training. Per affrontare questo problema, utilizziamo la recente proposta della Parametrizzazione ad Aggiornamento Massimale (muP), che consente la generalizzazione zero-shot degli iperparametri dell'ottimizzatore da modelli più piccoli a quelli più grandi. Estendiamo la teoria muP agli ottimizzatori appresi, trattando il problema del meta-training come la ricerca dell'ottimizzatore appreso sotto muP. La nostra valutazione mostra che i LOs meta-addestrati con muP migliorano sostanzialmente la meta-generalizzazione rispetto ai LOs addestrati con la parametrizzazione standard (SP). In particolare, quando applicati a modelli di grande ampiezza, il nostro miglior muLO, addestrato per 103 ore GPU, eguaglia o supera le prestazioni di VeLO, il più grande ottimizzatore appreso disponibile pubblicamente, meta-addestrato con 4000 mesi TPU di calcolo. Inoltre, i muLOs dimostrano una migliore generalizzazione rispetto alle loro controparti SP a reti più profonde e a orizzonti di addestramento molto più lunghi (25 volte più lunghi) rispetto a quelli visti durante il meta-training.

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.