$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers

TL;DR

This paper introduces $$LO, a meta-training approach for learned optimizers that enhances their ability to generalize to wider, deeper, and longer training tasks, significantly improving efficiency.

Contribution

The paper derives the Maximal Update Parametrization for learned optimizers and proposes a simple meta-training recipe that improves their meta-generalization capabilities.

Findings

Meta-trained $$LOs outperform standard parametrization LOs on unseen wider tasks.

$$LOs show improved generalization to deeper networks.

Enhanced generalization to longer training horizons.

Abstract

Learned optimizers (LOs) have the potential to significantly reduce the wall-clock training time of neural networks. However, they can struggle to optimize unseen tasks (meta-generalize), especially when training networks wider than those seen during meta-training. To address this, we derive the Maximal Update Parametrization ($\mu$P) for two state-of-the-art learned optimizer architectures and propose a simple meta-training recipe for $\mu$-parameterized LOs ($\mu$LOs). Our empirical evaluation demonstrates that LOs meta-trained with our recipe substantially improve meta-generalization to wider unseen tasks when compared to LOs trained under standard parametrization (SP) using the same compute budget. We also empirically observe that $\mu$LOs exhibit unexpectedly improved meta-generalization to deeper networks ($5\times$ meta-training) and surprising generalization to much longer training horizons ($25\times$ meta-training) when compared to SP LOs.