SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tiansheng Huang, Sihao Hu, Ka-Ho Chow, Fatih Ilhan, Selim Furkan Tekin, Ling Liu
Georgia Institute of Technology, Atlanta, USA
{thuang374, shu335, kchow35, filhan3, stekin6}@gatech.edu,ling.liu@cc.gatech.edu
Federated learning (FL) is vulnerable to backdoor attacks due to its distributed computing nature. Existing defense solution usually requires larger amount of computation in either the training or testing phase, which limits their practicality in the resource-constrain scenarios. A more practical defense, i.e., neural network (NN) pruning based defense has been proposed in centralized backdoor setting. However, our empirical study shows that traditional pruning-based solution suffers poison-coupling effect in FL, which significantly degrades the defense performance. This paper presents Lockdown, an isolated subspace training method to mitigate the poison-coupling effect. Lockdown follows three key procedures. First, it modifies the training protocol by isolating the training subspaces for different clients. Second, it utilizes randomness in initializing isolated subspaces, and performs subspace pruning and subspace recovery to segregate the subspaces between malicious and benign clients. Third, it introduces quorum consensus to cure the global model by purging malicious/dummy parameters. Empirical results show that Lockdown achieves superior and consistent defense performance compared to existing representative approaches against backdoor attacks. Another value-added property of Lockdown is the communication-efficiency and model complexity reduction, which are both critical for resource-constrain FL scenario. Our code is available at https://github.com/git-disl/Lockdown.
Federated Learning (FL) (McMahan et al., 2016) is a privacy-preserving machine learning paradigm that allows the training surrogates (clients) to collectively train a global model with data in local devices. However, because the training data and potentially the training process on the clients lacks supervision, it is possible that attackers can launch data poisoning attack on the global model (Tolpegin et al., 2020), so as to manipulate the prediction of the model.
Backdoor attack is one of the data poisoning attacks, which is stealthy and disruptive to the normal behavior of the model. Specifically, the prediction of the model can be manipulated such that it consistently predicts one (or some) chosen
Figure 1: Illustration of isolated subspace training. The blue circles and the shaded area respectively represent the parameters of the model and the training subspace. Left: vanilla FL where malicious clients can poison all the parameters. Right: isolated subspace training where malicious clients can only poison a subspace of them.
label whenever it is given samples with a backdoor trigger. Backdoor attack poses serious threats to
many security-critical applications, e.g., biometric authentication and autonomous driving (Chow & Liu, 2021), which vitiates the large-scale deployment of FL. Therefore, an effective defense that can mitigate such a risk is in an urgent need.
The main stream of existing research to defend backdoor attack in FL can be mainly classified into three genres, i) certified robustness (Xie et al., 2021), (Alfarra et al., 2022), ii) adversarial training (Zizzo et al., 2020; Shah et al., 2021), and iii) robust aggregation (Chow et al., 2023; Zhang et al., 2022a). Though existing defenses can mitigate attack, they are still far from their maturity. In particular, we highlight the in-efficiency of state-of-the-art defense solutions.
Existing defenses usually require extra computation in either the training or inference phases. For example, randomized smoothing, the key component for certified robustness defense, requires multiple forward pass of data in the inference phase for producing one effective prediction. Analogously, adversarial training requires to generate adversarial samples in the inner loop of training, and therefore demands extra data pass. This issue is particularly pronounced in the FL scenario, considering that training and the inference are conducted on edge devices with limited capacity (Wang et al., 2019b).
In this paper, we try to answer the following question:
Is there a solution that not only provides SOTA defense performance, but also requires equal or even less computation in both training and deployment phases?
To this end, we first apply the pruning-based defense solution in FL backdoor setting, but we observe a poison-coupling effect, which significantly degrades the defense performance. Driven by this observation, we propose Lockdown, a solution that utilizes isolated subspace training to de-couple the poison parameters, which are then pruned using the consensus from clients. Our experiments demonstrate that: i) Lockdown reduces the Attack Success Ratio (ASR) by up-to  $90\%$  compared to FedAvg without defense. ii) Lockdown consistently performs better than SOTA defense solutions in various attack settings. Specifically, Lockdown acquires  $83\%$  of ASR reduction with only  $2.7\%$  drop of benign accuracy when the data distribution is Non-IID (the other two defense baselines RLR and Krum respectively acquire  $14.2\%$  and  $86.1\%$  ASR reduction while sacrificing  $16.2\%$  and  $43.6\%$  benign accuracy in the same setting). iii) Lockdown reduces both downlink and uplink communication by at least  $0.75\mathrm{x}$ , and the number of parameters used in the model training and inference phases are also accordingly reduced by at least  $0.75\mathrm{x}$  thanks to the removal of malicious/dummy parameters.
To the end, we summarize our contribution as follows:
- We study the data-free pruning defense in FL setting, and find that the poison parameters tend to be statistically coupled with benign parameters, which we refer to as "poison-coupling" effect in FL backdoor. To the best of our knowledge, this finding is not available in the literature.
- We propose Lockdown, a backdoor defense that utilizes the idea of isolated subspace training to decouple the poison parameters. Notably, Lockdown i) enjoys communication reduction between server and clients, and ii) lowers the model's training/inference complexity.
- We conduct evaluations to show the efficacy of Lockdown. Results show that Lockdown consistently outperforms existing defense baselines under different attack settings (attack method, attacker number and poison ratio) and different data distributions (IID and Non-IID).
- Ablation study and hyper-parameter sensitivity analysis are conducted to verify the individual functionality of each component of Lockdown.
Federated Learning. Federated learning (McMahan et al., 2016) is a privacy-preserving distributed training paradigm that allows clients to collectively train a global model from distributed training data. Recent works on FL mostly lie in optimization (Karimireddy et al., 2020; Li et al., 2018; Sun et al., 2023a,c; Li et al., 2023b; Sun et al., 2023b), which study how to mitigate the Non-IID issue (Zhao et al., 2018) and system efficiency (Huang et al., 2020, 2022a; Li et al., 2021b,a).
Federated backdoor attack and defense. Classical backdoor attack in FL is empirically proven to be effective in (Bagdasaryan et al., 2020), and subsequently, new types of attack are developed, e.g., edge-case backdoor (Wang et al., 2020), stealthy model poisoning (Bhagoji et al., 2019), subnet replacement attack (Qi et al., 2022), and Distributed Backdoor Attack (Xie et al., 2019).
Backdoor defense are developed to counter the potential threat posed by the backdoor attacks (representative methods are neural cleanse (Wang et al., 2019a), meta analysis (Xu et al., 2021), etc). Under the FL setting, backdoor defenses can be classified into three main genres. The first genre is certified robustness. Certified robustness relies on randomized smoothing(Cohen et al., 2019), an approach with theoretical certification of the model robustness. Subsequent techniques, e.g., weight smoothing (Xie et al., 2021) and group ensemble (Cao et al., 2022) are studied for the defense of federated backdoor. The second genre is adversarial training (Zizzo et al., 2020; Shah et al., 2021), which generates adversarial samples to improve the model's robustness. However, both certified robustness and adversarial training require more computation in either the training or deployment stage. The third genre is robust aggregation, which studies how to identify and preclude the malicious update in the aggregation stage (Chow et al., 2023; Zhang et al., 2022a; Guo et al., 2021; Panda et al., 2022). However, the classification benign accuracy of this category of defense often perform worse when Non-IID extent increases.
Sparse training. Sparse training is originally designed to reduce the model complexity for both the training and deployment stage. SET (Mocanu et al., 2018) first proposes the idea of dynamic subspace searching with alternative pruning and recovery process, and are empirically studied with the concept of in-time over-parameterization (Liu et al., 2021b,a, 2022). (Evci et al., 2020) further introduces gradient information to guide the mask searching process. This model compression technique has recently been extended to FL (Bibikar et al., 2022; Huang et al., 2022b; Dai et al., 2022).
We utilize sparse training technique to develop a new genre of backdoor defense solution. To the best of our knowledge, this is the first attempt that connects sparse training with backdoor defense in FL.
3 Threat Models
We consider three threat models, named weak, medium and strong. All the models allow multiple attackers to co-exist in the system. We use  $N$  to denote the number of attackers out of  $M$  total clients.
Table 1: Permission of threat models.
Permission\Threat model	weak	medium	strong
Data manipulation	✓	✓	✓
Algorithm manipulation	✗	✓	✓
Global information Access	✗	✗	✓
Aggregation manipulation	✗	✗	✗
Permission. Different threats models are given different control permissions. Specifically, Weak model allows the attackers to arbitrarily manipulate its local data, but cannot control its local training process, while attackers in medium and strong model are allowed to do so. Only strong model allows the attackers to obtain other benign client's information, e.g., update from them. For control permission in the weak threat model, technique like trusted execution environment (Mo et al., 2021) can be applied to control each client to run the designated program, and thereby disabling algorithm manipulation. Existing literature usually adopt medium threat model, in which attackers can do whatever they like in their local devices, but has not extra information from server or from other benign clients. The permission of the threat models are summarized in Table 1.
Malicious objective. Denote the set of benign clients as  $\mathcal{M}$  and the set of malicious clients as  $\mathcal{N}$ . Formally, we characterize the objective of malicious clients as follows:
$$
\min  _ {\boldsymbol {w}} \frac {1}{M} \left(\sum_ {i \in \mathcal {N}} \tilde {f} _ {i} (\boldsymbol {w}) + \sum_ {i \in \mathcal {M} / \mathcal {N}} f _ {i} (\boldsymbol {w})\right) \tag {1}
$$
where  $\tilde{f}_i(\boldsymbol{w}) \triangleq \frac{1}{|\tilde{\mathcal{D}}_i|} \sum_{(\boldsymbol{x}_j, y_j) \in \tilde{\mathcal{D}}_i} CE(\boldsymbol{w}; \boldsymbol{x}_j, y_j)$  is the malicious objective of an attacker (to minimize the average cross-entropy of the backdoor dataset),  $f_i(\boldsymbol{w}) \triangleq \frac{1}{|\mathcal{D}_i|} \sum_{(\boldsymbol{x}_j, y_j) \in \mathcal{D}_i} CE(\boldsymbol{w}; \boldsymbol{x}_j, y_j)$  is the benign objective (to minimize the average CE loss for the dataset from benign clients), and  $\tilde{\mathcal{D}}_i$  and  $\mathcal{D}_i$  are respectively the backdoor dataset and pristine dataset.
Attack methods. Different attack methods are eligible in different threat models. For weak threat model, only data-level backdoor, e.g., BadNet (Gu et al., 2017), DBA (Xie et al., 2019), and Sinusoidal
(Barni et al., 2019) can be applied, since the attacker has no control to the local training, but only has access to the dataset. For medium threat model, the attacker can perform a wider range of attack methods, e.g., Scaling (Bagdasaryan et al., 2020), FixMask, etc, which requires the modification of the local training process. For the strong threat model, we test Neurontoxin (Zhang et al., 2022b), and Omniscient, which requires extra server-side information to aid the attack.
Defense goal. While solving the benign objective (see FL general objective in (McMahan et al., 2016)), we expect the global model  $\pmb{w}$  to be able to resist backdoor attack. In other words, the global model should minimize the benign objective while maximizing the malicious objective.
4 Case Study on Pruning-based Defense
Motivated by our goal to propose a computation-friendly defense, we extend the pruning-based defense originally proposed for centralized ML to federated learning setting. We study CLP defense proposed in (Zheng et al., 2022) towards two poisoned models trained with centralized SGD (centralized backdoor) and the global model produced by FedAvg (federated backdoor)<sup>1</sup>. Surprisingly, we observe that the pruning defense exhibits substantially different performance on the two models.
Figure 2: Properties of two models trained with centralized backdoor and federated backdoor. Left: ASR and benign accuracy with CLP defense in (Zheng et al., 2022). Middle: Channel lipschitzness of last convolutional layer of two models. Right: L2 norm of last convolutional layer of two models. The experiment is done on a ResNet-9 model on CIFAR10 dataset and under BadNet attack.
Observations. In the left of Figure 2, we observe that the pruning algorithm cannot efficiently eliminate the backdoor parameters for federated backdoor. The ASR can only drop to an acceptable range (say  $30\%$ ) when leveraging large pruning ratio, and accordingly, this also leads to a large drop of benign accuracy (say  $35\%$  drop). In the middle, we plot the channel lipschitzness (CL) of the two models, and find that the CL values for a federated backdoor couple in a more compact range, which makes it hard for the pruning method to find out the malicious channels to prune. This explains why CLP fails to work for federated backdoor. In the right, we plot the L2 norm of weights of the last convolutional layer, the same coupling effect for federated backdoor can be observed – the L2 norm of different channels are all within a small range.
**Poison-coupling effect.** In summary, our main takeaway is that the federated backdoor model is hard to be cured by classical pruning method, because the backdoor parameters does not exhibit substantial statistical difference with the benign ones, i.e., it tends to be coupled with the benign parameters. We refer to this phenomenon for federated backdoor as **poison-coupling** effect. Due to the existence of poison-coupling effect, it is hard, if not impossible, to accurately identify the poisoned parameters after the model has been fully trained by federated learning process, and thereby reducing the efficiency of the pure parameters pruning defense.
To mitigate the poison-coupling effect, we propose lockdown with isolated subspace training. Intuitively, we don't allow the malicious clients to get access to all the parameters, but only a subspace of them (See Figure 1), in order to prevent them from statistically coupling the poisoned parameters with the benign ones that are supposed to do normal function.
<sup>1</sup> Checkpoints of the centralized/federated backdoor models are available in https://www.dropbox.com/scl/fo/lhi0objkylgnyh299j5wz/h?rlkey=wrsa5msbl2m3kzuk2f8rcc8uw&dl=0
The high-level idea of Lockdown is as follows. i) we employ isolated subspace training to restrain the training of each client into its own subspace, such that the backdoor data cannot access and couple the poisoning function in all the parameters. ii) we employ dynamic subspace searching for the clients to search for local subspaces using their local datasets, which involve only the parameters that they deem important. iii) After the above local procedures, the server will aggregate the gradient updates into the corresponding subspaces of global model, and continues the next round of training. iv) Repeating  $T$  rounds of training, the server can identify the malicious parameters by consensus fusion. Under the intuition that the malicious/dummy parameters (at specific coordinates) should have less chance to
be involved in benign clients' subspaces (because they are not important parameters for the benign datasets). Considering the benign clients constitute the majority in the system, consensus fusion identify those parameters that have the least appearance among all the clients subspace as the malicious or dummy parameters, which are then pruned to mitigate backdoor behavior. As follows, we detail each step of Lockdown.
Subspace initialization. We enforce an overall sparsity  $s$  to the initial subspace of each client (denoted by a binary variable  $m_{i,0}$ ). For maintaining practical performance of sparsity, we follow (Evci et al., 2020) to use ERK for random subspace initialization, in which different layers in the corresponding subspace are designated different sparsity (See Appendix A.1 for details).
For each client, we enforce them to have the same initial subspace, which can be achieved via enforcing the same random seed to them for subspace generation. However, we do find that heterogeneous mask initialization can increase protection in some cases (See Section 6.2).
Isolated subspace training. In this stage, each client performs multiple local steps to train parameters within its isolated subspace (as enforced by its mask). Specifically, for local step  $k = 0,\dots ,K - 1$ , clients do this update:
$$
\boldsymbol {w} _ {i, t, k + 1} = \boldsymbol {w} _ {i, t, k} - \eta \boldsymbol {m} _ {i, t} \odot \nabla f _ {i} \left(\boldsymbol {w} _ {i, t, k}; \xi\right) \tag {2}
$$
where  $\xi$  is a piece of random sample within the local dataset,  $\eta$  is the learning rate, and  $\odot$  denotes Hadamard product. As shown, the binary mask  $m_{i,t}$  is applied to the gradient in every local step, which isolates the training into  $i$ -th client's own subspace. This process prevents the backdoor data from contaminating all the parameters within the parameter space.
Subspace searching. We conduct subspace searching for clients to search for their subspaces, i.e., the parameters that are important per the data they have. The searching process is decomposed to two phases, as follows:
Figure 3: Overview of Lockdown. First, clients receive subspace model from server, and apply isolated subspace training. Second, clients upload their subspace update. Third, server aggregates the updates into global model. After training, consensus fusion is applied to remove poison params.
Algorithm 1 Lockdown defense
input Training iteration  $T$ ; Local steps  $K$ ; Learning rate  $\eta$ ; Pruning/recovery rate  $\alpha_{t}$  decayed by cosine annealing (Specially,  $\alpha_{-1} = 0$ , i.e., no recovery for first round); Random seed seed; Initial model  $\boldsymbol{w}_0$ .  
output Clean model for deployment  $\tilde{\boldsymbol{w}}_T$
1: main Server's Main Loop
2:  $\pmb{m}_{i,0} = \mathbf{SubspaceInit}(\pmb{seed})$  for  $i \in \mathcal{M}$
3: for  $t = 0,1,\ldots ,T - 1$  do
4: for  $i\in \mathcal{M}$  do
5: Send  $\pmb{w}_{i,t,0} = \pmb{m}_{i,t}\odot \pmb{w}_t$  to client  $i$
6: Call Client  $i$  's main loop for training
7: Received  $\pmb{w}_{i,t,K}$  and  $\pmb{m}_{i,t+1}$
8: end for
9:  $\pmb{w}_{t + 1} = \pmb{w}_t - \frac{1}{M}\sum_{i = 1}^{M}\pmb{m}_{i,t}\odot (\pmb{w}_t - \pmb{w}_{i,t,K})$
10: end for
11:  $\tilde{\pmb{w}}_T = \pmb{w}_T\odot \mathcal{T}_{\theta}(\pmb{m}_{1,T},\dots,\pmb{m}_{M,T})$
12: Deploy  $\tilde{w}_T$  for serving/inference.
13: end main
14: main Client's Main Loop
15: Obtain local gradient  $\nabla f_{i}(\pmb{w}_{i,t,0})$
16:  $\pmb{m}_{i,t + \frac{1}{2}} = \pmb{m}_{i,t} + \mathrm{ArgTopK}_{\alpha_{t - 1}}(|\nabla f_i(\pmb{w}_{i,t,0})|)$
17: for  $k = 0, 1, \ldots, K - 1$  do
18:  $\pmb{w}_{i,t,k + 1} = \pmb{w}_{i,t,k} - \eta \pmb{m}_{i,t + \frac{1}{2}}\odot \nabla f_{i}(\pmb{w}_{i,t,k};\xi)$
19: end for
20:  $\pmb{m}_{i,t + 1} = \pmb{m}_{i,t + \frac{1}{2}} - \mathrm{ArgBottomK}_{\alpha_t}(|\pmb{w}_{i,t,K}|)$
21: Send  $\pmb{w}_{i,t,K}$  and  $\pmb{m}_{i,t+1}$  to server.
22: end main
1. **Subspace pruning.** In this phase, we prune the unimportant parameters within the client's current subspace that do not function at all. After  $K$  steps of local training, we can identify the unimportant parameters as those with the smallest magnitude within the subspace. To preclude them from future training, we prune those smallest  $\alpha_{t}$  parameters within each layer, and accordingly update
the mask to  $m_{i,t+1}$ . Formally, this process can be formulated as follows:
$$
\boldsymbol {m} _ {i, t + 1} = \boldsymbol {m} _ {i, t + \frac {1}{2}} - \operatorname {A r g B o t t o m K} _ {\alpha_ {t}} \left(\left| \boldsymbol {w} _ {i, t, K} \right|\right) \tag {3}
$$
where  $\mathrm{ArgBottomK}_{\alpha_t}(|w|)$  return the  $\alpha_{t}$  percentage of smallest coordinates of each layer absolute weights and mask them to 1, indicating that they will be pruned. We use cosine annealing to decay  $\alpha_{t}$  from the initial pruning rate  $\alpha_0$ . The decay process is postponed to the Appendix A.1.
2. **Subspace recovery.** After pruning, and before the start of next round local training, the client recovers the same amount of parameters (as indicated by  $\alpha_{t-1}$ ) to its subspace for exploration. To identify which parameters should be recovered, we use each client's data to extract the gradient w.r.t the weights after pruned, i.e., to extract  $\nabla f_i(\boldsymbol{w}_{i,t,0})$ . Then we recover  $\alpha_{t-1}$  percentage of parameters by identifying those with top- $\alpha_{t-1}$  percentage of gradient magnitude within each layer, and accordingly update the mask to  $\boldsymbol{m}_{i,t+\frac{1}{2}}$ . The intuition is that for an important parameter over the local data, its gradient magnitude should be larger than the unimportant ones, which should be included into the new subspace. Formally, the recovery process is formalized as follows:
$$
\boldsymbol {m} _ {i, t + \frac {1}{2}} = \boldsymbol {m} _ {i, t} + \operatorname {A r g T o p K} _ {\alpha_ {t - 1}} \left(\left| \nabla f _ {i} \left(\boldsymbol {w} _ {i, t, 0}\right) \right|\right) \tag {4}
$$
where  $\mathrm{ArgTopK}_{\alpha_{t-1}}(\boldsymbol{w})$  returns the  $\alpha_{t-1}$  percentage of largest coordinates of absolutely gradient and mask them to 1, indicating that they will be recovered. But we note here that we don't perform subspace recovery in the first round, since pruning is not done yet.
Aggregation. Once subspace training is done, the clients upload the gradient updates of the local subspace to the server for aggregation. To aggregate the knowledge into the global model, we average the gradient updates based on their coordinate-wise contributions, as follows:
$$
\boldsymbol {w} _ {t + 1} = \boldsymbol {w} _ {t} - \frac {1}{M} \sum_ {i = 1} ^ {M} \boldsymbol {m} _ {i, t + 1} \odot \left(\boldsymbol {w} _ {t} - \boldsymbol {w} _ {i, t, K}\right) \tag {5}
$$
where  $\pmb{w}_{i,t,K}$  is the weight from the client's local subspace, and  $\pmb{m}_{i,t+1}$  is the subspace after pruned.
Consensus fusion (CF). Given that those malicious parameters served to recognize backdoor triggers will be deemed unimportant for benign clients, they should not be contained in the subspace of benign clients, which accounts for the majority. This observation inspires us to eliminate the malicious parameters using consensus fusion after  $T$  rounds of global model training. Formally, the consensus fusion operator returns a vector that satisfies:
$$
\left[ \mathcal {T} _ {\theta} \left(\boldsymbol {m} _ {1, T}, \dots , \boldsymbol {m} _ {M, T} \right] _ {j} = \left\{ \begin{array}{l l} 1 & \sum_ {i = 1} ^ {M} \left[ \boldsymbol {m} _ {i, T} \right] _ {j} \geq \theta \\ 0 & \text {O t h e r w i s e} \end{array} \right. \right. \tag {6}
$$