Old Optimizer, New Norm: An Anthology
Authors: name: Jeremy Bernstein; name: Laker Newhouse
Published: 2024-09-30 14:26:12

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 R^{mtimes 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.