Celo2: Towards Learned Optimization Free Lunch

TL;DR

This paper introduces a simple, normalized learned optimizer that can be meta-trained efficiently and generalizes well across large-scale and diverse tasks, making learned optimization more practical.

Contribution

It demonstrates that a normalized optimizer architecture, combined with augmented meta-training, enables effective meta-learning with minimal compute and scales to billion-parameter tasks.

Findings

Meta-training requires only 4.5 GPU hours.

Scales effectively to GPT-3 XL 1.3B parameters.

Performs well on out-of-distribution tasks.

Abstract

Learned optimizers are powerful alternatives to hand-designed update rules like Adam, yet they have seen limited practical adoption since they often fail to meta-generalize beyond their training distribution and incur high meta-training cost. For instance, prior work, VeLO, scaled meta-training to 4,000 TPU months ($\sim$10$\times$ GPT-3 compute) to meta-train a general-purpose optimizer but it failed to generalize beyond 600M parameters tasks. In this work, we present a surprising finding: by crafting a simple normalized optimizer architecture and augmenting meta-training, it becomes feasible to meta-train a performant general-purpose learned update rule on a tiny fraction of VeLO compute, 4.5 GPU hours to be precise. Our learned update rule scales stably to a billion-scale pretraining task (GPT-3 XL 1.3B) which is six orders of magnitude larger than its meta-training distribution. Furthermore, it shows strong performance across diverse out-of-distribution tasks and is compatible with modern optimization harness that includes orthogonalization, distinct update rules for input-output and hidden weights, and decoupled weight decay. In all, this work paves the way for practically applicable learnable optimization algorithms, unlocking exploration of richer meta-training and data curation recipes to further improve performance.