Scheduled(Stable)-Weight-Decay-Regularization

February 3, 2024 ยท View on GitHub

The PyTorch Implementation of Scheduled (Stable) Weight Decay.

The algorithms were first proposed in our arxiv paper.

A formal version with major revision and theoretical mechanism "On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective" is accepted at NeurIPS 2023.

Why Scheduled (Stable) Weight Decay?

We proposed the Scheduled (Stable) Weight Decay (SWD) method to mitigate overlooked large-gradient-norm pitfalls of weight decay in modern deep learning libraries.

  • SWD can penalize the large gradient norms at the final phase of training.

  • SWD usually makes significant improvements over both L2 regularization and decoupled weight decay.

  • Simply fixing weight decay in Adam by SWD, with no extra hyperparameter, can usually outperform complex Adam variants, which have more hyperparameters.

The environment is as bellow:

Python 3.7.3

PyTorch >= 1.4.0

Usage

You may use it as a standard PyTorch optimizer.

import swd_optim

optimizer = swd_optim.AdamS(net.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False)

Test performance

DatasetModelAdamSSGD MAdamAMSGradAdamWAdaBoundPadamYogiRAdam
CIFAR-10ResNet184.910.045.010.036.530.036.160.185.080.075.650.085.120.045.870.126.010.10
VGG166.090.116.420.027.310.257.140.146.480.136.760.126.150.066.900.226.560.04
CIFAR-100DenseNet12120.520.2619.810.3325.110.1524.430.0921.550.1422.690.1521.100.2322.150.3622.270.22
GoogLeNet21.050.1821.210.2926.120.3325.530.1721.290.1723.180.3121.820.1724.240.1622.230.15

Citing

If you use Scheduled (Stable) Weight Decay in your work, please cite "On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective".

@inproceedings{xie2023onwd,
    title={On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective},
    author={Xie, Zeke and Xu, Zhiqiang and Zhang, Jingzhao and Sato, Issei and Sugiyama, Masashi},
    booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
    year={2023}
}