α-WEIGHTED FEDERATED ADVERSARIAL TRAINING

Abstract:
Federated Adversarial Training (FAT) helps us address the data privacy and governance issues, meanwhile maintains the model robustness to the adversarial attack. However, the inner-maximization optimization of Adversarial Training can exacerbate the data heterogeneity among local clients, which triggers the pain points of Federated Learning. This makes that the straightforward combination of two paradigms shows the performance deterioration as observed in previous works. In this paper, we introduce an α-Weighted Federated Adversarial Training (αWFAT) method to overcome this problem, which relaxes the inner-maximization of Adversarial Training into a lower bound friendly to Federated Learning. We present the theoretical analysis about this α-weighted mechanism and its effect on the convergence of FAT. Empirically, the extensive experiments are conducted to comprehensively understand the characteristics of α-WFAT, and the results on three benchmark datasets demonstrate α-WFAT significantly outperforms FAT under different adversarial learning methods and federated optimization methods.

heading: 1 INTRODUCTION
text: To handle the data privacy and governance issues, Federated Learning (McMahan et al., 2017) as one promising paradigm of distributed training has drawn the increasing attention (McMahan et al., 2017; Kairouz et al., 2019). However, training locally in Federated Learning also introduces the vulnerability from the adversarial attacks (Goodfellow et al., 2015; Kurakin et al., 2016), which drives us to consider the model robustness in this framework. Thus, recent advances (Kairouz et al., 2019) explore to apply the Adversarial Training methods (Madry et al., 2018) into Federated Learning.
However, the straightforward combination of Adversarial Training and Federated Learning presents some potential challenges due to the communication cost and the hardware requirement. For example, Shah et al. (2021) pointed out that the communication in Federated Learning may be a constraint to Adversarial Training, and proposed a dynamic schedule on the local training epochs to achieve the expected robustness in a short communication budget. Hong et al. (2021) considered the hardware constraint where some clients are not able to participate Adversarial Training, and they proposed a federated robust propagation method to share the adversarial robustness among the clients. Although these previous works indeed addressed some realistic problems, one critical issue in the way is the performance deterioration in the combination of two paradigms as observed in (Zizzo et al., 2020).
As shown in the left panel of Figure 1, one typical phenomenon in Federated Adversarial Training (FAT) is the robust accuracy of FAT (Zizzo et al., 2020) based on FedAvg (McMahan et al., 2017) will decrease significantly at the later stage of learning compared with the centralized Adversarial Training (Madry et al., 2018) that does not. Actually, it exists in many variants of Federated Learning methods (Shah et al., 2021) under different communication rounds and different local training epochs. However, it is still lack of the sufficient algorithmic breakthrough to overcome this issue, since almost all previous works (Zizzo et al., 2020; Shah et al., 2021; Hong et al., 2021) consistently apply the conventional update framework of Federated Learning with Adversarial Training.
We dive into this phenomenon and attribute it to the inner-maximization of Adversarial Training. Compared with the centralized Adversarial Training (Madry et al., 2018), the training data of FAT is distributed to each client, which leads to the Adversarial Training in each client unaware of the data in the others. Therefore, the adversarial examples generated by the inner-maximization of Adversarial Training tend to be highly biased to each local distribution, yielding the radical optimization to pursuit the model robustness (as shown in Figure 2, it has a severe local bias to the global optimum). In
short, the inner-maximization of Adversarial Training exacerbates the data heterogeneity among local clients, which actually triggers the pain points of Federated Learning that are exploring urgently.
To handle this problem, we propose a new learning framework based on a simple but effective reweighting mechanism, namely, α-Weighted Federated Adversarial Training (α-WFAT). Concretely, we relax the objective of the inner-maximization in Adversarial Training into a lower bound by an αweighting mechanism (as Eq. (3) in Section 4.2). Similar to the bias-variance trade-off (Kohavi et al., 1996), we introduce the small bias to the original objective via a low-bound relaxation, facilitating a friendly optimization in the combination of Adversarial Training and Federated Learning. This constructs a conservative optimization for Adversarial Training that weights different population based on their adversarial losses. Then, the harsh heterogeneous update can be down-weighted with this flexible weighting mechanism, and the convergence could be accelerated at the same time. Applying the similar idea to FAT, we propose α-WFAT that emphasizes the robust clients more compared with other non-robust clients to alleviate the heterogeneous bias caused by the local adversarial generation. The right panel of Figure 1 gives a simple comparison between α-WFAT and FAT. Empirically, we conducted extensive experiments to provide a comprehensive understanding of the proposed α-WFAT, and the results of α-WFAT in the context of different adversarial learning methods and federated optimization methods demonstrate its superiority to improve the model performance.
Main Contributions
• We derive an α-weighted relaxation for Adversarial Training to relax the inner-maximization by a lower bound, which builds a mediating function to alleviate the potential radical optimization in its straightforward combination with Federated Learning (in Section 4.2).
• We propose a new learning framework, i.e., α-Weighted Federated Adversarial Training (α-WFAT), to realize the relaxation of inner-maximization in FAT, which is simple and compatible with various Federated Learning or Adversarial Training methods (in Section 4.3).
• We conduct extensive experiments to comprehensively understand the characteristics of the α-WFAT, and confirm its effectiveness on improving the model performance for both IID and Non-IID settings in the context of several federated optimization methods (in Section 5).

