$μ$LO: Compute-Efficient Meta-Generalization of Learned Optimizers (2406.00153v3)

Abstract: Learned optimizers (LOs) can significantly reduce the wall-clock training time of neural networks, substantially reducing training costs. 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), as they are trained in existing work. We also empirically observe that $\mu$LOs trained with our recipe 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.

• The paper extends μP theory to learned optimizers, enabling robust zero-shot hyperparameter generalization from small to large models.
• The paper empirically validates μLO’s performance, showing that 103 GPU-hours of training rivals state-of-the-art results achieved with much higher compute.
• The paper demonstrates μLOs’ capacity to generalize to longer training durations and deeper architectures, underscoring their practical compute efficiency.

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

The paper "$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers" presents a substantive advancement in the field of learned optimizers (LOs). The focal point of the research is to bridge the gap in meta-generalization capabilities of LOs, especially when applied to larger neural network models than those encountered during meta-training. The authors employ the recently proposed Maximal Update Parametrization (denoted as $\mu$P), which facilitates zero-shot generalization of optimizer hyperparameters across varying model scales by aligning the gradient distribution between small and large neural networks.

Main Contributions

1. Extension of $\mu$P Theory to Learned Optimizers:
   • The researchers extrapolate the $\mu$P theory, initially formulated for adaptive optimizers, to the paradigm of learned optimizers. This theoretical augmentation is pivotal as it facilitates robust generalization of learned optimizer parameters from smaller to larger models.
2. Empirical Validation of $\mu$LO:
   • Through comprehensive empirical analysis, the paper establishes that LOs trained with $\mu$P (termed $\mu$LOs) manifest substantial improvements in meta-generalization over their standard parameterization (SP) counterparts.
   • Remarkably, the $\mu$LO trained for 103 GPU-hours matches or exceeds the performance of VeLO, a state-of-the-art LO meta-trained using 4000 TPU-months of compute, particularly on larger-width neural networks.
3. Generalization to Longer Training Horizons and Deeper Networks:
   • The paper further demonstrates that $\mu$LOs exhibit enhanced generalization capabilities to training durations extending 25 times longer than those observed during meta-training and to deeper neural network architectures compared to their SP-trained equivalents.

Empirical Findings and Analysis

The empirical results are segmented into various dimensions of generalization, delineating the following:

1. Generalization to Wider Networks:
   • The experiments showcase the superior performance of $\mu$LOs on tasks involving networks with a wide array of widths, including those significantly larger than those used during meta-training. For instance, $\mu$LOs not only converge faster but also achieve lower training losses compared to VeLO and other baselines on these wider networks.
2. Generalization to Larger Input Images:
   • Performance evaluation on classification tasks with larger image inputs reveals that $\mu$LOs maintain stability and achieve better performance compared to SP LOs and partially surpass VeLO's performance. This reinforces the potent generalization capacity of the $\mu$P framework.
3. Generalization to Different Datasets:
   • The analysis covers additional datasets such as CIFAR-10, demonstrating that $\mu$LOs generalize seamlessly to different data distributions, showcasing the general-purpose applicability of the $\mu$P approach.

Theoretical and Practical Implications

Theoretical Implications:

• The extension of $\mu$P to learned optimizers is a valuable theoretical contribution, underpinning the notion that parametric stability and gradient alignment can be extrapolated to meta-learning frameworks. This yields theoretical insights into the scaling properties of optimizer hyperparameters and offers a stable mathematical foundation for future explorations in the domain of learned optimization strategies.

Practical Implications:

• Practically, the research heralds a significant reduction in the computational expense required for training effective LOs. By ensuring that LOs generalize robustly to larger and more complex models without necessitating extensive compute resources, $\mu$P provides a cost-effective alternative to current methodologies.
• This is particularly relevant for scaling deep learning models in resource-constrained environments, making advanced optimization strategies accessible to a broader set of applications.

Speculations on Future Developments in AI:

• Building on the successes of $\mu$LO, future research could explore the integration of $\mu$P into other domains of machine learning such as reinforcement learning optimizers and GAN training. There is potential for $\mu$P to influence the development of universally generalizable optimization algorithms, ultimately contributing to more efficient and scalable AI systems.
• Additionally, extending $\mu$P principles to unsupervised and semi-supervised training regimes could be an intriguing avenue for research, potentially revolutionizing how models are trained across various data modalities and sparsity constraints.

Conclusion

The paper "Compute-Efficient Meta-Generalization of Learned Optimizers" delineates a significant stride in enhancing the efficacy and efficiency of learned optimizers. By leveraging the Maximal Update Parametrization, the authors provide a robust framework that facilitates superior meta-generalization capabilities, all while necessitating markedly fewer computational resources. The empirical evidence supporting the efficacy of $\mu$LOs underscores the potential for $\mu$P to redefine optimizer training paradigms and extend the applicability of advanced optimization techniques to broader, more computationally constrained settings.