Robuste Laaggewijze Schaalregels door Juiste Afstemming van Gewichtsverval

Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

October 17, 2025

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

cs.AI

Samenvatting

Empirische schaalwetten geven aan hoe parameters, data en rekenkracht moeten worden toegewezen, terwijl maximal-update-parameterisatie (muP) het overdragen van leerratio's over verschillende breedtes mogelijk maakt door de grootte van updates in de vroege fase gelijk te stellen. Echter, in moderne schaalinvariante architecturen komt de training snel in een door de optimizer beheerde stabiele toestand, waarbij normalisatielagen achterwaartse schaalgevoeligheid creëren en de effectieve leerratio afhankelijk wordt van de breedte, wat de muP-overdracht verslechtert. Wij pakken dit aan door een gewichtsvervalschaalregel voor AdamW in te voeren die de sublaagwinst over verschillende breedtes behoudt. Empirisch gezien schaalt het spectrum van singuliere waarden van elke matrixparameter in norm als eta/lambda met een ongeveer invariante vorm; bij breedteschaal d observeren we dat de top singuliere waarde ongeveer schaalt als eta/lambda * d^{0.75}. Door deze observatie te combineren met de muP-leerratieregel eta_2 ∝ d^{-1} voor matrixachtige parameters, impliceert dit een empirische gewichtsvervalschaalregel lambda_2 ∝ d die de sublaagwinst ongeveer breedte-invariant houdt. Samen met vectorachtige parameters die getraind worden bij eta_1 = Theta_d(1) en lambda_1 = 0, resulteert dit in zero-shot overdracht van zowel leerratio als gewichtsverval van proxy naar doelbreedtes, waardoor per-breedte zoekacties overbodig worden. We valideren de regel op LLaMA-style Transformers en in een minimale synthetische setting, en we bieden een eenvoudige diagnostische methode, het matchen van top singuliere waarden, om de sublaagwinst-invariantie te controleren. Onze resultaten breiden muP uit voorbij het near-init regime door expliciet de schalen in de stabiele toestand, ingesteld door de optimizer, te beheersen, en bieden een praktisch recept voor breedte-robuuste hyperparameteroverdracht onder AdamW.

English

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.