Saved in:
Bibliographic Details
Main Authors: Large, Tim, Liu, Yang, Huh, Minyoung, Bahng, Hyojin, Isola, Phillip, Bernstein, Jeremy
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2405.14813
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916257875361792
author Large, Tim
Liu, Yang
Huh, Minyoung
Bahng, Hyojin
Isola, Phillip
Bernstein, Jeremy
author_facet Large, Tim
Liu, Yang
Huh, Minyoung
Bahng, Hyojin
Isola, Phillip
Bernstein, Jeremy
contents To improve performance in contemporary deep learning, one is interested in scaling up the neural network in terms of both the number and the size of the layers. When ramping up the width of a single layer, graceful scaling of training has been linked to the need to normalize the weights and their updates in the "natural norm" particular to that layer. In this paper, we significantly generalize this idea by defining the modular norm, which is the natural norm on the full weight space of any neural network architecture. The modular norm is defined recursively in tandem with the network architecture itself. We show that the modular norm has several promising applications. On the practical side, the modular norm can be used to normalize the updates of any base optimizer so that the learning rate becomes transferable across width and depth. This means that the user does not need to compute optimizer-specific scale factors in order to scale training. On the theoretical side, we show that for any neural network built from "well-behaved" atomic modules, the gradient of the network is Lipschitz-continuous in the modular norm, with the Lipschitz constant admitting a simple recursive formula. This characterization opens the door to porting standard ideas in optimization theory over to deep learning. We have created a Python package called Modula that automatically normalizes weight updates in the modular norm of the architecture. The package is available via "pip install modula" with source code at https://github.com/jxbz/modula.
format Preprint
id arxiv_https___arxiv_org_abs_2405_14813
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Scalable Optimization in the Modular Norm
Large, Tim
Liu, Yang
Huh, Minyoung
Bahng, Hyojin
Isola, Phillip
Bernstein, Jeremy
Machine Learning
To improve performance in contemporary deep learning, one is interested in scaling up the neural network in terms of both the number and the size of the layers. When ramping up the width of a single layer, graceful scaling of training has been linked to the need to normalize the weights and their updates in the "natural norm" particular to that layer. In this paper, we significantly generalize this idea by defining the modular norm, which is the natural norm on the full weight space of any neural network architecture. The modular norm is defined recursively in tandem with the network architecture itself. We show that the modular norm has several promising applications. On the practical side, the modular norm can be used to normalize the updates of any base optimizer so that the learning rate becomes transferable across width and depth. This means that the user does not need to compute optimizer-specific scale factors in order to scale training. On the theoretical side, we show that for any neural network built from "well-behaved" atomic modules, the gradient of the network is Lipschitz-continuous in the modular norm, with the Lipschitz constant admitting a simple recursive formula. This characterization opens the door to porting standard ideas in optimization theory over to deep learning. We have created a Python package called Modula that automatically normalizes weight updates in the modular norm of the architecture. The package is available via "pip install modula" with source code at https://github.com/jxbz/modula.
title Scalable Optimization in the Modular Norm
topic Machine Learning
url https://arxiv.org/abs/2405.14813