Geometry-Aware Weight Perturbation for Adversarial Training

Abstract

Adversarial training is one of the most successful approaches to improve model robustness against maliciously crafted data. Instead of training on a clean dataset, the model is trained on adversarial data generated on the fly. Based on that, a group of geometry-aware methods are proposed to further enhance the model robustness by assigning higher weights to the data points that are closer to the decision boundary during training. Although the robustness against the adversarial attack seen in the training process is significantly improved, the model becomes more vulnerable to unseen attacks, and the reason for the issue remains unclear. In this paper, we investigate the cause of the issue and claim that such geometry-aware methods lead to a sharp minimum, which results in poor robustness generalization for unseen attacks. Furthermore, we propose a remedy for the issue by imposing the adversarial weight perturbation mechanism and further develop a novel weight perturbation strategy called Geometry-Aware Weight Perturbation (GAWP). Extensive results demonstrate that the proposed method alleviates the robustness generalization issue of geometry-aware methods while consistently improving model robustness compared to existing weight perturbation strategies.

1. Introduction

Deep neural networks (DNNs) are known to be vulnerable to maliciously generated adversarial examples, which severely limits their applications in safety-critical areas [1], for example, autonomous driving [2]. As shown in Figure 1a, we can find an adversarial example by marginally perturbing a natural image using a PGD attack [3]. Although it is easy for a person to correctly classify both images and ignore the minor change between them, a regularly trained model can be easily fooled by the adversarial example and make a wrong prediction. To improve model robustness, adversarial training (AT) [3] has been demonstrated to be highly effective where approximate worst-case adversarial examples are used to train the model.

Figure 1. An illustration of the problems under consideration in this paper. As shown in (a), a regularly trained model is fooled by an adversarial example generated with a PGD attack, although there is no obvious difference between the natural image and its adversarial counterpart. GAIRAT is proposed to improve the model robustness. However, it is found that, although the robustness against the PGD attack is improved, the model can be easily fooled by adversarial examples generated with unseen attacks during training, for example, the targeted PGD attack, as shown in (b).

Based on the vanilla AT, geometry-aware instance-reweighted adversarial training (GAIRAT) is developed [4] following a straightforward intuition that the data closer to the decision boundary are more “attackable” and should be assigned a higher weight during training. While models trained with GAIRAT demonstrate robustness improvement against the adversarial attack imposed during training, they unexpectedly become more vulnerable to other unseen attacks, such as the logit-scaling attack [5] and the targeted PGD attack [3], as illustrated in Figure 1b. This issue of poor robustness generalization makes GAIRAT less appealing than other AT variants, such as TRADES [6] and MART [7], because the model robustness against the PGD attack no longer reflects the overall robustness of the model for GAIRAT. Zhang et al. [4] suggested that robustness generalization can be improved by incorporating more adversarial data generated by various attacks. However, the cost of adversarial data simulation can be prohibitively high, making the suggested remedy impractical. Therefore, to develop a practical solution that improves the robustness generalization of GAIRAT, it requires further study on the root cause of the issue. Additionally, similar ideas to GAIRAT have been successfully applied in regular training settings, such as AdaBoost [8] and focal loss [9]. This suggests that GAIRAT still holds untapped potential for achieving overall improvements in model robustness by addressing its robustness generalization issue.

To explore the causes of the issue, we inspect the model’s weight loss landscape [10] based on the checkpoints saved during the training with GAIRAT. In Figure 2, we present a 2D visualization of the weight loss landscape for each checkpoint [11]. It is observed that the landscape roughness significantly increases as the training progresses, indicating that GAIRAT converges to a sharp local minimum. Since the flatness of the weight loss landscape is commonly used as an indicator of generalization in the standard training scenario [11,12], we claim that the robustness generalization issue of GAIRAT is led by the lack of regularization in the weight loss landscape during training. Adversarial weight perturbation (AWP) [10] is a powerful approach that achieves superior model robustness by explicitly regularizing the weight loss landscape flatness. It motivates us to combine GAIRAT and AWP to achieve a better generalization in terms of robustness against various types of adversarial attack.

We first substantiate our claim that the robustness generalization issue of GAIRAT arises from a rough weight loss landscape. Specifically, we apply the same criterion proposed in GAIRAT to determine the importance of each adversarial example. To emphasize the training of more important examples, we increase their losses by assigning higher weights to them during the adversarial perturbation of model weights, which is different from GAIRAT, which directly uses the importance scores as weights of adversarial risks. We empirically show that our method converges to a flatter local minimum and leads to a model that is more robust to unseen attacks, which supports our claim about the reason for the robustness generalization issue. Treating this approach as the baseline, we find that, although it achieves higher model robustness, it suffers from a more severe robust overfitting issue [13] when compared with the vanilla AWP [10]. Recent studies on robust overfitting suggest that increasing the weights of examples with low adversarial risk is crucial for eliminating robust overfitting [14,15]. Based on that, we propose a criterion that determines the importance of an adversarial example during weight perturbation by considering its adversarial risk and its distance to the decision boundary. In a nutshell, examples that are closer to the decision boundary and have a lower adversarial risk are assigned higher importance scores during weight perturbation. With this importance criterion, we develop a novel weight perturbation strategy, namely, the Geometry-Aware Weight Perturbation (GAWP). The proposed method not only significantly alleviates the robustness generalization issue of GAIRAT [4], but also outperforms existing weight perturbation strategies such as AWP [10] and RWP [14]. The main contributions of this paper are summarized as follows:

