Steepest Descent in the Modular Norm

Authors: Jeremy Bernstein, Laker Newhouse
Primary area: optimization
Keywords: Adam, Shampoo, Prodigy, optimizers, optimization, steepest descent, norms, modular norm, spectral norm

Abstract:
An old idea in optimization theory says that since the gradient is a dual vector it may not be subtracted from the weights without first being mapped to the primal space where the weights live. We take this idea seriously in this paper and construct such a gradient mapping for neural networks. We call our mapping modular dualization, and show how it may be used to derive steepest descent optimization methods in an architecture-aware norm called the modular norm. We prove that the Adam and Shampoo optimizers (without exponential moving averages) emerge as special cases of modular dualization. We go on to develop new special-purpose dualize functions for linear, embedding and convolution layers. And we contribute a novel, GPU-friendly algorithm for dualizing under the spectral norm. We hope that the technique of modular dualization might enable a next generation of fast and scalable optimizers that are well-matched to general neural architectures.

Decision: reject

Reviews and discussion:

Metareview: