Robust Layerwise Scaling Rules by Proper Weight Decay Tuning
Authors: name: Zhiyuan Fan, name: Yifeng Liu, name: Qingyue Zhao, name: Angela Yuan, name: Quanquan Gu
Organization: Massachusetts Institute of Technology

Empirical scaling laws prescribe how to allocate parameters, data, and
compute, while maximal-update parameterization (muP) 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 muP 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
eta/lambda with an approximately invariant shape; under width
scaling d, we observe that the top singular value scales approximately as
eta/lambdacdot d^{0.75}. Combining this observation with the muP
learning-rate rule eta_2propto d^{-1} for matrix-like parameters implies an
empirical weight-decay scaling rule lambda_2propto d that
approximately keeps sublayer gains width invariant. Together with vector-like
parameters trained at eta_1=Theta_d(1) and lambda_1=0, this yields
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 muP 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.