Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks (2305.17212v4)

Abstract: This study investigates how weight decay affects the update behavior of individual neurons in deep neural networks through a combination of applied analysis and experimentation. Weight decay can cause the expected magnitude and angular updates of a neuron's weight vector to converge to a steady state we call rotational equilibrium. These states can be highly homogeneous, effectively balancing the average rotation -- a proxy for the effective learning rate -- across different layers and neurons. Our work analyzes these dynamics across optimizers like Adam, Lion, and SGD with momentum, offering a new simple perspective on training that elucidates the efficacy of widely used but poorly understood methods in deep learning. We demonstrate how balanced rotation plays a key role in the effectiveness of normalization like Weight Standardization, as well as that of AdamW over Adam with L2-regularization. Finally, we show that explicitly controlling the rotation provides the benefits of weight decay while substantially reducing the need for learning rate warmup.

Rotational Equilibrium in Neural Network Optimization

The paper "Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks" by Kosson et al. presents an investigation into the role of weight decay—or $\ell_2$-regularization—in balancing updates across neural networks. The authors aim to demystify the dynamics that enable optimizations in neural network training, a topic that, despite its extensive use, remains poorly understood.

Main Contributions

1. Rotational Equilibrium: The primary conceptual framework introduced is "Rotational Equilibrium," which describes a steady state where weight decay aligns the learning dynamics across various neural network layers. This state emerges from the interplay between weight decay, which tends to shrink weight magnitudes, and gradient updates, which can increase them. The paper’s analysis, supported by experimental evidence, shows that achieving this equilibrium yields more homogeneous training dynamics across network layers.
2. Optimizer Analysis: The authors examine the equilibrium conditions for several common optimizers, including SGDM, AdamW, and Adam with $\ell_2$-regularization, providing insights into why certain configurations are naturally more effective. They elucidate that the homogeneity stems from establishing a balanced rotational update rate ($\eta_r$) across neurons.
3. Decoupled Weight Decay: One notable contribution is the explanation for why AdamW generally outperforms Adam with traditional $\ell_2$-regularization. The work argues that AdamW achieves a balanced rotation due to its decoupled nature, avoiding the inconsistent angular updates that $\ell_2$-regularization might cause when applied directly.
4. Rotational Variants of Optimizers: The paper introduces rotational optimizer variants that explicitly control the angular update sizes instead of applying weight decay, thus achieving similar benefits. The proposed variants remain competitive in performance while reducing the typical dependency on learning rate warmups.

Experimental Verification

The paper is supported by extensive empirical validation, notably across various neural network architectures including ResNet and GPT2-style models on tasks like image classification and LLMing. These experiments confirm that the theoretically derived rotational equilibrium conditions are observable in practice. For example, layers with batch normalization exhibit converging rotational dynamics as predicted.

Practical Implications and Future Directions

By clarifying the operational mechanics of weight decay, the research provides a foundation for refining hyperparameter tuning aimed at achieving equilibrium quicker, potentially leading to faster convergence or enhanced model performance. Practically, the insights into rotational dynamics offer a blueprint for designing new optimization strategies that circumvent the transient phases of training.

The theoretical speculations and empirical findings indicate that further adjustments in training protocols could harness better-balanced learning across layers and neurons. The introduction of rotational variants is particularly promising, showing potential to decouple weight management from learning rate scheduling, which could simplify training regimens and reduce computational overhead.

Conclusion

Kosson et al.'s work offers a compelling advancement in our understanding of weight decay's role in neural network training. By parsing weight dynamics into interpretable geometric models, the paper not only enriches the theoretical landscape of optimization in machine learning but also furnishes practical pathways for more efficient algorithm designs. Future research may build on this work by exploring alternative formulations for balance across other optimizer variants, potentially unlocking new efficiencies in large-scale neural network training.