Old Optimizer, New Norm: An Anthology

TL;DR

This paper reinterprets popular deep learning optimizers as first-order methods under specific norms, proposing a new perspective that could improve training stability and speed by customizing norms for different network components.

Contribution

It introduces a unified first-order framework for optimizers like Adam, Shampoo, and Prodigy, based on norm choices, and suggests a new design space for training algorithms.

Findings

Optimizers can be viewed as steepest descent under specific norms.

Different network layers should be assigned different operator norms.

This perspective may lead to more stable and scalable training methods.

Abstract

Deep learning optimizers are often motivated through a mix of convex and approximate second-order theory. We select three such methods -- Adam, Shampoo and Prodigy -- and argue that each method can instead be understood as a squarely first-order method without convexity assumptions. In fact, after switching off exponential moving averages, each method is equivalent to steepest descent under a particular norm. By generalizing this observation, we chart a new design space for training algorithms. Different operator norms should be assigned to different tensors based on the role that the tensor plays within the network. For example, while linear and embedding layers may have the same weight space of $\mathbb{R}^{m\times n}$, these layers play different roles and should be assigned different norms. We hope that this idea of carefully metrizing the neural architecture might lead to more stable, scalable and indeed faster training.