Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

TL;DR

This work addresses width-dependent steady-state dynamics in AdamW-trained, scale-invariant Transformer models, where normalization induces backward scale sensitivity and undermines μP transfer. By linking the steady-state weight norms to $\|\mW\|_{\mathrm{rms}} = \Theta_d\big(\sqrt{\eta/\lambda}\big)$, the authors derive a layerwise weight-decay scaling rule: matrix-like parameters should use $\lambda_2 = \Theta(\sqrt{d})$ with $\eta_2 = \Theta(d^{-1})$, while vector-like parameters keep $\eta_1 = \Theta_d(1)$ and $\lambda_1 = 0$, enabling zero-shot transfer of hyperparameters across widths. Empirical evidence from LLaMA-style Transformers and a synthetic model shows that the top singular values scale as $\sigma_{\max}(\mW) = \Theta\big(\sqrt{\eta/\lambda}\big) d^{0.75}$, so aligning spectra via $\sqrt{\eta_2/\lambda_2} \propto d^{-0.75}$ preserves sublayer gains across widths. This yields practical, width-robust hyperparameter transfer rules that extend μP beyond initialization, with a simple diagnostic based on singular-value spectra to verify sublayer-gain invariance. The approach provides a principled method for choosing learning-rate and weight-decay schedules in large-scale training under AdamW, reducing per-width sweeps and enabling scalable deployment of wide models.

Abstract

Empirical scaling laws prescribe how to allocate parameters, data, and compute, while maximal-update parameterization ($μ$P) enables learning-rate transfer across widths by equalizing early-time update magnitudes. However, in modern scale-invariant architectures, training quickly enters an optimizer-governed steady state where normalization layers create backward scale sensitivity and the effective learning rate becomes width dependent, degrading $μ$P transfer. We address this by introducing a weight-decay scaling rule for AdamW that preserves sublayer gain across widths. Empirically, the singular-value spectrum of each matrix parameter scales in norm as $\sqrt{η/λ}$ with an approximately invariant shape; under width scaling $d$, we observe that the top singular value scales approximately as $\sqrt{η/λ}\cdot d^{0.75}$. Combining this observation with the $μ$P learning-rate rule $η_2\propto d^{-1}$ for matrix-like parameters implies an empirical weight-decay scaling rule $λ_2\propto \sqrt{d}$ that approximately keeps sublayer gains width invariant. Together with vector-like parameters trained at $η_1=Θ_d(1)$ and $λ_1=0$, this yields \emph{zero-shot} transfer of both learning rate and weight decay from proxy to target widths, removing per-width sweeps. We validate the rule on LLaMA-style Transformers and in a minimal synthetic setting, and we provide a simple diagnostic, matching top singular values, to check sublayer-gain invariance. Our results extend $μ$P beyond the near-init regime by explicitly controlling steady-state scales set by the optimizer, offering a practical recipe for width-robust hyperparameter transfer under AdamW.

Figures (6)

• Figure 1: Statistics of the FFN weight matrices under AdamW with various learning rates $\eta$ and weight decay values $\lambda$. The plotted lines are averaged across matrix-like sublayers from all blocks. Left: Sublayer gain $\|\vy\|_{\rms}/\|\vx\|_{\rms}$ during training. Right: Singular value spectrum $\sigma_i(\mW)$ of the final weight matrices. The singular values are sorted in descending order; the horizontal axis shows the spectral index, and the vertical axis shows the corresponding singular value.
• Figure 2: Statistics of the FFN weight matrices under AdamW with various matrix-like weight decay scalings $\lambda_2$, and learning rate scaling specified in \ref{['tab:mup-scaling']}, where $\eta_1 = \Theta_d(1)$ is used for vector-like parameters and $\eta_2 \propto 1 / d$ for matrix-like parameters. The plotted lines are averaged across matrix-like sublayers from all blocks. Left: Sublayer gain $\|\vy\|_{\rms}/\|\vx\|_{\rms}$ during training. Right: Singular value spectrum $\sigma_i(\mW)$ of the final weight matrices. Alignment is observed when the weight decay scaling follows $\lambda_2 \propto \sqrt{d}$.
• Figure 3: Transfer of the optimal base learning rate $\eta_{\mathrm{base}}$(Left) and weight decay $\lambda_{\mathrm{base}}$(Right) across model widths. Each curve shows the loss landscape for a specific width $d$, with minima aligned after scaling according to \ref{['tab:new-scaling']}. For visualization, all curves are vertically shifted by constant offsets so that the losses are directly comparable across widths. The alignment of minima indicates that the proposed parameterization enables consistent hyperparameter transfer across scales.
• Figure 4: Validation loss differences across learning-rate and weight-decay pairs. Loss values are measured relative to the optimal configuration $(\eta_{\mathrm{base}} = 0.02,\; \lambda_{\mathrm{base}} = 0.075)$. A clear diagonal ridge of near-optimal points reveals that increasing $\lambda$ requires reducing $\eta$, confirming their strong correlation. Both axes are spaced approximately logarithmically.
• Figure 5: Statistics of the weight matrix $\mW_{\mathrm{in}}$ under AdamW in the synthetic run with learning rate scaling $\eta \propto 1 / d$ and various scalings of weight decay $\lambda$. Left: Singular value spectrum $\sigma_i(\mW)$ of the final weight matrices. When the weight decay follows $\sqrt{\eta / \lambda} \propto d^{-0.75}$, the top singular values are approximately aligned, matching our realistic training results in \ref{['sec:matching-spectra']}. Right: During training, the root-mean-square of the weight matrix $\|\mW\|_{\rms}$ converges to a stable value, approximately proportional to $\sqrt{\eta / \lambda}$.