Paper Overview
Research Field: Machine Learning (ML) Authors: Dayal Singh Kalra, Maissam Barkeshli Published: 2025-05-20 arXiv: 2505.15986
---
Introduction
Hyperparameter transfer enables extrapolation of optimal optimization hyperparameters from small-scale models to large-scale models. This is critical for training large language models (LLMs) because exhaustive hyperparameter search at full scale is computationally infeasible. Two main approaches exist:
1. Scaling-law fitting to hyperparameter data points across scales. 2. Scale-invariant parameterizations such as Maximal Update Parameterization (μP), which make optimal hyperparameters approximately invariant under width scaling.
Standard Parameterization (SP) lacks this property, often leading to brittle transfer.
---
Three-Metric Framework for Quantifying Transfer
The authors formalize hyperparameter transfer quality with three metrics:
1. Scaling-law fit quality — how well a power-law (or similar) model fits optimal hyperparameters measured at different widths. 2. Robustness to extrapolation errors — sensitivity of a large-scale run to small misestimations in the extrapolated hyperparameter. 3. Asymptotic loss penalty — the irreducible loss gap induced purely by the choice of parameterization in the infinite-width limit.
These metrics jointly distinguish parameterizations that merely fit clean curves from those that remain stable under noise and converge to competitive loss.
---
Why Does μP Outperform SP?
Existing theory does not fully explain the empirical advantage of μP for learning-rate transfer. Through a comprehensive series of ablations with AdamW, the authors trace μP's superiority to a single factor: maximizing the embedding layer learning rate.
Key empirical findings:
- In SP, the embedding layer learning rate acts as a bottleneck, causing training instability as width grows.
- Scaling the embedding-layer learning rate with width in SP — i.e., matching μP's prescription on this layer alone — significantly smooths training dynamics and improves hyperparameter transfer.
- Conversely, reverting this single layer in μP back to SP-style scaling largely erases μP's advantage.
- Weight decay improves scaling-law fits for the optimal learning rate.
- However, the same weight decay harms robustness to extrapolation errors, indicating a trade-off between fit quality and predictive reliability at scale.
- Hyperparameter transfer can be rigorously evaluated via three complementary metrics: fit quality, extrapolation robustness, and asymptotic loss penalty.
- μP's empirical superiority over SP under AdamW is largely attributable to a width-scaled embedding-layer learning rate.
- Modifying only the embedding layer's learning-rate scaling in SP reproduces most of μP's benefits.
- Weight decay introduces a tension between scaling-law fit quality and extrapolation robustness.
The result suggests μP's gains over SP are concentrated at the embedding layer rather than distributed across the network.
---
Role of Weight Decay
Under fixed tokens-per-parameter training:
---
Key Takeaways
Reference
Kalra, D. S., & Barkeshli, M. (2025). *Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate*. arXiv:2505.15986. Link: https://arxiv.org/abs/2505.15986