Weight Decay may matter more than muP for Learning Rate Transfer in Practice

TL;DR

This work challenges the view that Maximal Update Parameterization (muP) LR scaling alone enables reliable learning rate transfer across model widths. By developing a relative-update framework with update alignment and weight alignment measures, it shows independent weight decay ultimately overrides muP's scaling and stabilizes cross-width feature learning, explaining practical LR transfer in large-scale training. Importantly, the observed benefits of muP largely arise from an implicit warmup effect, which can be replicated with stronger explicit warmup schedules, suggesting practical alternatives to muP for robust transfer. The findings have implications for hyperparameter transfer in practice and point toward matrix-level optimizers as a promising avenue for further research.

Abstract

Transferring the optimal learning rate from small to large neural networks can enable efficient training at scales where hyperparameter tuning is otherwise prohibitively expensive. To this end, the Maximal Update Parameterization (muP) proposes a learning rate scaling designed to keep the update dynamics of internal representations stable across different model widths. However, the scaling rules of muP rely on strong assumptions, particularly about the geometric alignment of a layer's inputs with both its weights and gradient updates. In this large-scale empirical investigation, we show that these assumptions hold only briefly at the start of training in the practical setups where learning rate transfer is most valuable, such as LLM training. For the remainder of training it is weight decay rather than muP that correctly stabilizes the update dynamics of internal representations across widths, facilitating learning rate transfer. This suggests muP's scaling primarily acts as a form of implicit learning rate warmup, allowing us to largely replace it with modified warmup schedules. Together these findings fundamentally challenge prevailing beliefs about learning rate transfer and can explain empirical practice such as why muP requires the independent weight decay variant for successful transfer.

Figures (16)

• Figure 1: $\mu$P specifically requires independent weight decay for good learning rate transfer when pre-training LLaMA networks with AdamW (20B tokens, width 2048 and 1B params, see Appendix \ref{['details:llama_training']}). Independent weight decay achieves this by counteracting $\mu$P's update scaling later in training which is needed because the core assumptions $\mu$P's scaling is based on quickly break down. Standard weight decay does not counteract $\mu$P's scaling and consequently results in a poor learning rate transfer on longer experiments. No weight decay behaves somewhere in-between, see Appx. \ref{['appx:no_wd']}.
• Figure 2: Neural networks are commonly invariant to certain types of rescaling transformations. Examples where applying an in-place scaling factor $\kappa > 0$ does not affect the current block output ${\bm{y}}$. Across such transformations, an update with a fixed absolute size $\|{\bm{W}}\|$ has a varying impact on the output while a given relative change $\|\Delta {\bm{W}}\|/\|{\bm{W}}\|$ always has the same impact. This makes the relative change a more informative measure of the size of updates that better captures their impact.
• Figure 3: Standard $\mu$P scaling fails to maintain relative representation changes across widths. Comparing the relative representation change (RRC) for $\mu$P with standard weight decay scaling (first row) and independent weight decay scaling (second row) for LLaMA training with a 10% linear warmup followed by linear decay. Plots shown for a single representative layer with $C=128$ at $1\times$ width, see \ref{['details:fig:rrc_mup_std_vs_ind']} for details. Left: Standard scaling does not preserve the RRC for the same learning rate unlike independent scaling. Recall that the RRC is the product of the alignment ratio and the relative weight update (Equation \ref{['eq:rrc_to_rwu']}). Middle: The alignment ratio is similar for both widths and approximately 1. Right: To maintain the RRC across widths the relative weight update size should then be kept constant as achieved by independent weight decay rather than $\mu$P's prescribed scaling.
• Figure 4: The alignment assumptions of $\mu$P do not hold in practice. Measurements of the weight alignment, the update alignment and the alignment ratio for LLaMA training at different widths. The update alignment varies significantly over time, becoming width-dependent. As a result the alignment ratio loses its width dependence. This violates $\mu$P's core scaling assumptions which only hold very early in training. Plots shown for one representative MLP layer, see \ref{['details:fig:alignment_w_dw_ratio']} for details.
• Figure 5: $\mu$P with independent WD scaling has a warmup-like effect on relative updates. Measurements of the relative weight update (left) and relative representation change (middle) for different learning rate (LR) and weight decay (WD) combinations in a LLaMA layer with 2048 input features. Independent weight decay scaling achieves a given LR-WD product using a lower LR and higher WD (orange) compared to not performing $\mu$P scaling (blue). This results in smaller relative updates early in training, similar to an additional learning rate warmup factor. The right panel shows this effect through the ratio of the relative weight updates for the two runs, which goes from $\frac{1}{m}=\frac{1}{16}$ to asymptotically approaching 1 in a roughly exponential manner. See \ref{['details:fig:mup_ind_is_warmup']} for experimental details.