Sharpness-Aware Minimization with Dynamic Reweighting (2022)

Abstract: Deep neural networks are often overparameterized and may not easily achieve model generalization. Adversarial training has shown effectiveness in improving generalization by regularizing the change of loss on top of adversarially chosen perturbations. The recently proposed sharpness-aware minimization (SAM) algorithm conducts adversarial weight perturbation, encouraging the model to converge to a flat minima. SAM finds a common adversarial weight perturbation per-batch. Although per-instance adversarial weight perturbations are stronger adversaries and can potentially lead to better generalization performance, their computational cost is very high and thus it is impossible to use per-instance perturbations efficiently in SAM. In this paper, we tackle this efficiency bottleneck and propose sharpness-aware minimization with dynamic reweighting (delta-SAM). Our theoretical analysis motivates that it is possible to approach the stronger, per-instance adversarial weight perturbations using reweighted per-batch weight perturbations. delta-SAM dynamically reweights perturbation within each batch according to the theoretically principled weighting factors, serving as a good approximation to per-instance perturbation. Experiments on various natural language understanding tasks demonstrate the effectiveness of delta-SAM.

## introduction although deep neural networks (dnns) have demonstrated promising results in various fields such as natural language understanding @xcite and computer vision @xcite , they are often overparameterized and can easily overfit the training data @xcite . adversarial training has been proven effective in improving both model generalization @xcite and adversarial robustness @xcite . a general approach for adversarial training has been first to augment the inputs with small perturbations that lead to the maximum possible change of loss, and then optimize the model parameters to the direction where the changed amount is minimized. besides perturbing inputs, a recent work of sharpness-aware minimization (sam; @xcite ) has further proposed to adversarially perturb model weights. such a method works by first adversarially calculating a weight perturbation that maximizes the empirical risk and then minimizing the empirical risk on the perturbed model. this method demonstrates improved model generalizations across different datasets and models. in principle, each instance in a batch has its own worst-case weight perturbation and the weight perturbations of different instances need to be calculated separately and cannot be done in a single forward/backward pass. this leads to a significant increase in computational and memory cost. to allow a feasible algorithm, sam approximates per-instance perturbations by a single per-batch perturbation, where the weight perturbation is calculated on the averaged loss of the batch and shared by all instances in the batch. however, as the per-batch perturbation represents the average of perturbations yielded by different instances, it is a weaker adversary compared to per-instance perturbations, and may hinder the effectiveness of sam. in this paper, we study how to efficiently approximate per-instance weight perturbation for sharpness-aware minimization, while maintaining a similar computational cost to per-batch perturbation. we first theoretically analyze the gradient posed by the optimization of per-instance perturbation, and find that it can be effectively approximated with a weighted-batch perturbation under some assumptions, where the instances with a larger rate of gradient change are up-weighted. based on this motivation, we propose sharpnessaware minimization with dynamic reweighting (δ-sam). specifically, we first estimate the hessian and gradient norm of each instance by perturbing the loss with a random gaussian noise on model weights. next, δ-sam dynamically reweights the loss within each batch of training instances, and then calculate a shared weight perturbation that maximizes the reweighted batch loss. finally, we update the perturbed model on the original (unweighted) batch. compared to sam, δ-sam only requires extra computation cost in estimation of the rate of gradient change, which can be efficiently performed using three additional forward passes. we evaluate δ-sam on finetuning pretrained language models (plms). experiments on standard glue benchmark @xcite , selfsupervised semantic textual similarity (sts), and abstractive summarization tasks show that besides significantly outperforming base models, δ-sam also consistently outperforms sam with only 18% extra computational cost in average. the main contributions of this paper are threefold. first, we analyze the training objective of per-instance weight perturbation and find that under some assumptions, it can be approximated by a weighted-batch perturbation, where instances are efficiently reweighted according to their caused rates of gradient changes. second, we propose to use random perturbations as estimations to efficiently realize the weighting scheme. third, we evaluate δ-sam on a diverse set of datasets and find consistent improvements across the board. ## related work model generalization. deep neural networks are often overparameterized and may suffer from poor generalization @xcite . a lot of efforts have been devoted to improve the generalization of neural models, leading to methods including data augmentation @xcite @xcite @xcite , regularization @xcite @xcite , and improved optimization processes @xcite @xcite . these methods consider different aspects of generalization and may be combined to achieve better performance. among them, adversarial training @xcite has demonstrated its effectiveness in improving model generalization without the need of any extra data or external knowledge, and has been widely attempted to enhance nlp models @xcite @xcite . adversarial training works by adversarially perturbing the input embedding and either minimize the adversarial risk or regularizes the change of risk to be small. specifically, freelb @xcite uses projected gradient descent (pgd; @xcite to generate adversarial perturbations on input embedding, and recycles the computed gradients when updating model parameters in adjacent steps @xcite to reduce the computational costs. tat @xcite improves freelb by prioritizing the most frequently mispredicted classes in perturbation calculation. tavat @xcite uses a token-level accumulated perturbation vocabulary to guide the initialization in pgd. however, these works only consider the robustness on input feature representation, while we consider the robustness of all model weights. sharpness-aware minimization. @xcite leverage the correlation between flat minima and better model generalization and propose sam for training deep neural models that are robust to adversarial weight perturbations. it has demonstrated effectiveness in tasks on both vision @xcite and language @xcite modalities. several variants of sam have been proposed to improve its efficiency or effectiveness. for efficiency, @xcite propose to speed up sam by perturbing fewer instances in the batch; @xcite introduce stochastic weight perturbation and sharpness-sensitive data selection to reduce the computational overhead. for effectiveness; @xcite propose to adaptively set the sam's radius such that it is invariant to parameters' scales; zhuang et al. (2022) introduce a gradient ascent optimization step in the perturbed model's orthogonal direction to achieve better flat-algorithm 1: sam and δ-sam input: model f w , training set s ≜ {(x i , y i )} |s| @xmath0, loss function l : w × x × @xmath1, batch size n , neighborhood size ρ ∈ r + , optimizer h. output: a flat solution ŵ. initialize model weights w. while not converge do sample a batch @xmath2, y j )} n @xmath3. for δ-sam: reweigh b by eq. 5. compute gradient ∇l b (w) of the (reweighted) batch's empirical risk. perturb the model weights by ϵ * = ρ∇l b (w)/ ∥∇l b (w)∥ 2 . update w w.r.t. the unweighted empirical risk 1 n n @xmath4. ## sharpness-aware minimization (sam) in this section, we briefly review the principle of sam and discuss its limitations. literature has observed a direct correlation between flat minima and better model generalization, both empirically and theoretically @xcite @xcite . to find a flat loss landscape, sam @xcite ) adversarially perturbs the model weights and optimizes the following minmax objective on a batch of size n : the outer minimization can be performed with a standalone optimizer (e.g., adam; kingma and ba 2015). sam roughly doubles the computational cost of training the model, requiring two forward and two backward passes for each batch. the sam algorithm is outlined in alg. 1. besides perturbing by batches, weight perturbation can also be performed on individual instances: ## sam with dynamic reweighting in this paper, we seek to improve sam with a better adversary on weight perturbations. as the perbatch weight perturbation adopted by sam weakens the adversarial training, we propose a simple yet effective modification of sam, δ-sam (sam with dynamic reweighting), that can approximate per-instance weight perturbation without requiring much additional computational cost. our reweighting approach is motivated by a theoretical analysis on approximating the per-instance weight perturbation to justify its superior efficiency. based on this motivation, we then illustrate how δ-sam is realized in implementation. ## conclusion this paper presents a new sharpness-aware minimization method with dynamic reweighting (δ-sam). the proposed method represents the first successful attempt in realizing a per-instance weighting scheme. we achieve this by prioritizing instances with larger gradient change rate in adversarial weight perturbation, in comparison to previous approaches that adopt per-batch weight perturbation. we show that perturbation calculated on reweighted batch can serve as a better approximation to per-instance weight perturbation while requiring only similar computational cost to perbatch perturbation. we conduct extensive experiments on the glue, sts, and abstractive summarization benchmarks. across all 30 experimental setups that compares to sam, δ-sam achieves an consistent improvement over sam in 27 of them. when compared to a set of other competitive regularization methods, δ-sam achieves the best performance in 23 out of 33 of the setups. further, we quantitatively analyze δ-sam's impact on sharpness, finding that it indeed leads to flatter loss landscape. future work includes inventing new techniques to further reduce the computational cost of δ-sam and demonstrating its effectiveness on more tasks such as sequence tagging and question answering.