Quantitative Analysis of Hyperparameter Transfer in Large Model Training: The Critical Role of Embedding Layer Learning Rate

This study focuses on the hyperparameter transfer problem in large model training. The core finding is: The advantages of Maximal Update Parameterization (μP) mainly come from the increase in the embedding layer learning rate, rather than complex parameterization theories. Through quantitative analysis and ablation experiments, it is revealed that the embedding layer is a training bottleneck in standard parameterization. This article will provide a detailed interpretation from aspects such as background, methods, experiments, and suggestions.

Training a GPT-4-level large model is extremely costly (millions to tens of millions of dollars per run). Hyperparameter transfer has become a key strategy to reduce trial-and-error costs—searching for optimal hyperparameters in small-scale models and then extrapolating to large models. There are two main paths: 1. Fitting scaling laws to predict optimal values for large models; 2. Using parameterization designs (such as μP) to make hyperparameters approximately invariant across different scales. However, existing theories lack a sufficient explanation for the effectiveness of μP.

The study established a three-dimensional quantitative framework: 1. Scaling law fitting quality (measures the prediction accuracy of small-scale fitting rules for large models); 2. Extrapolation error robustness (evaluates the impact of small-scale errors on large model performance); 3. Parameterization asymptotic loss penalty (compares performance gaps between different parameterizations at the large-scale limit). These three together form a complete evaluation picture.

The core advantage of μP over Standard Parameterization (SP) lies in increasing the embedding layer learning rate. In SP, the embedding layer has a large number of parameters (proportional to vocabulary size and model dimension), leading to limited gradient updates, which causes problems such as unstable training, slow convergence, and hyperparameter sensitivity. μP removes this bottleneck by scaling the embedding layer learning rate by the model width (increasing by width times); simply increasing the embedding layer learning rate in SP to the μP level can significantly improve hyperparameter transfer effects.

Weight Decay has dual effects: The positive effect is improving the fitting quality of scaling laws, making small-scale extrapolation more reliable; the negative effect is impairing extrapolation robustness under fixed token-per-parameter settings (small experimental errors are easily amplified). Practical training needs to balance between fitting quality and robustness.

Verification through a series of ablation experiments found: 1. Embedding layer learning rate ablation: Adjusting the embedding layer learning rate alone in SP can achieve μP-like effects; 2. Component-level analysis: Identifying the embedding layer as a key component; 3. Scale extrapolation test: Extrapolating from small scale to large scale to verify transfer effects. The results strongly support the "embedding layer bottleneck" hypothesis.

Based on the research findings, four recommendations are proposed: 1. Prioritize the setting of the embedding layer learning rate; 2. Simplify μP implementation: Most benefits can be obtained by only increasing the embedding layer learning rate in SP; 3. Weight decay tuning: Adjust strategies according to training settings (fixed steps vs fixed token count); 4. Small experiment design: Ensure that the embedding layer behavior can represent large models to avoid extrapolation failure.

Research limitations: Verified only on AdamW optimizer and Transformer architecture; applicability to other optimizers (such as Adam, SGD) and architectures (such as Mamba, RWKV) remains to be verified. Future directions: 1. Deep mathematical principles of the embedding layer bottleneck; 2. Similar bottlenecks in multimodal models (such as vision-language); 3. Dynamically adjusting the embedding layer learning rate during training. Conclusion: This study reveals the source of μP's advantages through detailed exploration, provides concise and effective guidance for large model training, and emphasizes that insights come from details rather than complex theories.