Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

TL;DR

This paper introduces a framework to quantify hyperparameter transfer in large language models, highlighting the significance of embedding layer learning rate and the Maximal Update parameterization for improved transferability.

Contribution

It develops metrics to evaluate hyperparameter transfer quality and explains why Maximal Update parameterization enhances transferability by focusing on embedding layer learning rate.

Findings

Maximal Update parameterization improves hyperparameter transfer over standard parameterization.

Increasing embedding layer learning rate stabilizes training and enhances transfer.

Weight decay improves scaling law fit but can reduce robustness in extrapolation.

Abstract

Hyperparameter transfer allows extrapolating optimal optimization hyperparameters from small to large scales, making it critical for training large language models (LLMs). This is done either by fitting a scaling law to the hyperparameters or by a judicious choice of parameterization, such as Maximal Update ($\mu$P), that renders optimal hyperparameters approximately scale invariant. In this paper, we first develop a framework to quantify hyperparameter transfer through three metrics: (1) the quality of the scaling law fit, (2) the robustness to extrapolation errors, and (3) the asymptotic loss penalty due to choice of parameterization. Next, we investigate through a comprehensive series of ablations why $\mu$P appears to offer high-quality learning rate transfer relative to standard parameterization (SP), as existing theory is inadequate. We find that the overwhelming benefit of $\mu$P relative to SP when training with AdamW arises simply from maximizing the learning rate of the embedding layer. In SP, the embedding layer learning rate acts as a bottleneck that induces training instabilities; increasing it by a factor of width to match $\mu$P dramatically smooths out training while improving hyperparameter transfer. We also find that weight decay improves the scaling law fits, while, in the fixed token-per-parameter setting, it hurts the robustness of the extrapolation.