Celo: Training Versatile Learned Optimizers on a Compute Diet

TL;DR

Celo introduces a new learned optimizer that achieves strong meta-generalization and outperforms existing optimizers on diverse tasks with significantly less computational resources, enabling practical off-the-shelf use.

Contribution

The paper presents Celo, a learned optimizer architecture and training procedure that significantly improves meta-generalization with minimal meta-training time.

Findings

Celo outperforms state-of-the-art optimizers on diverse tasks.

Celo requires only 24 GPU hours for meta-training.

Celo demonstrates strong generalization to out-of-distribution tasks.

Abstract

Learned optimization has emerged as a promising alternative to hand-crafted optimizers, with the potential to discover stronger learned update rules that enable faster, hyperparameter-free training of neural networks. A critical element for practically useful learned optimizers, that can be used off-the-shelf after meta-training, is strong meta-generalization: the ability to apply the optimizers to new tasks. Recent state-of-the-art work in learned optimizers, VeLO (Metz et al., 2022), requires a large number of highly diverse meta-training tasks along with massive computational resources, 4000 TPU months, to achieve meta-generalization. This makes further improvements to such learned optimizers impractical. In this work, we identify several key elements in learned optimizer architectures and meta-training procedures that can lead to strong meta-generalization. We also propose evaluation metrics to reliably assess quantitative performance of an optimizer at scale on a set of evaluation tasks. Our proposed approach, Celo, makes a significant leap in improving the meta-generalization performance of learned optimizers and also outperforms tuned state-of-the-art optimizers on a diverse set of out-of-distribution tasks, despite being meta-trained for just 24 GPU hours.