$\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 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.