Federated Unlearning with Knowledge Distillation

In a general perspective, the problem includes two entities: the server $S$ and a group of clients $C$ participated in the FL. The server $S$ relies on the group of clients $C$ to help train a global model $M$ to be used by all the clients. According to the FL definition McMahan et al. (2017), training data is on the clients’ side and cannot be shared with others. Each client will train the global model $M$ on its local dataset and send its model updates to the server. The server will aggregate model updates from clients and generate an updated global model for the next round of training. Suppose a client $i\in C$ needs to revoke its previous contribution/updates to the global model $M$. In that case, the server should be able to update the current global model $M$ to an unlearned model $M^{C\setminus\{i\}}$, where $M^{C\setminus\{i\}}$ is a model that could be trained if client $i$ was never included in the training group $C$.

Assume that there are a total of $N$ clients to participate in the FL training at round $t$, and each client trains the global model $M_{t}$ with its private dataset and updates parameter changes to the server. The new global model $M_{t+1}$ is calculated by an averaged aggregation (FedAvg McMahan et al. (2017)) of the weight updates from these clients.

where $\Delta M_{t}^{i}$ is the parameter update contributed by client $i$ based on the model $M_{t}$. The server keeps running the training process and updates the global model from $M_{1}$ to $M_{F}$ when the termination criterion (e.g. test accuracy) has been satisfied at round $F-1$.

The federated unlearning task is defined as to completely remove the influence of all the updates $\Delta M_{t}^{i}$ from a target client $i$, with $t\in[1,F-1]$ from the model $M_{F}$. In other words, it removes the contribution of any client $i$ to the final global model $M_{F}$ and creates a new model $M_{F}^{C\setminus\{i\}}$ as if client $i$ has never participated in the training process.

With the non-deterministic property of FL, it is natural to ask the question of how to evaluate the unlearning model. Is there a way to measure the unlearning effect in FL system? Previous unlearning works on DNNs Bourtoule et al. (2021, 2021) compare the test accuracy between unlearning model and retrained model may not be convincing enough in FL systems. Therefore, we introduce the backdoor attacks Gu et al. (2017); Bagdasaryan et al. (2020) to help evaluate the effectiveness of unlearning methods. As one of the most powerful attacks to the FL system, the backdoor attack does not influence the global model’s performance on regular input but only distorts the predictions when triggered by specific inputs with the backdoor pattern. Such property makes it a perfect evaluation method to measure the effect of unlearning. A successful unlearning global model should perform reasonably well on the evaluation dataset but reduce backdoor attacks’ success rate when triggered by backdoored inputs.

To completely remove the contribution of any client $i$ to the final global model $M_{F}$, we want to erase all the historical updates $\Delta M_{t}^{i}$ from this client as long as $t\in[1,F-1]$. As we have shown in Equation 1, the update of the global model at round $t$ consists of averaged weight updates from participating clients. If we use $\Delta M_{t}$ to represent this update at round $t$, the finalized global model $M_{F}$ can be viewed as a composition of initial model weight $M_{1}$ and updates to the global model from round $1$ to round $F-1$.

For simplicity, we assume that each round, there are $N$ clients participating in the FL training, and client $N$ is the target client that wants to be unlearned from the global model. At this time, we can simplify the problem as to remove the contribution $\Delta M_{t}^{N}$ of the target client $N$ from the global model update $\Delta M_{t}$ at each round $t$.

There are two ways to calculate the new global model update $\Delta M_{t}^{\prime}$ at round $t$. The first one is to assume that only $N-1$ clients were participating in the FL at round $t$. In this way, the new global model updates $\Delta M_{t}^{\prime}$ at round $t$ becomes the following equation.

However, we can not directly accumulate the new updates to reconstruct the unlearning model because of the incremental learning property of FL, as discussed in the previous section. Any update to the global model $M_{t}$ will result in a requirement of updates to all the model updates that happened afterward. Hence, we use $\epsilon_{t}$ to represent the necessary amendment (skew) to the global model at each round $t$. After combining the above equation with Equation 2, we can get the unlearning version of the final global model $M_{F}^{\prime}$.

where $\epsilon_{t}$ is the necessary correction to amend the skew produced by change of the model at previous rounds. Because of this characteristic of the incremental learning process in FL, the skew $\epsilon_{t}$ will increase with more training rounds after updates to the global model. Thus, the above unlearning rule has a shortcoming that when the target client $N$ makes little contribution to the model at round $t$ (e.g., $\Delta M_{t}^{N}\approx 0$), the global model update $\Delta M_{t}$ will still change a lot by multiplying itself with a factor of $\frac{N}{N-1}$. This will bring more skew $\epsilon_{t}$ to the global model in the following rounds.

To mitigate this problem, we propose to use a lazy learning strategy to eliminate the influence of target client $N$. Specifically, we assume client $N$ still participated in the training process but set his updates $\Delta M_{t}^{N}=0$ for all rounds $t\in[1,F-1]$. The unlearning of the global model update can be simplified as follows.

A combination of the above formula with Equation 2 gives us the unlearning result of the final global model $M_{F}^{\prime}$.

Now the unlearning model update rule becomes surprisingly straightforward and easy to understand. We just need to subtract all the historical averaged updates from target client $N$ from the final global model $M_{F}$. Then, we remedy the skew $\epsilon_{t}$ caused by this process because of the incremental learning characteristic of the FL.

There is no existing method to calculate the skew $\epsilon_{t}$ without retraining the updated model on the original dataset again. However, as we discussed before, one of the challenges in FL is that we cannot rely on the clients to hold the dataset forever and prepare for the unlearning purpose. How to remedy this skew without relying on clients’ training data or the participation of clients becomes crucially important.

