μLO: Meta-Generalización Eficiente en Cómputo de Optimizadores Aprendidos

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

May 31, 2024

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

cs.AI

Resumen

Los optimizadores aprendidos (LOs, por sus siglas en inglés) pueden reducir significativamente el tiempo de entrenamiento en tiempo real de las redes neuronales, disminuyendo sustancialmente los costos de entrenamiento. Sin embargo, a menudo sufren de una mala meta-generalización, especialmente cuando se entrenan redes más grandes que las vistas durante el meta-entrenamiento. Para abordar esto, utilizamos la recientemente propuesta Parametrización de Actualización Máxima (muP, por sus siglas en inglés), que permite la generalización en cero disparos de los hiperparámetros del optimizador desde modelos más pequeños a modelos más grandes. Extendemos la teoría de muP a los optimizadores aprendidos, tratando el problema de meta-entrenamiento como la búsqueda del optimizador aprendido bajo muP. Nuestra evaluación muestra que los LOs meta-entrenados con muP mejoran sustancialmente la meta-generalización en comparación con los LOs entrenados bajo la parametrización estándar (SP, por sus siglas en inglés). Notablemente, cuando se aplican a modelos de gran anchura, nuestro mejor muLO, entrenado durante 103 horas de GPU, iguala o supera el rendimiento de VeLO, el optimizador aprendido disponible públicamente más grande, meta-entrenado con 4000 meses-TPU de cómputo. Además, los muLOs demuestran una mejor generalización que sus contrapartes SP en redes más profundas y en horizontes de entrenamiento mucho más largos (25 veces más largos) que los vistos durante el meta-entrenamiento.

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.