μLO: Compute-Efficient Meta-Generalization of Learned Optimizers
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.
665e929101ca1c80a0cef855
Benjamin Thérien
admin_assigned
2024-06-04T08:38:07.719Z
62d0e2a3fb896639b2a04ca1
https://cdn-avatars.huggingface.co/v1/production/uploads/62d0e2a3fb896639b2a04ca1/UBZzdldkvCAEHp-t6L2JJ.jpeg
Benjamin Therien
user
btherien
665e929101ca1c80a0cef856
Charles-Étienne Joseph
admin_assigned
2024-06-04T08:38:13.933Z
65c0016fdaa6f48256f0799f
/avatars/93658acaaaee76eca64df6b269172cc6.svg
Charles-Étienne Joseph
user
lefameuxbeding
665e929101ca1c80a0cef857
Boris Knyazev
admin_assigned
2024-06-04T08:38:19.906Z
63e16165f0039731dfdd442a
/avatars/37cf99dc016c291c800f60d260173482.svg
Boris Knyazev
user
bknyaz
665e929101ca1c80a0cef858
Edouard Oyallon
665e929101ca1c80a0cef859
Irina Rish
admin_assigned
2024-06-04T08:38:35.376Z
6227f764236b7b2eefa7f173
/avatars/0204e3da8e9ef12d88c98c6e955f8340.svg
Irina Rish
user
irinarish
665e929101ca1c80a0cef85a
Eugene Belilovsky
μLO: Compute-Efficient Meta-Generalization of Learned Optimizers
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.
Maximal Update Parametrization improves the meta-generalization of learned optimizers, allowing zero-shot generalization to larger and deeper models with reduced computational cost.
Learned optimizers Maximal Update Parametrization meta-generalization zero-shot generalization large-width models deeper networks training horizons