Learning Rate Transfer in Normalized Transformers
This work addresses the practical problem of hyperparameter tuning for large-scale Normalized Transformers, enabling more efficient training by allowing learning rates to be transferred across model sizes.
The authors identify that Normalized Transformers (nGPT) lack learning rate transfer across model dimensions and token horizons. They propose a new parameterization, νGPT, which achieves learning rate transfer across width, depth, and token horizon, validated through extensive experiments.
The Normalized Transformer, or nGPT (arXiv:2410.01131) achieves impressive training speedups and does not require weight decay or learning rate warmup. However, despite having hyperparameters that explicitly scale with model size, we observe that nGPT does not exhibit learning rate transfer across model dimension and token horizon. To rectify this, we combine numerical experiments with a principled use of alignment exponents (arXiv:2407.05872) to revisit and modify the $μ$P approach to hyperparameter transfer (arXiv:2011.14522). The result is a novel nGPT parameterization we call $ν$GPT. Through extensive empirical validation, we find $ν$GPT exhibits learning rate transfer across width, depth, and token horizon.