• We identify that the geometry-aware adversarial training method GAIRAT converges to a sharp local minimum. By incorporating GAIRAT with AWP, we can obtain a flatter optimization result that is shown to be robust against unseen attacks. It provides a better understanding of the robustness generalization issue of GAIRAT.
• We claim that the model weight perturbation of the data that are close to the decision boundary and have a low adversarial risk is essential for achieving higher model robustness and avoiding robust overfitting.
• A novel weight perturbation strategy, GAWP, is developed in this paper. Extensive experiments demonstrate that GAWP outperforms not only GAIRAT, but also existing weight perturbation strategies, achieving superior model robustness.

Figure 2. Contours of weight loss landscape for 4 checkpoints saved during the training with GAIRAT [4]. The analysis is based on a PreAct ResNet-18 [16] and CIFAR-10 [17] dataset. Two directions are used for the 2D visualization. One direction is determined by weight perturbation, and the other one is a random orthogonal direction.

Figure 2. Contours of weight loss landscape for 4 checkpoints saved during the training with GAIRAT [4]. The analysis is based on a PreAct ResNet-18 [16] and CIFAR-10 [17] dataset. Two directions are used for the 2D visualization. One direction is determined by weight perturbation, and the other one is a random orthogonal direction.

2. Related Work

2.1. Adversarial Attacks

Let $\mathcal{X}\subseteq {\mathbb{R}}^d$ denote the input space and ${\mathcal{B}}_p(x,ϵ)=\{\tilde{x}\in \mathcal{X}:{\parallel \tilde{x}-x\parallel }_p\le ϵ\}$ represent the $l_p-$ norm ball of radius $ϵ$ centered at the clean input data x. The associated label of x is denoted by y. Considering a $C-$class classifier $f:\mathcal{X}\to {\mathbb{R}}^C$, an adversarial attack method is designed to find an adversarial example $\tilde{x}$, such that

$$\underset{c=1,\cdots ,C}{\mathrm{argmax}}{f}_c\left(\tilde{x}\right)\ne y,\mathrm{where}\tilde{x}\in {\mathcal{B}}_p(x,ϵ).$$

To find $\tilde{x}$, it is common to solve a constrained optimization problem defined in terms of a surrogate function L. The model’s classification accuracy on the adversarial data is considered as the measure of robustness, making it essential to design a comprehensive adversarial attack method for an objective evaluation of the model robustness. Here, we introduce several widely used adversarial attack methods.

Projected Gradient Descent (PGD). PGD [3] finds adversarial data by approximately solving the optimization problem (1), where the surrogate loss is the cross-entropy loss denoted by $\mathtt{CE}(·)$:

$$\underset{\tilde{x}\in {\mathcal{B}}_p(x,ϵ)}{max}\mathtt{CE}(f\left(\tilde{x}\right),y).$$

(1)

Starting with a point uniformly sampled inside the ball ${\mathcal{B}}_p(x,ϵ)$, PGD perturbs x for K steps with a step size $\alpha$:

$${\tilde{x}}^k={\mathsf{\Pi }}_{{\mathcal{B}}_p(x,ϵ)}({\tilde{x}}^{k-1}+\alpha ·\mathtt{sign}\left({\nabla }_{{\tilde{x}}^{k-1}}\mathtt{CE}(f\left({\tilde{x}}^{k-1}\right),y)\right)).$$

(2)

Here, ${\mathsf{\Pi }}_{{\mathcal{B}}_p(x,ϵ)}$ represents the projection function that projects ${\tilde{x}}^k$ back into the ball ${\mathcal{B}}_p(x,ϵ)$ if necessary.

AutoAttack (AA). As shown in recent studies [18,19], cross-entropy loss can be easily manipulated by simply rescaling the output logits of the model, meaning that measuring the model test accuracy under the PGD attack alone can overestimate the model robustness. For more objective evaluation of the model robustness, AA [18] generates adversarial examples with an ensemble of complementary attacks consisting of three white-box attacks and a black-box attack. During the evaluation, AA considers the benchmark model as robust only if it correctly classifies all types of adversarial examples, making it one of the most reliable evaluation methods for model robustness.

2.2. Adversarial Training

To defend against adversarial attacks, it is intuitive to train the model on adversarial data, namely, adversarial training (AT) [3]. Given a training dataset with n samples, denoted by $\mathcal{D}={\left\{(x_i,y_i)\right\}}_{i=1}^n$, training with AT is formulated as a min–max optimization problem:

$$\underset{\theta }{min}{\displaystyle \frac{1}{n}}\sum _{i=1}^n\underset{{\tilde{x}}_i\in {\mathcal{B}}_p(x_i,ϵ)}{max}\mathtt{CE}(f({\tilde{x}}_i;\theta ),y_i),$$

(3)

where $f(·;\theta )$ is a DNN with trainable parameters $\theta$. The adversarial risk $\rho \left(\theta \right)$ is the objective function for the outer minimization. For AT,

$$\rho \left(\theta \right)={\displaystyle \frac{1}{n}}\sum _{i=1}^n\underset{{\tilde{x}}_i\in {\mathcal{B}}_p(x_i,ϵ)}{max}\mathtt{CE}(f({\tilde{x}}_i;\theta ),y_i).$$

(4)

Based on AT, numerous variants have been proposed to further enhance model robustness. One group of methods follows the same definition of adversarial risk as vanilla AT but further boosts the model robustness by assigning unequal weights to different adversarial examples, for example, geometry-aware instance-reweighted adversarial training (GAIRAT) [4]. Based on a straightforward intuition that a data point closer to the decision boundary is less robust, GAIRAT proposes to emphasize the training on those less robust examples. Specifically, the outer minimization is reformulated as

$$\underset{\theta }{min}{\displaystyle \frac{1}{n}}\sum _{i=1}^{n}w(x_i,y_i)\mathtt{CE}(f({\tilde{x}}_i;\theta ),y_i),$$

(5)

