ALLEVIATED FORGETTING IN LOCAL TRAINING
Year: 2022

Abstract:
Federated learning (FL) enables distributed optimization of machine learning models while protecting privacy by independently training local models on each client and then aggregating parameters on a central server, thereby producing an effective global model. Although a variety of FL algorithms have been proposed, their training efficiency remains low when the data are not independently and identically distributed (non-i.i.d.) across different clients. We observe that the slow convergence rates of the existing methods are (at least partially) caused by the catastrophic forgetting issue during the local training stage on each individual client, which leads to a large increase in the loss function concerning the previous training data at the other clients. Here, we propose FedReg, an algorithm to accelerate FL with alleviated knowledge forgetting in the local training stage by regularizing locally trained parameters with the loss on generated pseudo data, which encode the knowledge of previous training data learned by the global model. Our comprehensive experiments demonstrate that FedReg not only significantly improves the convergence rate of FL, especially when the neural network architecture is deep and the clients’ data are extremely non-i.i.d., but is also able to protect privacy better in classification problems and more robust against gradient inversion attacks. The code is available at: https://github.com/Zoesgithub/FedReg.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK

heading: 2.1 FEDERATED LEARNING
text: 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 cannot 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.

heading: 2.2 CATASTROPHIC FORGETTING
text: 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.

heading: 2.3 GRADIENT INVERSION ATTACKS
text: 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.

heading: 3 METHOD
text: 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.

heading: 3.1 GENERATION OF PSEUDO DATA AND PERTURBED DATA
text: With the above assumptions, the catastrophic forgetting issue can be alleviated by constraining the prediction values for data in N(D, δ). However, such a constraint is too strong and would hinder the learning of new knowledge from the local data. Moreover, it is also computationally inefficient to enumerate all such data. Observe that for a data point d′ = (x′,y′) ∈ N (Di, δ), after locally training θ(t,i) with Di, if fθ(t,i) (x′) = fθ(t−1) (x′) then the constraint on d′ will not contribute to addressing the forgetting issue. Therefore, to efficiently regularize the locally trained parameters, only data points with large changes in the predicted values should be used to alleviate knowledge forgetting. Such data points are usually within a small distance to the data points in Di, but their predicted values given by θ(t−1) could be very different from the labels in Di. Data satisfying these conditions can be easily obtained with the fast gradient sign method (Goodfellow et al., 2014) and we denote the data generated in this way as pseudo data Dsi . More precisely, each data point (xs,ys) ∈ Dsi is iteratively generated from a data point (x,y) ∈ Di as follows:
xs = xsE , y s = fθ(t−1) (x s) , (2)
xsj = x s j−1 + ηssign ( ∇xsj−1Lθ(t−1) (( xsj−1,y ))) , j = 1, ..., E, xs0 = x,
where ηs is the step size and E the number of steps. Despite the relaxation in constraints achieved by the pseudo data generated above, some information conflicting with the knowledge embedded in the local data Di could possibly be introduced in Dsi due to the inaccuracy of the global model during the training process. To further eliminate such conflicting information, perturbed data Dpi = {(xp,yp)} with small perturbations on Di are iteratively generated as follows:
xp = xpE , y p = y, xpj = x p j−1 + ηpsign ( ∇xpj−1Lθ(t−1) (( xpj−1,y ))) , j = 1, ..., E, xp0 = x,
(3) where the step size ηp satisfies ηp ≪ ηs in order to make sure that the perturbed data is much closer to the local data than the pseudo data. In our experiments, we take ηp = 0.01ηs and E = 10 to reduce the complexity of tuning hyper-parameters.

