Old Optimizer, New Norm: An Anthology

TL;DR

The paper reframes three popular deep-learning optimizers—Adam, Shampoo, and Prodigy—as fundamentally first-order steepest-descent methods under carefully chosen layerwise norms, discarding convexity or second-order assumptions. It shows Adam without EMA is equivalent to sign descent under a max-of-max norm, Shampoo without accumulation corresponds to spectral-norm steepest descent and a projection to semi-orthogonal matrices, and Prodigy without EMA behaves as sign descent with an adaptive step size that aims for escape velocity. A modular norm framework is introduced to assign different norms to different layers, enabling principled optimizer design that accounts for the role of each tensor in the network. The work also provides computational strategies for Shampoo and rigorous proofs of the core steepest-descent connections, highlighting a new design space where choosing the right norm and step size drives training stability, scalability, and speed.

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.

Figures (1)

• Figure 1: Steepest descent considers the problem of minimizing a linear functional under a quadratic penalty: $\mathop{\mathrm{arg\,min}}\limits_{\Delta {\bm{w}} \in \mathbb{R}^n} \left[{\bm{g}}^\top \Delta {\bm{w}} + \frac{\lambda}{2} \, \Vert {\Delta {\bm{w}}} \Vert^2 \right]$ for ${\bm{g}} \in \mathbb{R}^n$. Here we show how the solution varies with the sharpness $\lambda > 0$ and the choice of norm $\Vert {\cdot} \Vert$. We overlay different norm balls on top of a linear color gradient, and use arrows to denote the solution, meaning the member of the norm ball that "minimizes the color". a) Increasing the sharpness decreases the size of the solution vector. b) Changing the norm can change the direction of the solution vector. For different $\ell_p$ norms, the solution direction changes because the gradient is not axis-aligned. In practice, we should pick the sharpness and norm to fit the geometry of our loss.

Theorems & Definitions (9)

• Proposition 1: Steepest descent
• Definition 1: Induced operator norm
• Proposition 2: Sign descent as steepest descent under the infinity norm
• Proposition 3: Sign descent as steepest descent under the max-of-max norm
• Proposition 4: Projection to the closest semi-orthogonal matrix
• Proposition 5: Shampoo as steepest descent under the spectral norm
• Proposition 6: Bounding the square loss of a linear predictor
• Proposition 7: Steepest descent under the modular norm
• Proposition 8: $\ell_1\to\ell_p$ and $\ell_p \to \ell_\infty$ induced operator norms are tractable