where $w(x_i,y_i)$ is a non-increasing function w.r.t., the distance from the $x_i$ to the decision boundary. Additionally, it is required that $w(x_i,y_i)\ge 0$ and ${\displaystyle \frac{1}{n}}{\sum }_{i=1}^{n}w(x_i,y_i)=1$. One method of realization is proposed as follows: the geometric distance between $x_i$ and the decision boundary can be approximated by the least number of iterations $\kappa \left(x_i\right)$ that the PGD attack needs to generate an adversarial example ${\tilde{x}}_i$ to successfully fool the model given the maximum number of attack iterations K. Based on that, one possible design of $w(x_i,y_i)$ is proposed as follows:

$$w{(x_i,y_i;\lambda )}^{GAIRAT}=\mathtt{norm}({\displaystyle \frac{(1+\mathtt{tanh}(\lambda +5\times (1-2\times \kappa \left(x_i\right)/K)}{2}}),$$

(6)

where $\mathtt{norm}(·)$ represents the normalization operation by the batch sum, and $\lambda$ is a constant controlling the shape of the function. Note that, as $\lambda \to \infty$, Equation (6) is equivalent to assigning equal weights to all the adversarial examples, and GAIRAT degrades to vanilla AT. In essence, both AT and GAIRAT train a model based on adversarial data generated under the PGD attack. However, it is observed that, although GAIRAT achieves superior robustness against the PGD attack, the trained model is more vulnerable to adversarial unseen attacks during training [5] as compared to vanilla AT. As can be observed in Figure 3a, for GAIRAT, there is a more significant gap between model performances under the PGD attack and AA, indicating poor robustness generalization.

Figure 3. Comparison between GAIRAT and the baseline version of GAWP (GAWP_baseline). A comparison of model robustness is provided in (a). “PGD_10” represents the test accuracy against a 10-step PGD attack under the $L_{\infty }$ norm, which is seen during training. “AA” denotes the test accuracy against AutoAttack [18], representing the model robustness against unseen attacks. In (b), the Kullback–Leibler (KL) divergence between the adversarial distributions before and after a 5-step weight perturbation attack is shown. A smaller value indicates less sensitivity to the weight perturbation. A visualization of these loss distributions is provided based on the trained models after the 200-epoch training of GAIRAT (c) and GAWP_baseline (d).

Another group of AT variants is proposed to improve model robustness via novel designs of the adversarial risk $\rho \left(\theta \right)$. TRADES [6] develops a new design of adversarial risk through a theoretical analysis of the trade-off between the model robustness and its accuracy on the natural data. MART [7] reformulates the problem based on the intuition that misclassified and correctly classified examples should be differentiated during adversarial training. Technically, unlike the definition of adversarial risk (4) in vanilla AT, these methods propose alternative designs, summarized in Table 1.

Table 1. Adversarial risk definitions for AT and its variants. Here, $\mathtt{CE}(·)$ represents the cross-entropy loss, $\mathtt{BCE}(·)$ is the boosted cross-entropy loss [7] and $\mathtt{KL}(·\parallel ·)$ denotes the Kullback–Leibler divergence. $\lambda$ is a constant hyper-parameter. Note that, for MART, the adversarial examples $x^{\prime }$ are generated in the same manner as vanilla AT.

Method	$\mathit{\rho }\mathbf{\left(}\mathit{\theta }\mathbf{\right)}$
AT	${\displaystyle \frac{1}{n}}{\sum }_{i=1}^n{max}_{{\tilde{x}}_i\in {\mathcal{B}}_p(x_i,ϵ)}\mathtt{CE}(f({\tilde{x}}_i;\theta ),y_i)$
TRADES	${\displaystyle \frac{1}{n}}{\sum }_{i=1}^n\mathtt{CE}(f(x_i;\theta ),y_i)+\lambda ·{max}_{{\tilde{x}}_i\in {\mathcal{B}}_p(x_i,ϵ)}\mathtt{KL}(f(x_i;\theta )\parallel f({\tilde{x}}_i;\theta ))$
MART	${\displaystyle \frac{1}{n}}{\sum }_{i=1}^n\mathtt{BCE}(f({\tilde{x}}_i;\theta ),y_i)+\lambda ·\mathtt{KL}(f(x_i;\theta )\parallel f({\tilde{x}}_i;\theta ))·(1-f_{y_i}(x_i;\theta ))$

2.3. Adversarial Purification

As an alternative to adversarial training, adversarial purification refers to a group of defense methods that transform an adversarial example into its counterpart on the manifold of normal data using generative models. Meng and Chen proposed MagNet [20], which uses a collection of auto-encoders [21] for adversarial example purification. More recent methods such as defense-GAN [22] and DiffPure [23] apply more advanced generative models (GAN [24,25], Diffusion model [26]) to recover clean images through a reverse generation process. Although shown to be effective in defending against unseen attacks, these methods still suffer from the shortcomings of current generative models, for example, the mode collapse issue in GANs [24]. In addition to that, the extra purification process significantly increases the time complexity during inference, making them inapplicable for real-time tasks [23]. Considering these limitations, we focus on developing a more advanced adversarial training method in this paper.

3. Methodology

In this section, we first introduce the adversarial weight perturbation (AWP) mechanism [10] in detail. Then, we empirically investigate the underlying cause for the poor robustness generalization of GAIRAT [4] from the perspective of weight loss landscape smoothness. Based on this observation, we propose an effective remedy for the issue by incorporating GAIRAT with AWP. We further discuss how to utilize the geometric signal to guide weight perturbation for improving the model robustness, and a novel weight perturbation strategy, Geometry-Aware Weight Perturbation (GAWP), is proposed.

3.1. Preliminaries

Adversarial weight perturbation is proposed based on the key insight that, for a model with a smoother weight loss landscape, the gap between training robustness and test robustness is narrower [10]. Therefore, it is proposed to explicitly regularize the model’s weight loss landscape by imposing a double perturbation mechanism during adversarial training:

$$\underset{\theta }{min}\underset{v\in \mathcal{V}}{max}{\displaystyle \frac{1}{n}}\sum _{i=1}^n\underset{{\tilde{x}}_i\in {\mathcal{B}}_p(x_i,ϵ)}{max}\mathtt{CE}(f({\tilde{x}}_i;\theta +v),y_i).$$

(7)

Here, $\mathcal{V}$ is the feasible region of weight perturbation, i.e., $\{v\in \mathcal{V}:{\parallel v\parallel }_2\le \gamma {\parallel \theta \parallel }_2\}$, where $\gamma$ denotes the constraint on weight perturbation size. In essence, problem (7) is equivalent to ${min}_{\theta }\{\rho \left(\theta \right)+(\rho (\theta +v)-\rho \left(\theta \right))\}$, where the term $\left(\rho \right(\theta +v)-\rho (\theta \left)\right)$ explicitly regularizes the flatness of the weight loss landscape during training. Technically, for each step of weight perturbation, v is updated as follows:

$$v^{k+1}=v^k+{\nabla }_{v^k}{\displaystyle \frac{1}{n}}\sum _{i=1}^n\mathtt{CE}(f({\tilde{x}}_i;\theta +v^k),y_i).$$

(8)

Following AWP, it is discovered that, by assigning higher importance scores to the small-loss data during AWP, the model robustness is further enhanced, leading to advanced variants of AWP such as RWP [14] and MLCAT_WP [15]. Particularly, for RWP, v is updated based on the following:

$$\begin{array}{cc}\hfill & v^{k+1}=v^k+{\nabla }_{v^k}{\displaystyle \frac{1}{n}}\sum _{i=1}^n\mathbb{1}({\tilde{x}}_i,y_i)\mathtt{CE}(f({\tilde{x}}_i;\theta +v^k),y_i),\hfill \\ \hfill & \mathrm{where}\mathbb{1}({\tilde{x}}_i,y_i)=\left\{\begin{array}{cc}0\hfill & \mathrm{if}\mathtt{CE}(f({\tilde{x}}_i;\theta +v),y_i)>c_{min}\hfill \\ 1\hfill & \mathrm{otherwise}.\hfill \end{array}\right.\hfill \end{array}$$

(9)

Here, $c_{min}$ is the minimum loss value that is set as a hyper-parameter before training. MLCAT follows a similar update rule to compute v, where the difference lies in the design of indicator function $\mathbb{1}({\tilde{x}}_i,y_i)$.

3.2. Understanding the Robustness Generalization Issue of GAIRAT

We examine the model robustness of the “best” model returned by GAIRAT (the model with the highest test accuracy under the PGD-10 attack). In particular, a PreAct ResNet-18 [16] is trained with GAIRAT on CIFAR-10 [17] for 200 epochs following the same setting as in the original paper [4]. As shown in Figure 3a, although the test accuracy under the PGD-10 attack is as high as $63.28\%$, it drops significantly to $29.99\%$ under the AA, suggesting poor robustness generalization. Furthermore, we investigate the model weight loss landscape smoothness based on 4 checkpoints saved every 50 epochs of training. For quantification purposes, we measure the smoothness of the weight loss landscape with the Kullback–Leibler (KL) divergence between the adversarial risk distributions before and after a 5-step adversarial weight perturbation with the perturbation size $\gamma =0.01$. A lower KL divergence means the model is less sensitive to weight perturbation, implying a smoother weight loss landscape. According to Figure 3b, as the training of GAIRAT progresses, the model’s sensitivity to weight perturbation considerably increases, suggesting that the model gradually converges to a sharp local minimum. We further claim that this is the potential reason for the robustness generalization issue of GAIRAT.

To testify to our claim, we impose proper regularization on the smoothness of the weight loss landscape to improve the robustness generalization of GAIRAT. We apply the same rule as GAIRAT to determine the importance of each adversarial sample. However, we impose importance scores during weight perturbation rather than the outer minimization. Specifically, model weights are adversarially perturbed based on

$$v^{k+1}=v^k+{\nabla }_{v^k}{\displaystyle \frac{1}{n}}\sum _{i=1}^{n}w{(x_i,y_i;\lambda )}^{GAIRAT}\mathtt{CE}(f({\tilde{x}}_i;\theta +v^k),y_i).$$

(10)

In this manner, the training on the data that are closer to the decision boundary is implicitly emphasized through the weight perturbation mechanism. Regarding this approach as the baseline (GAWP_baseline), we observe that the performance gap between the PGD-10 attack and the AA obviously shrinks (PGD-10: $55.27\%$, AA: $50.17\%$), as shown in Figure 3a. Additionally, comparing against GAIRAT, the GAWP_baseline converges to a flatter local minimum, based on Figure 3b–d. In a nutshell, the experimental results support our claim about the cause of the poor robustness generalization of GAIRAT, which motivates us to develop a novel weight perturbation strategy that considers the geometric distance of each adversarial example to the decision boundary.

3.3. Geometry-Aware Weight Perturbation

Although the baseline version of GAWP effectively alleviates the robustness generalization issue of GAIRAT, it is observed that the proposed method suffers from the robust overfitting issue and only marginally improves the model robustness when compared against AWP; see Table 2. In particular, we observe that the model robustness starts degrading and never recovers after training for around 120 epochs with the GAWP_baseline, indicating robust overfitting. Therefore, we seek to further improve the model robustness by eliminating robust overfitting for the baseline method.

Table 2. Test robustness (%) on CIFAR-10 using PreAct ResNet-18 under the $L_{\infty }$ threat model. We focus on model performance on the clean data (“Natural”) and adversarial data under AutoAttack (“AA”) and PGD-10 attack (“PGD-10”). For PGD-10, we present the highest model robustness that is ever achieved (“Best”) and after the whole training process (“Last”). Their performance gap is also recorded (“Diff”), indicating the severity of the robust overfitting [13].

Method	Natural	AA	PGD-10
			Best	Last	Diff
GAIRAT	79.24	29.99	63.28	59.47	3.81
AT-AWP	83.54	50.16	55.65	54.91	0.74
AT-GAWP_baseline	82.80	50.17	55.27	53.97	1.30

For numerous adversarial training methods, it is a prevalent phenomenon that there is an obvious gap between the best and last model robustness, termed the robust overfitting issue [13]. To tackle the issue, various remedies have been proposed, such as semi-supervised learning [27], temporal ensembling [28], pruning [29] and label smoothing [30]. Moreover, recent studies [15,31] have sought an explanation for why robust overfitting occurs. It is discovered that, as AT progresses, a portion of adversarial data becomes easy to learn and fails to provide enough adversary strength. Additionally, the robust overfitting issue is found to be eliminated by simply increasing the adversarial risk of such small-loss data. It motivates us to introduce the adversarial risk as an additional factor when computing the importance score of each adversarial example.

Inspired by previous studies on robust overfitting [14,15,31], small-loss data should be assigned a higher importance score. Therefore, we propose to determine the importance of each example by its adversarial risk as follows:

$$w{(x_i,y_i)}^{loss}={\displaystyle \frac{1}{m-1}}-{\displaystyle \frac{\mathtt{exp}\left(\mathtt{CE}(f({\tilde{x}}_i;\theta ),y_i)\right)}{(m-1){\sum }_{j=1}^m\mathtt{exp}\left(\mathtt{CE}(f({\tilde{x}}_j;\theta ),y_j)\right)}}.$$

(11)

Here, m denotes the batch size. By combining it with Equation (6), the importance score is computed as follows:

$$w{(x_i,y_i;\lambda )}^{GAWP}={\displaystyle \frac{w{(x_i,y_i)}^{loss}+w{(x_i,y_i;\lambda )}^{GAIRAT}}{2}}.$$

(12)

Based on this, we propose Geometry-Aware Weight Perturbation (GAWP), summarized as Algorithm 1. Here, we take the combination of GAWP and vanilla AT as an example. Note that it can be extended to other classic methods such as TRADES [6] and MART [7]. For each optimization step, we sample one batch of data and generate corresponding adversarial examples with the PGD attack. The importance score of each adversarial sample is determined by its adversarial risk and geometric distance to the decision boundary, which is then applied to weight perturbation. In essence, during weight perturbation, GAWP focuses more on small-loss data that are close to the decision boundary.

Algorithm 1 Geometry-Aware Weight Perturbation (GAWP)

Input: DNN $f(·;\theta )$, training dataset $\mathcal{D}$, batch size m, learning rate $\eta$, PGD step size $\alpha$, PGD steps $K_1$, PGD constraint $ϵ$, GAWP stepes $K_2$, GAWP constraint $\gamma$, parameter $\lambda$. Outout: Adversarially robust model $f(·;\theta )$. Repeat  Sample a mini-batch $x_{\mathcal{B}}$ with a size m from training set $\mathcal{D}$.  Generate one batch of adversarial examples ${\tilde{x}}_{\mathcal{B}}$ based on Equation (2) for $K_1$ steps.  for  $i=1\cdots m$  do      Compute $w({x_{\mathcal{B}}}^{\left(i\right)},y^{\left(i\right)};{\lambda )}^{GAWP}$ based on Equation (12).  end for  Initialize $v\leftarrow 0$  for  $k=1\cdots K_2$  do      $v\leftarrow v+{\nabla }_v{\sum }_{i=1}^{m}w{({x_{\mathcal{B}}}^{\left(i\right)},y^{\left(i\right)};\lambda )}^{GAWP}\mathtt{CE}(f({{\tilde{x}}_{\mathcal{B}}}^{\left(i\right)};\theta +v),y^{\left(i\right)})$      $v\leftarrow \gamma {\displaystyle \frac{v}{\parallel v\parallel }}\parallel \theta \parallel$  end for  $\theta \leftarrow (\theta +v)-\eta {\nabla }_{\theta +v}{\displaystyle \frac{1}{m}}{\sum }_{i=1}^m\mathtt{CE}(f\left({{\tilde{x}}_{\mathcal{B}}}^{\left(i\right)}\right),y^{\left(i\right)})-v$ Until training converges

4. Experimental Results

In this section, we conduct extensive experiments to verify the effectiveness of GAWP, providing experimental settings, robustness evaluation and ablation studies.

4.1. Experimental Setup

For the model architecture, we use PreAct ResNet-18 [16] and Wide ResNet-34-10 [32]. Three datasets are considered: CIFAR-10, CIFAR-100 [17] and SVHN [33]. For CIFAR datasets, we apply random crops with four pixels of padding and random horizontal flips for the data augmentation. The proposed method GAWP is compared with a number of baseline methods including the vanilla AT [3], GAIRAT [4] and TRADES [6]. We also compare GAWP with advanced weight perturbation strategies, particularly AWP [10], RWP [14] and MLCAT_WP [15]. For adversary generation, a 10-step PGD attack with a random start is applied following the setting in [3]; under the $L_{\infty }$ threat model, the attack budget $ϵ=8/255$ and the perturbation step size $\alpha =2/255$ for the CIFAR datasets while $\alpha =1/255$ for SVHN; and, under the $L_2$ threat model, $ϵ=128/255$ and $\alpha =15/255$ for all datasets. All experiments are conducted on four NVIDIA GeForce 2080Ti GPUs. The program is developed based on PyTorch 1.8.1 [34].

Training parameters. We use an SGD optimizer with a momentum of $0.9$ and a weight decay of $5\times {10}^{-4}$. The model is trained for 200 epochs with a batch size of 128. The learning rate starts with $0.1$ ($0.05$ for SVHN) and decays by 0.1 at the ${100}_{th}$ and ${150}_{th}$ epoch, respectively. For GAWP, the number of weight perturbation steps is set as $K_2=10$, the perturbation constraint $\gamma =0.01$ and $\lambda =-1$ for all datasets. The hyper-parameter settings of the baseline methods follow the recommendations of their original papers.

4.2. Robustness Evaluation

In this section, we present the main results of robustness evaluation for GAWP and various benchmark methods. We showcase the effectiveness of GAWP from two aspects: firstly, it enhances the robustness generalization of GAIRAT and avoids overfitting issues; secondly, GAWP consistently improves the test robustness across different datasets and threat models compared to other weight perturbation strategies.

4.2.1. Mitigating Robustness Generalization and Overfitting Issues

To validate that GAWP improves the robustness generalization of GAIRAT, we compare vanilla AT, GAIRAT and AT-GAWP using PreAct ResNet-18 on CIFAR-10. The evaluation results are presented in Table 3. For the checkpoint that achieves the highest test accuracy against PGD-10 attack, “Natural” denotes the model performance on the clean test data, and “AA” denotes the test accuracy under AutoAttack. It is observed that AT-GAWP achieves the highest model robustness against AutoAttack under both $L_{\infty }$ and $L_2$ threat models when compared to the other two baselines. Therefore, the robustness generalization issue of GAIRAT is effectively addressed in AT-GAWP.

Table 3. Test robustness (%) on CIFAR-10 using PreAct ResNet-18. The highest accuracy or the smallest performance gap in each setting is marked in bold.

Threat Model	Method	Natural	AA	PGD-10
				Best	Last	Diff
$L_{\infty }$	AT	83.73	47.00	51.26	43.57	7.69
	GAIRAT	79.24	29.99	63.28	59.47	3.81
	AT-GAWP	81.07	50.01	55.77	54.93	0.84
	TRADES	82.86	49.37	53.82	51.32	2.50
	TRADES-GAWP	81.31	50.61	55.20	54.56	0.64
	MART	77.06	47.47	54.08	50.41	3.67
	MART-GAWP	78.04	49.15	56.21	55.81	0.40
$L_2$	AT	88.57	67.20	71.01	68.78	2.23
	GAIRAT	84.34	53.36	72.56	69.55	3.01
	AT-GAWP	89.26	70.06	74.61	74.17	0.44
	TRADES	86.47	68.06	71.64	69.70	1.94
	TRADES-GAWP	87.93	71.14	74.61	74.19	0.42
	MART	87.90	66.63	70.95	68.27	2.68
	MART-GAWP	88.07	68.89	74.09	73.57	0.52

As for the robust overfitting issue, the test accuracy against PGD-10 attack is selected as the measure of robustness. The extent of robust overfitting is quantified with the performance gap (“Diff”) between the highest robustness ever achieved (“Best”) and the robustness of the last epoch checkpoint (“Last”). A larger gap indicates more severe robust overfitting. As shown in Table 3, both AT and GAIRAT suffer from robust overfitting, with a gap of over $2\%$ across all settings. On the contrary, the performance gap significantly shrinks for AT-GAWP (<1%), indicating effective suppression of robust overfitting.

Along with exploring robust overfitting, we also examine the model’s generalization on natural data. The generalization gap, defined as the difference between training accuracy and test accuracy, is commonly chosen as an indicator of generalization. A smaller generalization gap implies a better generalization. According to Figure 4, when training with vanilla AT, the generalization gap surges after training for 100 epochs, showing a degradation of generalization to unseen natural data. With the help of GAWP, the generalization gap consistently stays below $10\%$ during the whole training process. For both $L_{\infty }$ and $L_2$ threat models, we obtain the same observation. Therefore, we claim that GAWP not only mitigates robust overfitting but also boosts the generalization to unseen natural data.

Figure 4. A comparison in terms of generalization gap (training accuracy - test accuracy) on the natural data. With the help of GAWP, the overfitting issues is obviously alleviated, and the trained model possesses better generalization to unseen data. Such an observation applies for both settings of $L_{\infty }$ norm (a) and $L_2$ norm (b).

In addition to vanilla AT, we also extend GAWP to TRADES and MART. The experimental results are summarized in Table 3 and Figure 4. Similar patterns have been observed when combing GAWP with TRADES and MART. In conclusion, GAWP consistently enhances robustness generalization, improves natural data generalization and mitigates robust overfitting across various baseline methods and threat models.

4.2.2. Imposing Geometric Distance Benefits Weight Perturbation

Unlike existing weight perturbation strategies, we propose to determine the importance score of each adversarial example based on not only its adversarial risk but also its geometric distance to the decision boundary. Recall that the vanilla AWP treats all the adversarial examples equally during weight perturbation, while RWP conducts weight perturbation only on the data with an adversarial risk lower than a threshold. To showcase that weight perturbation benefits from the introduced geometric information, we compare AT assisted by GAWP (AT-GAWP) to vanilla AT, AT-AWP, AT-RWP and AT-MLCAT_WP using Wide ResNet-34-10 across different datasets and threat models. Based on the checkpoint that achieves the highest PGD-10 robustness for each benchmark method, we conduct the robustness evaluation and summarize the results in Table 4.

Table 4. Test robustness (%) using Wide ResNet-34-10 across different datasets and threat models. The highest accuracy in each setting is marked in bold.

Threat Model	Method	CIFAR-10		CIFAR-100		SVHN
		Natural	AA	Natural	AA	Natural	AA
$L_{\infty }$	AT	86.50	51.67	60.88	27.71	92.89	50.23
	AT-AWP	86.80	54.64	60.73	29.99	93.78	55.83
	AT-RWP	86.18	55.53	62.95	30.30	92.74	54.28
	AT-MLCATWP	82.89	53.45	59.50	29.49	92.76	53.47
	AT-GAWP	87.29	56.28	63.97	30.91	93.82	56.01
$L_2$	AT	90.50	70.11	68.01	42.00	93.65	65.73
	AT-AWP	92.61	74.21	70.10	45.89	94.01	68.80
	AT-RWP	92.11	74.41	70.31	46.07	94.48	68.95
	AT-MLCATWP	90.38	72.45	70.30	45.97	94.40	68.99
	AT-GAWP	92.88	74.56	70.88	46.71	95.21	69.71

It is observed that GAWP consistently outperforms the other baseline weight perturbation strategies across different datasets and threat models in terms of test accuracy both on the clean dataset and under AutoAttack. This is due to GAWP providing a more comprehensive way to estimate the importance of each adversarial example during weight perturbation. We further claim that, to enhance the model robustness, weight perturbation benefits from assigning adversarial examples with unequal importance scores. In addition, an adversarial example with a lower adversarial risk and that is closer to the decision boundary should be assigned a higher importance score during weight perturbation.

4.3. Ablation Studies

4.3.1. Effect of Different Ways to Determine Importance Scores

GAWP determines the importance score of each example based on $w{(x,y;\lambda )}^{GAIRAT}$ and $w{(x,y)}^{loss}$. To observe the effect of each component, we conduct an ablation study and present the results in Table 5. When using $w{(x,y)}^{GAIRAT}$ only (GAWP_baseline), the model robustness against AutoAttack is $54.56\%$, which is comparable to AWP ($54.64\%$). Applying $w{(x,y)}^{loss}$ only, we can obtain a robustness of $55.94\%$ under AutoAttack, which is comparable with the state-of-the-art strategy RWP ($55.53\%$). Using an ensemble of $w{(x,y,\lambda )}^{GAIRAT}$ and $w{(x,y)}^{loss}$ (GAWP), model robustness is further improved. Similar patterns are observed on the natural data (“Natural”). Based on our observations, we assert that $w{(x,y,\lambda )}^{GAIRAT}$ and $w{(x,y)}^{loss}$ work complementarily in determining the importance of each example during weight perturbation, both providing essential training signals.

Table 5. Ablation study for examining the effect of each component based on CIFAR-10 using Wide ResNet-34-10 under the $L_{\infty }$ threat model. The highest accuracy is marked in bold.

$\mathit{w}{\mathbf{(}\mathit{x}\mathbf{,}\mathit{y}\mathbf{,}\mathit{\lambda }\mathbf{)}}^{\mathit{G}\mathit{A}\mathit{I}\mathit{R}\mathit{A}\mathit{T}}$	$\mathit{w}{\mathbf{(}\mathit{x}\mathbf{,}\mathit{y}\mathbf{)}}^{\mathit{l}\mathit{o}\mathit{s}\mathit{s}}$	Natural (%)	AA (%)
✓		85.21	54.56
	✓	86.95	55.94
✓	✓	87.29	56.28

For $w{(x,y)}^{loss}$, we testify another design following the same idea:

$$w{(x,y)}_{linear}^{loss}={\displaystyle \frac{1}{m-1}}-{\displaystyle \frac{\mathtt{CE}(f(x_i^{\prime };\theta ),y_i)}{(m-1){\sum }_{j=1}^m\mathtt{CE}(f(x_j^{\prime };\theta ),y_j)}}.$$

(13)

Unlike $w{(x,y)}^{loss}$, $w{(x,y)}_{linear}^{loss}$ removes the nonlinearity led by the $\mathtt{Softmax}(·)$ operation while assigning a higher importance score to the data with a lower adversarial risk. For $w{(x,y;\lambda )}^{GAIRAT}$, we investigate its reversed version, $\widehat{w}{(x,y;\lambda )}^{GAIRAT}$, defined as follows:

$$\widehat{w}{(x,y;\lambda )}^{GAIRAT}=1-w{(x,y;\lambda )}^{GAIRAT}.$$

(14)

Opposite to $w{(x,y;\lambda )}^{GAIRAT}$, $\widehat{w}{(x,y;\lambda )}^{GAIRAT}$ assigns a higher weight to the data that are more distant from the decision boundary. We conduct experiments based on the various strategies to determine the importance scores of adversarial examples during weight perturbation, and the results are summarized in Table 6.

Table 6. Ablation study for comparing different importance score determination strategies based on CIFAR-10 using Wide ResNet-34-10 under the $L_{\infty }$ threat model. The highest accuracy in each setting is marked in bold.

Strategy	Natural (%)	AA (%)
$w{(x,y)}^{loss}$	86.95	55.94
$w{(x,y)}_{linear}^{loss}$	86.83	55.88
$w{(x,y)}^{loss}+w{(x,y)}^{GAIRAT}$	87.29	56.28
$w{(x,y)}^{loss}+\widehat{w}{(x,y;\lambda )}^{GAIRAT}$	86.55	55.19

It is observed that applying $w{(x,y)}^{loss}$ leads to a more robust model due to the additional $\mathtt{Softmax}(·)$ operation. Additionally, assigning a higher importance score to the data that are further away from the decision boundary results in degradation of the model robustness. This observation supports our claim that, during weight perturbation, the data closer to the decision boundary deserve a higher importance score.

4.3.2. Hyper-Parameter Sensitivity

In this part, we investigate how sensitive the performance of GAWP is to the values of three main hyper-parameters: $\lambda$, $\gamma$ and $K_2$. Recall that $\lambda$ defines the shape of $w{(x,y,\lambda )}^{GAIRAT}$. $\gamma$ controls the size of the feasible region during the weight perturbation. $K_2$ is the number of steps for weight perturbation. We focus on the combination of AT and GAWP in this part. The analysis is based on CIFAR-10 under the $L_{\infty }$ threat model using PreAct ResNet-18.

We first explore the performance of GAWP under different settings of $\lambda$. Recall that $\kappa$ is the number of attack steps required for an example to fool the model. A smaller $\kappa$ indicates that the adversarial example is closer to the decision boundary. $w{(x,y;\lambda )}^{GAIRAT}$ is defined to assign a higher weight to the example that is closer to the decision boundary, while the hyper-parameter $\lambda$ controls the shape. As shown in Figure 5a, as $\lambda$ increases, $w{(x,y;\lambda )}^{GAIRAT}$ tends to assign equal weights to all the examples, which degrades to vanilla AT as $\lambda \to \infty$. To showcase the effect of $\lambda$, we train models under different values of $\lambda$ while fixing $K_2=10$ and $\gamma =0.01$. Four values of $\lambda$ are examined in the ablation study, and the results are summarized in Figure 5b. It is observed that the model achieves the highest robustness with $\lambda =-1.0$, which is the same as the default setting proposed in GAIRAT [4].

Figure 5. The ablation study investigating the importance of three hyper-parameters: $\lambda$, $\gamma$ and $K_2$. The parameter $\lambda$ determines the shape of $w^{GAIRAT}$, as shown in (a). Measuring the model robustness in terms of the test accuracy under AutoAttack, the results of ablation study in terms of $\lambda$, $\gamma$ and $K_2$ are presented in (b), (c) and (d), respectively. With different values of $K_2$, the time cost per epoch of training is presented in (e).

We also investigate the effect of $\gamma$. Fixing $K_2=10$ and $\lambda =-1$, the performance of GAWP is examined with a $\gamma$ that varies from 0.0025 to 0.0125. It can be observed in Figure 5c that the model achieves the highest robustness when $\gamma =0.01$. Model robustness decreases rapidly when $\gamma <0.005$, indicating that the regularization is insufficient. When $\gamma >0.01$, the model robustness starts decreasing, showing that over-regularization is also detrimental to the training process.

As for studying the effect of $K_2$, based on the results of previous studies, we fix $\lambda =-1.0$ and $\gamma =0.01$. The parameter $K_2$ of GAWP varies from 1 to 10. It can be observed in Figure 5d that the model achieves higher robustness as $K_2$ increases, which meets our expectation. Using RWP ($K_2=10$, $\gamma =0.01$) as the baseline, GAWP obviously outperforms RWP with the same $K_2$. Such an observation further reflects the benefits of imposing geometric information to determine the importance of each example during weight perturbation. We also summarize the average time cost of one training epoch for each setting in Figure 5e. As $K_2$ increases, it takes longer for training due to the additional time complexity caused by the weight perturbation.

5. Limitations and Future Work

To provide a balanced view of our work, the limitations of the proposed method and future research directions are discussed in this section. There are two main limitations that await further study.

Extension to other data types. Note that the proposed method is only discussed in the context of image data. This is because the proposed method is based on a well-defined continuous optimization problem (7), which can be significantly different for text data or audio data. For text data, the generation of adversarial examples [35,36] is a discrete problem, and it requires the elimination of grammatical and semantic problems perceptible to humans [37,38]. For audio data, although the adversary generation process is similar to that in the image setting [39,40], the training process may be entangled with regular training [41] or introduce additional models [42]. Since the problem formulations in these settings are different from (7), whether the insights provided in this paper can be applied to text data or audio data remains unclear, which suggests an interesting research direction.

Theoretical proof. Although the effectiveness of AWP [10] has been theoretically justified, the reason why treating examples in an unequal manner enhances model robustness is still unaddressed. Similar to related works such as RWP [14] and MLCAT_WP [15], this paper develops a weight perturbation strategy based on empirical intuitions. Although the effectiveness of the proposed method has been substantiated by extensive experiment results, discovering how to understand the robustness improvement from a theoretical perspective remains an open question.

6. Conclusions

In this paper, we investigate the robustness generalization issue of GAIRAT from the perspective of the weight loss landscape. We identify that GAIRAT leads to a sharp local minimum and mitigate the robustness generalization issue by imposing regularization on the weight loss landscape through weight perturbation. Inspired by this, we propose Geometry-Aware Weight Perturbation (GAWP). Unlike existing weight perturbation strategies, we impose additional geometric information of each adversarial example along with its adversarial risk to determine its importance score during weight perturbation. Comprehensive experiments show that GAWP eliminates the robustness generalization issue of GAIRAT and improves adversarial robustness across various network architectures and datasets (Table 3) as compared to existing weight perturbation strategies (Table 4).

Note that GAWP is a general adversarial training method that can be applied to various real-world computer vision applications, especially those requiring model robustness against input perturbations. For instance, the intelligent power line inspection system detects external obstacles of transmission lines in the power grid [43,44]. Technically, the system acquires image data using visual sensors and leverages object detection techniques to identify obstacles in the images. In real-world scenarios, the acquired images may be distorted by the sensor noise or adverse weather conditions. Therefore, it is essential for the object detection model to maintain its performance despite these perturbations. In the near future, we plan to apply GAWP to train the model of an intelligent drone-based power asset management system to further enhance its robustness.