heading: 3.2 ALLEVIATION OF CATASTROPHIC FORGETTING
text: With the pseudo data Dsi and perturbed data D p i generated above, we regularize θ (t,i) by requiring∑ ds∈Dsi Lθ(t,i) (d s) ≤ ∑ ds∈Dsi Lθ(t−1) (d s) , (4)
∑ dp∈Dpi Lθ(t,i) (d p) ≤ ∑ dp∈Dpi Lθ(t−1) (d p) , (5)
where constraint (4) alleviates the catastrophic forgetting issue and constraint (5) eliminates conflicting information introduced in (4). Constraint (5) also helps improve the robustness of the resultant model (Madry et al., 2018). Expanding both sides of (4) and of (5) respectively with θ (t,i) β = 0.5 ( θ(t,i) + θ(t−1) ) so that the second-order term can be eliminated and ignoring higherorder terms (see Appendix A.1 for the details), we obtain( θ(t−1) − θ(t,i) )T gβ (Dsi ) ≥ 0, gβ (Dsi ) = ∇θ(t,i)β 1 |Dsi | ∑
ds∈Dsi
L θ (t,i) β
(ds) (6)
( θ(t−1) − θ(t,i) )T gβ (Dpi ) ≥ 0, gβ (D
p i ) = ∇θ(t,i)β
1 |Dpi | ∑
dp∈Dpi
L θ (t,i) β
(dp) (7)
Due to the existence of potentially conflicting information in Dsi and D p i , directly finding θ (t,i) to satisfy the above two inequalities is mathematically ill-defined. Besides, given the complex architectures of deep neural networks, it is computationally too challenging to solve either of the above inequalities. Therefore, we approximate the optimal parameters θ(t,i) ∗ by solving the following constrained optimization problems:
θ(t,i) ′ = argminθ ∥θ − θ(t,i)∥2 s.t. ( θ(t−1) − θ )T gβ (Dsi ) ≥ 0, (8)
θ(t,i) ∗ = argminθ ∥θ − θ(t,i) ′∥2 s.t. ( θ(t−1) − θ )T gβ(Dpi ) ≥ 0. (9)
By solving (8) and (9) (see Appendix A.2 for the detailed derivation), we obtain
θ(t,i) ∗ = θ(t,i) − wsgβ (Dsi )− wpgβ(D p i ), (10)
ws = max
( (θ(t,i)−θ(t−1)) T gβ(Dsi )
gβ(Dsi ) T gβ(Dsi )
, 0 ) , wp = max ( (θ(t,i)−wsgβ(Dsi )−θ (t−1)) T gβ(Dpi )
gβ(Dpi ) T gβ(Dpi )
, 0 ) .
Based on the generation process of the pseudo data, it is easy to note that ∑
ds∈Dsi Lθ(t−1)(d
s) = minθ ∑
ds∈Dsi Lθ(d s), and thus the inequality sign in (4) can in fact be replaced by an equality sign. In practice, however, we have observed that the training process would be more stable and result in non-negative ws and wp with the inequality formulation, probably due to errors introduced in the above approximation. Finally, to prevent the gradient from converging to zero in a very small number of steps, slowly-updated parameters (Zhang et al., 2019b) θ(t,i)γ = γθ(t,i) + (1− γ)θ(t−1) with γ ∈ [0, 1] are used to compute the gradient for Di, i.e., gγ(Di) = ∇θ(t,i)γ 1 |Di| ∑ d∈Di Lθ(t,i)γ (d).
In summary, taking into acount the above improvements, the local parameters are updated in each training step as:
θ(t,i) ← θ(t,i) − ηθgγ (Di) , θ(t,i) ← θ(t,i) − wsgβ (Dsi )− wpgβ (D p i ) . (11)

heading: 3.3 MODIFICATION OF GRADIENTS TO PROTECT PRIVACY
text: We further notice that in classification problems, the pseudo data can be used to modify the gradient gγ(Di) to enhance the protection of privacy. If we denote Dsi ′ = {(x, 1n ∑n j=1 e (j))|(x,y) ∈ Dsi} for an n-classification problem, where e(j) is the standard basis vector with a 1 at position j, then the data points in Dsi ′ are similar to those in Di, but they may have totally different labels. Hence,
it is reasonable to assume that gγ(Di) and gγ(Dsi ′) contain similar semantic information but different classification information. Since in the training of FL, only the classification information contributes to the model performance, removing the semantic information in gγ(Di) will enhance the privacy protection capability of FL without severely degrading the performance of the resultant model. Based on this intuition, we presume that the components of gγ(Di) in the same (or roughly the same) direction of gγ(Dsi
′) contain the semantic information and the other (orthogonal) components contain the classification information. Thus, we compute a modified gradient (MG) to enhance the protection of privacy as
g̃γ (Di) = gγ (Di)− vgγ ( Dsi ′) , v = gγ (Di)T gγ (Dsi ′) gγ ( Dsi
′)T gγ (Dsi ′) (12) We refer to the version of FedReg in which the gγ (Di) term in (11) is replaced by g̃γ (Di) as FedReg with MG.

