Learning Rate Transfer in Normalized Transformers

TL;DR

This paper introduces GPT, a new parameterization of normalized transformers that enables effective transfer of learning rates across various model dimensions and token horizons, improving training efficiency.

Contribution

It proposes GPT, a novel hyperparameter transfer method for normalized transformers, extending the P approach with alignment exponents for better transferability.

Findings

GPT exhibits learning rate transfer across width, depth, and token horizon.

It achieves training speedups without requiring weight decay or warmup.

Extensive empirical validation confirms improved hyperparameter transferability.

Abstract

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 $\mu$P approach to hyperparameter transfer (arXiv:2011.14522). The result is a novel nGPT parameterization we call $\nu$GPT. Through extensive empirical validation, we find $\nu$GPT exhibits learning rate transfer across width, depth, and token horizon.