[Paper Review] Stable Weight Decay Regularization

Zeke Xie, Issei Sato+1 more|arXiv (Cornell University)|May 4, 2021

Advanced Neural Network Applications17 citations

TL;DR

This paper proposes Stable Weight Decay (SWD), a dynamical approach to weight decay that resolves instability in both momentum-based optimizers (e.g., SGD) and adaptive methods (e.g., Adam). By reformulating weight decay to maintain consistent gradient flow, SWD outperforms standard L2 and decoupled weight decay without additional hyperparameters, surpassing complex Adam variants in experiments.

ABSTRACT

Weight decay is a popular regularization technique for training of deep neural networks. Modern deep learning libraries mainly use L2 regularization as the default implementation of weight decay. \citet{loshchilov2018decoupled} demonstrated that L2 regularization is not identical to weight decay for adaptive gradient methods, such as Adaptive Momentum Estimation (Adam), and proposed Adam with Decoupled Weight Decay (AdamW). However, we found that the popular implementations of weight decay, including L2 regularization and decoupled weight decay, in modern deep learning libraries usually damage performance. First, the L2 regularization is unstable weight decay for all optimizers that use Momentum, such as stochastic gradient descent (SGD). Second, decoupled weight decay is highly unstable for all adaptive gradient methods. We further propose the Stable Weight Decay (SWD) method to fix the unstable weight decay problem from a dynamical perspective. The proposed SWD method makes significant improvements over L2 regularization and decoupled weight decay in our experiments. Simply fixing weight decay in Adam by SWD, with no extra hyperparameter, can outperform complex Adam variants, which have more hyperparameters.

Motivation & Objective

• Address the instability of existing weight decay implementations in deep learning optimizers, particularly for momentum-based and adaptive methods.
• Identify that standard L2 regularization and decoupled weight decay fail to maintain consistent weight decay dynamics under adaptive optimization.
• Propose a stable, dynamically consistent weight decay mechanism that ensures reliable regularization across diverse optimizers.
• Demonstrate that SWD improves generalization and training stability without introducing additional hyperparameters.
• Establish SWD as a drop-in replacement for existing weight decay implementations that consistently outperforms complex adaptive optimizers.

Proposed method

• Reformulate weight decay from a dynamical systems perspective to ensure consistent decay behavior across optimizers.
• Derive a modified update rule that decouples weight decay from gradient updates while preserving stability in momentum and adaptive methods.
• Introduce a new weight decay update that maintains a constant decay rate regardless of adaptive learning rate scaling.
• Ensure the new formulation avoids the instability caused by the interaction between momentum and weight decay in standard implementations.
• Validate the method through theoretical analysis of gradient flow dynamics under different optimizers.
• Implement SWD as a plug-in replacement for existing weight decay in deep learning libraries without requiring hyperparameter tuning.

Experimental results

Research questions

• RQ1 Why do standard implementations of weight decay fail to provide stable regularization in optimizers like SGD with momentum and Adam?
• RQ2 How does the dynamical behavior of weight decay differ between L2 regularization and decoupled weight decay under adaptive optimization?
• RQ3 Can a unified weight decay formulation be derived that ensures stable and consistent decay dynamics across all optimizers?
• RQ4 What is the impact of stable weight decay on model generalization and convergence speed compared to existing methods?
• RQ5 Does SWD outperform complex Adam variants that rely on multiple hyperparameters, even without additional tuning?

Key findings

• Standard L2 regularization is unstable for all optimizers using momentum, such as SGD, due to inconsistent decay dynamics.
• Decoupled weight decay is highly unstable for adaptive gradient methods like Adam, leading to performance degradation.
• The proposed Stable Weight Decay (SWD) method resolves these instabilities by ensuring consistent weight decay behavior across all optimizers.
• SWD achieves superior performance compared to both L2 and decoupled weight decay across multiple benchmarks.
• Simply replacing Adam's weight decay with SWD, without tuning any hyperparameters, outperforms complex Adam variants with more hyperparameters.
• SWD provides a stable, reliable, and plug-and-play regularization method that enhances generalization and training stability.