Research map / Large language models / Language, safety and interpretability
AlphaDecay: Module-wise Weight Decay for Heavy-Tailed Balancing in LLMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AlphaDecay, which assigns module-wise weight decay per LLM module by scaling it with each module's Hill-estimated heavy-tailed exponent from HT-SR theory, improving C4 pre-training perplexity for LLaMa models from 60M to 1B.
Problem
Uniform weight decay across all layers and modules in LLM pre-training ignores structural diversity and differing spectral properties of modules (attention vs. MLP), which may be suboptimal.
Approach
Estimate each module's heavy-tailed self-regularization exponent (PL_Alpha_Hill of the weight correlation matrix eigenvalues) periodically during training and linearly map it to a per-module weight decay in [s1*eta, s2*eta], giving larger decay to modules with higher alpha, to balance module-wise spectral properties.
Key results
- On C4 pre-training, AlphaDecay reaches validation perplexity 15.13 on LLaMa-1B at weight decay 1e-5 versus 15.36 for Uniform, and 22.55 on LLaMa-135M at 5e-6 versus 23.14 for Uniform.
- On seven zero-shot commonsense tasks with LLaMa-1B, AlphaDecay averages 43.35 versus 41.50 for Uniform and is best on 6 of 7 benchmarks.
- On GLUE fine-tuning with roberta-base, AlphaDecay averages 85.12 versus 83.83 for Uniform.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace uniform weight decay with module-wise Hill-guided weight decay | LLaMa pretraining · C4 validation, LLaMa-135M/350M/1B | perplexity: 22.99 -> 22.76; 23.14 -> 22.55; 24.14 -> 23.49; 17.12 -> 17.00; 16.74 -> 16.66; 17.50 -> 16.88; 15.36 -> 15.13; 14.66 -> | Helped |
| Replace uniform weight decay with module-wise Hill-guided weight decay | LLaMa-1B pretrained model · Zero-shot commonsense average | score: 41.50 -> 43.35 | Helped |
| Replace uniform weight decay with module-wise Hill-guided weight decay | roberta-base finetuning · GLUE average | score: 83.83 -> 85.12 | Helped |
| Replace Uniform assignment with Linear weight-decay assignment (ablation) | AlphaDecay · C4, LLaMa-60M | perplexity: 32.386 -> 31.565 | Helped |
| Replace Linear assignment with Sqrt, Log2, or Sigmoid-like assignment (ablation) | AlphaDecay · C4, LLaMa-60M | perplexity: 31.565 vs 32.326, 32.301, 32.201 | Hurt |
| Use Hill rather than alternative HT-SR metrics to guide decay (ablation) | AlphaDecay · C4, LLaMa-135M | perplexity: 22.993 -> 22.763; 23.138 -> 22.551; 24.142 -> 23.488 | Helped |
| Replace Hill with GradNorm, FrobeniusNorm, or SpectralNorm (ablation) | AlphaDecay · C4, LLaMa-135M | perplexity: 22.763 vs 22.855, 23.265, 23.348; 22.551 vs 22.714, 22.91, 23.006; 23.488 vs 23.801, 23.596, 23.601 | Hurt |
| Use Median rather than Goodness-of-fit or Fix-finger PL fitting (ablation) | AlphaDecay · C4, LLaMa-60M and LLaMa-135M | computation time: 8.73±0.30 -> 1.67±0.01; 7.90±0.33 -> 1.67±0.01; 24.86±0.11 -> 4.67±0.02; 21.53±0.83 -> 4.67±0.02 | Helped |
| Use Median rather than alternative PL fitting methods (ablation) | AlphaDecay · C4, LLaMa-60M, weight decay 1e-5 | perplexity: 32.166 -> 31.628; 32.231 -> 31.628 | Helped |
| Adjust module-wise weight decay every 500 steps rather than using Uniform decay (ablation) | AlphaDecay · C4, LLaMa-135M | perplexity: 22.994 -> 22.763; 23.138 -> 22.551; 24.142 -> 23.488 | Helped |
| Increase update frequency from every 500 steps to every step (ablation) | AlphaDecay · C4, LLaMa-135M | computation time: 5.7h -> 74.5h | Hurt |
| Vary the module-wise weight-decay update interval from 500 steps to 1, 50, or 100 steps (ablation) | AlphaDecay · C4, LLaMa-135M | perplexity: 22.763; 22.809; 22.758; 22.779; 22.551; 22.581; 22.539; 22.569; 23.488; 23.488; 23.468; 23.479 | No clear effect |
Surprising findings
- Existing adaptive weight decay methods AWD and AdaDecay, designed for non-attention architectures, often fail to beat the Uniform baseline on LLM pre-training (e.g., AWD is worse than Uniform at all LLaMa-135M settings).
Limitations
- Evaluation on Mixture-of-Experts models is left for future work.
- Interactions with other regularization and optimization techniques are not systematically assessed.
- Gains are modest in several settings and depend on tuned hyperparameters (s1, s2) and learning-rate grid search.
Experimental setup
Datasets: C4 GLUE ImageNet-1K lm-eval-harness commonsense tasks
Benchmarks: C4 validation perplexity ARC-c ARC-e PIQA Hellaswag OBQA Winogrande BoolQ GLUE ImageNet-1K
Metrics: validation perplexity zero-shot accuracy GLUE accuracy top-1 accuracy p-value (paired t-test)
Baselines: Uniform AWD AdaDecay WD=0 GradNorm FrobeniusNorm SpectralNorm
Builds on: Heavy-Tailed Self-Regularization (HT-SR) theory AdamW AWD AdaDecay
Closest papers by meaning
- AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning of Large Language Models (NeurIPS 2024 · similarity 0.90)
- Treasures in Discarded Weights for LLM Quantization (AAAI 2025 · similarity 0.88)
- OWQ: Outlier-Aware Weight Quantization for Efficient Fine-Tuning and Inference of Large Language Models (AAAI 2024 · similarity 0.88)
- Sign-Based Optimizers Are Effective Under Heavy-Tailed Noise (NeurIPS 2026 · similarity 0.87)
- OSTQuant: Refining Large Language Model Quantization with Orthogonal and Scaling Transformations for Better Distribution Fitting (ICLR 2025 · similarity 0.87)
- Adaptive Layer Sparsity for Large Language Models via Activation Correlation Assessment (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.