Research map / Architectures and efficiency / RNNs and normalization
Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks
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.
Derives steady-state (rotational equilibrium) neuronal update dynamics induced by weight decay for AdamW, Adam+L2, Lion and SGDM under a random-walk model, and introduces Rotational Variants of optimizers that directly control angular updates without weight decay.
Problem
Understanding why weight decay helps modern neural network training with normalization layers, and why AdamW outperforms Adam with L2 regularization, at the level of individual neuron update dynamics.
Approach
Weight decay drives scale-invariant weight norms to an equilibrium where the expected angular update (effective update size) becomes stable and, for AdamW, identical across neurons and layers (balanced rotation). Rotational Variants fix the angular update directly, removing weight decay while retaining the benefits and reducing warmup need.
Key results
- Rotational Variant of AdamW (RV-AdamW) matches AdamW baseline across tasks, e.g. CIFAR-10 ResNet-20 batch 128: 92.4±0.3 vs 92.2±0.2 baseline; GPT2-124M OpenWebText loss 3.18 vs 3.18.
- Disabling weight decay in AdamW degrades results (e.g. CIFAR-10 ResNet-20 batch 128: 91.1±0.3 vs 92.2±0.2; GPT training diverges without weight decay).
- RV-AdamW substantially reduces the need for learning rate warmup on GPT2-124M OpenWebText, where AdamW shows a significant gap with vs without warmup while RV shows minimal difference.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Disable weight decay (λ=0) | AdamW on ResNet-20 · CIFAR-10 | Top-1 accuracy (%): 92.2$\pm$0.2 -> 91.1$\pm$0.3 | Hurt |
| Replace weight decay with controlled angular updates (zero-shot RV) | AdamW on ResNet-20 · CIFAR-10 | Top-1 accuracy (%): 92.2$\pm$0.2 -> 92.4$\pm$0.3 | No clear effect |
| Disable weight decay (λ=0) | AdamW on ResNet-20 · CIFAR-10 | Top-1 accuracy (%): 91.5$\pm$0.3 -> 90.4$\pm$0.3 | Hurt |
| Replace weight decay with controlled angular updates (zero-shot RV) | AdamW on DeiT tiny · ImageNet-1k | Top-1 accuracy (%): 72.1 -> 71.3 | No clear effect |
| Disable weight decay (λ=0) | AdamW on DeiT tiny · ImageNet-1k | Top-1 accuracy (%): 72.1 -> 69.3 | Hurt |
| Replace weight decay with controlled angular updates (zero-shot RV) | AdamW on Transformer-S · IWSLT2014 de-en | BLEU: 34.6$\pm$0.1 -> 20.0$\pm$0.1 | Hurt |
| Tune weight decay for controlled angular updates (few-shot RV) | AdamW on Transformer-S · IWSLT2014 de-en | BLEU: 34.6$\pm$0.1 -> 34.7$\pm$0.1 | No clear effect |
| Disable weight decay (λ=0) | AdamW on GPT2-124M · OpenWebText | loss: 3.18$\pm$0.02 -> 3.21$\pm$0.01 | Hurt |
| Replace weight decay with controlled angular updates (zero-shot RV) | AdamW on GPT2-124M · OpenWebText | loss: 3.18$\pm$0.02 -> 3.18$\pm$0.01 | No clear effect |
| Disable weight decay (λ=0) | AdamW on GPT2-124M · WikiText | perplexity: 19.6$\pm$0.1 -> unstable | Hurt |
| Replace weight decay with controlled angular updates | AdamW on GPT2-124M · WikiText | perplexity: 19.6$\pm$0.1 -> 19.0$\pm$0.2 | Helped |
| Add Weight Standardization | Layer-Normalized ResNet-18 · CIFAR-100 | accuracy | Helped |
Surprising findings
- Without weight decay, GPT training eventually diverges (at least in bfloat16).
- Even small imbalances in neuronal rotation speed can significantly degrade performance.
Limitations
- Analysis relies on a random-walk approximation where noise dominates gradients; predictions are approximate for real training.
- Rotational variants are presented as research tools and are not shown to outperform baselines.
- Some layers (non-scale-invariant, e.g. final FC) deviate from predicted equilibrium due to radial gradient components.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet-1k IWSLT2014 de-en WikiText OpenWebText
Benchmarks: CIFAR-10 CIFAR-100 ImageNet-1k IWSLT2014 de-en WikiText OpenWebText
Metrics: top-1 accuracy BLEU perplexity validation loss expected angular update (rotation)
Baselines: AdamW Adam+L2 SGDM Lion LARS LAMB Nero AdamP
Builds on: Spherical Motion Dynamics Online Normalization AdamW Lion LARS
Closest papers by meaning
- Why Do We Need Weight Decay in Modern Deep Learning? (NeurIPS 2024 · similarity 0.88)
- Wide Neural Networks Trained with Weight Decay Provably Exhibit Neural Collapse (ICLR 2025 · similarity 0.87)
- Improving Deep Learning Speed and Performance Through Synaptic Neural Balance (AAAI 2025 · similarity 0.87)
- Understanding Adam Requires Better Rotation Dependent Assumptions (NeurIPS 2025 · similarity 0.87)
- Improving equilibrium propagation without weight symmetry through Jacobian homeostasis (ICLR 2024 · similarity 0.87)
- From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks (ICLR 2025 · 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.