Robust Layerwise Scaling Rules by Proper Weight Decay Tuning
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.
68f5aac18589920bf4d32153
Zhiyuan Fan
68f5aac18589920bf4d32154
Yifeng Liu
claimed_verified
2025-10-20T12:10:32.620Z
653d276681f52ceb4d12bd85
/avatars/56601a25e5f883a8f6dc15f6fd9dcc57.svg
Yifeng Liu
user
Lewis-Lau
68f5aac18589920bf4d32155
Qingyue Zhao
68f5aac18589920bf4d32156
Angela Yuan
68f5aac18589920bf4d32157
Quanquan Gu
Robust Layerwise Scaling Rules by Proper Weight Decay Tuning
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.
A new weight-decay scaling rule for AdamW is introduced to preserve sublayer gain across widths in modern scale-invariant architectures, enabling zero-shot transfer of learning rate and weight decay.
maximal-update parameterization $\mu$P optimizer-governed steady state normalization layers weight-decay scaling rule AdamW singular-value spectrum sublayer gain LLaMA-style Transformers