PaperlayerView exact revision on arXiv

Source-linked AI summary

Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

Zhiyuan Fan, Yifeng Liu, Qingyue Zhao, Angela Yuan, Quanquan Gu

arXiv:2510.15262v1cs.LGcs.AIstat.ML

TL;DR

μP supports learning-rate transfer across widths through early-time update matching, but its near-initialization focus leaves optimizer-governed steady-state scales unresolved. The paper proposes width-scaled AdamW weight decay to preserve sublayer gain, reporting zero-shot transfer of learning rate and weight decay across widths.

  • Problem

    μP analyses primarily cover near-initialization behavior, leaving its transfer behavior in optimizer-governed steady states insufficiently addressed.

    Source: §1 · p. 1, §2.1 · p. 3

  • Method

    The paper matches matrix-parameter singular-value scales through width-dependent AdamW weight decay, alongside the standard layerwise μP learning-rate split.

    Source: §4.2 · p. 8, §6 · p. 11

  • Results

    The layerwise rule enables zero-shot transfer of learning rate and weight decay across widths, with approximately invariant sublayer gains.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    Matching top singular values provides a diagnostic for checking sublayer-gain invariance under the proposed scaling rule.

    Source: §6 · p. 11

    The reported scope is AdamW and LLaMA architectures with fixed head count and FFN ratio; universality across other architectures is uncertain.

    Source: §6 · p. 11

Abstract

from arXiv · show

