English

Depth Dependence of $\mu$P Learning Rates in ReLU MLPs

Machine Learning 2023-05-16 v1 Machine Learning

Abstract

In this short note we consider random fully connected ReLU networks of width nn and depth LL equipped with a mean-field weight initialization. Our purpose is to study the dependence on nn and LL of the maximal update (μ\muP) learning rate, the largest learning rate for which the mean squared change in pre-activations after one step of gradient descent remains uniformly bounded at large n,Ln,L. As in prior work on μ\muP of Yang et. al., we find that this maximal update learning rate is independent of nn for all but the first and last layer weights. However, we find that it has a non-trivial dependence of LL, scaling like L3/2.L^{-3/2}.

Keywords

Cite

@article{arxiv.2305.07810,
  title  = {Depth Dependence of $\mu$P Learning Rates in ReLU MLPs},
  author = {Samy Jelassi and Boris Hanin and Ziwei Ji and Sashank J. Reddi and Srinadh Bhojanapalli and Sanjiv Kumar},
  journal= {arXiv preprint arXiv:2305.07810},
  year   = {2023}
}