$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers
<|reference_start|>$\mu$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 ($\mu$P), which allows zero-shot generalization of optimizer hyperparameters from smaller to larger models. We extend $\mu$P theory to learned optimizers, treating the meta-training problem as finding the learned optimizer under $\mu$P. Our evaluation shows that LOs meta-trained with $\mu$P substantially improve meta-generalization as compared to LOs trained under standard parametrization (SP). Notably, when applied to large-width models, our best $\mu$LO, 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, $\mu$LOs 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.<|reference_end|>
@article{thérien2024$\mu$lo:,
    title={$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers},
    author={Benjamin Th'erien, Charles-'Etienne Joseph, Boris Knyazev, Edouard
  Oyallon, Irina Rish, Eugene Belilovsky},
    journal={arXiv preprint arXiv:2406.00153},
    year={2024},
    archivePrefix={arXiv},
    eprint={2406.00153},
    primaryClass={cs.LG}
}