Empirical scaling laws prescribe how to allocate parameters, data, and compute, while maximal-update parameterization ($\mu$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 $\mu$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{\eta/\lambda}$ with an approximately invariant shape; under width scaling $d$, we observe that the top singular value scales approximately as $\sqrt{\eta/\lambda}\cdot d^{0.75}$. Combining this observation with the $\mu$P learning-rate rule $\eta_2\propto d^{-1}$ for matrix-like parameters implies an empirical weight-decay scaling rule $\lambda_2\propto \sqrt{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 \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 $\mu$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.

1 Introduction

The paper extends μP transfer into optimizer-governed steady states by tuning AdamW weight decay to preserve sublayer gain across widths. It links width-dependent spectral scaling to a layerwise rule and reports transfer of learning rate and weight decay together.

1 Introduction

  • μP prescribes width-invariant update magnitudes, but its early-time analysis does not cover the optimizer-governed steady state reached as norms stabilize.Normalization preserves approximate forward scale invariance while making backward gradients scale-sensitive, so effective learning rates can vary with width.Source: §1 · pp. 1–2
  • Matrix spectra grow in norm with η/λ while their shape remains nearly unchanged, linking spectrum magnitude to sublayer gain.The paper presents this observation for weight matrices at the steady state of AdamW training.Source: §1 · p. 2
  • The top singular value magnitude approximately scales as η/λ · d^0.75; the passages connect this scaling to maintaining sublayer gain invariance through weight-decay scaling for matrix-like parameters.The stated learning-rate scaling for matrix-like parameters is η2 ∝ d−1.Source: §1 · p. 2
  • The proposed layerwise scheme transfers learning rate and weight decay together, using constant-scale learning rates and zero decay for vector-like parameters.The matrix-like parameter decay is described as width-scaled, while vector-like parameters use η1 = Θ_d(1) and λ1 = 0.Source: §1 · p. 2
  • A synthetic random-data example also exhibits the square-root weight-decay rule, which the authors suggest may reflect an architectural property.They present this as a clue for studying more fine-grained layerwise scaling rules.Source: §1 · p. 2

2 Related Work

Prior work offers global scaling laws and μP-based hyperparameter transfer, while its core analyses mainly address near-initialization behavior. Work on weight decay motivates width scaling, but this paper targets the rule that preserves sublayer gains under AdamW.

2.1 Hyperparameter Transfer

  • μP keeps vector-like learning rates constant and scales matrix-like learning rates inversely with width, supporting transfer from proxy to larger target models.The related-work discussion notes empirical µTransfer across architectures, optimizers, and LLM scale.Source: §2.1 · pp. 2–3
  • Existing Tensor Program and μP analyses primarily describe near-initialization dynamics, while long-horizon training leaves terminal implicit bias and generalization insufficiently characterized.The discussion contrasts effective update counts small relative to width with modern pretraining horizons much larger than width.Source: §2.1 · p. 3
  • Normalization preserves forward-scale invariance but makes gradients depend inversely on activation magnitude, motivating width-aware decay to preserve sublayer gain.As norms stabilize, optimizer dynamics set the effective scale and can make layerwise output scales width-dependent.Source: §2.1 · p. 3

2.2 Weight Decay Scaling Rule

  • Prior proposals include linear AdamW decay scaling and evidence that increasing decay can recover transfer as width grows.The discussion also notes related linear-scaling variants and layerwise arguments for hidden matrices.Source: §2.2 · p. 3
  • Prior steady-state analyses relate parameter norms to η/λ, while this paper asks which width-dependent decay keeps AdamW sublayer gains invariant.The authors report that λ2 ∝ √d, alongside η2 ∝ d^-1 for matrix-like parameters and η1 = Θ_d(1), λ1 = 0 for vector-like parameters, keeps singular-value scales and sublayer gains width-invariant.Source: §2.2 · p. 3

3 Preliminaries

The preliminaries establish the notation and AdamW setup, distinguish parameter classes by how their sizes scale with width, and review the µP scaling principle that motivates studying weight decay.

3.1 General Notations

  • The notation section defines vector and matrix RMS norms and specifies conventions for asymptotic order, elementwise operations, and logarithms.Θ_d(1) emphasizes that the bounding constants are independent of width d.Source: §3.1 · p. 4

3.2 AdamW Optimizer

  • The AdamW analysis treats the model statically at a fixed step, uses bias-corrected gradient moments, and neglects the small stabilizer ε.AdamW is described as using decoupled weight decay.Source: §3.2 · p. 4

3.3 Parameter Classes

  • Parameters are grouped by width scaling: vector-like parameters have linearly many entries, while matrix-like parameters have quadratically many.The paper holds depth, attention-head count, and FFN expansion coefficient fixed while varying width d.Source: §3.3 · p. 4

3.4 Maximal-update Parameterization

  • µP targets width-invariant per-step functional changes, with different hyperparameter scalings for parameter types based on their shapes.The paper studies how the weight-decay coefficient λ should scale with model width d.Source: §3.4 · p. 5

4 Robust Scaling Rule by Propoer Weight Decay Scaling

The section develops a weight-decay scaling rule to preserve sublayer gain across widths, using steady-state spectra under AdamW. Experiments report consistent transfer of optimal base learning rate and weight decay across model widths, alongside a tradeoff between the two hyperparameters.

4 Robust Scaling Rule by Propoer Weight Decay Scaling

  • The method targets width-invariant gain by tuning weight decay rather than initial variance, in a steady state where AdamW weight norms stabilize.In the described regime, updates behave like rotations on a sphere.Source: §4 · p. 5

4.1 Matching Weight Matrix Spectra under Fixed Model Width

  • At fixed width, proportional scaling of learning rate and weight decay leaves sublayer gains and singular values nearly invariant.The study measures gain during training and computes matrix-parameter spectra at its end.Source: §4.1 · p. 6, Figure 1 · p. 6

4.2 Matching Spectra Induces Sublayer Gain Invariance

  • When η2/λ2 ∝ d^-0.75, spectra align across widths and sublayer amplification is preserved; alternative weight-decay rules produce spectral deviations.The reported alignment factor ρ(d) approximately follows d^0.75 in the LLaMA configuration studied.Source: §4.2 · p. 7, Figure 2 · p. 7

4.3 Tuning Weight Decay Enables Hyperparameter Transfer

  • The proposed parameterization enables consistent transfer of the optimal base learning rate and weight decay across model widths.The loss landscapes’ minima align after scaling according to the proposed rules.Source: Figure 3 · p. 8, §4.3 · p. 9

4.4 Tradeoff between Learning Rate and Weight Decay

  • Learning rate and weight decay trade off along a near-optimal diagonal ridge, where models achieve almost identical final losses despite different weight decays.The authors suggest choosing a weight decay, sweeping learning rate, then transferring the configuration with µP scaling.Source: §4.4 · pp. 9–10, Figure 4 · p. 9
  • In a synthetic AdamW run with η ∝ 1/d, top singular values approximately align when η/λ ∝ d^-0.75.This alignment is reported as matching the realistic-training results; the figure also tracks weight-matrix RMS norm during training.Source: Figure 5 · p. 10

5 Illustrative Example of Sqrt Weight Decay Scaling

A synthetic two-layer FFN illustrates the square-root weight-decay rule under AdamW. Across widths, the scaled models’ top singular-value spectra align, consistent with the LLaMA experiments.

5 Illustrative Example of Sqrt Weight Decay Scaling

  • The authors present the synthetic example as evidence that the square-root rule arises from architecture rather than data distribution.They suggest this may inform finer-grained layerwise scaling rules.Source: §5 · p. 10
  • The illustrative model is a two-layer ReLU FFN with independently sampled standard-normal inputs and upstream gradients.Its weight matrices are both d×d.Source: §5 · p. 10
  • Across widths 256–2048, top singular-value spectra align under the square-root weight-decay rule, consistent with the LLaMA results.The synthetic model is trained for 20,000 steps without learning-rate annealing, reaching a rotational-equilibrium steady state.Source: §5 · p. 10

6 Conclusion

The conclusion proposes a layerwise AdamW scaling rule that aligns matrix top singular values and supports zero-shot transfer across widths. Its evidence is scoped to AdamW and LLaMA architectures with fixed head count and FFN ratio.

6 Conclusion

  • The rule is presented as complementary to μP’s early-stage guarantees, matching sublayer gain during steady-state training.The paper also frames the work as a step toward predictive models of steady-state dynamics.Source: §6 · p. 11
  • For matrix-like linear-layer parameters, η/λ ∝ d^-0.75 aligns top singular values, while η2 ∝ d^-1 and λ2 ∝ √d approximately preserve sublayer gain across widths.Vector-like parameters use η1 = Θ_d(1) and λ1 = 0; together these rules enable zero-shot hyperparameter transfer instead of per-width learning-rate sweeps.Source: §6 · p. 11
  • A concurrent approach matches operator norms by adjusting batch size, whereas this work uses explicit weight decay for layerwise prescriptions.The paper contrasts batch size as an implicit regularizer with weight decay as explicit regularization.Source: §6 · p. 11
  • The demonstrated scope is AdamW and LLaMA architectures with fixed head count and FFN ratio; universality across other architectures is uncertain.The authors identify mixture-of-experts and alternative attention designs as settings that may change the scaling factor.Source: §6 · p. 11
  • Future work includes other optimizers, mixture-of-experts and structured-sparse models, width-varying batch sizes or token counts, and depth scaling.The authors also propose linking data distributions, optimizer statistics, and spectral shape to explain η/λ.Source: §6 · p. 11

A Hyperparameter Table

Table 3 is identified as the default experimental setup, with width ratio m_d defined relative to d_base = 256.

A Hyperparameter Table

  • The table caption defines m_d as d/d_base, with d_base = 256.The caption describes the table as the default experimental setup.Source: Table 3 · p. 16

B Other Figures

Figure 6 reports statistics for Wout in the synthetic run and complements Figure 5.

B Other Figures

  • Figure 6 presents weight-matrix Wout statistics from the synthetic run, complementary to Figure 5.Source: Figure 6 · p. 17

Related papers

Loading 2510.15262v1…