heading: 4 EXPERIMENTS
text: We compare the performance of FedFeg against the above-mentioned baseline FL algorithms including FedAvg, FedProx, FedCurv and SCAFFOLD as well as the well-known SGD algorithm (Bottou, 2012). Note that SGD corresponds to a special case of FedAvg where the local training stage consists of only one step and all data on a client are used in a large single batch.

heading: 4.1 DATA PREPARATION
text: FedReg is evaluated on MNIST (Deng, 2012), EMNIST (Cohen et al., 2017), CIFAR-10 and CIFAR100 (Krizhevsky et al., 2009), and CT images of COVID-19 (He, 2020). To simulate a scenario for FL, the training data in each dataset are split into multiple clients in different ways and the performance of the trained model is evaluated on the test data. The data preparation steps for each dataset are described below and more experimental details are provided in Appendix B.
MNIST The training data are split into 5,000 clients. Each client has images either from only one class, named MNIST (one class), or from two classes, named MNIST (two classes). The number of images per client follows a power law distribution (Li et al., 2018). A 5-layer neural network with three layers of convolutional neural networks (CNN) and two layers of full connections (FC) is used.
EMNIST The training data of the digits in EMNIST are split into 10,000 clients. Each client owns 24 images from the same class. The same model architecture used on MNIST is used here.
CIFAR-10 The training data are split into 10,000 clients. Each client owns 5 images from random classes, named CIFAR-10 (uniform), or from only one class, named CIFAR-10 (one class). A ResNet-9 network with the Fixup initialization (Zhang et al., 2019a) is trained from scratch.
CIFAR-100 The training data are split into 50,000 clients and each client owns 1 image. A ResNet9 network with the Fixup initialization is trained from scratch, named CIFAR-100 (ResNet-9), and a pre-trained transformer (Dosovitskiy et al., 2020) is fine-tuned, named CIFAR-100 (Transformer).
CT images related to COVID-19 The dataset contains 349 CT images of COVID-19 patients with source information and 397 non-COVID-19 CT images without source information. The data are split into training and test data in the same way provided in (He, 2020). Then, the CT images of COVID-19 patients in the training data are split to make sure that the images from the same source are assigned to the same client, and 32 clients are obtained in this way. The non-COVID-19 images in the training data are distributed to each client following the proportion of COVID-19 images. A DenseNet-121 network (Huang et al., 2017) and a 10-layer neural network with 9 layers of CNNs and one layer of FC are trained on these CT images to distinguish the COVID-19 images from the non-COVID-19 ones.

heading: 4.2 COMPARISON OF CONVERGENCE RATES
text: The results on MNIST and EMNIST are shown in Table 1. FedReg required fewer communication rounds to converge compared with the baseline methods and achieved a higher final accuracy (i.e.,
accuracy of the respective models after finishing all training rounds) under all the experimental settings, especially on EMNIST, where the number of clients is twice compared to MNIST, showing clearly the superiority of FedReg in dealing with non-i.i.d. data when a shallow model is used and its robust scalability. FedCurv also aims to allieviate the forgetting issue, but its performance is not significantly better than FedProx. Note that due to its heavy reliance on the optimization history when estimating the correction term, in our experiments, SCAFFOLD performed worse than FedAvg. Table 2 illustrates that, on CIFAR-10, when the data are randomly distributed, FedReg significantly outperformed all the baseline methods in both convergence speed (measured in communication rounds) and final model accuracy, demonstrating the superiority of FedReg in training deep neural networks. When data across clients become more heterogeneous and the images at each client are from the same class, the multiple local training steps in the baseline algorithms (other than SGD) failed to improve the performance of the locally trained model, resulting in a worse performance than that of SGD, while FedReg was still able to save about a half of the communication costs compared with SGD. A similar conclusion can be drawn when training the ResNet-9 network from scratch and fine-tuning the transformer model on CIFAR-100, as shown in Table 3, indicating the robustness of FedReg on various datasets and model architectures. The results in distinguishing COVID-19 CT images from non-COVID-19 images on a DenseNet-121 network, as shown in Table 3, further exhibits the superiority of FedReg in medical image processing applications when the model architecture is deep and complex. To assess the advantage of FedReg on real-world data, it is compared with the other methods on a popular real-world dataset Landmarks-User-160k (Hsu et al., 2020; Lai et al., 2021), where the data at each client consists of many classes, and as shown in supplementary Table C.1, FedReg still outperformed the baseline methods. The wall-clock time usages of the methods in each local training round on MNIST and EMNIST are given in supplementary Table C.2. Note that since FedReg performs additional computation in each round, its local training requires a bit more time.

