A Proof of Learning Rate Transfer under $μ$P (2511.01734v1)

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.

• The paper demonstrates that the optimal learning rate under μP converges to a constant, enabling consistent hyperparameter transfer from small to large models.
• It employs a polynomial analysis of training loss in linear MLPs, revealing a dominant linear term that simplifies learning rate tuning as width increases.
• Empirical results confirm that μP preserves feature learning and stable training dynamics, unlike standard and neural tangent parametrizations.

Proof and Implications of Learning Rate Transfer under $\mu$P

Introduction and Motivation

The paper establishes a rigorous theoretical foundation for learning rate (LR) transfer in neural networks parametrized with Maximal Update Parametrization ($\mu$P). The central result is that, for deep linear Multi-Layer Perceptrons (MLPs) trained with gradient descent, the optimal learning rate converges to a non-zero constant as the model width $n \to \infty$. This property, termed "learning rate transfer," enables hyperparameter tuning on small models and direct transfer to larger models, significantly reducing computational cost and complexity in large-scale neural network training. The analysis contrasts $\mu$P with Standard Parametrization (SP) and Neural Tangent Parametrization (NTP), showing that only $\mu$P supports robust LR transfer.

Figure 1: Optimal learning rate as a function of model width, demonstrating convergence to a non-zero constant under $\mu$P and vanishing optimal LR under SP.

Theoretical Framework

Neural Parametrizations

The paper formalizes neural parametrizations as explicit scaling rules for initialization and learning rate with respect to model width. For $\mu$P, the initialization of the output layer $V$ uses variance $n^{-2}$, and the learning rate is held constant for gradient descent (or scaled as $\eta n^{-1}$ for Adam). In contrast, SP uses $n^{-1}$ variance for $V$ and does not scale the learning rate, leading to suboptimal training dynamics as width increases.

Definition of Learning Rate Transfer

Learning rate transfer is defined as the convergence in probability of the optimal learning rate $\eta_n$ for width $n$ to a deterministic constant $\eta_\infty > 0$ as $n \to \infty$. This property is crucial for practical scalability, as it allows hyperparameter settings to remain stable across model sizes.

Main Results

Polynomial Structure of Training Loss

The analysis leverages the fact that, for linear MLPs, the training loss after $t$ steps can be expressed as a polynomial in the learning rate $\eta$, with coefficients determined by the random initialization. Under $\mu$P, higher-order coefficients vanish as $n$ increases, leaving a dominant linear term. This leads to a quasi-quadratic loss landscape in $\eta$ and enables precise characterization of the optimal learning rate.

Convergence of Optimal Learning Rate

The paper proves that, under $\mu$P, the optimal learning rate $\eta_n^{(t)}$ converges to a non-zero constant $\eta_\infty^{(t)}$ at each training step $t$, with convergence rate $O(n^{-1/2})$. The limiting value is given by a closed-form expression involving the input Gram matrix $K$ and target vector $y$:

$$\eta_\infty^{(1)} = \frac{m}{L} \frac{y^\top K y}{\|K y\|^2}$$

where $m$ is the dataset size and $L$ is the network depth.

Figure 2: Train loss as a function of learning rate at $t=5$ and $t=10$ for both $\mu$P and SP, illustrating stable optimal LR under $\mu$P and shifting optimum under SP.

Failure of LR Transfer in SP and NTP

For SP, the optimal learning rate vanishes as $n \to \infty$, and for NTP, it diverges. This is attributed to the scaling of the output layer and the lack of feature learning in the infinite-width limit, resulting in either vanishing or exploding feature updates.

Generalization to Arbitrary Training Steps

The proof is extended to arbitrary training steps $t$, showing that the polynomial structure persists and the optimal learning rate continues to converge to a non-zero constant under $\mu$P, provided the limiting loss has a unique minimizer.

Empirical Validation

Extensive experiments confirm the theoretical predictions. Linear and non-linear MLPs (with ReLU activation), varying depths, and different optimizers (SGD, Adam) all exhibit robust LR transfer under $\mu$P. The convergence rate of the optimal LR matches theoretical expectations for moderate widths and accelerates for larger widths.

Figure 3: Train loss as a function of learning rate at $t=20$ for linear and ReLU MLPs of varying depth, showing consistent LR transfer under $\mu$P.

Figure 4: Train loss as a function of learning rate at $t=100$ for a deep ReLU MLP trained with Adam, demonstrating LR transfer near convergence.

Practical and Theoretical Implications

Hyperparameter Tuning Efficiency

The results provide a rigorous justification for the empirical practice of tuning learning rates on small models and transferring them to large models when using $\mu$P. This dramatically reduces the computational burden of hyperparameter search in large-scale neural network training.

Feature Learning in the Infinite-Width Limit

The analysis clarifies that $\mu$P is unique in preserving feature learning as width increases, in contrast to SP and NTP, which either suppress or exaggerate feature updates. This has implications for the design of scalable neural architectures and the theoretical understanding of training dynamics.

Extension to Non-Linear Networks and Other Optimizers

While the proofs are restricted to linear MLPs and gradient descent, empirical results suggest that LR transfer under $\mu$P extends to non-linear architectures and optimizers such as Adam. Future work may generalize the theoretical framework to these cases.

Scaling Laws and Depth Dependence

The limiting optimal learning rate depends on network depth, suggesting that depth-aware versions of $\mu$P (e.g., Depth-$\mu$P) may be necessary for optimal scaling in very deep networks.

Limitations and Future Directions

The current analysis is limited to linear networks and does not address the full complexity of non-linear activations or adaptive optimizers. Extending the proof techniques to these settings will require new mathematical tools, particularly for controlling large-width deviations and non-quadratic loss landscapes. Additionally, the empirical observation of accelerated convergence rates for very large widths warrants further investigation.

Conclusion

This work provides the first rigorous proof of learning rate transfer under $\mu$P, establishing that the optimal learning rate converges to a non-zero constant as model width increases. The results have significant practical implications for scalable neural network training and deepen the theoretical understanding of feature learning in the infinite-width regime. The findings motivate further research into generalizing LR transfer to non-linear networks, alternative optimizers, and more complex architectures.