Acceleration of Federated Learning with Alleviated Forgetting in Local Training

1 INTRODUCTION . Federated learning ( FL ) has emerged as a paradigm to train a global machine learning model in a distributed manner while taking privacy concerns and data protection regulations into consideration by keeping data on clients ( Voigt & Von dem Bussche , 2017 ) . The main challenge faced in FL is how to reduce the communication costs ( in training ) without degrading the performance of the final resultant model , especially when the data on different clients are not independently and identically distributed ( non-i.i.d . ) ( Yuan & Ma , 2020 ; Wang et al. , 2020b ) . The most popular FL algorithm is FedAvg ( McMahan et al. , 2017a ) . In each training round of FedAvg , local training steps are separately performed at every client and the locally trained models are transferred to the server . Then , the server aggregates the local models into a global model by simply averaging their parameters . Although FedAvg succeeds in many applications , its training processes often diverge when the data are non-i.i.d. , also known as heterogeneous , across the clients ( Li et al. , 2019 ; Zhao et al. , 2018 ) . Several FL algorithms have been designed to improve FedAvg and tackle the heterogeneity issue mainly by reducing the difference between locally trained parameters ( Li et al. , 2018 ; Karimireddy et al. , 2020 ) or aggregating these parameters into different groups ( Wang et al. , 2020a ; Yurochkin et al. , 2019 ) . However , the performance of these methods is still far from satisfactory when a deep neural network architecture is employed ( Rothchild et al. , 2020 ; Li et al. , 2021 ) . On the other hand , recent work in the literature ( Geiping et al. , 2020 ) shows that the transmission of trained model parameters does not ensure the protection of privacy . Privacy-sensitive information can be recovered ∗Minlie Huang and Tao Jiang are the co-corresponding authors . with gradient inversion attacks ( Zhu et al. , 2019 ) . Differential privacy ( DP ) ( McMahan et al. , 2017b ; Abadi et al. , 2016 ) is one of the most widely used strategies to prevent the leakage of private information . However , when DP is incorporated into FL , the performance of the resulting model decays significantly ( Jayaraman & Evans , 2019 ) . We observe that when the data are non-i.i.d . across the clients , the locally trained models suffer from severe forgetting of the knowledge of previous training data at the other clients ( i.e. , the well-known catastrophic forgetting issue ) , perhaps due to the discrepancy between local data distributions and the global data distribution . As shown in Figure 1 and supplementary Figure C.1 , this forgetting issue leads to a large increase in the loss concerning these training data under the non-i.i.d . setting , thereby slowing down the convergence speed . In this work , we propose a novel algorithm , FedReg , that reduces the communication costs in training by alleviating the catastrophic forgetting issue in the local training stage . FedReg reduces knowledge forgetting by regularizing the locally trained parameters with generated pseudo data , which are obtained by using modified local data to encode the knowledge of previous training data as learned by the global model . The potential conflict with the knowledge in the local data introduced by the pseudo data is dampened by the use of perturbed data , which are generated by making small perturbations to the local data , whose predictive values they help ensure . The generation of the pseudo data and perturbed data only relies on the global model received from the server and the local data on the current client . Thus , compared with FedAvg , no extra communication costs concerning the data on the other clients are incurred . Our extensive experiments demonstrate the superiority of FedReg in accelerating the FL training process , especially when the neural network architecture is deep and the clients ’ data are extremely heterogeneous . Furthermore , the pseudo data can be further utilized in classification problems to defend against gradient inversion attacks . More specifically , we show that with similar degree of protection of private information , the degradation in the performance of the global model learned by our method is much less than that learned by using FedAvg combined with DP . Our contributions . We demonstrate that when the data across clients are non-i.i.d. , catastrophic forgetting in the local training stage is an important factor that slows down the FL training process . We therefore propose a novel algorithm , FedReg , that accelerates FL by alleviating catastrophic forgetting with generated pseudo data . We also perform extensive experiments to establish the superiority of FedReg . We further show that in classification problems , the pseudo data generated in FedReg can help protect private information from gradient inversion attacks with much less impact on the performance of the resultant model compared with DP . 2 RELATED WORK . 2.1 FEDERATED LEARNING . FL is proposed to address privacy concerns in a distributed learning environment . In the FL paradigm , a client i keeps its local data Di on the client machine during the whole training process so that the server and other clients can not directly access Di . The objective is to find the parameter θ∗ that minimizes the loss on the global data ∪i∈CDi , where C is the set of clients , i.e. , θ∗ = argmin θ 1∑ i∈C |Di| ∑ d∈∪i∈CDi Lθ ( d ) , ( 1 ) In this equation , Lθ is the loss function with parameters θ , which can be cross-entropy or in some other custom-defined form . FedAvg ( McMahan et al. , 2017a ) and FedProx ( Li et al. , 2018 ) are the most widely used FL algorithms . In FedAvg , in a training round t , the server sends the initial parameters θ ( t−1 ) to a set of sampled clients C ( t ) . Then , the parameters are updated independently for S epochs on each of these clients to minimize the loss on the local data , and the locally trained parameters θ ( t , i ) are then sent to the server . The server aggregates the parameters by simply averaging over them and obtains the parameters θ ( t ) , i.e. , θ ( t ) = 1K ∑ i∈C ( t ) θ ( t , i ) , where K is the number of sampled clients in round t. FedAvg is unstable and diverge when the data are non-i.i.d . across different clients ( Li et al. , 2018 ) . FedProx stabilizes FedAvg by including a proximal term in the loss function to limit the distance between θ ( t , i ) and θ ( t−1 ) . Although FedProx provides a theoretical proof of convergence , empirically it fails to achieve good performances when deep neural network architectures are used ( Li et al. , 2021 ) . SCAFFOLD ( Karimireddy et al. , 2020 ) assumes that the heterogeneity of data leads to a client-drift in gradients and correlates the drift with gradient correction terms . As clients need to send extra information to the server , SCAFFOLD increases the risk of privacy leakage and doubles the communication burden compared with FedAvg . Furthermore , the accuracy of the correction terms is highly correlated with the training history of the client . Thus , the performance of SCAFFOLD decreases significantly when the number of clients is large , in which case each client is only sampled for few times and the estimation of the gradient correction terms is often inaccurate . FedCurv ( Shoham et al. , 2019 ) aims to tackle the forgetting issue on non-i.i.d . data with elastic weight consolidation ( Kirkpatrick et al. , 2017 ) . To achieve this , FedCurv needs to transfer Fisher information between the server and clients , which increases the communication costs to 2.5 times compared with FedAvg . Multiple methods ( Luo et al. , 2021 ; Hao et al. , 2021 ; Goetz & Tewari , 2020 ) have also tried to introduce synthesized data to help reduce the effect of heterogeneity , but they either rely on some specific architectures of neural networks ( Luo et al. , 2021 ) such as batch normalization ( Ioffe & Szegedy , 2015 ) or require the synthesized data to be shared with the server ( Hao et al. , 2021 ; Goetz & Tewari , 2020 ) , which contradicts the privacy protection objectives of FL . 2.2 CATASTROPHIC FORGETTING . Catastrophic forgetting occurs specifically when the neural network is trained sequentially on multiple tasks . In this case , the optimal parameters for the current task might perform poorly on the objectives of previous tasks . Many algorithms have been proposed to alleviate the forgetting issue . Memory-based methods ( Parisi et al. , 2019 ) have achieved excellent performances in accommodating new knowledge while retaining previously learned experience . Such memory-based methods , such as gradient episodic memory ( GEM ) ( Lopez-Paz & Ranzato , 2017 ) and averaged gradient episodic memory ( A-GEM ) ( Chaudhry et al. , 2018 ) , store a subset of data from previous tasks and replay the memorized data when training on the current task . For instance , A-GEM treats the losses on the episodic memories of previous tasks as inequality constraints in optimizing the objectives of current tasks and changes the model updates from the plain gradient g to g − wgref , where gref is the gradient computed from the loss on the memorized data and w is a non-negative weight . Unfortunately , these memory-based techniques are not suitable for FL due to privacy concerns . 2.3 GRADIENT INVERSION ATTACKS . FL provides a privacy guarantee by keeping the users ’ data on individual client machines locally and only sharing the model updates . However , recent research has shown that data information can be recovered from the parameter updates ( Geiping et al. , 2020 ; Zhu et al. , 2019 ) in the FedAvg framework by simply finding data with updates similar to the values returned from the client . DP , which avoids privacy leakage by adding noise to training data ( Sun et al. , 2019 ) or model updates ( Abadi et al. , 2016 ; McMahan et al. , 2017b ) , is the most widely used strategy to protect private information . When adding noise to model updates , such as differentially private SGD ( DPSGD ) ( Abadi et al. , 2016 ) , the noise level in DP is controlled by the gradient norm bound C and the noise scale σ. DP often causes a large performance degradation in the resultant model . 3 METHOD . Our main challenge is how to alleviate the forgetting of previously learned knowledge at each client without having to access data at the other clients in the local training stage . For any data set D , we denote the set of data near D within the Euclidean distance of δ as N ( D , δ ) = ∪d= ( x , y ) ∈D { ( x′ , y′ ) |0 < ∥x − x′∥ ≤ δ } . We assume that in round t on client i , ( 1 ) for all data d− ∈ ∪j∈C/ { i } Dj local to the other clients , the global model with parameter θ ( t−1 ) has a better feature representation than the local model with parameter θ ( t , i ) ; and ( 2 ) for any c > 0 , ∃ϵ , δ > 0 such that if ∀d′ = ( x′ , y′ ) ∈ N ( Di , δ ) , Lθ ( t , i ) ( ( x′ , fθ ( t−1 ) ( x′ ) ) ) − Lθ ( t−1 ) ( ( x′ , fθ ( t−1 ) ( x′ ) ) ) ≤ ϵ , then ∀d− = ( x− , y− ) ∈ ∪j∈C/ { i } Dj , Lθ ( t , i ) ( ( x− , fθ ( t−1 ) ( x− ) ) ) − Lθ ( t−1 ) ( ( x− , fθ ( t−1 ) ( x− ) ) ) ≤ c , where fθ is the prediction function with parameters θ . The assumption ( 2 ) guarantees that , in the local training stage , the range of changes in the predicted values of previous training data at the other clients can be controlled by constraining the changes in the predicted values of the local data near Di . Based on these assumptions , we first generate pseudo data and then alleviate the catastrophic forgetting issue by regularizing the locally trained parameters with the loss on the pseudo data . In other words , we hope that the pseudo data would achieve the same effect as the previous training data in constraining the local model . Note that FedReg does not assume any particular form of the loss function , but the method of modifying gradients to enhance privacy protection ( to be discussed below in Section 3.3 ) is currently only applicable to classification problems . The pseudo code of FedReg is provided in Appendix D .

Summary: This paper aims to alleviate forgetting in federated learning on non iid data. The method proposed FedReg focuses on regularizing locally trained parameters with the loss on generated pseudo data, which are based on adversarial examples of the global model in the previous step and the adversarial examples of the local model at current step. This paper experiments on several small scale dataset with large number of clients. It shows improved performance and less forgetting during the training.