$\alpha$-Weighted Federated Adversarial Training

1 INTRODUCTION . 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 ) . 2 RELATED WORK . 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 . 3 PRELIMINARY . 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 ) . 3.1 ADVERSARIAL TRAINING . 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 . 3.2 FEDERATED LEARNING . 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 .

Summary: 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.