Rethinking SWATS (Switches from Adam to SGD) Optimizer. Performing the switch Locally than Globally.

• PyTorch == 1.5.1

Elementwise scaling of learning rate adopted by Adaptive Optimizers such as Adam, RMSProp etc often generalise poorly due to unstable and non-uniform learning rates at the end of training, although they scale well during the initial part of training. Hence SGD is the go-to for SOTA results since it generalizes better than adaptive methods.

SWATS is a method from Keskar and Socher which switches from Adam to SGD when the difference between the bias corrected projected learning rate and the projected learning rate is less than a threshold ϵ.The projected learning rate is found by projecting the SGD's update onto Adam's update. The switch is global i.e. if one of the layers of the network switches to SGD, all the layers are switched to SGD.

Switching all the layers to SGD just because a particular layer switched to SGD is something which has to be investigated. Why should a Dense or Batch Norm Layer impact Convolutional layer's training? So I went ahead and made the switch local i.e a layer's switch to SGD is independent of other layer switching to SGD i.e. at a given time some layers will be in SGD Phase while the rest of them are still using Adam.

To perform the analysis I use ResNet-18 on Cifar-10 from this repository.
Since paper's default value of ϵ (threshold to switch) wasn't working for me, I've reduced ϵ value from 10^-9 to 10^-5. Since ϵ value is left as a hyperparameter in the paper, I guess this is okay to do :)

• batch-size = 128
• epochs = 200

1. SWATS (GLOBAL SWITCH)

   • initial step size (adam) 0.001
   • step size decay by a factor of 10 at 100 epochs.

2. New SWATS (LOCAL SWITCH)

   • step size decay by a factor of 10 at 75,150 epochs for layers which are in SGD Phase.
   • step size decay by a factor of 10 at 100 epochs for layers which are in Adam.

3. Adabound

   • step size decay by a factor of 10 at 100 epochs for AdaBound.

1. SWATS (GLOBAL SWITCH)

   • initial step size (adam) 0.001
   • step size decay by a factor of 10 at 150 epochs.

2. New SWATS (LOCAL SWITCH)

   • step size decay by a factor of 10 at 75,150 epochs for layers which are in SGD Phase.
   • step size decay by a factor of 10 at 100 epochs for layers which are in Adam.

3. Adabound

   • step size decay by a factor of 10 at 150 epochs for AdaBound.

Since AdaBound is the only paper which I know which changes smoothly or in other words continuously transforms from Adam to SGD, rather than a hard switch like SWATS, it is fair to compare SWATS with AdaBound.

The idea of having some layers stay back in Adam phase and some of them moving forward to SGD(M) phase show that you need not switch all the layers to generalize better. Morever in Vanilla SWATS if the switch happens too early, Adaptive estimates cannot be utilised to the full extent. But if the switch is Local, some layers which are in Adam phase will help the initial part of training although the remaining layers have moved on to SGD.

As seen from the table above some layers(convolutional) in DenseNet have learning rates 4.96,6.37,8.56, 28.56. One may wonder how is the model converging at such high lr. I think this is due to Batch Normalization and it working well with exotic exponential learning rate schedulers Ref

Since there is no official implementation of SWATS from the authors, the code for it is borrowed from here . (SwatsVanillaGlobal.py).

If you find any bug in my implementation of switching it locally or find the idea of switching it locally "obscure", kindly feel free to drop an issue or email me at [email protected].