heading: 4.3 ALLEVIATION OF CATASTROPHIC FORGETTING
text: To prove that FedReg indeed alleviates the catastrophic forgetting, increases of the loss values concerning previous training data at the other clients are compared between FedReg and FedAvg. As
shown in Figure 2, the increases of the loss are significantly lower in FedReg than those in FedAvg, suggesting that although FedAvg and FedReg both forget some of the learned knowledge, the forgetting issue is less severe in FedReg. To further explore the role of the generated pseudo data in the alleviation of forgetting, the values of the empirical Fisher information (Ly et al., 2017) in the pseudo data and previous training data are compared. As shown in Figure 2, the values of the pseudo data and previous training data are significantly correlated. Following the Laplace approximation (MacKay, 1992; Kirkpatrick et al., 2017), if we approximate the posterior distribution of the optimal parameters on a dataset by a normal distribution centered at θ(t−1) with a diagonal precision matrix, then the precision matrix can be approximated by the empirical Fisher information. The similarity in the empirical Fisher information suggests that a model with high performance on the pseudo data has a relatively high probability to perform well on previous training data. Therefore, regularizing the parameters with the pseudo data can help reduce the performance degradation on previous training data and alleviate the knowledge forgetting issue.

heading: 4.4 THE IMPACT OF STEP SIZE IN GENERATING PSEUDO DATA
text: The impact of the hyper-parameter ηs, which determines the distance between the pseudo data and the local data, is explored here. Intuitively, when ηs is too large, the distance between the generated pseudo data and the local data will be too large to effectively regularize the model parameters. On the other hand, when ηs is too small, the conflicting information from an inaccurate global model will slow down the convergence speed. This intuition is empirically confirmed by the experiments presented in supplementary Figure C.2. On the EMNIST and CIFAR-10 (one class) datasets, when ηs is too small, it takes more communication rounds to achieve the same accuracy. As ηs increases, the number of communication rounds decreases to a minimum and then bounces back, indicating that a less effective regularization on the model parameters in the local training stage has been reached.

heading: 4.5 PROTECTION OF PRIVACY
text: The protection of private information is one of the most important concerns in the FL paradigm, especially in medical applications. To illustrate the privacy security of FedReg with MG, gradient inversion attacks are used to recover information from updated parameters in classification problems. Under a comparable degree of privacy protection, the resultant model performance is compared between FedReg with MG and the baseline methods with DPSGD (Abadi et al., 2016). Using gradient inversion attacks, the quality of images recovered from the updated parameters in a local step is examined. Note that when only one local training step is performed in each round, the updated parameters in FedAvg, FedProx, FedCurv, and SGD are all the same, leading to the same images recovered from these methods. We do not include SCAFFOLD here since its performance was significantly inferior to the other methods in the above experiments. As shown in supplementary Figures C.3 and C.4, for some images sampled from the EMNIST, CIFAR-10 and COVID-19 CT datasets, the quality of the images recovered from FedReg with MG is significantly worse than those recovered from the baseline methods, exhibiting a better privacy protection capability of FedReg with MG. To compare the impact of privacy protection on model performance, we tune the noise level in DPSGD to generate images with similar quality as measured by PSNR (Hore & Ziou, 2010) and compare the performance of the resultant models. The results are shown in Table 4. It can be seen that the protection of private information with DP costs a large performance degradation in the baseline methods, while the performance of FedReg with MG is degraded much less when maintaining a similar privacy protection level. Note that the CT images used in the above experiments contain some privacy-sensitive information, including the date of birth of the patient and the date when the image was taken. In the basic FedAvg, when a 10-layer neural network with 9 layers of CNNs and one layer of FC is used, such information can be easily recovered with gradient inversion attacks. When DP is applied, although such time information could be protected, the resultant model suffers from severe performance degradation, as shown in Table 4. In contrast, FedReg with MG is able to protect the sensitive time information with only mild model performance decay.

