Règles robustes de mise à l'échelle par couches grâce à un réglage approprié de la décroissance des poids

Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

October 17, 2025

papers.authors: Zhiyuan Fan, Yifeng Liu, Qingyue Zhao, Angela Yuan, Quanquan Gu

cs.AI

papers.abstract

Les lois d'échelle empiriques prescrivent comment allouer les paramètres, les données et les ressources de calcul, tandis que la paramétrisation à mise à jour maximale (muP) permet le transfert du taux d'apprentissage entre différentes largeurs en égalisant les magnitudes des mises à jour en début d'entraînement. Cependant, dans les architectures modernes invariantes à l'échelle, l'entraînement entre rapidement dans un état d'équilibre gouverné par l'optimiseur, où les couches de normalisation créent une sensibilité à l'échelle en rétropropagation et le taux d'apprentissage effectif devient dépendant de la largeur, dégradant ainsi le transfert muP. Nous abordons ce problème en introduisant une règle de mise à l'échelle de la décroissance de poids pour AdamW qui préserve le gain des sous-couches à travers les largeurs. Empiriquement, le spectre des valeurs singulières de chaque paramètre matriciel se met à l'échelle en norme comme eta/lambda avec une forme approximativement invariante ; sous une mise à l'échelle de largeur d, nous observons que la valeur singulière supérieure se met à l'échelle approximativement comme eta/lambda * d^{0.75}. En combinant cette observation avec la règle muP du taux d'apprentissage eta_2 ∝ d^{-1} pour les paramètres de type matriciel, cela implique une règle empirique de mise à l'échelle de la décroissance de poids lambda_2 ∝ d qui maintient approximativement les gains des sous-couches invariants à la largeur. Avec les paramètres de type vecteur entraînés à eta_1 = Theta_d(1) et lambda_1 = 0, cela permet un transfert à zéro coup des taux d'apprentissage et de la décroissance de poids des largeurs proxy aux largeurs cibles, éliminant ainsi les balayages par largeur. Nous validons cette règle sur des Transformers de style LLaMA et dans un cadre synthétique minimal, et nous fournissons un diagnostic simple, en faisant correspondre les valeurs singulières supérieures, pour vérifier l'invariance des gains des sous-couches. Nos résultats étendent muP au-delà du régime proche de l'initialisation en contrôlant explicitement les échelles d'état d'équilibre définies par l'optimiseur, offrant une recette pratique pour le transfert robuste des hyperparamètres sous 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.