Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate
Why one simple tweak to embedding layer training speeds up AI model scaling
Training large language models requires finding the right hyperparameters—settings like learning rates—at small scale and then scaling them up. This paper reveals that a popular technique called Maximal Update Parameterization (μP) works so well primarily because it increases the learning rate for one specific component: the embedding layer. Simply boosting the embedding layer's learning rate in standard training setups by a factor equal to model width produces the same scaling benefits, suggesting the real advantage isn't deep theory but rather fixing a training bottleneck.
Training large language models is expensive and time-consuming. If you can nail hyperparameters on a small, cheap model and confidently scale them to a massive one, you save weeks of computation and millions in hardware costs. This work shows practitioners exactly which knob to turn—the embedding layer learning rate—to make that transfer reliable, potentially cutting wasted training runs and accelerating AI development timelines.