A thorough reproduction and evaluation of µP

This paper is an independent empirical reproduction of the claimed benefits of the µP parametrization proposed in Yang & Hu (2020) and Yang et al. (2021). Under the so-called Standard Parametrization (SP), the weights of neural networks are initialized from the Gaussian distribution with variance scaling as the inverse of “fan-in”, with the learning rate being the same for every layer. While this guarantees that (pre)activations are $\mathcal{O}$(1) at initialization with respect to width, it causes their scale to be width-dependent during training. To address this, Yang & Hu (2020) and Yang et al. (2021) proposed the Maximal Update Parametrization (µP), which is also claimed to make the optimal value of various hyperparameters independent of width. However, despite its alleged benefits, µP has not gained much traction among practitioners. Possibly, this could stem from a lack of thorough independent evaluation of µP against SP. We address this by independently reproducing the empirical claims of the original works. At the same time, we substantially increase the scale of the experiments, by training 16000 neural networks of sizes from 500 to 1B parameters, and empirically investigate µP’s effect on outputs, gradient updates, weights, training loss and validation loss. We find that generally µP indeed delivers on its promises, even though this does not always translate to improved generalization.