To tackle this problem, we propose to leverage the knowledge distillation method to train the unlearning model using the original global model. The distillation technique is first motivated by the goal of reducing the size of DNN architectures or ensembles of models Hinton et al. (2015). It uses the prediction results of class probabilities produced by an ensemble of models or a complex DNN to train another DNN of a reduced number of parameters without much loss of accuracy. Later on, Papernot et al. proposed to use this method as a defensive mechanism to reduce the effectiveness of adversarial samples on DNNs Papernot et al. (2016a). They showed that the defensive distillation can improve the generalizability and robustness of the trained DNNs. The intuition is based on the fact that the knowledge acquired by DNNs during the training process is not only encoded in the weight parameters but also can be reflected from the class probability prediction output of the model. Distillation training uses these soft probabilities instead of hard ground truth labels to provide additional information about each class and the prediction logic to the training process. For example, given an image of “7” from the MNIST dataset, the distillation training process can train the new model with more information (except from the ground truth label that the image belongs to class “7”), such as this image is close to “1” and “9” but a lot different from “3” and “8”.

To perform knowledge distillation in the unlearning problem, we treat the original global model as the teacher model and the skewed unlearning model as the student model. Then, the server can use any unlabeled data to train the unlearning model and remedy the skew $\epsilon_{t}$ caused by the previous subtraction process. Specifically, the original global model produces class prediction probabilities through a “softmax” output layer that converts the logit, $z_{i}$, computed for each class into a probability, $q_{i}$, by comparing $z_{i}$ with the other logits.

where $T$ is a parameter named temperature and shared across the softmax layer. The value of $T$ is normally set to 1 for traditional ML training and predictions. A higher temperature $T$ makes the DNN produce a softer probability distribution over classes. In other words, the probability output will be forced to produce relatively large values for each class, and logits $z_{i}$ become negligible compared to temperature $T$. An example is that the output probability for each class $i$ converge to $1/Z$ (assume there are $Z$ classes for prediction) as $T\rightarrow\infty$. In summary, higher temperature $T$ produces probability distribution more ambiguously while lower temperature $T$ produces probability distribution more discretely.

We use this soft class prediction probability produced by the original global model $M_{F}$ to label the dataset. The skewed unlearning model is then trained with this dataset with soft labels (with high temperature $T$). On the other hand, if the server possesses a labeled dataset, we can also leverage a combination of hard labels (ground truth with temperature $T=1$) and soft labels produced by the global model. As suggested by Hinton et al., using a weighted average of these two objective functions with a considerably lower weight on the objective function of the hard labels can produce the best results Hinton et al. (2015). The temperature will be set back to $1$ after distillation training, so the unlearning model $M_{F}^{\prime}$ can produce more discrete class prediction probabilities during test time.

In this section, we evaluate the performance of the unlearning model under different steps during unlearning. To begin with, we show, in Figure 2, the behavior of the model after directly erasing all the historical parameter updates from the target client (attacker). Subtracting the parameter updates from the global model will create a non-negligible skew. As we observed from the figure, the unlearning model’s accuracy has never exceeded 60% and has a trend of slightly decreasing with more FL training rounds. The advantage of this process is that it can thoroughly remove the influence of the target client (attacker) from the global model. Compared with the original global model with an almost 100% backdoor attack success rate, the unlearning model keeps the attack success rate 0% all the time.

Then, we can look at the following knowledge distillation process used to recover the skew caused by the subtraction process. From Figure 3, we notice that the test accuracy of the model is quickly recovered within five epochs of distillation training. The loss on the test dataset keeps getting closer to the original global model while the backdoor attack success rate keeps as low as 0% all the time. The attacker’s influence on the global model does not transmit to the unlearning model after the knowledge distillation training process. The reason is that the distillation training method does not use any data from the target client (attacker), so the backdoor in the original model has not been triggered and thus will not be learned by the unlearning model. This also demonstrates the privacy protection of a client in the case of right to be forgotten. because the influence of his contributions will also be removed from the new global model.

In the end, Table 1 reports the integrated experiment results with different datasets and model architectures. The “Training” column reports the performance of global model $M_{F}$ on the evaluation dataset, including the test accuracy and backdoor attack success rate. The “UL-Subtraction” column reports the unlearning model’s performance after merely subtracting the target client’s historical parameter updates from the global model. The “UL-Distillation” column reports the performance of unlearning model after the knowledge distillation process on the server’s side. The “Post-Training” column represents how much we can further improve the unlearning model if putting it back to the FL system and continue training without the participation of the target client (attacker). The “Re-Training” column is the widely used golden standard for the unlearning problem. We retrain the model from scratch excluding the participation of the attacker.

From the results, we can conclude that subtracting the attacker’s historical parameter updates eliminates his influence on the global model. In all cases, the attack success rate drops to zero. The knowledge distillation training process can help remedy the skew caused by subtraction and recover the model’s performance back to an acceptable standard. The test accuracy of the unlearning model after distillation is almost identical to that of the model retraining from scratch (with differences less than 1%). In the post-training results, we find the distillation can even help improve the model’s accuracy with follow-up training with clients. What is more, we are delighted to observe that the distillation process does not pass the unlearning target’s attributes from the original global model to the unlearning model. It supports our assumption that as long as we are not using the same dataset to activate the global model, the original global model’s private information/attributes will not be learned by the unlearning model. One may ask why the attack success rate is larger than zero for even the retraining model on the CIFAR-10 dataset. It is caused by the prediction errors of the model (the test accuracy is only around 80%). In other words, the original images (without backdoor patterns) will be misclassified as the backdoor target label too. Thus, such errors will be automatically counted as backdoor success rates here. Compared with the attack success rate of almost 100% in the original global model, under this error rate (less than 7%), we may consider it a successful unlearning of the target client’s contribution.