Source-linked AI summary
Old Optimizer, New Norm: An Anthology
Jeremy Bernstein, Laker Newhouse
TL;DR
Deep learning optimizers are commonly motivated by convex or approximate second-order theory, while the paper explores a first-order account without convexity assumptions. It interprets Adam, Shampoo, and Prodigy as steepest descent under particular norms after disabling EMA, then uses this separation of norm and step size to propose a broader norm-based optimizer design space.
Problem
Deep learning optimizers are often motivated through convex or approximate second-order theory, leaving untapped algorithmic opportunity in exact first-order theory without convexity assumptions.
Method
The paper analyzes Adam, Shampoo, and Prodigy after disabling EMA, expressing each as steepest descent under a particular norm and separating norm choice from step-size selection.
Results
Adam and Prodigy use sign descent under the infinity or max-of-max norm, while Shampoo uses steepest descent under the spectral norm.
Takeaways & Limitations
The norm-based view opens a design space in which optimizer norms can be chosen more intentionally for the structure and role of tensors in a neural network.
Takeaways & Limitations
Choosing a norm and sharpness suited to the optimization problem remains an open design challenge beyond the anthology’s scope, and Prodigy’s monotone step-size increase may be sub-optimal.
Abstract
from arXiv · showhide
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.
Prologue
The paper reframes Adam, Shampoo, and Prodigy as first-order steepest-descent methods under particular norms after disabling exponential moving averages. This norm-based view separates step direction from step size and motivates designing optimizers around the roles of tensors within neural architectures.
- Prologue: Adam, Shampoo, and Prodigy can each be understood as steepest descent under a particular norm after their exponential moving averages are disabled.The paper presents this as a first-order explanation without convexity or approximate Hessian assumptions.
- Prologue: Steepest descent chooses a norm and sharpness parameter in advance, without using an approximate Hessian during training.The norm controls the geometry of the update, while sharpness controls its scale.
- Prologue: Varying vector norms induces a large family of matrix norms and corresponding steepest-descent optimizers.The authors argue that optimizer design should become more intentional about the norm assigned to the optimization problem.
- Prologue: Adam without EMA is sign descent, which is steepest descent under the infinity norm and also respects neural-network matrix structure through a max-of-max norm.This connects a flattened-vector interpretation with a structured, layer-aware interpretation.
- Prologue: Shampoo without accumulation performs a semi-orthogonal gradient update and corresponds to steepest descent under the maximum spectral norm across layers.The paper also connects this update to minimizing a spectral-norm quadratic upper bound for linear-predictor square loss.
- Prologue: Prodigy without EMA uses the same sign-descent direction as Adam but selects a dynamically increasing step size through a heuristic warm-up procedure.The authors interpret this as separating norm choice from step-size selection, while noting that Prodigy’s monotone increase and other line-search choices may be sub-optimal.
Epilogue
The epilogue frames optimizer design as choosing a norm and a step size, then generalizes this view to assign norms according to each layer’s role. It presents this as a route toward more intentional training algorithms and potentially improved speed and scalability.
- Adam, Shampoo, and Prodigy fit a framework that separates optimizer design into choosing a norm and choosing a step size.
- Modular norm: The modular norm assigns each layer its own norm and combines the resulting scaled layer norms through a maximum.
- Modular norm: Under the modular norm, each layer’s update direction reflects its assigned norm, while the global step size uses a weighted sum of gradient dual norms.
- Norm assignment: Induced operator norms such as ℓ1 → ℓp and ℓp → ℓ8 are computationally tractable, with interpretations based on maximum column or row norms.
- Norm assignment: Norms should be selected according to a layer’s role, since linear and embedding layers can share a weight space while serving different network functions.
- The paper suggests that better norms may improve training speed and scalability, while presenting this as a direction for future optimizer design.
Appendix A. Computational Strategies for Shampoo
The appendix describes several ways to compute or approximate Shampoo’s update, which equals the orthogonal factor UV^T of a gradient’s reduced SVD. It also discusses iterative computation, convergence conditions, alternative polynomial iterations, and practical trade-offs.
- With EMA disabled, Shampoo’s update for a gradient G with reduced SVD G = UΣV^T is -ηUV^T.
- Computational strategies: The update can be computed by SVD, randomized sketching, Newton iteration for inverse matrix roots, or Newton–Schulz iteration.
- Newton–Schulz iteration: Newton–Schulz iteration initializes a normalized gradient matrix and converges toward UV^T as the iteration count tends to infinity.
- Newton–Schulz iteration: The cubic iteration acts independently on singular values, and convergence requires the initial singular values to lie in a suitable positive bounded interval.
- Polynomial alternatives: Higher-degree odd polynomial iterations can approximate sign(x), with coefficients tunable for faster convergence.
- The practical usefulness of these computational methods may depend on the gradient matrix’s condition number and available computational resources.
Appendix B. Proofs
The appendix derives steepest-descent solutions by separating update magnitude from direction, then applies this framework to modular and matrix operator norms. It also characterizes semi-orthogonal projections and tractable induced operator norms.
- Steepest descent: Steepest-descent updates decompose into a nonnegative magnitude and a unit-norm direction, whose separate minimizers combine into the overall update.The direction and magnitude are optimized independently before being multiplied to obtain the minimizing update.
- Steepest descent: For the ℓ∞ direction constraint, the maximizing direction is sign(g), and the corresponding dual norm is ||g||1.This follows from maximizing the gradient inner product over unit ℓ∞ vectors.
- Semi-orthogonal projection: For a matrix G, minimizing Frobenius distance over semi-orthogonal matrices is equivalent to maximizing Frobenius alignment; the solution A⋆ = UV^T attains the upper bound.With a reduced SVD of G, the semi-orthogonal optimizer is expressed using the corresponding singular-vector matrices.
- Semi-orthogonal projection: The semi-orthogonal solution is unique when G is full rank, but not unique when G is rank-deficient.Full rank requires all singular values to be positive; rank deficiency leaves multiple equally good solutions.
- Modular norms: The modular-norm minimizer assigns each layer an update proportional to its alignment-maximizing direction, scaled by η/sl.The common scalar η is determined by the layerwise magnitude constraints and yields the overall minimizer.
- Induced operator norms: The induced norms satisfy ||M||ℓ1→ℓp = maxj ||colj(M)||p, while the ℓp→ℓ∞ case follows from vector-norm duality.The ℓ1→ℓp maximum is attained by selecting the column with largest ℓp norm; the conjugate exponents obey 1/p + 1/q = 1.