Understanding AdamW through Proximal Methods and Scale-Freeness

TL;DR

The paper analyzes why AdamW often outperforms Adam-$\ell_2$ by linking AdamW to proximal updates that fully utilize the regularizer and by identifying a scale-free update property that makes AdamW robust to gradient rescaling. It establishes a theoretical connection to proximal methods and shows that scale-free updates can effectively reduce conditioning, with empirical evidence that the advantage of AdamW grows in deep, BN-free networks. Through extensive CIFAR experiments, the work correlates AdamW's performance with gradient scale diversity and demonstrates the practical similarity between AdamW and its proximal counterpart, AdamProx. The findings illuminate when and why decoupling weight decay from the gradient benefits optimization, and they guide future exploration of scale-free optimization in deep learning and distributed settings.

Abstract

Adam has been widely adopted for training deep neural networks due to less hyperparameter tuning and remarkable performance. To improve generalization, Adam is typically used in tandem with a squared $\ell_2$ regularizer (referred to as Adam-$\ell_2$). However, even better performance can be obtained with AdamW, which decouples the gradient of the regularizer from the update rule of Adam-$\ell_2$. Yet, we are still lacking a complete explanation of the advantages of AdamW. In this paper, we tackle this question from both an optimization and an empirical point of view. First, we show how to re-interpret AdamW as an approximation of a proximal gradient method, which takes advantage of the closed-form proximal mapping of the regularizer instead of only utilizing its gradient information as in Adam-$\ell_2$. Next, we consider the property of "scale-freeness" enjoyed by AdamW and by its proximal counterpart: their updates are invariant to component-wise rescaling of the gradients. We provide empirical evidence across a wide range of deep learning experiments showing a correlation between the problems in which AdamW exhibits an advantage over Adam-$\ell_2$ and the degree to which we expect the gradients of the network to exhibit multiple scales, thus motivating the hypothesis that the advantage of AdamW could be due to the scale-free updates.

Figures (12)

• Figure 1: Non-scale-free GD v.s. scale-free AdamW on quadratic functions with different condition numbers.
• Figure 2: The final Top-1 test error on using AdamW vs. Adam-$\ell_2$ on training a Resnet/DenseNet with Batch Normalization on CIFAR10/100 (the black circle denotes the best setting).
• Figure 3: On using AdamW vs. Adam-$\ell_2$ on training a Resnet/DenseNet without Batch Normalization on CIFAR10/100. (Left two) The final Top-1 test error (the black circle denotes the best setting). (Middle two) The training loss and test accuracy curve when employing the initial step size and the weight decay parameter that gives the smallest test error. (Right two) The histogram of the magnitude of corresponding updates of all coordinates of the network near the end of the training when employing the initial step size and the weight decay parameter that gives the smallest test error.
• Figure 4: The final top-1 test error of AdamW vs. Adam-$\ell_2$ on optimizing a 110-layer Resnet with BN removed on CIFAR-10 with the loss function multiplied by 10 (left two figures) and 100 (right two figures).
• Figure 5: The final Top-1 test error of using AdamW vs. AdamProx on training (the black circle denotes the best setting). (Top row) a 110-layer ResNet with BN removed on CIFAR-10 (trained for 300 epochs). (Bottom row) a 100-layer DenseNet-BC with BN removed on CIFAR-100 (trained for 100 epochs).

Theorems & Definitions (6)

• Theorem 3.1
• proof
• Theorem B.1
• Theorem B.2
• proof : Proof of Theorem \ref{['thm:restart_adagrad']}
• proof : Proof of Theorem \ref{['thm:AdaGrad']}