AlphaDecay: Module-wise Weight Decay for Heavy-Tailed Balancing in LLMs

Li Shen (Sun Yat-Sen University) · Songjun Tu (Institute of Automation Chinese Academy of Sciences) · Ganzhao Yuan (Shenzhen University of Advanced Technology) · Lu Yin (University of Surrey) · Shiwei Liu (University of Oxford) · Di He (Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences, Chinese Academy of Sciences) · Ajay Jaiswal (UT Austin)

adaptive decay.alphadecayempirical spectral densityfeature learning generalizationheavy-tailed self-regularizationheavy-tailednessmodule-wise differencesperplexity spectral properties structural diversitytailored weight decayweight correlation matricesweight decay

Weight decay is a standard regularization technique for training large language models (LLMs). While it is common to assign a uniform decay rate to every layer, this approach overlooks the structural diversity of LLMs and the varying spectral properties across modules. In this paper, we introduce AlphaDecay, a simple yet effective method that adaptively assigns different weight decay strengths to each module of an LLM. Our approach is guided by Heavy-Tailed Self-Regularization (HT-SR) theory, which analyzes the empirical spectral density (ESD) of weight correlation matrices to quantify “heavy-tailedness.” Modules exhibiting more pronounced heavy-tailed ESDs, reflecting stronger feature learning, are assigned weaker decay, while modules with lighter-tailed spectra receive stronger decay. Our method leverages tailored weight decay assignments to balance the module-wise differences in spectral properties, leading to improved performance. Extensive pre-training tasks with various model sizes from 60M to 1B demonstrate that AlphaDecay achieves better perplexity and generalization than conventional uniform decay and other adaptive decay baselines. The code is available at https://github.com/hed-ucas/AlphaDecay.