How to Set the Learning Rate for Large-Scale Pre-training? (2601.05049v1)

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.

• The paper introduces an empirical scaling law that reduces the learning rate search complexity from cubic to nearly linear.
• It compares two paradigms and demonstrates that fitting laws yield superior downstream performance over μTransfer.
• The study shows that modern normalization techniques reduce the need for module-specific learning rate tuning in large-scale MoE LLMs.

Optimal Learning Rate Selection in Large-Scale Pre-training: Fitting Laws versus $\mu$Transfer

Problem Formulation and Context

The paper "How to Set the Learning Rate for Large-Scale Pre-training?" (2601.05049) rigorously investigates strategies to determine optimal learning rates (LR) in large-scale LLM pre-training, focusing on both empirical complexity and performance. As LLMs scale in parameter count and data regimes, the sensitivity of learning rates becomes acute, especially given constraints in compute that preclude exhaustive hyperparameter sweeps. The study organizes approaches into two methodological paradigms: the Fitting Paradigm, based on regression and scaling laws from sub-scale models, and the Transfer Paradigm, specifically maximal update parameterization ($\mu$P, $\mu$Transfer), which seeks to transfer optimal settings from proxy to large target models.

Fitting Paradigm: Scaling Laws for LR

The Fitting Paradigm models validation loss as an explicit function of learning rate, model size ($N$), and data size ($D$). The principal contribution is a reduction in the search complexity for optimal LR from cubic to nearly linear by exploiting the empirical observation that, for fixed $N$ and $D$, loss versus log LR forms an invex (close to quadratic) curve. This enables efficient quadratic fitting on a sparse grid:

Figure 1: Results of fitting loss as a quadratic function of log LR, capturing the relationship for efficient extrapolation.

Additionally, leveraging the known power law $L(D) = L_0 + A D^{-\gamma}$ (as observed in prior works such as Chinchilla [hoffmann2022trainingcomputeoptimallargelanguage]), the search along the $D$ axis is further compressed. The composite scaling law for the global optimal LR is fit as:

$$\text{LR}(N, D) = 38.46 \cdot N^{-0.2219} \cdot D^{-0.3509}$$

Figure 2: Visualization of the fitted optimal LR surface and its behavior with $N$ and $D$.

The paradigm is formally validated using Qwen3-MoE architectures over scales from 0.5B to 12B parameters under a Warmup-Stable-Decay (WSD) schedule, with downstream performance measured on both MMLU and CMMLU.

Transfer Paradigm: Scaling $\mu$Transfer to Ultra-Large MoEs

The Transfer Paradigm implements and extends $\mu$Transfer for Mixture-of-Experts (MoE) models, including scaling rules across width, depth, weight decay, and total training tokens. The implementation draws on recent refinements ([mup], [depthmup], [complete-d-p]) and integrates weight decay proportionality as highlighted in [wang2025setadamwsweightdecay] and [fan2025robustlayerwisescalingrules]. Depth transfer is realized via depth-dependent scaling in residual connections:

$$H^{i+1} = H^i + m_L^{-\alpha} \mathcal{F}(H^i),\quad \alpha = 1,$$

where $m_L$ denotes layer count scaling. The transfer procedure is rigorously compared to fitting on target models up to 12B parameters and 500B tokens.

Empirical Evaluation and Main Results

Results unambiguously indicate that the Fitting Paradigm, instantiated via the derived scaling law for LR, yields consistently superior downstream performance over $\mu$Transfer in large-scale MoE pre-training. The difference in both in-domain and out-of-domain metrics is non-trivial:

Figure 3: Downstream performance (MMLU/CMMLU) for 4B models, showing clear advantage for the Fitting paradigm.

Figure 4: Downstream performance for 12B models, with the gap between fitting and $\mu$Transfer persisting and widening.

Module-wise LR search reveals that no single module (e.g., LM Head, Hidden, Embedding, Router) significantly benefits from specialized LR, contradicting the intuition motivating $\mu$Parameterization for feature learning balance at this scale:

Figure 5: Loss versus LR curves for individual modules, showing optimal LRs colinear with the global optimum.

Furthermore, full model training with global versus module-specific optima produces nearly identical curves, showing $\Delta L_\mathrm{val} < 0.01$ typically:

Figure 6: Validation losses for 4B models with global and module-wise optimal LRs overlap during stable training.

Mechanistic Analysis and LR Robustness

Layerwise update magnitudes (RMS norm) are stable and nearly constant ($\sim 0.2$) throughout training for both standard and $\mu$P parameterizations (with or without QK-Norm), indicating convergence to a regime where feature learning is uniformly distributed across layers:

Figure 7: RMS norm of updates for several model scales and parameterizations, all converging to 0.2.

Analysis of internal activations (embeddings, logits) across width shows no evidence of instability ("blow up") under standard parameterization when QK-Norm is present, in contrast with original $\mu$P findings. Removing QK-Norm immediately recovers the instability, highlighting the role of architectural details in scaling laws:

Figure 8: Embedding/logit variance trajectories with width, showing stability in modern architectures.

Figure 9: Ablation—removal of QK-Norm reproduces classic scaling instability for attention logits.

Data size proves more critical than model size for update scaling and stability, as reflected in the empirical exponents of the global scaling law. As training data increases, activation variations increase more rapidly than with parameter count:

Figure 10: Variation of key internal states through training as $D$ increases, illustrating heightened sensitivity.

Decay-Phase Training and Practical Implications

Extending training into a decay phase using high-quality data and reduced LR, as per WSD schedules, the performance advantage of the fitting-based LR selection persists and, in some cases, grows (up to +2.23% on CMMLU compared to $\mu$Transfer):

Figure 11: Downstream performance after decay; fitting paradigm continues to dominate in both stable and decay phase.

Theoretical and Practical Implications

This work demonstrates that, at present scales and architectures, the balance of feature learning and stability is largely achieved by modern normalization techniques and optimizer dynamics (AdamW), rather than module-wise LR calibration. The result that global LR scaling laws sharply outperform $\mu$Transfer at scale suggests a shift in hyperparameter protocol for industrial and research LLM pre-training: empirical fitting paired with scaling laws is more robust, less brittle with respect to architecture changes, and crucially bypasses several subtleties required for effective $\mu$Transfer. The strong predictive power of the empirically fit law ($R^2 \approx 0.96$) makes it a practical drop-in for planning ultra-large training runs, and reduces resource wastage typical of brute-force hyperparameter searches.

From a theoretical viewpoint, the findings prompt deeper investigation into the changing nature of optimization landscapes and feature-learning regimes as normalization and architectural practices evolve.

Conclusion

The study provides compelling evidence that in contemporary, ultra-large MoE LLM pre-training, learning rates optimized via scaling-law-guided fitting deliver performance superior to $\mu$Transfer-derived settings, both in stability and downstream task quality. These results have both immediate engineering value—streamlining resource allocation and hyperparameter selection—and broad theoretical implications, especially regarding the shifting relevance of parameterization methods at scale. Future improvements to scaling law methodology or normalization design may further alter best practices for hyperparameter transfer and optimization in LLMs.