按模块特性动态调整权重衰减，提升大模型训练效果
AlphaDecay: Module-wise Weight Decay for Heavy-Tailed Balancing in LLMs
- 根据权重相关矩阵的重尾程度，为不同模块分配不同衰减强度
- 在60M到1B参数模型上，困惑度更低且泛化能力更强
- 适合追求训练稳定性和性能优化的大模型研究者
权重衰减是训练大语言模型的标准正则化方法。尽管通常对所有层采用统一衰减率，但这种方法忽略了大模型的结构多样性及各模块间谱性质的差异。本文提出AlphaDecay，一种基于重尾自正则化（HT-SR）理论的简单而有效的方法，可自适应地为每个模块分配不同的权重衰减强度。该方法通过分析权重相关矩阵的经验谱密度（ESD）来量化“重尾性”：重尾特征越明显（反映更强的特征学习能力）的模块，赋予更弱的衰减；轻尾谱的模块则接受更强衰减。通过这种模块级谱特性平衡，显著提升了性能。在多个从60M到1B参数规模的预训练任务中，AlphaDecay在困惑度和泛化能力上均优于传统均匀衰减及其他自适应衰减基线。代码已开源：https://github.com/hed-ucas/AlphaDecay。
原文摘要 · Abstract (English)
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. Our code is available at https://github.com/hed-ucas/AlphaDecay.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。