Recently we proposed a new optimization algorithm called Adaptive Sharpness-Aware Minimization (ASAM), which pushes the limit of deep learning via PAC-Bayesian theory. ASAM has been improving generalization performance for various tasks leveraging geometry of loss landscape, which is highly correlated with generalization. In recognition of its theoretical contribution and practicality, our paper was published at the 38th International Conference on Machine Learning (ICML 2021).

Motivation

Generalization means how well a trained deep neural network can classify or forecast unseen data. Nowadays, deep learning has made paradigm shifts in many academic or industrial areas utilizing generalization power of deep neural networks. Thus, boosting generalization power of deep neural networks via introduction of proper training techniques has been a main goal in many studies on deep learning.

To achieve this goal, deep learning finds a (local) minimum for a loss surface of training data by using gradient-based optimizers e.g. stochastic gradient descent (SGD) or Adam, which indirectly induces minimization of generalization loss.

To guarantee this indirect minimization of generalization loss, however, sufficiently large training data and large model size are required. In many practical cases, sufficiently large data and model size are not available, which limits the generalization power of deep neural networks.

To go one step further beyond only minimizing the training loss, here we design an optimizer which enhances the generalization ability of trained neural networks. This improvement of generalization power can contribute to data efficiency and reduction of required model size.

Background

A description of the correlation between the generalization gap and the sharpness of loss function is a good starting point for the discussion of the ASAM algorithm. The generalization gap and sharpness of training loss can be defined as follows:

where

As can be seen in the definition above, the generalization gap is defined as the difference between the generalization loss and the training loss. Thus, a large generalization gap means that a neural network has weak generalization power. The sharpness of loss function can be defined as the difference between the maximum training loss in an \[{\ell}^p\] ball with a fixed radius \[\mathbf{ρ}.\] and the training loss at \[\mathbf{w}.\]

The paper [1] shows the tendency that a sharp minimum has a larger generalization gap than a flat minimum does. The following figure is a conceptual sketch showing the positive correlation between the sharpness and generalization gap.

This tendency has also been investigated and verified via works based on PAC-Bayesian theory in the literature ([2], [3]). Inspired by the sharpness and the generalization gap, [4] suggests sharpness-aware minimization (SAM). SAM aims to minimize the local maximum of the training loss in epsilon-ball of \[\parallel{ϵ}\parallel_2 ≤ ρ\], which can be decomposed into the sharpness and training loss, by solving the following minimax problem to minimize the generalization loss indirectly.

However, [5] points out that this kind of sharpness with a fixed radius could have a weak correlation with the generalization gap due to vulnerability to weight scaling. This issue becomes more pronounced when a rectifier activation function such as ReLU is used. As in the following example, if the weight parameters before and after a ReLU activation function is multiplied by a scaling factor a and \[\frac{1}{a}\], respectively, the generalization power is not affected by the scaling factor, whereas the sharpness changes with the scaling factor.

EX) \[f\left(x;a\right)=\frac{1}{a}w_1ReLU(aw_2x)\]

This vulnerability to weight scaling of sharpness is the main cause that sharpness has a weak correlation with generalization, which prevents SAM from guaranteeing sufficiently strong generalization power of neural networks.

Adaptive Sharpness: A Scale-Invariant Generalization Measure

To strengthen the correlation by eliminating the vulnerability to weight scaling, we introduce the concept of adaptive sharpness defined as

where \[{T}_\mathbf{w}^{-\mathbf{1}}\] is a normalization operator which can cancel out the weight scaling effect. Since the maximization region \[\parallel{T}_\mathbf{w}^{-\mathbf{1}}ϵ\parallel ≤ ρ\] is adjusted with respect to positions in weight space, in contrast to the sharpness, the adaptive sharpness can be independent of weight scaling like generalization power as can be seen in Table 1. Through correlation analysis including granulated coefficients ψ and Kendall rank coefficients τ which are correlation measures suggested by [3] as well as scatter plots, we can demonstrate that the adaptive sharpness has stronger correlation compared to the sharpness.

Adaptive Sharpness-Aware Minimization (ASAM)

Using the concept of adaptive sharpness, we can formulate the Adaptive Sharpness-Aware Minimization (ASAM) problem as follows.

By minimizing the local maximum of the training loss in \[\parallel{T}_\mathbf{w}^{-\mathbf{1}}ϵ\parallel_2 ≤ ρ\], which can be decomposed into the adaptive sharpness and training loss, ASAM induces the minimization of generalization loss without the vulnerability to weight scaling. Also due to the negligible calculation cost of \[T_\mathbf{w}^{-1}\], ASAM can be considered as efficient as SAM.

The following figure shows behaviors of SAM and ASAM on loss contours of a simple neural network.

As can be seen in this figure, ASAM does not depend on the initial weight by benefiting from adjusting the region, whereas SAM may fail to converge to the valley entirely.

Experimental Results

We also empirically show that ASAM consistently outperforms SGD and SAM in the tests on CIFAR [6] and ImageNet [7] datasets using various neural network architectures. Moreover, for CIFAR-100 dataset, Pyramidnet-272 [8] trained with ASAM shows the state-of-the-art performance among the models using no extra data.

In the test of IWSLT2014 [9] which is a dataset of machine translation task, we could observe that generalization performance of transformer [10] can be improved by ASAM combined with Adam optimizer [11].

These results imply that ASAM can enhance generalization performance of deep neural network architectures in various tasks beyond the image classification.

Usability of ASAM

Another strength of ASAM is that it is an easy-to-apply, general-purpose optimizer like SGD or Adam. The code on the left shows a PyTorch example of typical training codes using SGD and we can improve the performance just by adding 4 lines to it as can be seen in the right code.

Conclusions

In the field of deep learning, optimization techniques (SGD, Adam, etc.) are one of the key factors influencing the performance of deep neural network models. As shown in this study, the development of such optimization techniques can generally improve the performance of models for various tasks, and the performance can be improved just by replacing the optimization technique without modifying the model structure, so we expect significant positive impact from ASAM. Contributing to the development of core technologies in the field of deep learning, ASAM will serve as a stepping stone to raise the level of quality of our products equipped with various AI models.