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

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 3888 - 3903 December 7-11, 2022 ©2022 Association for Computational Linguistics GA-SAM: Gradient-Strength based Adaptive Sharpness-Aware Minimization for Improved Generalization Zhiyuan Zhang1, Ruixuan Luo2, Qi Su3, 1, Xu Sun1 1MOE Key Laboratory of Computational Linguistics, School of Computer Science, Peking University 2Center for Data Science, Peking University 3School of Foreign Languages, Peking University {zzy1210,luoruixuan97,sukia,xusun}@pku.edu.cn Abstract Recently, Sharpness-Aware Minimization (SAM) algorithm has shown state-of-the-art generalization abilities in vision tasks. demonstrates that flat minima tend to implyIt 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. 1 Introduction Recently, researchers (Wu et al., 2020; Weng et al., 2020; Sun et al., 2021; Zhang et al., 2021; Foret et al., 2020; Liu et al., 2021) propose that for better generalization ability, learning algorithms should find flat minima that have better robust- ness resistant to parameter corruptions or pertur- bations. Many learning algorithms that take the flatness or sharpness of the parameters into con- sideration are motivated by the observation that flat minima tend to imply better generalization abilities. Among them, Sharpness-Aware Mini- mization (SAM) (Foret et al., 2020) algorithm has achieved state-of-the-art generalization abilities in vision tasks. It adopts virtual adversarial parame- ter corruptions or perturbations during training and lowers the risk after parameter corruptions. How- ever, traditional SAM algorithms usually adopt fixed strengths of parameter corruptions and con- straint the corruptions with L2-norm or L+∞-norm balls. It cannot conduct flexible strengths of param- eter corruptions for different parameters, or during different stages of training. Thus, it is difficult to apply SAM to some natural language tasks, espe- cially to models with drastic gradient changes, such as RNNs. To settle this issue, many adaptive SAM algorithms (Kwon et al., 2021; Liu et al., 2021) are proposed empirically. In this work, we propose a gradient-strength based adaptive solution based on our theoretical framework. Existing studies (Wu et al., 2020; Zhang et al., 2021) try to explain the relation between the flatness of the local minimum and its general- ization ability according to Probably Approxi- mately Correct (PAC) Bayesian generalization bounds (Neyshabur et al., 2017). In this work, we propose a novel theoretical framework to ana- lyze this relation from a more intuitive and direct perspective. In the Distributionally Robust …

Content truncated to 500 words