heading: 5 CONCLUSIONS AND FUTURE WORK
text: In this work, we proposed a novel algorithm, FedReg, to accelerate the convergence speed of FL by alleviating the catastrophic forgetting issue in the local training stage. Pseudo data are generated to carry knowledge about previous training data learned by the global model without incurring extra communication costs or accessing data provided at the other clients. Our extensive experiments show that the generated pseudo data contain similar Fisher information as the previous training data at the other clients, and thus the forgetting issue can be alleviated by regularizing the locally trained parameters with the pseudo data. The pseudo data can also be used to defend against gradient inversion attacks in classification problems with only mild performance decay on the resultant model compared with DP. Although FedReg exhibits competitive performance in accelerating convergence speeds, its performance is influenced by hyper-parameters. Automatic tuning of the hyperparameters could make FedReg easier to use. Meanwhile, the question of how to modify gradients in FedReg to effectively protect privacy in regression problems remains to be explored. Moreover, a rigorous proof of the convergence rate of FedReg is of theoretical interest. We leave these questions to future work.

heading: 6 ACKNOWLEDGEMENT
text: This work has been supported in part by the US National Institute of Health grant 1R01NS125018, the National Key Research and Development Program of China grant 2018YFC0910404 and the Guoqiang Institute of Tsinghua University with grant no. 2019GQG1.

heading: A SUPPLEMENTARY METHOD
text: A.1 TAYLOR EXPANSION
Expanding both sides of the inequality (4) with θ(t,i)β = θ(t,i)+θ(t−1)
2 , we obtain∑ ds∈Dsi Lθ(t,i) (d s) = ∑ ds∈Dsi L θ (t,i) β (ds) + ( θ(t,i) − θ(t,i)β )T ∇ θ (t,i) β ∑ ds∈Dsi L θ (t,i) β (ds)
+ 12 ( θ(t,i) − θ(t,i)β )T ∇2
θ (t,i) β
(∑ ds∈Dsi L θ (t,i) β (ds) )( θ(t,i) − θ(t,i)β )
+O ( ∥θ(t,i)β − θ(t,i)∥3 ) ,∑
ds∈Dsi
Lθ(t−1) (d s) = ∑ ds∈Dsi L θ (t,i) β (ds) + ( θ(t−1) − θ(t,i)β )T ∇ θ (t,i) β ∑ ds∈Dsi L θ (t,i) β (ds)
+ 12 ( θ(t−1) − θ(t,i)β )T ∇2
θ (t,i) β
(∑ ds∈Dsi L θ (t,i) β (ds) )( θ(t−1) − θ(t,i)β )
+O ( ∥θ(t,i)β − θ(t−1)∥3 ) .
Substituting both sides of (4), dividing them by |Dsi | and ignoring all high order terms, then only the first order terms are left and the thus inequality (6) is obtained, i.e.,(
θ(t−1) − θ(t,i) )T ∇
θ (t,i) β
1 |Dsi | ∑
ds∈Dsi
L θ (t,i) β
(ds) ≥ 0.
The inequality (7) can be derived from (5) in a similar way.
A.2 SOLVING THE CONSTRAINT OPTIMIZATION
The constraint optimization problems in (8) and (9) can be solved by the Lagrange multiplier method. More specifically, for the problem (8), i.e.,
θ(t,i) ′ = argminθ ∥θ − θ(t,i)∥2 s.t. ( θ(t−1) − θ )T gβ(Dsi ) ≥ 0,
its Lagrangian can be written as L(θ, α) = ∥θ − θ(t,i)∥2 + α ( θ − θ(t−1) )T gβ (Dsi ) ,
where α ≥ 0. The dual problem is D(α) = min
θ L(θ, α).
Setting the derivatives of L(θ, α) w.r.t. θ to zero, the value of θ(t,i)′ minimizing L(θ, α) is obtained as
2 ( θ(t,i) ′ − θ(t,i) ) + αgβ (Dsi ) = 0,
θ(t,i) ′ = θ(t,i) − 12αgβ (D s i ) .
Thus, D(α) can be simplified as
D(α) = −1 4 α2gβ(Dsi )Tgβ(Dsi ) + α(θ(t,i) − θ(t−1))Tgβ (Dsi ) .
Solving α∗ = argmaxD(α), the value of α∗ is
α∗ = max
( 2 ( θ(t,i) − θ(t−1) )T gβ (Dsi )
gβ (Dsi ) T gβ (Dsi )
, 0 ) .
Thus, the value of θ(t,i) ′ is
θ(t,i) ′ = θ(t,i) −max
(( θ(t,i) − θ(t−1) )T gβ (Dsi )
gβ (Dsi ) T gβ (Dsi )
, 0 ) gβ (Dsi ) .
The constraint optimization problem 9 can be solved similarly.

