Regras Robustas de Escalonamento por Camada através da Ajuste Adequado da Decaimento de Peso

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

Resumo

As leis de escalonamento empíricas prescrevem como alocar parâmetros, dados e computação, enquanto a parametrização de atualização máxima (muP) permite a transferência da taxa de aprendizado entre larguras ao equalizar as magnitudes das atualizações no início do treinamento. No entanto, em arquiteturas modernas invariantes à escala, o treinamento rapidamente entra em um estado estacionário governado pelo otimizador, onde camadas de normalização criam sensibilidade à escala no retropropagação e a taxa de aprendizado efetiva se torna dependente da largura, degradando a transferência do muP. Nós abordamos isso introduzindo uma regra de escalonamento de decaimento de peso para o AdamW que preserva o ganho das subcamadas entre larguras. Empiricamente, o espectro de valores singulares de cada parâmetro matricial escala em norma como eta/lambda com uma forma aproximadamente invariante; sob escalonamento de largura d, observamos que o valor singular máximo escala aproximadamente como eta/lambda * d^{0.75}. Combinando essa observação com a regra de taxa de aprendizado do muP eta_2 proporcional a d^{-1} para parâmetros do tipo matriz implica uma regra empírica de escalonamento de decaimento de peso lambda_2 proporcional a d que aproximadamente mantém os ganhos das subcamadas invariantes à largura. Juntamente com parâmetros do tipo vetor treinados em eta_1 = Theta_d(1) e lambda_1 = 0, isso resulta na transferência zero-shot tanto da taxa de aprendizado quanto do decaimento de peso de larguras proxy para alvo, eliminando varreduras por largura. Validamos a regra em Transformers no estilo LLaMA e em um cenário sintético mínimo, e fornecemos um diagnóstico simples, correspondendo os valores singulares máximos, para verificar a invariância do ganho das subcamadas. Nossos resultados estendem o muP além do regime próximo à inicialização ao controlar explicitamente as escalas de estado estacionário definidas pelo otimizador, oferecendo uma receita prática para a transferência robusta de hiperparâmetros em largura sob o 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.