heading: 2 RELATED WORK
text: Federated Learning The representative work in Federated Learning is FedAvg (McMahan et al., 2017), which has been proved effectiveness during the distributed training to maintain the data privacy. To further address the heterogeneous issues, several optimization approaches have been proposed e.g., FedProx (Li et al., 2018), FedNova (Wang et al., 2020b) and Scaffold (Karimireddy et al., 2020). FedProx introduced a proximal term for FedAvg to constrain the model drift cause by heterogeneity; FedNova proposed a general framework that eliminated the objective inconsistency and preserved the fast convergence; Scaffold utilized the control variates to reduce the gradient variance in the
local updates and accelerate the convergence. MOON (Li et al., 2021a) alleviated the heterogeneity by maximizing the agreement between the representation of the local model and that of the global model, which helps correct the local training of individual parties. Reisizadeh et al. (2020) developed a robust federated learning algorithm to against distribution shifts in clients samples. Our α-WFAT introduces the relaxation into federated adversarial training, which is orthogonal to and compatible with the previous optimization methods.
Adversarial Training As one of the defensive methods (Papernot et al., 2016; Gao et al., 2021), Adversarial Training (Madry et al., 2018; Zhang et al., 2019; Jiang et al., 2020; Chen et al., 2021) is to improve the robustness of machine learning models. The classical AT (Madry et al., 2018) is built upon on a min-max formula to optimize the worst case, e.g., the adversarial example near the natural example (Goodfellow et al., 2015). Zhang et al. (2019) decomposed the prediction error for adversarial examples as the sum of the natural error and the boundary error, and proposed TRADES to balance the classification performance between the natural and adversarial examples. Wang et al. (2020c) further explored the influence of the misclassified examples on the robustness, and proposed MART that emphasizes the minimization of the misclassified examples to boost the AT. Zhang et al. (2020) investigated the “benign adversarial examples” in AT and further improve the natural performance of robust model. In this work, our α-WFAT framework leverages the client-level measure to alleviate the heterogeneous issue in the straightforward combination of adversarial training and federated learning. It is compatible to further incorporate those centralized adversarial training methods to improve the model performance.
Federated Adversarial Training. Recently, several works have made the exploration on the Adversarial Training in the context of Federated Learning, which consider the data privacy and the robustness in one framework. To our best knowledge, Zizzo et al. (2020) take the first trial to study the feasibility of extending Federated Learning (McMahan et al., 2017) with the standard AT on both IID and Non-IID settings. Empirically, they found that there was a large performance gap existing between the distributed and the centralized adversarial training, especially on the Non-IID data. Shah et al. (2021) designed a dynamic schedule for the local training to pursue a larger robustness under the constrained communication budget of Federated Learning. Hong et al. (2021) explored how to effectively propagate the adversarial robustness when only limited clients in Federated Learning have the sufficient computational budget to afford AT. Although previous works have investigated to solve the challenges about the constrained communication or computational budget, one critical issue that affects the performance when combined Adversarial Training with Federated Learning has received only few discussion. In this work, we consider such a basic issue (see Figure 1) in the straightforward combination of AT with Federated Learning and introduce our solution to this problem.

heading: 3 PRELIMINARY
text: In this section, we will briefly formalize the notations of Adversarial Training (Goodfellow et al., 2015; Madry et al., 2018) and Federated Learning as well as FedAvg (McMahan et al., 2017).

