{'fraction_non_alphanumeric': 0.04549420013570713, 'fraction_numerical': 0.017674238262948724, 'mean_word_length': 4.513985391056476, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
GA-SAM: Gradient-Strength based Adaptive Sharpness-Aware Minimization for Improved Generalization

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 (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 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) (Foret et al., 2020) 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 (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 generalization ability according to Probably Approximately Correct (PAC) Bayesian generalization bounds (Neyshabur et al., 2017). 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) (Rahimian and Mehrotra, 2019) 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. arXiv:2210.06895v1 [cs.LG] 13 Oct 2022 To validate the effectiveness of the proposed GA-SAM algorithm, we choose several natural language models and benchmarks, including Convolution Neural Networks (CNN) (Kim, 2014) on text classification, Long Short-term Memory (LSTM) (Merity et al., 2017) networks on language modeling, and Transformer (Vaswani et al., 2017) on neural machine translation. We also compare our proposed GA-SAM algorithm with the traditional SAM algorithm (Foret et al., 2020) and its multiple variants, including multi-step adversarial parameter defense algorithm (Zhang et al., 2021), adaptive SAM (Kwon et al., 2021), layerwise SAM (Liu et al., 2021) 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: • We propose a novel theoretical framework to analyze the relation between the flatness of the local minimum and its generalization ability. Under our proposed theoretical framework, the shift of the training and test distributions can be equivalently seen as a virtual parameter corruption or perturbation. Thus, the flatness or the robustness against parameter corruptions can indicate the generalization ability.
• On the basis of our novel framework, we further propose a Gradient-Strength based Adaptive Sharpness-Aware Minimization (GA-SAM) algorithm to set flexible strengths of parameter corruptions for different parameter groups, during different stages of training for an improvement over generalization ability.
• Experimental results show the effectiveness of the GA-SAM algorithm compared to the traditional SAM algorithm and its variants.

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.

Preliminary
Let us consider a neural network with the parameter vector w ∈ R n . For a data instance z = (x, y), denote (w; z) as the loss of the data instance, L(w; D) as the average loss of a dataset D, and p(z) as the probability distribution of D, we have: Denote θ as the optimal parameter: and the Hessian matrix on θ is H = ∇ 2 θ L(θ; D). Similarly, denote D * and p * (z) as the test set and its distribution, θ * as its optimal parameter, and the Hessian matrix on θ * is H * . Define the parameter shift of the test and training minima as Suppose n parameters are divided into l groups and the i-th group has n (i) parameters (e.g., l = 1, n = n (1) when the whole model adopt the same strength and we call it model-wise, n = l, n (i) = 1 when element-wise, l is the layer number when layer-wise, l is the filter number when filter-wise, etc.

The Distribution Shift between the Training and Test Sets
The elementary assumption in the Distributionally Robust Optimization (DRO) (Rahimian and Mehrotra, 2019) field is that there exists a small distributional shift between the training and test sets. Previous studies usually assume that the divergence or the distance of the training and test distributions is bounded by a constant, e.g., the Kullback-Leibler divergence (Kullback and Leibler, 1951), KL(p(z)||p * (z)) ≤ Constant. In this work, more generally, we assume that the f -divergence (Rényi, 1961) D f of the distributions is bounded by C f : where the function f is convex and f (1) = 0, its Taylor expansion should also satisfy f (1 +

Parameter Corruptions as Results of Distribution Shifts
We propose a novel theoretical framework to analyze this relation from a more intuitive and direct perspective. The main theoretical motivation is Theorem 1. Proofs and details are in Appendix. Theorem 1. The distribution shifts of datasets D and D * can be equivalently treated as a parameter corruption near the corresponding minima, where Constant=L(θ * ; D * ) − L(θ; D). Let a = −δ, when w is near θ and θ * , we have It shows that the distribution shifts will cause a parameter corruption or parameter shift. Therefore, optimizing the parameter corruption risk L(w + a; D) can help optimize the loss on the test set L(w; D * ). Define S as the possible corruption constraint set of potential corruptions a. Since a = −δ is determined by the invisible distribution shifts, we optimize the risk under potential corruptions a instead, which is exactly the SAM optimization, Thus, we reveal why flat minima that are robust against potential parameter corruptions or perturbations have better generalization performances in our theoretical framework. Traditional SAM algorithms adopt fixed strengths of parameter corruptions and constraint the corruptions with L 2 -norm or L +∞ -norm balls, namely S = {a : a 2 ≤ } or S = {a : a +∞ ≤ }. However, in Theorem 2, it reveals that the potential parameter corruption δ is determined by the distribution shifts and the local geometry near the local minimum in the loss basin. Based on this, we have Proposition 1. Theorem 2. Define r(z) = p * (z)/p(z) − 1. When the distribution shift is small enough, namely r(z) is small, we can estimate the parameter shift δ as, Proposition 1. Suppose the loss L(w; D) is µstrongly convex 1 , and D f (p * ||p) ≤ C f , there exists such that δ 2 ≤ C δ , namely C δ is a upper bound.
1 Note that L is only required to be µ-strongly convex in the neighborhood of the loss basin including θ and θ * , instead of the entire R n .

Verification of the Theoretical Framework
In this section, we conduct several analytic trials 2 to verify our theoretical explanations and analyze the corruption strength. Results show that our theoretical framework fits the simulation. Visualization of Distribution Shift Effects. Fig. 1 visualizes the training loss, test loss and the shifted training loss. The shifted training loss is similar to the test loss near the local minimum, namely we have L(w; D * ) ≈ L(w − δ; D)+Constant, which validates Theorem 1. This phenomenon can also be observed in visualizations of training and test loss landscapes in other studies.
Relation between Corruption Strength and Distribution Shift Strength. We conduct analytic trials to reveal the relation between the corruption strength and the distribution shift to verify our theoretical framework. Suppose D is the training set and D * is the test set. We can construct a mixed dataset D mix , mixed with (1 − η) of the training data from D and η of the test data from D * . We have p mix = (1 − η)p + ηp * approximately. Define η as the relatively distribution shift strength between D mix and D. Proposition 2 reveals that the corruption strength is proportional to the distribution shift strength, which fits both the simulation results of analytic trails in Fig. 2

and the intuition.
Proposition 2. Suppose the mixed distribution of D mix is p mix = (1 − η)p + ηp * , then we have D f (p mix ||p) ≤ C mix f = η 2 C f . Denote θ mix as the optimal parameter on D mix , δ mix = θ mix − θ, then we have: 3 Gradient-Strength based Adaptive Sharpness-Aware Minimization In this section, we propose a Gradient-Strength based Adaptive Sharpness-Aware Minimization (GA-SAM) algorithm based on the proposed theoretical framework, which conducts flexible strengths of parameter corruptions for different parameter groups for better generalization abilities.

Adaptive Sharpness-Aware Minimization
As illustrated in Section 2.3, the SAM (Foret et al., 2020) optimization objective is exactly the risk under potential corruption as a result of distribution shift, where the constraint S is S = {a : a p ≤ }. Adaptive SAM algorithms (Kwon et al., 2021;Liu et al., 2021) are proposed to set flexible strengths of parameter corruptions, which set the constraint S as S = {a : T −1 a p ≤ }, where T is the transformation matrix (usually diagonal) controlling the corruption strengths of corresponding parameter groups. Define T = diag{T (1) I n (1) , · · · , T (l) I n (l) }, where T (i) controls the corruption strengths of group i. Adaptive SAM (ASAM) (Kwon et al., 2021) empirically adopts T (i) = w (i) 2 element-wisely or filterwisely in CNNs, and layer-wise SAM (Layer-SAM) (Liu et al., 2021) empirically adopts T (i) = w (i) 2 / g (i) 2 layer-wisely.

The Relation between the Corruption Strength and the Gradient Strength
We hope to derive T from the theoretical framework instead of setting T empirically. In Theorem 2, besides the term r(z) determined by the distribution shift, corruption strength is also determined by the local geometry near the local minimum (H, ∇ θ (θ; z)). Suppose n is the parameter number, Suppose the Fisher information matrix assumption (Pascanu and Bengio, 2013) , the local geometry is determined by the gradient strength G.
In Proposition 3, we analyze the relation between the corruption strength and the gradient strength. We have δ 2 ∝ √ n/G and δ (i) 2 ∝ √ n (i) /G (i) , corruption strengths and scales T should be determined by gradient strengths G.
Proposition 3. Define the average gradient strength as G = E p(z) [ ∇ θ (θ; z) 2 ]/ √ n, and the average gradient strength of group i as G

Proposed Algorithm
Our proposed algorithm adopts the constraint S = {a : T −1 a p ≤ } in SAM learning, where T are adaptive scales derived from the theoretical framework, and we adopt the multi-step implementation.
and we can set , then our proposal is derived, Existing algorithms are mainly single-step based, we adopt the multi-step implementation inspired by multi-step adversarial parameter defense algorithm (Zhang et al., 2021), which optimizes where a 0 = 0. Suppose the k-th update of the corruption is u k = arg max T −1 u p≤η u T g k−1 , which is generated based on a k−1 and g k−1 = ∇ w L(w + a k−1 ; B), where η is the step size and following Zhang et al. (2021), we set η = 1.5 /K, then To get the k-th corruption a k , we project the updated corruption a k−1 + u k into the set S, and the solutions to the commonly adopted L 2norm and L +∞ -norm constraints are: To summarize, we adopt the adaptive scale T (i) according to the gradient strength and a multi-step implementation. It should be noted that even when K = 1, our multi-step implementation, which optimizes (L(w)+L(w +a 1 ))/2, is different from the single-step SAM implementation, which optimizes L(w + a 1 ). However, when K = 1, they have similar speeds since they both require to generate a 1 and need two backward propagation processes. Our multi-step implementation, however, allows setting a larger K for better generalization.
We name the proposed algorithm as the Gradient-Strength based Adaptive SAM (GA-SAM). The algorithm is shown in Algorithm 1. The proofs and theoretical details are in Appendix.

Experiments
In this section, we report the tasks, datasets, and implementation details. Main results are in Table 1.

Algorithm 1 GA-SAM Algorithm
Require: Parameters w; loss L and dataset D; steps K; training iterations; batch size |B|. 1: Prepare batches {B} and initialize w.

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 (

Implementations
We implement our proposed GA-SAM algorithm with a multi-step risk minimization, and set layerwise adaptive scales The multi-step adversarial parameter defense algorithm (Zhang et al., 2021) adopts a multi-step risk implementation and sets fixed scales. We also try to combine these techniques and implement other possible variants of GA-SAM, and we conduct an ablation study to compare GA-SAM with these variants. For a fair comparison, the settings of different SAM algorithms and variants are the same as the base models except for the SAM settings (K, , L p , T (i) ). We grid search the optimal SAM hyper-parameters for each algorithm. The details of base models and hyper-parameters are in Appendix.

Main Results
The main results are shown in Table 1. Evidently, our proposed GA-SAM leads to significant performance gains over the base models.
Single-step SAM algorithms cannot improve the performance of the LSTM base model on the language modeling task. SAM algorithms that set adaptive scales cannot improve the performance of the traditional SAM algorithm consistently. But it may help deal with drastic changes in gradient scales of different parameters or different learning phases in NLP tasks. The multi-step risk minimization generally outperforms SAM and helps improve the stability of learning and the generalization abilities of models. GA-SAM can both achieve the adaptive scales deduced from the theoretical framework and help improve the stability and the generalization via the multi-step implementation. Thus, GA-SAM outperforms base models  and other SAM algorithms.

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.

Ablation Study
We implement the single-step variants, elementwise or model-wise variants, and variants with other scales T (i) on IMDB and PTB-LM. The results of the ablation study are reported in Table 2.  Experimental results show that layer-wise implementation outperforms element-wise or modelwise implementations and GA-SAM with the multistep implementation (K = 1) outperforms GA-SAM with the single-step implementation.
For adaptive scales, the multi-step adversarial parameter defense algorithm (Zhang et al., 2021) adopts T (i) = 1. We also adopt T (i) = w (i) 2 / g (i) 2 following Liu et al. (2021), and T (i) = w (i) 2 following Kwon et al. (2021). We also try other variants with similar formulas. Experimental results show that GA-SAM outperforms other variants. Gradient-strength ( g (i) 2 ) based adaptive scales usually outperform other adaptive scales, and adaptive scales can enhance the multistep adversarial parameter defense algorithm generally, which validates our theoretical framework.

Analysis of Hyper-parameters
We analyze the influence of hyper-parameters and settings in GA-SAM learning in Table 3.
On IMDB, adopting larger K cannot improve the accuracy, while on PTB-LM, larger K can achieve better PPL. The reason might be that on PTB-LM, the LSTM model with more drastic gradient changes needs more steps for better stability.   Both under L 2 and L +∞ constraints, on both datasets, the performance can often be improved substantially with small . However, when grows too large, it may harm the learning and the performance will drop. We also find that in this work, the best performance is achieved under L +∞ .

Sharpness and Parameter Robustness
As analyzed in our theoretical framework, flat minima tend to imply better generalization abilities. In this section, we validate that GA-SAM can help find flat minima, which tends to help improve the generalization abilities of models.
The sharpness near the local minima can be evaluated by the parameter robustness against parameter corruptions via the multi-step adversarial parameter corruption algorithm (Zhang et al., 2021). In Table 4, we evaluate the robustness of the baselines and models with GA-SAM against the L 2 or L +∞ constrained parameter corruptions on En-Vi. We can see that models with GA-SAM are more robust than baselines, which implies that GA-SAM can help find flat minima. We also adopt the Fisher  information matrix assumption (Pascanu and Bengio, 2013) to estimate the top 50 eigenvalues of the Hessian matrix to evaluate the sharpness of minima. As shown in Fig. 3, the eigenvalues of models with GA-SAM are lower, which illustrates that GA-SAM can help find flat minima.

Comparisons to Adversarial Training.
Some adversarial training algorithms improve the generalization ability of neural networks by optimizing the loss of virtual adversarial examples. In this section, we compare GA-SAM with FreeLB (Zhu et al., 2019), an existing highperformance algorithm for NLP tasks, on En-Vi and De-En. Detailed settings are in Appendix.
In Table 5, we can see that FreeLB (Zhu et al., 2019) can improve the accuracy of NLP models. The reasons that FreeLB (Zhu et al., 2019) works may lie in two aspects: (1) FreeLB adopts a multistep minimization that is helpful for training stability; and (2) FreeLB only involves attacks on word embeddings while the ideal attack strengths for parameters in different layers vary a lot due to gradient vanishing and explosion in NLP models. Therefore, FreeLB does not need flexible scales as necessarily as SAM. However, GA-SAM can still outperform FreeLB.

Why NLP Tasks Need GA-SAM
Compared with the traditional SAM algorithm, GA-SAM adopts the multi-step risk minimization and gradient-strength based adaptive corruption strengths. From the ablation study, we can see that the multi-step risk minimization can enhance the traditional SAM algorithm on NLP tasks. It shows that NLP tasks do need GA-SAM for better stability and generalization with the multi-step risk minimization algorithm. The gradient-strength based adaptive corruption strengths can also enhance SAM algorithms since gradient strengths change drastically during different learning phases and gradient strengths vary in different layers.  In this section, we further visualize the gradient strengths during different learning phases and different layers on En-Vi to illustrate why NLP tasks need GA-SAM. As shown in Fig. 4, during the early phase of learning, the gradient norms are much larger and we should conduct parameter corruptions with smaller strengths. As shown in Fig. 5, we can see that higher layers of the Transformer encoder or decoder have larger gradients and need smaller parameter corruptions. To conclude, we need gradient-strength based adaptive corruption strengths since gradient strengths vary in different learning phases and different layers, while the traditional SAM algorithm fails to conduct flexible strengths of parameter corruptions for different parameters, or during different stages of training.

Computational Complexity
Single-step implementations usually involve two forward and backward propagations, including one for generating the corruption a and another for optimizing the L(w + a) loss function or the L(w) + λa T ∇ w L(w) loss function). Multi-step implementations involve K + 1 forward and backward propagations, including K + 1 times for forwarding and backwarding {L(w + a k )} 0≤k≤K . Therefore, when K = 1, the neural network forward and backward propagation cost (involving two forward and backward propagations) of GA-SAM is the same as SAM or GA-SAM with the single-step implementation. In Table 2, we conduct a fair comparison between GA-SAM (K = 1) and SAM or GA-SAM with the single-step implementation (all involve two forward and backward propagations). GA-SAM still outperforms SAM and GA-SAM with the single-step implementation, which illustrates the effectiveness of our proposed GA-SAM and the multi-step implementation. Existing studies research parameter corruptions or perturbations mainly for better generalization ability (Zheng et al., 2020;Foret et al., 2020;Kwon et al., 2021;Liu et al., 2021;Du et al., 2021;Sun et al., 2021;Zhang et al., 2021), safety issue (Garg et al., 2020;Rakin et al., 2020)

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.
= ρC Consider the gradients, Consider the function f , Therefore, the change in the term f (v) in the loss function can be omitted compared to the change in the loss function, Let a = −δ, w = θ * + v, we have, A.2 Proofs of Theorem 2 Theorem 2. Define r(z) = p * (z)/p(z) − 1. When the distribution shift is small enough, namely r(z) is small, we can estimate the parameter shift δ as, Proof. With the change-of-measure technique, we have According to the definition, Conduct Taylor Expansion, we have, Solve it, we have, where the term r(z)∇ 2 θ (θ; z)δ = o( δ 2 ) since the distribution shift r(z) = o(1). To conclude,

A.3 Proofs of Propositions
Proposition 1. Suppose the loss L(w; D) is µstrongly convex 3 , and D f (p * ||p) ≤ C f , there exists such that δ 2 ≤ C δ , namely C δ is a upper bound.
3 Note that L is only required to be µ-strongly convex in the neighborhood of the loss basin including θ and θ * , instead of the entire R n .
Proof. With the change-of-measure technique, we have: Then Therefore, According to the upper bound of E p(z) |r| 2 , There exists Denote θ mix as the optimal parameter on D mix , δ mix = θ mix − θ, then we have: Proof.
To get the k-th corruption a k , we project the updated corruption a k−1 +u k into the set S. Zhang et al. (2021) give the projecting functions, similarly, when S = {u : Therefore,

B Datasets and Baseline Implementations
In this section, we introduce the datasets and baseline implementations. The models enhanced with SAM algorithms adopt the same hyper-parameters to baselines except for the SAM hyper-parameter settings (K, T (i) , , L p ). All experiments are conducted on NVIDIA TITAN RTX GPUs. We conduct every experiment for 4 runs and report the mean and validation results.

B.1 IMDB
We implement the base model TextCNN (Kim, 2014) on the IMDb movie reviews dataset (IMDB) (Maas et al., 2011). The reviews are classified into 3 classes, namely the positive reviews, negative reviews, and neutral reviews. The training size is 25000, the validation size is 25000, and the test size is 50000. In the preprocessing of the text, the vocab size is 30000, and the text length is 200. The embedding size is 500. In TextCNN, the filter window sizes are 3, 4, and 5, with 500 feature maps each. The optimizer is Adam with a learning rate of 10 −3 and a batch size of 64. We train models for 10 epochs and test the accuracy on the checkpoint with the best valid accuracy.

B.2 PTB-LM
We implement a 2-layer LSTM as a language model following Merity et al. (2017) on the wordlevel Penn TreeBank dataset (PTB) 4 (Marcus et al., 1993). In the preprocessing of the text, the vocab size is 10000. We use the SGD optimizer with an initial learning rate of 50 and a gradient norm clip of 0.25. We adopt a learning rate decay of 0.5 after 10 epochs. The input and output embeddings are tied. The embedding size is 500, and the hidden size is 500. We train models for 20 epochs and test the perplexity on the checkpoint with the lowest valid perplexity.

B.3 En-Vi
We implement the transformer model following the fairseq Ott et al. (2019) "transformer_wmt_en_de" implementation as our baseline model on the En-Vi dataset, which is provided by the IWSLT 2015 Evaluation Campaign (Cettolo et al., 2015). The training size is 133K, the validation set is TED tst2012 with a size of 1553 sentences, and the test set is TED tst2013 with a size of 1268 sentences. In the preprocessing of the text, the English and Vietnamese vocabulary sizes are 17200 and 7800 respectively. We use the same hyper-parameters following fairseq (Ott et al., 2019). We train the model for 52 epochs and adopt a checkpoint average of 10. In testing, We adopt the BLEU metric and the beam size of the model inference is 5.

B.4 De-En
We implement the transformer model following the fairseq Ott et al. (2019)

C Experimental Settings
In this section, we report experimental settings in the paper, including analytic trial details and hyperparameters of different algorithms.

C.1 Analytic Trial Settings
We conduct a 3-layer Multi-Layer Perceptrons (MLP, sizes are 784, 100, 100, 10) on the MNIST dataset (LeCun et al., 1998). The test minimum θ * is fine-tuned from the training minimum θ, and δ = θ * − θ. We utilize linear interpolation to plot αθ + (1 − α)θ * for different α, for visualizing the training loss, test loss and the shifted training loss. It can be concluded that the shifted training loss is similar to the test loss near the local minimum. To visualize the relation between δ 2 and the distribution shift. We conduct 100 trials with the same training set D and different mixed test sets D mix (mixed with (1 − η) of the training data from D and η of the test data from D * ). Here the test minimum is fine-tuned from the training minimum and δ = θ * − θ. η can be utilized to measure the strength of the distribution shift between D and D mix . η = 1%, 2%, 3%, · · · , 99%, 100%. Results show that there exists an approximately linear relationship between δ 2 and distribution shift strengths.

C.2 Main Result Settings
We try both L 2 and L +∞ . For , we first binary search proper order of magnitude, for example 10 −5 to 10 −4 , then we grid search , for example, {1 × 10 −5 , 2 × 10 −5 , 5 × 10 −5 , 8 × 10 −5 , 1 × 10 −4 , 2 × 10 −4 , 5 × 10 −4 , 8 × 10 −4 }. SAM on Transformer Models. The Transformer models need larger K and more detailed hyper-parameter tuning since SAM learning on the Transformer is unstable. Besides, Zhang et al. (2021) propose that SAM learning in the early stage may harm the learning. In our GA-SAM, g 2 is large and we can omit the SAM learning in the early stage. Therefore, we adopt K = 2 in multistep implementations, and following Zhang et al. (2021), we adopt a start epoch of 30 for Transformer models.

C.3 Ablation Study Settings
In this section, we report hyper-parameter settings in the ablation study.

C.4 Adversarial Training Settings
In this section, we report hyper-parameter settings in the study of adversarial training.