How to Set the Learning Rate for Large-Scale Pre-training?

TL;DR

The paper tackles the challenge of setting the learning rate for large-scale pre-training under costly compute. It presents two paradigms: a Fitting approach using scaling laws to predict the optimal LR across model size $N$ and data size $D$, and a Transfer approach (μTransfer) applied to MoE proxy models with width, depth, weight decay, and token-horizon considerations. Empirically, the Fitting paradigm outperforms μTransfer at 4B and 12B scales, and its module-level LR tuning yields negligible gains over a well-chosen global LR; analysis of training stability and feature learning provides mechanistic explanations, such as uniform update magnitudes and architectural effects like QK-Norm mitigating instability. The work delivers practical guidelines for industrial pre-training and offers a theoretical lens for understanding hyperparameter transfer and stability in ultra-large models, while acknowledging limitations related to LR schedules, MoE-centric focus, and extrapolation bounds.

Abstract

Optimal configuration of the learning rate (LR) is a fundamental yet formidable challenge in large-scale pre-training. Given the stringent trade-off between training costs and model performance, the pivotal question is whether the optimal LR can be accurately extrapolated from low-cost experiments. In this paper, we formalize this investigation into two distinct research paradigms: Fitting and Transfer. Within the Fitting Paradigm, we innovatively introduce a Scaling Law for search factor, effectively reducing the search complexity from O(n^3) to O(n*C_D*C_η) via predictive modeling. Within the Transfer Paradigm, we extend the principles of $μ$Transfer to the Mixture of Experts (MoE) architecture, broadening its applicability to encompass model depth, weight decay, and token horizons. By pushing the boundaries of existing hyperparameter research in terms of scale, we conduct a comprehensive comparison between these two paradigms. Our empirical results challenge the scalability of the widely adopted $μ$ Transfer in large-scale pre-training scenarios. Furthermore, we provide a rigorous analysis through the dual lenses of training stability and feature learning to elucidate the underlying reasons why module-wise parameter tuning underperforms in large-scale settings. This work offers systematic practical guidelines and a fresh theoretical perspective for optimizing industrial-level pre-training.

Figures (23)

• Figure 1: Results of Equation \ref{['eq:ld']} and \ref{['eq:loss-n-d']}. These approaches allow for a substantial reduction in the time and storage cost of the search process.
• Figure 2: Middle: Visualization of the optimal learning rate relative to model size $N$ and data size $D$. Left: The relationship between the optimal learning rate and data size $D$ with model size fixed at $N=4\text{B}$. Right: The relationship between the optimal learning rate and model size $N$ with data size fixed at $D=140\text{B}$.
• Figure 3: Downstream task performances of 4B model with global optimal LR and $\mu$P respectively.
• Figure 4: Downstream task performances of 12B model with global optimal LR and $\mu$P respectively.
• Figure 5: The relationship between loss and learning rate for (a)LM Head, (b)Router, (c)Hidden and (d)Embedding parameters during the module-level learning rate search. Each curve corresponds to a model of a specific size. The dashed lines indicate the loss achieved by the corresponding model size under the global optimal learning rate setting. Triangle markers denote the optimal learning rate for the current module, while star markers represent the global optimal learning rate.