Title: AlphaDecay: Module-wise Weight Decay for Heavy-Tailed Balancing in LLMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: Weight decay

Large Language Models

Heavy-Tailed Self-Regularization
Abstract: 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAfter establishing that there are variations in the spectral properties of different components in transformer blocks of LLMs, i.e., the weight matrices of attention modules have singular values that follow power law distributions with heavier tails as opposed to the singular values of weight matrices attached to MLP modules, which follow power law distributions with lighter tails, the authors of this work propose a weight decay regularization scheme that aids in closing the gap between the aforementioned spectral properties of attention and MLP components. Namely, the AlphaDecay regularization strategy results in making heavier tails lighter (attention) and lighter tails heavier (MLP).\n\nExperimental evidence in the paper suggests that AlphaDecay outperforms, in terms of achieving lower perplexity score, other baselines such as the Uniform weight decay assignment, Adaptive Weigth Decay (ADW), and AdaDecay.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper is well-written, the scope and contributions clearly articulated by the authors, and the methodology elegant. Although, I would like to point out to what I see as potential weaknesses and ground for substantial improvement.\n\nWeaknesses: \n\n- A downside is that the experimental results are presented for a single family of models, i.e., LLaMa, on a single dataset, that is the C4 dataset. While C4 is a widely used corpus for pre-training, other datasets could be used to assess the capabilities of LLaMa models post-training when applying the AlphaDecay regularization strategy.\n- Besides using only the perplexity score for validation, additional metrics like accuracy, BLEU or F1 could be reported if the model was fine-tuned or evaluated on downstream tasks derived from C4.\n- When describing the hyperparameter settings, it is mentioned that all the models were trained using the Adam optimizer. Would the results remain stable under using other optimizers such as Adam, AdamW, AdamS or even SGD?\n- Finally, all the results in Figures 5-8 are discussed when fixing the scale of the LLaMa model with 135M parameter count (besides Fig. 6 which considered the LLaMa model with 60M parameters). It will be good to demonstrate the universality of the results presented in those figures when considering all the LLaMa models (with 60M, 135M, 350M, 1B parameters).\n\nQUESTIONS:\n1. lines 41-42: the phrasing is slightly ambiguous. What should be understood from the context is that the model uses a single, time-dependent weight decay value that is updated over training steps. At each training time-step, this value is applied uniformly across all layers.\n\n\n2. lines 52-54: In listing the first contribution of their work, the authors claim that the substantial variations in the spectral properties of attention and MLP layers, are primarily responsible for harming the generalization ability, referring to Fig. 4. However, it seems that the left subfigure in Fig. 4 does not readily yield that link. Instead, it can be more accurately said, that the AlphaDecay scheme that the authors propose, balances the variations in the module-wise spectral properties, and in comparison with applying a module-wise uniform weight decay, the proposed scheme establishes superior generalization in terms of the perplexity score.\n\n\n3. lines 55-57: The spectral alignment between the various modules is witnessed in Fig. 1 as well as in Fig. 3. There is an erroneous figure number referenced in line 57. Instead of Fig. 3, the improved training performance is reported in Fig. 4. It will be also good to include in the caption of Fig. 4 whether perplexity refers to training or validation perplexity.\n\n\n4. line 70: While the first part of the Related Work section is related to weight decay and LLMs, reference [20] is mentioned as showing the importance of weight decay in scale-invariant systems. Although, LLMs are not globally scale-invariant, one can especially refer to some of their components (more specifically layer normalization) as scale-invariant. And since this work is focusing on weight decay schemes applied solely to attention and MLP components, it might be out of scope to refer to the relevance of weight decay in scale-invariant systems.\n\n\n5. There is an overlap of information between lines 74-79 and 88-92.\n\n\n6. line 115: The reference to Fig. 3 might be unnecessary and misleading since the paragraph describes the findings of Fig. 2.\n\n\n7. line 134: L instead of N modules.\n\n\n8. line 135: Are those the correlation or covariance matrices? Are the weight matrices normalized?\n\n\n9. The last sentence in Fig. 3 is misleading. What the figure demonstrates is that the authors’ method forces balance between “overtrained” and “undertrained” modules which according to references [24, 46] contributes to overall model quality, which is rightly mentioned in l. 151. Besides that, Fig. 3 does not say anything about model performance so far.\n\nLIMITATIONS:\nThe authors have addressed the limitations properly. See also weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes AlphaDecay, an adaptive, module-wise weight decay method for LLM training. AlphaDecay dynamically assigns different decay strengths to each module of a transformer based on the heavy-tailedness of their weight spectra, quantified via the Hill estimator of power-law exponents. The authors justify their approach using heavy-tailed self-regularization (HT-SR) theory and show empirically, on LLaMa-based models of various sizes (60M–1B), that AlphaDecay improves pre-training perplexity over standard and other adaptive baselines (AWD, Adadecay). Extensive ablations demonstrate the robustness of the method to different hyperparameters and computation intervals.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The motivation for the proposed method is simple, straightforward, and easy to grasp.  the use of HT-SR theory to analyze module-specific training behavior and justify differential regularization is well-grounded.\n- The experimental setup is solid and provides clear evidence of the method’s effectiveness and robustness.\n\nWeakness:\n- The number of compared baselines is limited. Notably, the strongest baseline cited was published in 2019, which raises concerns about whether there are more recent or advanced approaches in this research area that should be included for comparison. \n- The largest model evaluated in this work has 1B parameters, which is relatively small by current LLM standards. Evaluating on larger models, such as 7B or above, would make the results more compelling. Additionally, the experiments are limited to the LLaMa series of models, which restricts the generalizability of the findings.\n\nQUESTIONS:\nQ1: In Figure 8, the setting with GAP=50 achieves the best performance. Could the authors provide details on the computational cost of AlphaDecay when the interval size is set to 50?\n\nQ2: The paper focuses on the pretraining stage. Is it possible to extend AlphaDecay to the finetuning stage, especially for specific downstream tasks such as question answering or text classification? Any insights or preliminary results in this direction would be valuable.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates adaptive layer-wise weight decay method, an interesting area in the LLMs training. The authors utilize heavy-tailed self-regularization theory to guide the research direction, which analyzes the empirical spectral density of the weight correlation matrices to indicate the heavy-tailedness of the module. The proposed method assign weaker weight decay to modules with more obvious heavy-tailed ESDs reflects and stronger weight decay otherwise, since the heavy-tailed ESDs reflect the degree of feature learning. Simulation results demonstrate the effectiveness and robustness of the proposed method.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n1.\tThis paper introduces an adaptive weight decay method, which differs from the traditional methods that utilizes a uniform or time-wise decay rate.\n2.\tThe authors adopt HT-SR theory to guide the weight decay rate selection. The theory provides a strong theoretical basis for the proposed method.\n3.\tThe simulation results and performance improvements across various model sizes show the effectiveness and robustness of the proposed method.\n\nWeaknesses\n1.\tThe rationale and theoretical framework is vague. The authors only provide high-level description of the theory and there is no detailed mathematical formulation.\n2.\tThe clarity of the paper requires improvement. There are some typos and poorly-defined variables. Some important notion such as heavy-tailedness and some figures are not well explained.\n\nQUESTIONS:\n1.\tThe authors should provide detailed description of the theory and mathematical formulation.\n2.\tThere are some typos and poorly-defined variables. HT-SR seems to be used as HR-SR in this work, which is a big misleading. Is the Hill estimator equivalent to \\alpha? If so, it is suggested to replace PL_ALPHA_HILL with variable like \\hat{\\alpha} to enhance clarity. \n3.\tSome notations and figures require further explanation and clarification. The notion of heavy-tailedness requires further clarification. Is this notion defined by metric singular values, i.e., Fig. 2, or ESD, i.e., Fig. 3. The attention layers (att.q and att.k) exhibits less heavy-tailedness in Fig.2 while more heavy-tailedness in Fig. 3.In Fig.2, why are some layers, e.g., layer 1, 2 and 39, 40, drawn with different color?  In Fig.4, the PL_ALPHA_HILL does not change significantly with weight decay, please explain why. Same in Fig. 4, is the weight decay equivalent to \\eta in the context? In Fig.8, it seems more frequent weight decay adjustment does not yield performance improvement, which requires further explanation.\n4.\tThe performance of the proposed method is only evaluated on one dataset and metric. More datasets and metrics are needed to demonstrate robustness.\n5.\tCertain parameter setting is not motivated. For example, the selection of the range of scaling ratios should be explained.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper is concerned with the problem  weight decay in LLMs to  improve generalization.  In particular the  authors propose a scheme based on HT-SR theory to adaptively choose the weight day for different modules. The introduced method called AlphaDecay first calculates the PL_Alpha_Hill values for all modules, and then assign larger weight decay to modules with higher\nPL_Alpha_Hill values, while assigning smaller weight decay to those with lower PL_Alpha_Hill\nvalues. This scheme balances the modules.\n\nSTRENGTHS AND WEAKNESSES:\nThe strength of this paper is that it is  tackling the weight decay for improving generalization of LLMs which is awidely sought topic and interesting to the  community. The weakness is:\n\n- The effect of  promoting uniformity in PL_Alpha_Hill across modules on model quality is only shown emprically and no mathematical proof is provided to prove it. This is while  the model quality is measure only by perplexity. I am not convinced this  strategy will improve other metrics and generalization.\n\nQUESTIONS:\nIn figure 3, the performance of model is not shown and only the ESD is  demostrated, therefore I am not inferring the following from figure 3:\n\n“We empirically demonstrate that promoting uniformity in PL_Alpha_Hill across modules (e.g.,\nattention and MLP components) can further enhance overall model quality (see figure 3).”\n\nLIMITATIONS:\n-The paper is purely empirical and reporting observations on balancing weight day  on models qulity.  The terms such as model quality is  very generally used and it is not clear if the methods enhances performance in terms of ID and OOD  generalization gaps or other metrics.\n\n-The method is only implemented on LLama model and it is not clear if the same improvement will be achieved on other architetures.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-21"}