GA - SAM : Gradient-Strength based Adaptive Sharpness-Aware Minimization for Improved Generalization (2022)

Abstract: Recently, Sharpness-Aware Minimization (SAM) algorithm has shown state-of-the-art generalization abilities in vision tasks. It demonstrates that flat minima tend to imply better generalization abilities. However, it has some difficulty implying SAM to some natural language tasks, especially to models with drastic gradient changes, such as RNNs. In this work, we analyze the relation between the flatness of the local minimum and its generalization ability from a novel and straightforward theoretical perspective. We propose that the shift of the training and test distributions can be equivalently seen as a virtual parameter corruption or perturbation, which can explain why flat minima that are robust against parameter corruptions or perturbations have better generalization performances. On its basis, we propose a Gradient-Strength based Adaptive Sharpness-Aware Minimization (GA-SAM) algorithm to help to learn algorithms find flat minima that generalize better. Results in various language benchmarks validate the effectiveness of the proposed GA-SAM algorithm on natural language tasks.

## introduction recently, researchers @xcite @xcite @xcite propose that for better generalization ability, learning algorithms should find flat minima that have better robustness resistant to parameter corruptions or perturbations. many learning algorithms that take the flatness or sharpness of the parameters into consideration are motivated by the observation that flat minima tend to imply better generalization abilities. among them, sharpness-aware minimization (sam) @xcite algorithm has achieved state-of-the-art generalization abilities in vision tasks. it adopts virtual adversarial parameter corruptions or perturbations during training and lowers the risk after parameter corruptions. however, traditional sam algorithms usually adopt fixed strengths of parameter corruptions and constraint the corruptions with l 2 -norm or l +∞ -norm balls. it cannot conduct flexible strengths of parameter corruptions for different parameters, or during different stages of training. thus, it is difficult to apply sam to some natural language tasks, especially to models with drastic gradient changes, such as rnns. to settle this issue, many adaptive sam algorithms @xcite are proposed empirically. in this work, we propose a gradient-strength based adaptive solution based on our theoretical framework. existing studies @xcite try to explain the relation between the flatness of the local minimum and its generalization ability according to probably approximately correct (pac) bayesian generalization bounds @xcite . in this work, we propose a novel theoretical framework to analyze this relation from a more intuitive and direct perspective. in the distributionally robust optimization (dro) @xcite field, the elementary assumption is that there exists a shift between the distributions of the training set and the test set. we propose that a small distribution shift can be equivalently seen as a virtual parameter corruption or perturbation on the loss function. we conduct analytic trials to verify our theoretical account and the results show that it fits the simulation well and can therefore explain why flat minima that are robust against parameter corruptions or perturbations have better generalization performances. we also analyze the strength of the parameter corruption within this framework, based on which we propose a gradient-strength based adaptive sharpness-aware minimization (ga-sam) algorithm, which can set flexible strengths of parameter corruptions for different parameter groups, during different training stages. to validate the effectiveness of the proposed ga-sam algorithm, we choose several natural language models and benchmarks, including convolution neural networks (cnn) @xcite ) on text classification, long short-term memory (lstm) @xcite networks on language modeling, and transformer @xcite on neural machine translation. we also compare our proposed ga-sam algorithm with the traditional sam algorithm @xcite and its multiple variants, including multi-step adversarial parameter defense algorithm @xcite , adaptive sam @xcite , layerwise sam @xcite and other possible variants of our proposed algorithm. experimental results show that our proposed ga-sam gains better generalization compared to the traditional sam algorithm and other variants. our contributions can be summarized as follows: ## proposed theoretical framework in this section, we propose a novel theoretical framework to reveal the relation between distribution shifts and parameter corruptions from an intuitive and direct theoretical perspective. ## experiments in this section, we report the tasks, datasets, and implementation details. main results are in update w as minimizing eq. ( 14 ). 12: end while 4.1 tasks and datasets we adopt three typical neural networks and natural language tasks to validate the effectiveness of the proposed ga-sam algorithm on nlp tasks. on the text classification task, we adopt convolution neural networks (cnn) @xcite on the imdb movie reviews dataset (imdb) @xcite with the accuracy (acc) evaluation metric. on the language modeling task, we adopt long short-term memory (lstm) @xcite networks on the english penn treebank (ptb-lm) @xcite dataset with the perplexity (ppl) evaluation metric. on the neural machine translation task, we adopt the transformer @xcite model based on the fairseq implementation @xcite on iwslt 15 english-vietnamese (en-vi) @xcite and iwslt 14 german-english (de-en) @xcite datasets with the bleu score evaluation metric. compared with the classification and language modeling tasks, the datasets of the machine translation task are relatively large. other details are in appendix. ## analysis in this section, we first conduct an ablation study and analyze the hyper-parameters. then we illustrate the sharpness and hessian spectra with ga-sam and analyze the difference between cv and nlp learning to explain why nlp tasks need ga-sam. ## limitation and broader impact limitation. one limitation of our work is that, similar to other sam learning, the hyper-parameters tuning, especially ϵ, involves many numerical experiments, which is time costly and environmentally unfriendly. to settle this issue, we recommend researchers binary search the proper order of magnitude of ϵ first, and then grid search ϵ in a small range for a faster hyper-parameter search, instead of directly grid searching ϵ in a large range. broader impact. our work is beneficial for the security of nlp models since our work can help improve the robustness of nlp models against parameter corruptions, which can occur as random noises at the hardware level, quantization, or model compression. our work also has negative social impacts. our proposed ga-sam can be utilized to enhance nlp models and improve the accuracy of base models. however, the hyper-parameters tuning, especially ϵ, involves many numerical experiments, which is also a limitation of our work, and it is environmentally unfriendly. ## conclusion in this paper, we propose a novel theoretical framework to analyze the relation between parameter corruptions and generalization abilities. based on our proposed framework, we propose a gradient-strength based adaptive sharpness-aware minimization (ga-sam) algorithm. experimental results validate the effectiveness of ga-sam compared to the traditional sam algorithm and its variants. further analyses also show that ga-sam can help find flat minima and improve the generalization ability of neural networks. yaowei zheng, richong zhang, and yongyi mao. 2020. regularizing neural networks via adversarial model perturbation. corr, abs/2010.04925. chen zhu, yu cheng, zhe gan, siqi sun, tom goldstein, and jingjing liu. 2019. freelb: enhanced adversarial training for language understanding. corr, abs/1909.11764.