heading: B EXPERIMENTAL DETAILS
text: In all the experiments, the learning rates for different methods are optimized by grid search, the optimal weight for the proximal term in FedProx is searched among {1.0, 0.1, 0.01, 0.001}, the optimal λ for FedCurv is searched among {0.001, 0.0001, 0.00001}, and the hyper-parameters (γ and ηs) in FedReg are optimized by grid search as well. The number of epochs in the local training stage is optimized in FedAvg and applied to the other methods.
B.1 COMPUTATION OF loss(t−1) AND loss(t)
Formally, the value of loss(t−1) for a client i ∈ C(t−1) is computed as
loss(t−1)(i) = 1 |Di| ∑ d∈Di Lθ(t−1)(d),
and the value of loss(t) is computed as
loss(t)(i) = 1 |Di| ∑ d∈Di
∑ j∈C(t)(Lθ(t,j)(d))
|C(t)| ,
where θ(t,j) is the parameters sent from client j in round t.
B.2 COMPARISON OF CONVERGENCE RATES
The hyper-parameters in the experiments are listed below:
• MNIST. Ten clients are sampled in each training round and 500 rounds are run in total with a batch size of 10 so that all training data are utilized once on average. The learning rate in all methods is 0.1. On MNIST (two classes), in the local training stage, the local data are processed in 40 epochs. The weight for the proximal term in FedProx is 0.001 and λ in FedCurv is 10−4. In FedReg, γ = 0.4 and ηs = 0.2. On MNIST (one classes), the local data are processed in 20 epochs. The weight for the proximal term in FedProx is 0.01 and λ in FedCurv is 10−4. In FedReg, γ = 0.3 and ηs = 0.2.
• EMNIST. 20 clients are sampled in each training round, and in the local training stage, the local data are processed in 20 epochs with a batch size of 24. 500 rounds are run in total so that all the training data are utilized once on average. The learning rate in SCAFFOLD is 0.1 and in other methods, it is 0.2. The weight for the proximal term in FedProx is 0.001 and λ in FedCurv is 10−4. In FedReg, γ = 0.4 and ηs = 0.05.
• CIFAR-10. Following the experimental settings in (Rothchild et al., 2020), 240 rounds are run in total with 100 clients sampled in each training round, and in the local training stage, the batch size is 5. On CIFAR-10 (uniform), the local data are processed in 30 epochs. The learning rate is 0.05 in FedAvg, FedCurv and FedProx, 0.1 in SGD and FedReg, and 0.01 in SCAFFOLD respectively. The weight for the proximal term in FedProx is 0.001 and λ in FedCurv is 10−5. In FedReg, γ = 0.5 and ηs = 0.1. On CIFAR-10 (one class), the local data are processed in 20 epochs. The learning rate is 0.05 in FedAvg, FedCurv, FedProx, SCAFFOLD and FedReg, and 0.1 in SGD, respectively. The weight for the proximal term in FedProx is 0.01 and λ in FedCurv is 10−3. In FedReg, γ = 0.5 and ηs = 0.1.
• CIFAR-100. When ResNet-9 is adopted, 1,200 rounds are run in total with 100 clients sampled in each training round, and in the local training stage, the local data are processed in 10 epochs with a batch size of 1. The learning rate is 0.05 in FedAvg, FedCurv and FedProx, and 0.1 in SGD and FedReg, respectively. The weight for the proximal term in FedProx is 0.01 and λ = 10−3 in FedCurv. In FedReg, γ = 0.25 and ηs = 0.1. When the pre-trained transformer is adopted, 100 rounds are run in total with 500 clients sampled in each training round, and in the local training stage, the local data are processed in 5 epochs with a batch size of 1. The learning rate is 0.1 in all methods. The weight for the proximal term in FedProx is 0.001 and λ = 10−5 in FedCurv. In FedReg, γ = 0.02 and ηs = 10−2. Note that as SCAFFOLD is too memory-consuming to run on CIFAR100 given our computational resources, the results of SCAFFOLD on CIFAR-100 are not reported.
• CT images related to COVID-19. In our experiments, ten rounds are run in total and in each round, 10 clients are sampled. In the local training stage, the local data are processed in 20 epochs with a batch-size of 10. The learning rate is 1×10−3 in FedAvg, FedCurv and FedProx, 5×10−4 in SCAFFOLD, 1×10−4 in SGD, and 5×10−3 in FedReg, respectively. The weight for the proximal term in FedProx is 1.0 and λ in FedCurv is 10−4. In FedReg, γ = 0.5 and ηs = 10−6. Note that group-normalization (Wu & He, 2018) is used in the DenseNet-121 network as the batch size is small in the local training stage.
• Landmarks-User-160k. The training data in this dataset has 164,172 images of 2,028 landmarks from 1,262 users, and the test data contains 19,526 images. Each user has images of 29 classes on average. A DenseNet-121 network is trained on this dataset. 2,000 rounds are run in total, and in each round 100 clients are sampled. In the local training stage, the local data is processed in 40 epochs. The learning rate is set as 0.1 in SGD, FedAvg, FedReg, and pFedGP (Achituve et al., 2021). As pFedGP is designed for personalized federated learning scenarios and trains an independent GP-tree for each client, to apply it to this dataset, a small global dataset is built by sampling 2,028 images from all training data, which covers all the classes contained in Landmarks-User-160k. In each round, the deep kernel, which is the DenseNet-121 network without the classification layer, is trained on each client and aggregated in the server. Then, a global GP-tree is trained on the sampled global sub-dataset to classify the images in the test data.
B.3 ALLEVIATION OF CATASTROPHIC FORGETTING
The hyper-parameters used here, except γ in FedReg, are the same as the values used in the previous section.
• Computation of loss increment. To fairly compare the forgetting issue in FedReg and FedAvg, γ is set to be 1.0 in FedReg, and in each round, the client parameters are initialized with the parameters obtained in FedReg to avoid a large discrepancy between the values of loss(t−1) computed in FedReg and FedAvg. For a client i ∈ C(t−1), the loss increase is computed as loss(t)(i)− loss(t−1)(i).
• The computation of Fisher information. Following (Kunstner et al., 2019), in round t of client i ∈ C(t), for the previous training data D(t−) = ∪j∈C(t−)Dj , where C(t−) = ∪s∈[t−10,t−1]C(s), the empirical Fisher information is computed as
fisherθ(t,i)(D(t−)) = 1 |D(t−)| ∑
(x,y)∈D(t−)
( ∇θ(t,i)yT log pθ(t,i) (y|x) )2 ,
Similarly for the pseudo data Dsi , the empirical Fisher information is computed as
fisherθ(t,i)(Dsi ) = 1 |Dsi | ∑
(xs,ys)∈Dsi
( ∇θ(t,i)ysT log pθ(t,i) (ys|xs) )2 .
Note that considering the generation process of the pseudo data, the parameter values θ(t,i) used to compute Fisher information are the same local model parameters after a local training step.
B.4 PROTECTION OF PRIVATE INFORMATION
The target of gradient inversion attacks in (Geiping et al., 2020) and (Zhu et al., 2019) is to find
d∗ = argmin d
Distance ( ∆θ ( θ(t),dtrue ) ,∆θ ( θ(t),d )) + wTV (d)
where ∆θ ( θ(t), · ) denotes the amounts (updates) that the parameters have changed from θ(t) with
the corresponding data, TV (d) the total variation of d, and dtrue a data point on the client. The distance function could be cosine distance (Geiping et al., 2020) or L2-distance (Zhu et al., 2019). Note that the method proposed in (Zhu et al., 2019) sets w = 0. To consider the worst case, in our attacking experiments, the parameters are only trained for one step with each data point, in
which case the values of ∆θ ( θ(t),d ) obtained in FedAvg, FedCurv, FedProx and SGD are exactly the same. 32,000 attacking iterations are used in all experiments and both distance functions (i.e., cosine and L2) are considered. Adam (Kingma & Ba, 2014) optimizer is used, which performs better than L-BFGS empirically. The learning rate and the weight for the total variation are optimized in (the basic) FedAvg and applied to the other methods. The learning rate decaying scheme follows the method provided in (Geiping et al., 2020).
More details concerning each dataset are discussed below.
• EMNIST. The images to be recovered are randomly sampled. The model architecture and hyper-parameters used in FedAvg, FedProx, FedCurv, SGD, and FedReg are the same as the ones used in section B.2. In the baseline methods with DPSGD, C = 1.0 and σ = 0.01, and other hyper-parameters are the same as the ones used in the corresponding baseline methods (without DPSGD). In FedReg with MG, the learning rate and the value of γ are the same as the ones used in FedReg, and ηs = 0.03. In the attacking algorithm, the initial learning rate is 1.0, the distance function is L2-distance, and the weight for the total variation term is 10−8.
• CIFAR-10. The images to be recovered are randomly sampled. The model architecture and hyper-parameters used in FedAvg, FedProx, FedCurv, SGD, and FedReg are the same as the ones used in section B.2. In the baseline methods with DPSGD, C = 1.0 and σ = 0.05, and the other hyper-parameters are the same as the ones used in the corresponding baseline methods (without DPSGD). In FedReg with MG, the learning rate and the value of γ are the same as the one used in FedReg, and ηs = 0.001. In the attacking algorithm, the initial learning rate is 1.0, the distance function is cosine distance, and the weight for the total variation term is 10−6.
• CT images. Only the images containing time information are used here. As DenseNet-121 is too complex to be attacked, we use a ten-layer neural network with 9 layers of CNNs and one layer of FC instead. 10 rounds are run in total and in each round, 10 clients are sampled. In the local training stage, the local data are processed in 20 epochs. The learning rate is 5 × 10−4 in SGD, FedAvg, FedProx and FedReg, and 1 × 10−3 in FedCurv, respectively. The weight for the proximal term in FedProx is 0.1 and λ in FedCurv is 0.001. In FedReg, γ = 0.5 and ηs = 10−6. In the baseline methods with DPSGD, C = 1.0 and σ = 0.001, and other hyper-parameters are the same as the ones used in the corresponding methods (without DPSGD). In FedReg with MG, the learning rate and the values of γ and ηs are the same as those in FedReg. In the attacking algorithm, the initial learning rate is 1.0, the distance function is L2-distance, and the weight for the total variation term is 10−8.

