English

An Empirical Study of $\mu$P Learning Rate Transfer

Machine Learning 2025-02-17 v6

Abstract

Deep learning models have become a cornerstone of modern AI research, yet their initializations and learning rates may at times be set in an opaque or ad-hoc fashion due to the high cost of hyperparameter sweeps. The μ\mu-Parameterization (μ\muP) offers a possible solution to this challenge, yielding scaling rules for model initialization and learning rates while reportedly enabling zero-shot hyperparameter transfer from small to large models. Despite its evident promise, the μ\muP method is not yet widely adopted, perhaps due to higher implementation complexity, many variations, or complex theoretical background. This work considers μ\muP empirically, focusing on the popular transformer architecture, and aims to answer a simple question: does μ\mu-Transfer yield near-optimal learning rates in practice? Studying over a dozen ablations with up to 1.2B parameters and 33B tokens and a large-scale experiment with up to 10B parameters and 190B tokens, we observe a positive answer for most settings, and discuss improvements otherwise.

Keywords

Cite

@article{arxiv.2404.05728,
  title  = {An Empirical Study of $\mu$P Learning Rate Transfer},
  author = {Lucas Lingle},
  journal= {arXiv preprint arXiv:2404.05728},
  year   = {2025}
}

Comments

Improved exposition, added backlog of old experiments