Reglas Robustas de Escalado por Capas mediante el Ajuste Adecuado de la Decaída de Pesos

Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

October 17, 2025

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

cs.AI

Resumen

Las leyes de escalado empírico prescriben cómo asignar parámetros, datos y recursos computacionales, mientras que la parametrización de actualización máxima (muP) permite la transferencia de la tasa de aprendizaje a través de diferentes anchos al igualar las magnitudes de las actualizaciones en etapas tempranas. Sin embargo, en arquitecturas modernas invariantes a escala, el entrenamiento rápidamente entra en un estado estacionario gobernado por el optimizador, donde las capas de normalización crean sensibilidad a la escala en la retropropagación y la tasa de aprendizaje efectiva se vuelve dependiente del ancho, degradando la transferencia de muP. Abordamos esto introduciendo una regla de escalado de decaimiento de pesos para AdamW que preserva la ganancia de las subcapas a través de diferentes anchos. Empíricamente, el espectro de valores singulares de cada parámetro matricial escala en norma como eta/lambda con una forma aproximadamente invariante; bajo escalado de ancho d, observamos que el valor singular superior escala aproximadamente como eta/lambda * d^{0.75}. Combinando esta observación con la regla de tasa de aprendizaje de muP eta_2 ∝ d^{-1} para parámetros de tipo matricial, se implica una regla empírica de escalado de decaimiento de pesos lambda_2 ∝ d que aproximadamente mantiene las ganancias de las subcapas invariantes al ancho. Junto con parámetros de tipo vectorial entrenados a eta_1 = Theta_d(1) y lambda_1 = 0, esto permite la transferencia de cero disparos tanto de la tasa de aprendizaje como del decaimiento de pesos desde un ancho proxy a un ancho objetivo, eliminando la necesidad de barridos por ancho. Validamos la regla en Transformers de estilo LLaMA y en un entorno sintético mínimo, y proporcionamos un diagnóstico simple, la coincidencia de los valores singulares superiores, para verificar la invariancia de la ganancia de las subcapas. Nuestros resultados extienden muP más allá del régimen cercano a la inicialización al controlar explícitamente las escalas de estado estacionario establecidas por el optimizador, ofreciendo una receta práctica para la transferencia robusta de hiperparámetros en función del ancho bajo 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.