A Proof of Learning Rate Transfer under $\mu$P

TL;DR

This paper proves that in $$-parametrized neural networks, the optimal learning rate remains non-zero as width increases, explaining learning rate transfer, unlike other parametrizations.

Contribution

It provides the first proof of learning rate transfer under $$-parametrization in wide neural networks, supported by empirical validation.

Findings

Optimal learning rate converges to a non-zero constant under $$-parametrization.

Learning rate transfer fails under Standard and Neural Tangent Parametrizations.

Theoretical results are supported by extensive empirical experiments.

Abstract

We provide the first proof of learning rate transfer with width in a linear multi-layer perceptron (MLP) parametrized with $\mu$P, a neural network parameterization designed to ``maximize'' feature learning in the infinite-width limit. We show that under $\mu P$, the optimal learning rate converges to a \emph{non-zero constant} as width goes to infinity, providing a theoretical explanation to learning rate transfer. In contrast, we show that this property fails to hold under alternative parametrizations such as Standard Parametrization (SP) and Neural Tangent Parametrization (NTP). We provide intuitive proofs and support the theoretical findings with extensive empirical results.