heading: 3.1 ADVERSARIAL TRAINING
text: Let (X , d∞) denote the input feature space X with the infinity distance metric d∞(x, x̃) = ‖x− x̃‖∞, and B [x] = {x̃ ∈ X | d∞(x, x̃) ≤ } be the closed ball of radius > 0 centered at x in X . Dataset S = {(xn, yn)}Nn=1, where xn ∈ X and yn ∈ Y = {0, 1, ..., C − 1}. The objective function of the standard adversarial training (AT) (Madry et al., 2018) is defined as follows,
min fθ∈F
1
N N∑ n=1 max x̃n∈B [xn] `(fθ(x̃n), yn), (1)
where x̃ is the most adversarial data within the -ball centered at x, fθ(·) : X → RC is a score function, ` : RC × Y → R is a composition of a base loss `B : ∆C−1 × Y → R (e.g., the CrossEntropy loss) and an inverse link function `L : RC → ∆C−1 (e.g., the Softmax). Here, ∆C−1 is the corresponding probability simplex that yields `(fθ(·), y) = `B(`L(fθ(·)), y). For the innermaximization of Eq. (1), the multi-step projected gradient descent (PGD) (Madry et al., 2018) is usually applied to find the most adversarial samples, which are then used for the outer-minimization.

heading: 3.2 FEDERATED LEARNING
text: Let Dk denotes a finite set of samples from the k-th client, and in each round, a set of datasets {Dk}Kk=1 from K clients are involved into the training. The objective of Federated Learning is to learn a machine learning model without any exchange of the training data between the clients and the server. The current popular strategy, namely FedAvg, is introduced by McMahan et al. (2017), where the clients collaboratively send the locally trained model weights θk to the server for the global average aggregation. Concretely, each client runs on a local copy of the global model (parameterized by θt in the t-th round) with its local data to optimize the objective like Eq. (1). Then, the server receives their updated model weights {θtk}Kk=1 of all clients and performs the following aggregation
θt+1 = 1
N K∑ k=1 Nkθ t k, (2)
where Nk denotes the number of the samples in Dk and N = ∑K k=1Nk. Then, the weights θ
t+1 for the global model will be sent back to each client for another lifecycle. After the sufficient rounds of such a periodic synchronization and aggregation, we expect the stationary point of Federated Learning will approximately approach to or have a small gap with that from the centralized counterpart.
4 α-WEIGHTED FEDERATED ADVERSARIAL TRAINING
In the following sections, we will first introduce our motivation for Federated Adversarial Training through analyzing the challenges brought by the straightforward combination. Then, we relax the inner-maximization problem of Adversarial Training as a lower bound by an α-weighted decomposition, and present the theoretical understanding of such a relaxation. Finally, we will present the α-Weighted Federated Adversarial Training and the corresponding analysis.

heading: 4.1 MOTIVATION
text: For the general FAT (Zizzo et al., 2020), in each round, the clients will receive the latest model from the server. Then, they conduct Adversarial Training with the local data and send their optimized model parameters to the server. The server will aggregate the parameters of all clients into a global model with FedAvg or other methods in Federated Learning. However, considering the characteristics of Federated Learning and Adversarial Training, it naturally brings the following critical challenge.
The inner-maximization can enlarge the data heterogeneity in Federated Learning. In terms of the combination of adversarial training and federated learning, the key point is that such adversarial examples generated by the inner-maximization exacerbate the heterogeneity, which induces the performance deterioration under the model aggregation methods. As the Figure 2 shows, the client A with large adversarial loss might make a dominating effect on the convergence of the global model, given the current model aggregation methods treat all client models indiscriminately. Thus, it might be better to selectively weight the client models in the model aggregation.
4.2 α-WEIGHTED RELAXATION OF DECOMPOSED ADVERSARIAL TRAINING
As previous analysis, the inner-maximization of Adversarial Training is not very compatible with Federated Learning due to the exacerbation of heterogeneity when combing them together. An idea to alleviate this problem is building a mediating function that keeps the original goal but is friendly to two paradigms. In this section, we present one possible solution to this intuition, which tries to find a relaxation of the inner-maximization in Adversarial Training. Formally, we decompose the inner-maximization objective into the independent K populations that corresponds to the K clients in Federated Learning, and relaxes it into a lower bound by the α-weighted mechanism as follows,
LAT = 1
N N∑ n=1 max x̃n∈B [xn] `(f(x̃n), yn) = K∑ k=1 Nk N
( 1
Nk Nk∑ n=1 max x̃n∈B [xn] `(f(x̃kn), y k n) ) ︸ ︷︷ ︸
Lk
≥ (1 + α) K̂∑ k=1 Nφ(k) N Lφ(k) + (1− α) K∑ k=K̂+1 Nφ(k) N Lφ(k) s.t. α ∈ [0, 1), K̂ ≤ K 2 . = Lα(K̂), (3)
where φ(·) is a function which maps the index to the original population sorted by {NkN Lk} in an ascending order. The following theorems provide us more analysis about the α-weighted relaxation.
Theorem 4.1. Lα(K̂) is monotonically decreasing w.r.t. both α and K̂, i.e., Lα1(K̂) < Lα2(K̂) if α1 > α2 and Lα(K̂1) < Lα(K̂2) if K̂1 > K̂2. Specifically, Lα(K̂) recovers L of adversarial training when α achieves 0, and Lα(K̂) relaxes L to a lower bound objective by increasing K̂ and α.
We can flexibly emphasize the importance of partial populations by setting the proper hyperparameters, alleviating the evenly averaging of harsh heterogeneous updates in FedAvg (McMahan et al., 2017).
Theorem 4.2. Assume the loss function `(·, ·) in Eq. (3) satisfies the Lipschitzian smoothness condition w.r.t. the model parameter θ and the training sample x, and is λ-strongly concave for all x, and E [ ||∇Lα(K̂)−∇θ`(f(x̃), y)||22 ] ≤ δ2, where x̃ is the adversarial example. Then, after the
sufficient T -step optimization i.e., T ≥ L∆δ2 , for α-weighted relaxation of decomposed Adversarial Training with the constant stepsize √ ∆
LTσ2 in PGD, we have the following convergence property,
1
T T∑ t E [∣∣∣∣∣∣∣∣∇Lα(K̂)∣∣θt∣∣∣∣∣∣∣∣2 2 ] ≤ ( 1 + α 1 T ∑T t ξ (t) N )( 4L2θx λ + 4δ √ L∆ T ) , (4)
where L = Lθθ + LθxLxθλ defined by the Lipschitzian constraints, ∆ ≥ L α(K̂)|θ0 − infθ Lα(K̂) and ξ(t) = ∑K̂ k N (t) φ(k) − ∑K K̂+1N (t) φ(k) meaning the accumulative counting difference of the t-th step.
When α = 0, Eq.(3) recovers the original loss of Adversarial Training, and the first part in the RHS of Eq. (4) goes to 1 that recovers the convergence rate of Adversarial Training (Sinha et al., 2018). When α → 1, Eq.(3) becomes more biased, while simultaneously the straightforward benefit is that we can achieve a faster convergence in Eq. (4) if 1T ∑T t ξ (t) < 0, i.e., ( 1 + α 1 T ∑T t ξ (t) N ) < 1. Actually, this is possible when the sample number is approximately similar among all clients and the top-1 choice easily has −N < ξ(t) = N (t)φ(1) − ∑K k=2N (t) φ(k) < 0 in each optimization step. Especially, a larger α has a faster convergence in this case. Therefore, the α-weighted relaxation of decomposed Adversarial Training provides us a trade-off between maintaining the robustness from the standard robust training and achieving the faster convergence with some biased approximation by α and K̂.
4.3 α-WEIGHTED FEDERATED ADVERSARIAL TRAINING
Inspired by the previous analysis of α-weighted relaxation, we propose an α-Weighted Federated Adversarial Training to combine Adversarial Training and Federated Learning. The intuition is
Algorithm 1 α-Weighted Federated Adversarial Training (α-WFAT) Input: number of clients: K, number of communication rounds: T , number of client training epochs
per round: E, initial server’s model parameter: θ0, hyper-parameter for aggregation: α, number of enhanced clients: K̂;
Output: a globally robust model with parameter θT ; 1: for t = 0, . . . , T − 1 do 2: Clients: [ perform adversarial training] 3: for client k = 1, . . . , K do 4: θtk,Lk = AT(θtk, E) (Madry et al., 2018) 5: end for 6: Server: [ performs aggregation over weight updates] 7: Lall ← [N1N L1, N2 N L2, . . . , NK N LK ]
8: Lsorted ← Ascending_Sort(Lall) 9: ∀k, Pk ← (1 + α) · 1(NkN Lk ≤ Lsorted[K̂]) + (1− α) · 1( Nk N Lk > Lsorted[K̂]);
10: θt+1 = 1∑K k=1 PkNk ∑K k=1 PkNkθ t k; . α-weighted mechanism 11: end for
applying the α-weighted mechanism into the inner-maximization in FAT, formalized as follows,
minLα−WFAT = min fθ∈F 1∑K k PkNk K∑ k=1 PkNk
( 1
Nk Nk∑ n=1 max x̃n∈B [xn] `(fθ(x̃ k n), y k n) ) ︸ ︷︷ ︸
Lk
, (5)
where Pk = (1+α) ·1(NkN Lk ≤ Lsorted[K̂])+(1−α) ·1( Nk N Lk > Lsorted[K̂]) denotes the weight assigned to the k-th client based on the ascending sort of weighted client losses compared with the K̂-th one. We summarize the procedure of α-WFAT in Algorithm 1, which consists of multi-round iterations between the local training on the client side and the global aggregation on the server side.
Concretely, on the client side, after downloading the global model parameter from the server, each client will perform the Adversarial Training on its local data. At the same time, the client loss on the adversarial examples is also recorded, which acts as the soft-indicator of the local bias induced by the radical adversarial generation. Then, when the training steps reach to the condition, the client will upload its model parameter and the loss to the server. On the server side, after collecting the model parameters {θk}Kk=1 and the losses {Lk}Kk=1 of all clients, it will first sort the losses in an ascending order to find the top-K̂ clients. Based on that, the global model parameters will be aggregated by the α-weighted mechanism in which the model parameters of the top-K̂ clients are upweighted with (1 + α) and the remaining is downweighted by (1− α). For some atypical layers (Li et al., 2021b) e.g., BN, it is outside the scope of this paper and we keep the aggregation same as FedAvg. Note that, one interesting point in α-WFAT is the top-K̂ clients with the higher weights are not fixed, and they can be routing among all clients. In Figure 3, we trace this dynamic of α-WFAT in one experiment.
In the following, we provide the theoretical analysis of α-WFAT on the convergence in the context of Federated Learning (Li et al., 2019), which is slightly different from previous centralized counterpart. Theorem 4.3. Assume the loss function `(·, ·) in Eq. (5) is L-smooth and λ-strongly concave w.r.t. the model parameter θ, and the expected norm and the variance of the stochastic gradient in each client respectively satisfy E [ ||∇θ`(f(x̃k), yk)||22 ] ≤ ς2 and E [ ||∇θ`(f(x̃k), yk)−∇θLk||22 ] ≤ δ2k. Let κ = Lλ , γ = max{8κ,E} where E is the iteration number of the local Adversarial Training with the learning rate ηt = 2λ(γ+t) . Then, after the sufficient T -step communication rounds for α-WFAT, we have the following asymptotics to the optimal point,
E[Lα−WFAT ]− L∗ ≤ κ
γ + T − 1
( 2B
λ + λγ 2 E [ ||θ0 − θ∗||2 ]) , (6)
where L∗ is the minimum value of Lα−WFAT , θ∗ is the optimal model parameter, and
B = K∑ k
( P
(T ) k
1 + α ξ (T )
N
)2( Nk N δk )2 + 6L ( L∗ − K∑ k P (T ) k 1 + α ξ (T ) N Nk N L∗k ) + 8(E − 1)2ς2.
When α = 0, we have P (t)k = 1 and Eq.( 6) becomes the convergence rate of FedAvg on non-IID data (Li et al., 2019). Different from Theorem 4.2 that concludes in the centralized training setting, when α→ 1, the convergence is indefinite compared to the standard FAT, since the emerging terms in
B, i.e., ( P (T ) k
1+α ξ (T )
N
)2 and P (T ) k
1+α ξ (T )
N
, are acted as the scalar timing by the personalized variance bound
δ2k and the local optimum L∗k of each client. One possible case is when the optimization approaches to the optimal parameter θ∗, δ2k can be in a smaller scale relative to the scale of L∗k. In this case, the increment of the first term of B can be totally counteracted by the loss of second term of B so that in sum B becomes smaller. Then, we can have a tighter upper bound for α-WFAT in Eq.( 6) to achieve a faster convergence than FAT. The completed proof of Theorem 4.3 can refer to Appendix C and the experiments in the following section will confirm α-WFAT can reach to a more robust optimum.

heading: 5 EXPERIMENTS
text: In this section, we will provide an in-depth analysis of α-WFAT and empirically verify its efficiency compared with the current methods on a range of IID and non-IID datasets.

heading: 5.1 EXPERIMENTAL SETUP
text: Dataset. We conduct the experiments on three benchmark datasets, i.e., CIFAR-10, CIFAR100 (Krizhevsky, 2009) and SVHN (Netzer et al., 2011) for Federated Adversarial Training. For the IID scenario, we just randomly and evenly distribute the samples to each client. For the NonIID scenario, we follow McMahan et al. (2017); Shah et al. (2021) to partition the training data based on their labels. To be specific, a skew parameter s is utilized in the data partition introduced by Shah et al. (2021), which enables K clients to get a majority of the data samples from a subset of classes. We denote the set of all classes in a dataset as Y and create Yk by dividing all the class labels equally among K clients. Accordingly, we split the data across K clients that each client has (100 − (K − 1) × s)% of data for the class in Yk and s% of data in other split sets. The detailed training and evaluation settings can refer to Appendix E.

heading: 5.2 ABLATION STUDY
text: In this part, we conduct various experiments on CIFAR-10 to visualize the characteristics of α-WFAT.
Non-AT vs. AT. In the left two panels of Figure 4, we respectively apply our α-weighted mechanism to Federated Standard Training (α-WFST) and Federated Adversarial Training (α-WFAT) under K̂ = 1. We also consider both FedAvg and FedProx in this experiment to guarantee the universality. From the curves, we can see that α-WFST has the negative effect on the natural accuracy, while α-WFAT consistently improve the robust accuracy both based on FedAvg or FedProx. This indicates our α-weighted mechanism is tailored for the inner-maximization of Federated Adversarial Training instead of the outer-minimization considered by other federated optimization methods.
Impact of α and K̂. To study the effect of hyperparameter in α-WFAT, we conduct several ablation experiments to verify the model performance. Regarding the experiments of α, we set the client
number K = 5 and K̂ = 1 to upweight/downweight the client models in each communication rounds. The right middle panel of Figure 4 shows α ∈ (0, 1/6] can significantly improve the robust accuracy and the natural accuracy, while a larger αmight be inappropriate to the natural accuracy. Regarding the choice of K̂, we specially set K = 8 in this experiment to span the range of K̂ due to the constraint K̂ <= K/2. The right panel of Figure 4 presents the accuracy of α-WFAT with increasing K̂.
According to the plot, both the natural accuracy and the robust accuracy are improved even with larger K̂, which shows the effect of K̂ on the relaxation of inner-maximization in the proposed α-WFAT. In addition, we conduct the experiments about emphasizing/deemphasizing the client with smallest adversarial loss by adjusting the α, the results (in Appendix E.1) confirmed that the relaxation introduced by our α-WFAT is needed to improve performance, and is consistent with our previous analysis.
Different AT methods. In Table 1, we validate the combination of α-weighted mechanism and different Adversarial Training methods (i.e., TRADES (Zhang et al., 2019) and MART (Wang et al., 2020c)), where we switch different local Adversarial Training methods on the client side. Through the comparison with FAT, the results show that α-WFAT can consistently improve both the natural performance and the robust performance, and is general to the state-of-the-art Adversarial Training works under Federated Learning scenarios.

heading: 5.3 PERFORMANCE EVALUATION
text: In this section, we compare our α-WFAT with FAT on various benchmark datasets to demonstrate its effectiveness. Specifically, we validate both the non-IID and IID settings with three represenative Federated optimization methods i.e., FedAvg, FedProx and Scalffold. Besides, a centralized AT baseline is provided as the reference of the IID setting. Note that, there is no such a baseline for the non-IID setting, since the centralized case is one distribution which is incomparable and meaningless. Considering the sensitivity of data selection in Non-IID settings, we also report the results with Mean and Std values in Appendix E.2 after running experiments for multiple times.
According to Table 2 on CIFAR-10, we can find that α-WFAT significantly outperforms FAT on the Non-IID data in terms of both the natural accuracy (∼2%-6%) and the robust accuracy (∼2%-5%). For the IID data, our method acquires a similar improvement on the robust accuracy without the deterioration of the natural accuracy. The reason might be because even the data is IID, Adversarial Training can still drive the independently-initialized overparameterized network (Allen-Zhu & Li, 2020) on each client side towards at the robust overfitting of different directions, yielding the model heterogeneity. Thus, the proper relaxation to the inner-maximization makes Adversarial Training more compatible with Federated Learning. Another interesting observation is that Federated Adversarial Training shows better performance than centralized Adversarial Training in the case of the IID setting. This gain could be from the distributed training paradigm that helps Adversarial Training converge
to the more robust optimum by the divide-and-conquer mechanism. This might enlighten the more explore in Adversarial Training to improve the model robustness via Federated Learning.
On CIFAR-100 and SVHN, we can find the similar improvement in Table 2 as that of CIFAR-10 under three types of federated optimization methods. Nevertheless, α-WFAT only becomes superior in terms of the robust accuracy but comparable with FAT in terms of the natural accuracy especially on CIFAR-100. It indicates that the inner-maximization of Adversarial Training when combined with Federated Learning mainly affects the model robustness, and thus the α-weighted relaxation correspondingly helps the model converge to a more robust optimum.

heading: 6 CONCLUSION
text: In this work, we explore the performance deterioration in the straightforward combination of Adversarial Training with Federated Learning. To alleviate the potential radical optimization, we apply an α-weighted relaxation into Adversarial Training to relax the inner-maximization. Based on this α-weighted mechanism, we further propose α-Weighted Federated Adversarial Training (α-WFAT). We provide the theoretical analysis and empirical evidences to understand the proposed simple but effective method. The experimental results under different settings confirm the effectiveness of α-WFAT. Nevertheless, we only move a small step on the heterogeneous issue in the combination of two paradigms and more issues in their cross field could be further explored in the future.

heading: 7 ETHICS STATEMENT
text: This paper does not raise any ethics concerns. This study does not involve any human subjects, practices to data set releases, potentially harmful insights, methodologies and applications, potential conflicts of interest and sponsorship, discrimination/bias/fairness concerns, privacy and security issues, legal compliance, and research integrity issues.

heading: 8 REPRODUCIBILITY STATEMENT
text: To ensure the reproducibility of experimental results, we will provide a link for an anonymous repository about the source codes of this paper in the discussion forums.

heading: A PROOF OF EQ. (3) AND THEOREM 4.1
text: We proof the Eq. (3) and Theorem 4.1 in this section.
Recall the α-weighted relaxation for the inner-maximization objective decomposition with K independent populations as follows,
LAT = 1
N N∑ n=1 max x̃n∈B [xn] `(f(x̃n), yn) = K∑ k=1 Nk N
( 1
Nk Nk∑ n=1 max x̃n∈B [xn] `(f(x̃kn), y k n) ) ︸ ︷︷ ︸
Lk
≥ (1 + α) K̂∑ k=1 Nφ(k) N Lφ(k) + (1− α) K∑ k=K̂+1 Nφ(k) N Lφ(k) s.t. α ∈ [0, 1), K̂ ≤ K 2 . = Lα(K̂), (7)
where φ(·) is a function which maps the index to the original population group sorted by {NkN Lk} in an ascending order.
proof of Eq. (3). The deduction of the inequality in Eq. (7) can be formulated in the following. Given α ∈ [0, 1) and K̂ ≤ K2 with the population sorted by {
Nk N Lk} in an ascending order, we have∑K̂
k=1 Nφ(k) N Lφ(k) ≤ ∑K k=K̂+1 Nφ(k) N Lφ(k). Then, we have the following relationship by subtraction,
K̂∑ k=1 Nφ(k) N Lφ(k) + K∑ k=K̂+1 Nφ(k) N Lφ(k) − (1 + α) K̂∑ k=1 Nφ(k) N Lφ(k) − (1− α) K∑ k=K̂+1 Nφ(k) N Lφ(k)
= α ·  K∑ k=K̂+1 Nφ(k) N Lφ(k) − K̂∑ k=1 Nφ(k) N Lφ(k)  ≥ 0. (8)
proof of Theorem 4.1. It can be naturally proved by Eq. (8). If α1 > α2, then we have,
Lα1(K̂)− Lα2(K̂) = (1 + α1) K̂∑ k=1 Nφ(k) N Lφ(k) + (1− α1) K∑ k=K̂+1 Nφ(k) N Lφ(k)
−(1 + α2) K̂∑ k=1 Nφ(k) N Lφ(k) − (1− α2) K∑ k=K̂+1 Nφ(k) N Lφ(k) (9)
= (α1 − α2)  K̂∑ k=1 Nφ(k) N Lφ(k) − K∑ k=K̂+1 Nφ(k) N Lφ(k)  ≤ 0 Similarly, we can prove Lα(K̂1) < Lα(K̂2) if K̂1 > K̂2.

heading: B PROOF OF THEOREM 4.2
text: Based on the convergence of Adversarial Training (Sinha et al., 2018), we proof Theorem 4.2 in this section.
proof of Theorem 4.2. Let c : X ×X → R+ ∪ {∞}, where c(x, x0) is the “cost” for an adversary to perturb x0 to x. Let f(θ, x;x0) = `(θ;x)− γc(x, x0), noting that the gradient steps is preformed as
gt = ∇θf(θt, x̂;xt), where x̂ is an approximate maximizer of f(θ, x;xt) in x, and θt+1 = θt−µtgt. We assume µt ≤ 1Lψ in the rest of the proof, which is satisfied for the constant step size µ =
√ ∆L
LψTδ2
and T ≥ Lψ∆Lδ2 . By a Taylor expansion using the Lψ-smoothness of the objective Lk for k-th client, we have
Lk ∣∣ θt+1 ≤ Lk ∣∣ θt + 〈 ∇Lk ∣∣ θt , θt+1 − θt 〉 + Lψ 2 ||θt+1 − θt||22 (10)
= Lk ∣∣ θt − µt||∇Lk ∣∣ θt ||22 + Lψµ 2 2 ||gt||22 + µt 〈 ∇Lk ∣∣ θt ,∇Lk ∣∣ θt − gt 〉 = Lk ∣∣ θt − µt ( 1− 1
2 Lψµ
2 ) ||∇Lk ∣∣ θt ||22
+ µt(1− Lψµ) 〈 ∇Lk ∣∣ θt ,∇Lk ∣∣ θt − gt 〉 + Lψµ 2 2 ||gt −∇Lk ∣∣ θt ||22
Consider the function φγ(θ;x0) = supx∈Z f(θ, x;x0), we define the potentially biased errors ζt = gt −∇θφγ(θt;xt). Then we have the following relationship,
Lk ∣∣ θt+1 ≤Lk ∣∣ θt − µt ( 1− 1
2 Lψµ
2 ) ||∇Lk ∣∣ θt ||22 (11)
+ µt(1− Lψµ) 〈 ∇Lk ∣∣ θt ,∇Lk ∣∣ θt −∇θφγ(θt;xt) 〉 − µt(1− Lψµt) 〈 ∇Lk ∣∣ θt , ζt 〉 + Lψµ 2 2 ||∇θφγ(θt;xt) + ζ2 −∇Lk ∣∣ θt ||22
= Lk ∣∣ θt − µt ( 1− 1
2 Lψµ
2 ) ||∇Lk ∣∣ θt ||22
+ µt(1− Lψµ) 〈 ∇Lk ∣∣ θt ,∇Lk ∣∣ θt −∇θφγ(θt;xt) 〉 − µt(1− Lψµt) 〈 ∇Lk ∣∣ θt , ζt 〉
+ Lψµ
2 t
2
( ||ζt||22 + ||∇θφγ(θt;xt)−∇Lk ∣∣ θt ||22 + 2 〈 ∇θφγ(θt;xt)−∇Lk ∣∣ θt , ζt 〉) .
Since ±〈a, b〉 ≤ 12 ( ||a||22 + ||b||22 ) , we have
Lk ∣∣ θt+1 ≤Lk ∣∣ θt − µt
2 ||∇Lk
∣∣ θt ||22 + µt((1− Lψα)) 〈 ∇Lk ∣∣ θt ,∇Lk ∣∣ θt −∇θφγ(θt;xt) 〉 (12)
+ µt((1 + Lψµ))
2 ||ζ||22 + Lψµ2t ||∇θφγ(θt;xt)−∇Lk ∣∣ θt ||22.
Then, letting xt∗ = arg maxx f(θ t, x;xt), the error ζt satisfies,
||ζ||22 =||∇θφγ(θt;xt)−∇f(θ, x̂t;xt)||22 = ||∇θ`(θ, xt∗)−∇θ`(θ, x̂t)||22 (13)
≤ Lθx ||x̂t − xt∗||22 ≤ 2L2θx λ , (14)
where the final inequality utilize the λ = γ − Lxx strong-concavity of x 7→ f(θ, x;x0). For convenience, let ̂ =
2L2θx γ−Lxx . Taking conditional expectations in Eq. (12) and using E [∇θφγ(θ
t;xt)|θt] = ∇Lk ∣∣ θt , we have,
E [ Lk ∣∣ θt+1 − Lk ∣∣ θt |θt ] ≤ −µt
2 ||∇Lk
∣∣ θt ||22 + µt((1 + Lψµ))
2 ̂+ Lψµ
2 t ||∇θφγ(θt;xt)−∇Lk ∣∣ θt ||22
(15)
≤ −µt 2 ||∇Lk ∣∣ θt ||22 + µt̂+ Lψµ2t ||∇θφγ(θt;xt)−∇Lk ∣∣ θt ||22
Since µt ≤ 1Lψ , taking a fixed step size µ, we have,
E [ ||∇Lk ∣∣ θt ||22 ] − 2̂ ≤ 2 µ E [ Lk ∣∣ θt − Lk ∣∣ θt+1 ] + 2Lψµδ 2 (16)
Because E [ ||∇θφγ(θ;Z)−∇Lk ∣∣ ( θ)||22 ] ≤ δ2, summing over t, we have,
1
T T∑ t E [ ||∇Lk ∣∣ θt ||22 ] − 2̂ ≤ 2 µT (Lk ∣∣ θ0 − E[Lk ∣∣ θT ]) + 2Lψµδ 2
≤ 2∆ µT + 2Lψµδ 2 (17)
Since µ = √
Delta LψTδ2 , and λ = γ − Lxx, we can get the following result,
1
T T∑ t E [ ||∇Lk ∣∣ θt ||22 ] ≤ 4L 2 θx λ + 4δ
√ L∆
T (18)
Adopting our α-weighted relaxation, we have,
1
T T∑ t E [∣∣∣∣∣∣∣∣∇Lα(K̂)∣∣θt∣∣∣∣∣∣∣∣2 2 ]
= 1
T T∑ t E ∣∣∣∣∣∣∣∣(1 + α) K̂∑ k=1 N (t) φ(k) N ∇Lφ(k) ∣∣ θt + (1− α) K∑
k=K̂+1
N (t) φ(k)
N ∇Lφ(k)
∣∣ θt ∣∣∣∣∣∣∣∣2 2  ≤ 1 T T∑ t (1 + α) K̂∑ k=1 N (t) φ(k) N E [∣∣∣∣∣∣∣∣∇Lφ(k)∣∣θt∣∣∣∣∣∣∣∣2 2 ] + (1− α) K∑ k=K̂+1 N (t) φ(k) N E [∣∣∣∣∣∣∣∣∇Lφ(k)∣∣θt∣∣∣∣∣∣∣∣2 2
] ≤ 1 T T∑ t (1 + α) K̂∑ k=1 N (t) φ(k) N + (1− α) K∑ k=K̂+1 N (t) φ(k) N (4L2θx λ + 4δ √ L∆ T )
= 1
T T∑ t 1 + α∑K̂k=1N (t)φ(k) −∑Kk=K̂+1N (t)φ(k) N (4L2θx λ + 4δ √ L∆ T )
= ( 1 + α 1 T ∑T t ξ (t)
N
)( 4L2θx
λ + 4δ
√ L∆
T
) ,
(19)
where ξ(t) = ∑K̂ k=1N (t) φ(k) − ∑K k=K̂+1N (t) φ(k) to simplify the notations.

heading: C PROOF OF THEOREM 4.3
text: Based on the convergence of FedAvg (Li et al., 2019), we proof Theorem 4.3 in this section.
First, we make the following assumptions and present some useful lemmas. Specifically, we make the following assumptions. Assumption C.1 and C.2 are standard (typical examples are the `2-norm regularized linear regression, logistic regression, or softmax classifier). Assumption C.3 and C.4 have been made by the previous works (Zhang et al., 2013; Li et al., 2019). Assumption C.1. L1, . . . , LK are all L-smooth: for all v and w, Lk(v) ≤ Lk(w) + (v − w)T∇Lk(w) + L2 ||v − w|| 2 2. Assumption C.2. L1, . . . , LK are all λ-strongly convex: for all v and w, Lk(v) ≥ Lk(w) + (v − w)T∇Lk(w) + λ2 ||v − w|| 2 2.
Assumption C.3. Let ξkt be sampled from the k-th device’s local data uniformly at random. The variance of stochastic gradients in each device is bounded: E||∇Lk(wkt , ξkt )−∇Lk(wkt )||2 ≤ δ2k for k = 1, · · · ,K. Assumption C.4. The expected squared norm of stochastic gradients is uniformly bounded, i.e., E||∇Lk(wkt , ξkt )||2 ≤ ς2 for all k = 1, · · · ,K and t = 1, · · · , T − 1.
We use the following lemmas proved by Li et al. (2019). Let θkt denotes the model parameter maintained in the k-th client at t-th step, Θ represents an immediate result of one step SGD update from θkt . For convenience, we define Θ̄t = ∑K k=1 Nk N Θt, θ̄t == ∑K k=1 Nk N θt, ḡt = ∑K k=1 Nk N ∇Lk(θ k t )
and gt = ∑K k=1 Nk N ∇Lk(θ k t , ξ k t ). Therefore, Egt = ḡt. Lemma C.1 (Results of one step SGD). Assume Assumption C.1 and C.2. If ηt ≤ 14L , we have
E||Θ̄t+1 − θ∗||2 ≤ (1− ηtλ)E||θ̄t − θ∗||2 (20)
+ η2tE||gt − ḡt||2 + 6Lη2tΓ + 2E K∑ k=1 Nk N ||θ̄t − θkt ||2,
where Γ = L∗ − ∑K k Et Nk N L ∗ k ≥ 0 Lemma C.2 (Bounding the variance). Assume Assumption C.3. It follows that
E||gt − ḡt||2 ≤ K∑ k=1 ( Nk N )2 δ2k, (21)
Lemma C.3 (Bounding the divergence of θkt ). Assume Assumption C.4, that ηt is non-increasing and η ≤ 2ηt+E for all t > 0. It follows that
E [ K∑ k=1 Nk N ||θ̄t − θkt ||2 ] ≤ 4η2t (E − 1)2ς2 (22)
proof of Theorem 4.3. Let ∆t = E||θt − θ∗||2. From Lemma C.1, Lemma C.2 and Lemma C.3, it follows that
∆t+1 ≤ (1− ηtλ)∆t + η2tB, (23) where,
B = K∑ k
( P
(T ) k
1 + α ξ (T )
N
)2( Nk N δk )2 + 6L ( L∗ − K∑ k P (T ) k 1 + α ξ (T ) N Nk N L∗k ) + 8(E − 1)2ς2, (24)
and ( P (T ) k
1+α ξ (T )
N
)2 and P (T ) k
1+α ξ (T )
N , are acted as the scalar timing by the personalized variance bound δ2k and the local optimum L∗k of each client.
For a diminishing stepsize, ηt = βγ+t for some β > 1 λ and γ > 0 such that η1 ≤ min{ 1 λ , 1 4L} = 1 4L and ηt ≤ 2ηt+E . We will prove that ∆t ≤ νγ+t , where ν = max{ β2B βλ−1 , (γ + 1)∆1}. The above can be proved by induction. Firstly, the definition of ν ensures that it holds for t = 1. Assume the conclusion holds for some t, it follows that,
∆t+1 ≤ (1− ηtλ)∆t + η2tB
≤ (1− βλ t+ γ ) ν t+ γ +
β2B
(t+ γ)2
= t+ γ − 1 (t+ γ)2 ν + [ β2B (t+ γ)2 − βλ− 1 (t+ γ)2 ν]
≤ ν t+ γ + 1 .
(25)
Then by the L-smoothness of L(·),
E[Lt]− L∗ ≤ L
2 ∆t ≤
L
2
ν
γ + t (26)
Specifically, if we choose β = 2λ , γ = max{8 L λ , E} − 1 and denote κ = L λ , then ηt = 2 λ 1 γ+t . One can verify that the choice of ηt satisfies ηt ≤ 2ηt+E for t ≥ 1. Then we have
ν = max
{ β2B
βλ− 1 , (γ + 1)∆1
} ≤ β 2B
βλ− 1 + (γ + 1)∆1 ≤
4B λ2 + (γ + 1)∆1, (27)
and
E[Lt]− L∗ ≤ L
2
ν γ + t ≤ κ γ + t
( 2B
λ + λ(γ + 1) 2 ∆1
) (28)
D LEARNING FRAMEWORK AND ALGORITHM
Here, we provide an intuitive illustration of our proposed α-WFAT in Figure 5. Based on the αweighted mechanism, we provide a new flexible framework for the combination of adversarial training with federated learning. It is orthogonal to a variety of different adversarial training (Zhang et al., 2019; Wu et al., 2020; Zhang et al., 2020; Chen et al., 2020; 2021; Wang et al., 2020a; Jiang et al., 2020; Wang et al., 2020c) methods and federated learning algorithms (Li et al., 2018; Kairouz et al., 2019; Hong et al., 2021) which gain the adversarial robustness or alleviate the data heterogeneity on the client side, and can be simply but effectively combined with other approach.

heading: E EXPERIMENTAL DETAILS
text: Dataset. We conduct the experiments on three benchmark datasets, i.e., SVHN (Netzer et al., 2011), CIFAR-10 and CIFAR-100 (Krizhevsky, 2009) for federated adversarial training. For the IID scenario, we randomly distribute these datasets to each client. For simulating the Non-IID scenario, we follow McMahan et al. (2017); Shah et al. (2021) to distribute the training data based on their labels. To be specific, a skew parameter s is utilized in the data partition introduced by Shah et al. (2021), which enables K clients to get a majority of the data samples from a subset of classes. We denote the set of all classes in a dataset as Y and create Yk by dividing all the class labels equally among K clients. Accordingly, we split the data across K clients that each client has (100− (K − 1)× s)% of data for the class in Yk and s% of data in other split sets. In our experiments, we set s = 2 for simulating the Non-IID partition with 5 clients as Shah et al. (2021) recommended.
Training and Evaluation. In the experiments, we follow the previous works to leverage the same architectures, i.e., NIN (Lin et al., 2014) for CIFAR-10, ResNet-18 (He et al., 2016) for CIFAR-100 and Small CNN (Zhang et al., 2020) for SVHN. For the local training batch size, we set 32 for CIFAR-10, 128 for CIFAR-100 and SVHN. For the training schedule, SGD is adopted with 0.9 momentum for 100 communication rounds under 5 clients as in (Hong et al., 2021; Shah et al., 2021), and the weight decay = 0.0001. For adversarial training, we set the configurations of PGD respectively (Madry et al., 2018) for different datasets. On CIFAR-10/CIFAR-100, we set the perturbation bound = 8/255, the PGD step size 2/255 and set the PGD step number 10. On SVHN, we set the perturbation bound = 4/255, the PGD step size 1/255, and keep the same step number 10. Regarding the evaluation, the accuracy for the natural test data and that for the adversarial test data are computed
following Wang et al. (2019); Wu et al. (2020). Note that, the adversarial test data are generated by FGSM, PGD-20, C&W (Carlini & Wagner, 2017) attack with the same perturbation bound and step size as the training. All the adversarial generations have a random start, i.e, the uniformly random perturbation of [− , ] added to the natural data before attacking iterations. Besides, we also report the robustness under a stronger AutoAttack, termed as AA for simplicity.
As for our α-WFAT, different training tasks adopt different α-weighted mechanism considering different characteristic of local training data, we set α = 1/6 (i.e., 1+α1−α = 1.4) for the experiments on CIFAR-10, and α = 1/41 (i.e., 1+α1−α = 1.05) for the experiments on CIFAR-100 and α = 1/11 (i.e., 1+α1−α = 1.2) for the experiments on SVHN. As for FedProx, we set µ = 0.01 for each dataset and its α for α-weighted mechanism are 1/11, 1/41, and 1/11. As for Scaffold, the α adopted for previous datasets are 1/11, 1/101 and 1/11 respectively.
As for the choice of the hyper-parameter α, one useful way to set it might be progressively probing its effect in a value-growth manner. When the α is very small, the objective will approximately degenerate the original objective of FAT, so does the performance with no harm. Slightly enlarging α can improve the performance due to the benefit of the bias-variance trade-off, and then make a stop in one point where the performance becomes drop.
E.1 EMPHASIZE/DE-EMPHASIZE IN OUR α-WFAT
We conduct an empirical comparison between α-WFAT that emphasizes (relatively de-emphasize the other clients) the client model with the smallest loss and a contrary variant that de-emphasizes it (relatively emphasize the other clients) as follows. We find that de-emphasizing the client with smallest loss (relatively emphasize those with larger loss) consistently harm the model performance across these evaluations.
E.2 MEAN AND STD RESULTS OF THE NON-IID SETTINGS
Considering that the Non-IID results are sensitive to the selection of data in each client, we conduct our experiments on Non-IID settings for multiple times and conclude the results as follows. In summary, our α-WFAT can consistently improve the model robustness with comparable or even better natural performance than previous federated optimization methods.

Review summaries:
This paper tackles an important problem in federated adversarial training: robustness accuracy significantly drops at the later stage of training. The authors first raise their assumption for the cause of this phenomenon: Adversarial training amplifies the heterogeneity of data distributions across different clients, and overfitted local robustness can not well generalized to other clients. Based on this assumption, the authors proposed \alpha-weighted federated adversarial training, which essentially up-weights the model trained on benign distributions and down-weights those on harsh distributions when averaging them up at the cloud center. Results show the proposed method outperforms previous state-of-the-arts under different adversarial training and federated learning settings. 

This work studied the limitation of conventional Federated Adversarial Training approach, and proposed an \alpha-weighted relaxation for Adversarial Training in the federated learning setting. Then it proposed a novel \alpha-Weighted Federated Adversarial Training for minimizing a lower bound of the inner-maximization in Federated Adversarial Training. The performance of the proposed \alpha-Weighted Federated Adversarial Training were validated for both IID and Non-IID federated learning settings.

This paper introduces the alpha Weighted Federated Adversarial Training algorithm. The key of the idea is that in the aggregate step, the center prefers the local machine that yields smaller lost. Some theoretical results are delivered with numerical experiments. The paper claims that the alpha-weighted mechanism is tailored for the inner-maximization of Federated Adversarial Training, which is the rationale of the whole work.

The authors explore the adversarial robustness of federated learning. They claim that the inner-maximization optimization of AT can exacerbate the data heterogeneity among local clients. They propose an algorithm, $\alpha$--WFAT, which relaxes the inner-maximization of Adversarial Training into a lower bound friendly to Federated Learning.. The authors also experimentally establish that federated learning models are most susceptible to attacks when clients are using non-IID training sets. The experiments are performed over the CIFAR-10 , SVHN and CIFAR-100datasets.