Training Deep Learning Models with Norm-Constrained LMOs

TL;DR

A new optimization framework using norm-ball-based LMO that unifies several existing methods and accelerates deep learning training without Adam.

Problem

Deep learning optimization primarily relies on adaptive methods like Adam, which have high memory usage and difficult hyperparameter tuning. Additionally, there is a lack of a general framework that unifies various optimization methods.

Approach

We propose a family of stochastic optimization algorithms that use the constrained LMO to adapt to the geometry of the problem. We extend it to unconstrained problems and propose an explicit choice of norm for deep architectures, enabling hyperparameter transferability.

Results & Contribution

The proposed algorithm, Scion, achieves up to 2x speedup over Adam in nanoGPT training, is memory-efficient (requiring only one set of weights and one set of gradients), and demonstrates hyperparameter transferability across model sizes.