Improving Sharpness-Aware Minimization with Fisher Mask for Better Generalization on Language Models (2022)

Abstract: Fine-tuning large pretrained language models on a limited training corpus usually suffers from poor generalization. Prior works show that the recently-proposed sharpness-aware minimization (SAM) optimization method can improve the model generalization. However, SAM adds a perturbation to each model parameter equally (but not all parameters contribute equally to the optimization of training), which we argue is sub-optimal and will lead to excessive computation. In this paper, we propose a novel optimization procedure, namely FSAM, which introduces a Fisher mask to improve the efficiency and performance of SAM. In short, instead of adding perturbation to all parameters, FSAM uses the Fisher information to identity the important parameters and formulates a Fisher mask to obtain the sparse perturbation, i.e., making the optimizer focus on these important parameters. Experiments on various tasks in GLUE and SuperGLUE benchmarks show that FSAM consistently outperforms the vanilla SAM by 0.67 1.98 average score among four different pretrained models. We also empirically show that FSAM works well in other complex scenarios, e.g., fine-tuning on generation tasks or limited training data. Encouragingly, when training data is limited, FSAM improves the SAM by a large margin, i.e., up to 15.1.

## introduction the "pretraining-finetuning" paradigm has become the de facto standard for the community of natural language processing (nlp) @xcite @xcite @xcite . given a pretrained language model (plm), the dominant fine-tuning manner is tuning the entire pretrained parameters for each downstream task @xcite . while fine-tuning the entire plm can improve performance on a wide range of nlp tasks, it usually suffers from overfitting and poorer generalization ability @xcite , especially in the largescale plms and limited training data scenarios. hence, some existing efforts attempt to provide more regularization in the fine-tuning stage @xcite @xcite , among which the optimization of the training loss is an intuitive and effective method. specifically, motivated by the finding @xcite that the smoother loss landscape refers to the better model generalization, @xcite propose the "sharpness-aware minimization" (sam) to simultaneously minimize loss value and loss sharpness, where the sharpness can be quantified as the maximized difference of loss when a perturbation is added to the current weights. in practice, sam performs two forwardbackward computations for each optimization step, where the first forward-backward is to obtain the perturbation for each model parameter and the second one is to update the parameters. many prior works @xcite show the effectiveness of sam in the vision domain, motivated by this, @xcite first apply the sam to the language domain, more recently. although @xcite empirically show the remarkable performance of sam on several language understanding tasks, sam calculates perturbations indiscriminately for all parameters, which is time-consuming and hinders the application of sam. furthermore, inspired by the finding @xcite ) that only about 5% of parameters are sharp and rise steeply during optimization, we notice that not all parameters contribute equally to the optimization of training. hence, this raises a question that whether we can calculate perturbations for only some individual parameters, and thus make the optimizer focus on these important parameters. to this end, we propose a novel optimization approach, fisher sam (fsam), which introduces a fisher mask to improve the efficiency and effectiveness of sam. in short, fsam first uses the fisher information @xcite as the metric to identify the sharper parameters foot_0 and formulates a binary fisher mask correspondingly. then, the fisher mask is multiplied with the perturbations to obtain the sparse perturbations, which are lastly used to perform regularization in the parameter update. in this way, only parts of sharper parameters will be added into the perturbations, and the optimizer can thus focus more on these important parameters. also, the sparse perturbations could ensure the training acceleration via sparse backpropagation foot_1 . moreover, one may concern that the sparse fisher mask would affect the convergence rate of fsam @xcite . hence, we theoretically provide the convergence analysis of fsam, ensuring that the convergence of fsam is irrelevant to the fisher mask. we conduct a large-scale and systematic study to evaluate the performance and effectiveness of fsam. firstly, we apply sam and fsam to finetune various plms on parts of glue and super-glue benchmarks, where the results show that fsam consistently outperforms the vanilla sam by 0.67∼1.98 average score among these plms, and surpasses the adam (kingma and ba, 2015) optimizer by 1.41∼1.91 points. secondly, we conduct experiments on two popular generation tasks (i.e., xsum and conll2014) and prove that fsam can deliver promising results against sam. lastly, quantitative analysis and in-depth discussion demonstrate the universality and effectiveness of fsam in various complex scenarios, and prove that fsam indeed brings better model generalization. specifically, we show that our fisher mask strategy not only works well in the sam, but also can be applied to other sam variants. to summarize, our contributions are two-fold: (1) we propose a novel optimization approach (namely fsam) with theoretical convergence guar-antee for plms. specifically, fsam improves the performance and efficiency of recently-proposed sam via a fisher mask strategy, which can also be applied to more sam variants. (2) extensive experiments show that fsam consistently outperforms the sam by a large margin on both language understanding and generation tasks. the systematic study demonstrates the effectiveness and universality of fsam on improving model generalization. improving generalization. recently, we have witnessed numerous plms that achieved tremendous success in the community of nlp @xcite @xcite @xcite @xcite @xcite . the current dominant fine-tuning approach needs to tune all pretrained parameters for each downstream task, which makes the plm easily memorize the training data and thus leads to overfitting. to tackle this issue, some works attempt to provide implicit and explicit regularization into the training of models, such as dropout @xcite , label smoothing @xcite , mixup @xcite and other data-augmentation methods @xcite @xcite @xcite . on the other hand, motivated by the successful applications of sam in the vision domain, @xcite involve applying sam to optimize the t5 @xcite model on multiple language tasks and show that sam can improve the generalization of plms effectively. we depart from the prior work @xcite and ours as follows: 1) different motivations: instead of verifying the effect of vanilla sam on several language understanding tasks, we aim to improve the efficiency and effectiveness of sam. ## methodology in this section, we first review the sharpness-aware minimization, and then propose our sharpness-aware minimization with fisher mask, coined as fsam. finally, we theoretically analyze the convergence of fsam with adaptive learning rate. ## main results fsam outperforms vanilla sam by a large margin across different plms. ## analysis and discussion in this section, we examine whether our approach works in more complicated scenarios, and provide a more intuitive comparison between different optimizers towards the generalization. more analysis and results can be found in appendix. ## conclusion in this paper, we improve the recently-proposed sam optimization method with a novel fisher mask strategy, and propose a new approach fsam. different from the vanilla sam that adds a constant perturbation to all parameters, fsam uses the fisher information to calculate the fisher mask and further obtains the sparse perturbation. such a method can not only reduce the computation cost of optimization potentially, but also make the optimizer focus on the optimization of the important sharper parameters. extensive experiments on five plms and various language understanding and generation tasks show that our fsam consistently improves the performance of sam by a large margin across all plms and tasks. additionally, in-depth analysis and discussion demonstrate the robustness and universality of fsam on improving the generalization of language models.