TLDR: A new research paper provides the first theoretical proof for learning rate transfer in deep linear neural networks under Maximal Update Parametrization (µP). It demonstrates that µP allows the optimal learning rate to converge to a stable, non-zero constant as network width increases, unlike Standard Parametrization (SP) and Neural Tangent Parametrization (NTP) where it shifts or diverges. This stability enables efficient hyperparameter tuning on smaller models for use in larger ones, a property supported by extensive empirical results across various network configurations, including non-linear models and different optimizers.
A recent research paper sheds light on a crucial aspect of training large neural networks: the stability of optimal learning rates as models grow in size. This phenomenon, known as learning rate transfer, is particularly important because it allows developers to tune hyperparameters on smaller models and apply them to much larger, more complex networks without extensive re-tuning, saving significant computational resources and time.
The paper, titled A Proof of Learning Rate Transfer under µP, provides the first rigorous theoretical proof for learning rate transfer in deep linear networks. The focus is on a specific neural network parameterization called Maximal Update Parametrization (µP). µP is designed to maximize feature learning in the infinite-width limit, a characteristic that the authors suggest is key to its superior performance in hyperparameter transfer.
The core finding is that under µP, the optimal learning rate converges to a non-zero constant as the network’s width increases indefinitely. This means that beyond a certain size, the ideal learning rate for a µP-parametrized model remains stable, making it predictable and transferable. This is a significant advantage over other common parameterizations, such as Standard Parametrization (SP) and Neural Tangent Parametrization (NTP), where the optimal learning rate either shifts towards zero or diverges as width grows, necessitating costly re-tuning for larger models.
The researchers achieved this proof by demonstrating that the loss function in linear Multi-Layer Perceptrons (MLPs) can be expressed as a polynomial function of the learning rate at any given training step. By analyzing the convergence dynamics of these polynomials and their roots as network width approaches infinity, they were able to show that the optimal learning rate indeed converges to a stable, non-zero value under µP. In contrast, for SP and NTP, the coefficients of these polynomials behave differently, leading to unstable optimal learning rates.
Empirical results strongly support these theoretical findings. Experiments with linear MLPs of varying widths showed that under µP, the optimal learning rate quickly stabilizes, while under SP, it consistently decreases towards zero with increasing width. The paper also extends its analysis beyond the initial training step, proving learning rate transfer for general training steps under mild conditions, further solidifying the practical utility of µP.
Beyond linear networks, the research explored more challenging setups, including non-linear MLPs with ReLU activation functions and different optimizers like Adam. Even in these cases, the empirical evidence suggests that learning rate transfer holds true for µP, indicating its broad applicability. The depth of the network also appears to influence the optimal learning rate, with deeper networks showing a decrease in the optimal learning rate, but the transferability property remains.
Also Read:
- Decoding In-Context Learning: How Induction Heads Emerge in Transformers
- Balancing Act: How Efficient Fine-Tuning Shapes LLM Safety and Fairness
This work represents a foundational step in understanding why µP enables efficient hyperparameter tuning across different model scales. While the theoretical proofs are currently limited to linear networks trained with gradient descent, the empirical evidence suggests that the principles extend to more complex, non-linear architectures and advanced optimizers. This research paves the way for future work to generalize these proofs, potentially leading to more efficient and predictable training of even larger and more sophisticated AI models.
