Regole di Scalatura Robuste per Livelli tramite Ottimizzazione del Decadimento dei Pesi

Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

October 17, 2025

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

cs.AI

Abstract

Le leggi di scaling empiriche prescrivono come allocare parametri, dati e risorse computazionali, mentre la parametrizzazione a massimo aggiornamento (muP) consente il trasferimento del tasso di apprendimento tra diverse larghezze bilanciando le magnitudini degli aggiornamenti nelle fasi iniziali. Tuttavia, nelle moderne architetture invarianti alla scala, l'allenamento entra rapidamente in uno stato stazionario governato dall'ottimizzatore, dove i livelli di normalizzazione creano una sensibilità alla scala nel backward e il tasso di apprendimento effettivo diventa dipendente dalla larghezza, degradando il trasferimento muP. Affrontiamo questo problema introducendo una regola di scaling del decadimento del peso per AdamW che preserva il guadagno dei sottolivelli tra diverse larghezze. Empiricamente, lo spettro dei valori singolari di ciascun parametro matriciale scala in norma come eta/lambda con una forma approssimativamente invariante; sotto scaling di larghezza d, osserviamo che il valore singolare più alto scala approssimativamente come eta/lambda * d^{0.75}. Combinando questa osservazione con la regola muP del tasso di apprendimento eta_2 ∝ d^{-1} per parametri di tipo matrice implica una regola empirica di scaling del decadimento del peso lambda_2 ∝ d che mantiene approssimativamente invarianti i guadagni dei sottolivelli rispetto alla larghezza. Insieme a parametri di tipo vettore allenati a eta_1 = Theta_d(1) e lambda_1 = 0, ciò permette il trasferimento zero-shot sia del tasso di apprendimento che del decadimento del peso dalle larghezze proxy a quelle target, eliminando la necessità di sweep per ogni larghezza. Validiamo la regola su Transformer in stile LLaMA e in un contesto sintetico minimale, e forniamo una semplice diagnostica, confrontando i valori singolari più alti, per verificare l'invarianza del guadagno dei sottolivelli. I nostri risultati estendono muP oltre il regime vicino all'inizializzazione controllando esplicitamente le scale dello stato stazionario impostate dall'ottimizzatore, offrendo una ricetta pratica per il trasferimento robusto agli iperparametri rispetto alla larghezza sotto 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.