Reproduce “Learning Rate Transfer in Normalized Transformers” on OpenResearch

About this paper

Boris Shigida, Boris Hanin, Andrey Gromov

Researchers from Meta Superintelligence Labs and Princeton University introduced νGPT, a new parameterization for Normalized Transformers (nGPT), enabling robust learning rate transfer across model width, depth, and token horizon. This approach adapts existing hyperparameter transfer techniques by integrating empirically measured alignment exponents specific to nGPT, making large-scale model training more efficient by reducing hyperparameter tuning costs.

Key insights

• Empirical measurements revealed that the 'alignment exponents' ω_hidden and ω_output in nGPT are consistently 1/2, contradicting the ω_output = 1 assumption of classical μP and supporting order-1 logits at initialization.
• The optimal learning rate for hidden layers in nGPT is independent of model depth; applying 'Depth-μP' style depth scaling (e.g., `m_depth^-1/2`) to hidden learning rates breaks transfer in this architecture.
• The optimal global learning rate for nGPT models empirically follows a power law, scaling approximately as `(iteration count)^-1/3`.

Summary by alphaXiv.