heading: C SUPPLEMENTARY RESULTS

heading: D ALGORITHM
text: Algorithm D.1 Training Procedure of FedReg. Input: K, T , {Di}i∈C, θ(0), γ, ηθ, ηs, ηp, S, E Output: {θ(t)}Tt=1
1: for t = 1, 2, ..., T do 2: Randomly sample K clients C(t) 3: for i ∈ C(t) do 4: θ(t,i) ← θ(t−1), generate Dsi and D p i with ηs, ηp and E 5: for s = 1, 2, ..., S do 6: Compute gγ (Di), θ(t,i) ← θ(t,i) − ηθgγ (Di), θ(t,i)β ← 0.5 ( θ(t,i) + θ(t−1)
) 7: Compute gβ (Dsi ), gβ(D p i ), ws and wp, θ (t,i) ← θ(t,i) − wsgβ (Dsi )− wpgβ (D p i ) 8: end for 9: Send θ(t,i) to the server
10: end for 11: θ(t)← 1K ∑ i∈C(t) θ (t,i) 12: end for 13: return {θ(t)}Tt=1

Review summaries:
The authors claim that in non-iid federated learning setups models experience catastrophic forgetting of previously encountered data. As a result, the performance of the model degrades and the convergence is slower. As a solution, the authors propose FedReg, a method to regularize the model with pseudo-data generated from the local data of the client. Two datasets are generated using the fast gradient sign method. The first dataset, termed pseudo-data, aims at simulating data of other clients in the system. The second dataset, termed perturbed data, aims at retaining the model power on the local dataset of the current client. These datasets are used along with the original data of the client to update the local model parameters. The authors further introduce a method to modify the gradients to enhance the protection of privacy. The authors compared their method against several baselines on MNIST, EMNIST, CIFAR-10, CIFAR-100, and CT images. The comparisons show that FedReg achieves significant improvements in accuracy and convergence, reduction in catastrophic forgetting, and better protection of privacy compared to baseline methods.

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.

This paper considers the catastrophic forgetting issue in federated learning. The authors observe that this issue is (at least partially) responsible for slow convergence of existing FL methods when the data are not independently and identically distributed (non-i.i.d.) across different clients. This paper proposes FedReg, an algorithm to alleviate the catastrophic forgetting issue by regularizing the local model parameters on the generated pseudo data.