SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kangyang Luo $^{1}$ , Shuai Wang $^{1}$ , Yexuan Fu $^{1}$ , Xiang Li $^{1*}$ , Yunshi Lan $^{1}$ , Ming Gao $^{1,2}$
School of Data Science & Engineering
KLATASDS-MOE in School of Statistics2
{52205901003, 51215903058, 51215903042}@stu.ecnu.edu.cn
{xiangli, yslan, mgao}@dase.ecnu.edu.cn
Federated Learning (FL) is a privacy-constrained decentralized machine learning paradigm in which clients enable collaborative training without compromising private data. However, how to learn a robust global model in the data-heterogeneous and model-heterogeneous FL scenarios is challenging. To address it, we resort to data-free knowledge distillation to propose a new FL method (namely DFRD). DFRD equips a conditional generator on the server to approximate the training space of the local models uploaded by clients, and systematically investigates its training in terms of fidelity, transferability and diversity. To overcome the catastrophic forgetting of the global model caused by the distribution shifts of the generator across communication rounds, we maintain an exponential moving average copy of the generator on the server. Additionally, we propose dynamic weighting and label sampling to accurately extract knowledge from local models. Finally, our extensive experiments on various image classification tasks illustrate that DFRD achieves significant performance gains compared to SOTA baselines.
Our code is here: https://anonymous.4open.science/r/DFRD-0C83/.
With the surge of data, deep learning algorithms have made significant progress in both established and emerging fields [1-3]. However, in many real-world applications (e.g., mobile devices [4], IoT [5], and autonomous driving [6], etc.), data is generally dispersed across different clients (i.e., data silos). Owing to the high cost of data collection and strict privacy protection regulations, the centralized training that integrates data together is prohibited [7, 8]. Driven by this reality, Federated Learning (FL) [9, 10] has gained considerable attention in industry and academia as a promising distributed learning paradigm that allows multiple clients to participate in the collaborative training of a global model without access to their private data, thereby ensuring the basic privacy.
Despite its remarkable success, the inevitable hurdle that plagues FL research is the vast heterogeneity among real-world clients [11]. Specifically, the distribution of data among clients may be non-IID (identical and independently distributed), resulting in data heterogeneity [42, 12, 13, 8, 41]. It has been confirmed that the vanilla FL method FedAvg [14] suffers from client drift in this case, which leads to severe performance degradation. To ameliorate this issue, a plethora of modifications [15-21, 34, 39] for FedAvg focus on regularizing the objectives of the local models to align the global optimization objective. All of the above methods follow the widely accepted assumption of model homogeneity, where local models have to share the same architecture as the global model. Nevertheless,
when deploying FL systems, different clients may have distinct hardware and computing resources, and can only train the model architecture matching their own resource budgets [22, 24], resulting in model heterogeneity. In this case, to enable the FL systems with model homogeneity, on the one hand, clients with low resource budgets, which may be critical for enhancing the FL systems, will be discarded at the expense of training bias [43, 8, 23]. On the other hand, keeping a low complexity for the global model to accommodate all clients leads to performance drop due to the limited model capacity [30]. Therefore, the primary challenge of model-heterogeneous FL is how to conduct model aggregation of heterogeneous architectures among clients to enhance the inclusiveness of federated systems. To solve this challenge, existing efforts fall broadly into two categories: knowledge distillation (KD)-based methods [40, 45, 44, 25-27] and partial training (PT)-based methods [30, 28, 29, 31], yet each of them has its own limitations. Concretely, KD-based methods require additional public data to align the logits outputs between the global model (student) and local models (teachers). But the desired public data is not always available in practice and the performance may decrease dramatically if the apparent disparity in distributions exists between public data and clients' private data [33]. PT-based methods send width-based sub-models to clients, which are extracted by the server from a larger global model according to each client's resource budget, and then aggregate these trained sub-models to update the global model. PT-based methods can be considered as an extension of FedAvg to model-heterogeneous scenarios, which means they are implementation-friendly and computationally efficient, but they also suffer from the same adverse effects from data heterogeneity as FedAvg or even more severely. In a word, how to learn a robust global model in FL with both data and model heterogeneity is a highly meaningful and urgent problem.
Figure 1: The full pipeline for DFRD combined with a PT-based method. DFRD works on the server and contains two phases, training generator and robust model distillation, where  $\mathcal{L}_{tran},\mathcal{L}_{div},\mathcal{L}_{fid}$  and  $\mathcal{L}_{kl},\widetilde{\mathcal{L}}_{kl}$  are the loss objectives of the conditional generator and the global model, respectively.
To this end, we systematically investigate the training of a robust global model in FL with both data and model heterogeneity with the aid of data-free knowledge distillation (DFKD) [35, 36, 53, 59]. See the related work in Appendix A for more DFKD methods. Note that the strategy of integrating DFKD to FL is not unique to us. Recently, FedFTG [48] leverages DFKD to fine-tune the global model in model-homogeneous FL to overcome data heterogeneity, and DENSE [49] aggregates knowledge from heterogeneous local models based on DFKD to train a global model for one-shot FL. They all equip the server, which possesses powerful hardware and computing resources, with a generator to approximate the training space of the local models (teachers), and train the generator and the global model (student) in an adversarial manner. However, the local models uploaded per communication round are not only architecturally heterogeneous but also trained on non-IID distributed private data in the situation of both data and model heterogeneity. In this case, the generator tends to deviate from the real data distribution. Also, its output distribution may undergo large shifts (i.e., distribution shifts) across communication rounds, causing the global model to catastrophically forget useful knowledge learned in previous rounds and suffer from performance degradation. To confront the mentioned issues, we propose a novel Data-Free Robust Distillation FL method called DFRD, which utilizes a conditional generator to generate synthetic data and thoroughly studies how to effectively and accurately simulate the local models' training space in terms of fidelity, transferability and
diversity [48, 49]. To mitigate catastrophic forgetting of the global model, an exponential moving average (EMA) copy of the conditional generator is maintained on the server to store previous knowledge learned from the local models. The EMA generator, along with the current generator, provides training data for updates of the global model. Also, we propose dynamic weighting and label sampling to aggregate the logits outputs of the local models and sample labels respectively, thereby properly exploring the knowledge of the local models. We revisit FedFTG and DENSE, and argue that DFRD as a fine-tuning method (similar to FedFTG) can significantly enhance the global model. So, we readily associate the PT-based methods in model-heterogeneous FL with the ability to rapidly provide a preliminary global model, which will be fine-tuned by DFRD. We illustrate the schematic for DFRD as a fine-tuning method based on a PT-based method in Fig. 1. Although FedFTG and DENSE can also be applied to fine-tune the global model from the PT-based methods after simple extensions, we empirically find that they do not perform as well, and the performance of the global model is even inferior to that of local models tailored to clients' resource budgets.
Our main contributions of this work are summarized as follows. First, we propose a new FL method termed DFRD that enables a robust global model in both data and model heterogeneity settings with the help of DFKD. Second, we systematically study the training of the conditional generator w.r.t. fidelity, transferability and diversity to ensure the generation of high-quality synthetic data. Additionally, we maintain an EMA generator on the server to overcome the global model's catastrophic forgetting caused by the distribution shifts of the generator. Third, we propose dynamic weighting and label sampling to accurately extract the knowledge of local models. At last, our extensive experiments on six real-world image classification datasets verify the superiority of DFRD.
2 Notations and Preliminaries
Notations. In this paper, we focus on the centralized setup that consists of a central server and  $N$  clients owning private labeled datasets  $\{(\mathbf{X}_i,\mathbf{Y}_i)\}_{i = 1}^N$ , where  $\mathbf{X}_i = \{\pmb {x}_i^b\}_{b = 1}^{n_i}$  follows the data distribution  $\mathcal{D}_i$  over feature space  $\mathcal{X}_i$ , i.e.,  $\pmb{x}_i^b\sim \mathcal{D}_i$ , and  $\mathbf{Y}_i = \{y_i^b\}_{b = 1}^{n_i}\subset [C]\coloneqq \{1,\dots ,C\}$  denotes the ground-truth labels of  $\mathbf{X}_i$ . And  $C$  refers to the total number of labels. Remarkably, the heterogeneity for FL in our focus includes both data heterogeneity and model heterogeneity. For the former, we consider the same feature space, yet the data distribution may be different among clients, that is, label distribution skewness in clients (i.e.,  $\mathcal{X}_i = \mathcal{X}_j$  and  $\mathcal{D}_i\neq \mathcal{D}_j,\forall i\neq j,i,j\in [N])$ . For the latter, each client  $i$  holds an on-demand local model  $f_{i}$  parameterized by  $\theta_{i}$ . Due to the difference in resource budgets, the model capacity of each client may vary, i.e.,  $|\pmb {\theta}_i|\neq |\pmb {\theta}_j|,\exists i\neq j,i,j\in [N]$ . In PT-based methods, we define a confined width capability  $R_{i}\in (0,1]$  according to the resource budget of client  $i$ , which is the proportion of nodes extracted from each layer in the global model  $f$ . Note that  $f$  is parameterized by  $\pmb{\theta}$ , and  $|\pmb {a}|$  denotes the number of elements in vector  $\pmb{a}$ .
PT-based method is a solution for model-heterogeneous FL, which strives to extract a matching width-based slimmed-down sub-model from the global model as a local model according to each client's budget. As with FedAvg, it requires the server to periodically communicate with the clients. In each round, there are two phases: local training and server aggregation. In local training, each client trains the sub-model received from the server utilizing the local optimizer. In server aggregation, the server collects the heterogeneous sub-models and aggregates them by straightforward selective averaging to update the global model, as follows [28-31]:
$$
\boldsymbol {\theta} _ {[ l, k ]} ^ {t} = \frac {1}{\sum_ {j \in \mathcal {S} _ {t}} p _ {j}} \sum_ {i \in \mathcal {S} _ {t}} p _ {i} \boldsymbol {\theta} _ {i, [ l, k ]} ^ {t}, \tag {1}
$$
where  $S_{t}$  is a subset sampled from  $[N]$  and  $p_i$  is the weight of client  $i$ , which generally indicates the size of data held by client  $i$ . At round  $t$ ,  $\pmb{\theta}_{[l,k]}^{t}$  denotes the  $k^{th}$  parameter of layer  $l$  of the global model and  $\pmb{\theta}_{i,[l,k]}^{t}$  denotes the parameter  $\pmb{\theta}_{[l,k]}^{t}$  updated by client  $i$ . We can clearly see that Eq. (1) independently calculates the average of each parameter for the global model according to how many clients update that parameter in round  $t$ . Instead, the parameter remains unchanged if no clients update it. Notably, if  $|\pmb{\theta}_i^t| = |\pmb{\theta}^t|$  for any  $i \in [N]$ , PT-based method becomes FedAvg. The key to PT-based method is to select  $\pmb{\theta}_i^t$  from the global model  $\pmb{\theta}^t$  when given  $R_i$ . And existing sub-model extraction schemes fall into three categories: static [28, 29], random [30] and rolling [31].
In this section, we detail the proposed method DFRD. We mainly work on considering DFRD as a fine-tuning method to enhance the PT-based methods, thus enabling a robust global model in FL with both data and model heterogeneity. Fig. 1 visualizes the training procedure of DFRD combined with a PT-based method, consisting of four stages on the server side: training generator, robust model distillation, sub-model extraction and model aggregation. Note that sub-model extraction and model aggregation are consistent with that in the PT-based methods, so we detail the other two stages. Moreover, we present pseudocode for DFRD in Appendix C.
3.1 Training Generator
At this stage, we aim to train a well-behaved generator to capture the training space of local models uploaded from active clients. Specifically, we consider a conditional generator  $G(\cdot)$  parameterized by  $w$ . It takes as input a random noise  $z \in \mathbb{R}^d$  sampled from standard normal distribution  $\mathcal{N}(0, I)$ , and a random label  $y \in [C]$  sampled from label distribution  $p(y)$ , i.e., the probability of sampling  $y$ , thus generating the synthetic data  $s = G(h = o(z, y), w)$ . Note that  $o(z, y)$  represents the merge operator of  $z$  and  $y$ . To the best of our knowledge, synthetic data generated by a well-trained generator should satisfy several key characteristics: fidelity, transferability, and diversity [48, 49]. Therefore, in this section, we construct the loss objective from the referred aspects to ensure the quality and utility of  $G(\cdot)$ .
Fidelity. To commence, we study the fidelity of the synthetic data. Specifically, we expect  $G(\cdot)$  to simulate the training space of the local models to generate the synthetic dataset with a similar distribution to the original dataset. To put it differently, we want the synthetic data  $s$  to approximate the training data with label  $y$  without access to clients' training data. To achieve it, we form the fidelity loss  $\mathcal{L}_{fid}$  at logits level:
$$
\mathcal {L} _ {f i d} = C E \left(\sum_ {i \in \mathcal {S} _ {t}} \tau_ {i, y} f _ {i} \left(\boldsymbol {s}, \boldsymbol {\theta} _ {i}\right), y\right), \tag {2}
$$
where  $CE$  denotes the cross-entropy function,  $f_{i}(\boldsymbol{s},\boldsymbol{\theta}_{i})$  is the logits output of the local model from client  $i$  when  $s$  is given,  $\tau_{i,y}$  dominates the weight of logits from different clients  $\{i|i\in S_t\}$  when  $y$  is given. And  $\mathcal{L}_{fid}$  is the cross-entropy loss between the weighted average logits  $\sum_{i\in S_t}\tau_{i,y}f_i(\boldsymbol{s},\boldsymbol{\theta}_i)$  and the label  $y$ . By minimizing  $\mathcal{L}_{fid}$ ,  $s$  is enforced to be classified into label  $y$  with a high probability, thus facilitating the fidelity of  $s$ .
In reality, the conditional generator  $G(\cdot)$  easily generates synthetic data with low classification errors (i.e.  $\mathcal{L}_{fid}$  close to 0) as the training proceeds. This may cause the synthetic data to fall into a space far from the decision boundary of the ensemble model (i.e.,  $\sum_{i\in S_t}\tau_{i,y}f_i(\cdot ,\pmb {\theta}_i)$ ) if only  $\mathcal{L}_{fid}$  is optimized, as shown in the synthetic data represented by red circles in Fig. 2 (a). Note that  $d_S$  and  $d_T$  denote the decision boundaries of the global model (student) and the ensemble model (teacher), respectively. An obvious observation is that the red circles are correctly classified on the same side of the two decision boundaries (i.e.,  $d_S$  and  $d_T$ ), making it difficult to transfer teacher's knowledge to student. We next explore how to augment the transferability of the synthetic data to ameliorate this pitfall.
Transferability is intended to guide  $G(\cdot)$  in generating synthetic data that moves the decision
boundary of the global model towards that of the ensemble model, such as synthetic data with black circles in Fig. 2 (b). However, during the training of  $d_{S}$  to approach  $d_{T}$ , we find that  $G(\cdot)$  can easily generate two other types of synthetic data, the yellow and purple circles in Fig. 2 (c). Both of them are misclassified by the ensemble model  $(d_{T})$ , while the yellow circles are correctly classified and
Figure 2: The visualization of synthetic data and decision boundaries of global model (student) and ensemble model (teacher). Left panel: synthetic data (red circles) are far away from the decision boundary  $d_{T}$ . Middle panel: synthetic data (black circles) near the decision boundaries  $d_{T}$ . Right panel: synthetic data (yellow and purple circles) cross over the decision boundary  $d_{T}$ .
the purple circles are misclassified by the global model  $(d_S)$ . For the conditional generator  $G(\cdot)$  that takes label information as one of the inputs, yellow and purple circles can mislead the generator, thereby leading to  $d_S$  approximating  $d_T$  with a large deviation, as shown in Fig. 2 (c). Based on the above observation, we reckon that the synthetic data  $\boldsymbol{s} = G(\boldsymbol{h} = o(\boldsymbol{z}, y), \boldsymbol{w})$  is useful if it is classified as  $y$  by the ensemble model but classified not as  $y$  by the global model. To realize it, we maximize the logits discrepancy between the global model and the ensemble model on synthetic data with black circles by leveraging Kullback-Leibler divergence loss, which takes the form:
$$
\mathcal {L} _ {t r a n} = - \varepsilon \cdot K L \left(\sum_ {i \in S _ {t}} \tau_ {i, y} f _ {i} (\boldsymbol {s}, \boldsymbol {\theta} _ {i}), f (\boldsymbol {s}, \boldsymbol {\theta})\right), \tag {3}
$$
where  $KL$  is Kullback-Leibler divergence function and  $f(\boldsymbol{s},\boldsymbol{\theta})$  denotes the logits output of the global model on  $\boldsymbol{s}$  with label  $y$ . Note that  $\varepsilon = 1$  if  $\arg \max f(\boldsymbol{s},\boldsymbol{\theta}) \neq y$  and  $\arg \max \sum_{i\in S_t}\tau_{i,y}f_i(\boldsymbol{s},\boldsymbol{\theta}_i) = y$  hold, otherwise  $\varepsilon = 0$ .
We would like to point out that the existing works [48] and [49] are in line with our research perspective on the transferability of generator, which aims to generate more synthetic data with black circles. However, they do not carefully frame learning objective for enhancing the transferability of generator. Concretely, [48] does not consider the type of synthetic data, i.e.,  $\varepsilon = 1$  always holds, thus inducing the generation of synthetic data with yellow and purple circles. ( $\triangle$ ) [49] focuses on synthetic data satisfying  $\arg \max f(s,\theta)\neq \arg \max \sum_{i\in S_t}\tau_{i,y}f_i(s,\theta_i)$ , but enables the generation of synthetic data with purple circles yet.  $(\nabla)^2$
Diversity. Although we enable  $G(\cdot)$  to generate synthetic data that falls around the real data by optimizing  $\mathcal{L}_{fid}$  and  $\mathcal{L}_{tran}$ , the diversity of synthetic data is insufficient. Due to the fact that the generator may get stuck in local equilibria as the training proceeds, model collapse occurs [57, 58]. In this case, the generator may produce similar data points for each class with little diversity. Also, the synthetic data points may not differ significantly among classes. This causes the empirical distribution estimated by  $G(\cdot)$  to cover only a small manifold in the real data space, and thus only partial knowledge of the ensemble model is extracted. To alleviate this issue, we introduce a diversity loss  $\mathcal{L}_{div}$  with label information to increase the diversity of synthetic data as follows:
$$
\mathcal {L} _ {d i v} = \exp \left(- \sum_ {j, k \in [ B ]} \| \boldsymbol {s} _ {j} - \boldsymbol {s} _ {k} \| _ {2} * \| \boldsymbol {h} _ {j} - \boldsymbol {h} _ {k} \| _ {2} / B ^ {2}\right), \tag {4}
$$
where  $B$  denotes the batch size and  $\pmb{s}_{j / k} = G(\pmb{h}_{j / k} = o(\pmb{z}_{j / k},y_{j / k}),\pmb {w})$ . Intuitively,  $\mathcal{L}_{div}$  takes  $\| h_j - h_k\| _2$  as a weight, and then multiplies it by the corresponding  $\| s_j - s_k\| _2$  in each batch  $B$ , thus imposing a larger weight on the synthetic data points pair  $(s_j$  and  $s_k)$  at the more distant input pair  $(h_j$  and  $h_k)$ . Notably, we merge the random noise  $\textbf{z}$  with label  $y$  as the input of  $G(\cdot)$  to overcome spurious solutions [53]. Further, we propose a multiplicative merge operator, i.e.,  $o(z,y) = z\times \mathcal{E}(y)$  where  $\mathcal{E}$  is a trainable embedding and  $\times$  means vector element-wise product. We find that our merge operator enables synthetic data with more diversity compared to others, possibly because the label information is effectively absorbed into the stochasticity of  $z$  by multiplying them when updating  $\mathcal{E}$ . See Section 4.3 for more details and empirical justification.
Combining  $\mathcal{L}_{fid}$ ,  $\mathcal{L}_{tran}$  and  $\mathcal{L}_{div}$ , the overall objective of the generator can be formalized as follows:
$$
\mathcal {L} _ {\text {g e n}} = \mathcal {L} _ {\text {f i d}} + \beta_ {\text {t r a n}} \cdot \mathcal {L} _ {\text {t r a n}} + \beta_ {\text {d i v}} \cdot \mathcal {L} _ {\text {d i v}}, \tag {5}
$$
where  $\beta_{tran}$  and  $\beta_{div}$  are tunable hyper-parameters. Of note, the synthetic data generated by a well-trained generator should be visually distinct from the real data for privacy protection, while it can capture the common knowledge of the local models to ensure similarity to the real data distribution for utility. More privacy protection is discussed in Appendices I and J.
3.2 Robust Model Distillation
Now we update the global model. Normally, the global model attempts to learn as much as possible logits outputs of the ensemble model on the synthetic data generated by the generator based on knowledge distillation [40, 46, 48, 49]. The updated global model and the ensemble model are then served to train  $G(\cdot)$  with the goal of generating synthetic data that maximizes the mismatch
Note that  $\triangle, \bigtriangledown$  and  $\diamond$  denote the strategies for the transferability of generator in [48], in [49] and in this paper, respectively.
between them in terms of logits outputs (see transferability discussed in the previous section). This adversarial game enables the generator to rapidly explore the training space of the local models to help knowledge transfer from them to the global model. However, it also leads to dramatic shifts in the output distribution of  $G(\cdot)$  across communication rounds under heterogeneous FL scenario (i.e., distribution shifts), causing the global model to catastrophically forget useful knowledge gained in previous rounds. To tackle the deficiency, we propose to equip the server with a generator  $\widetilde{G}(\cdot)$  parameterized by  $\widetilde{w}$  that is an exponential moving average (EMA) copy of  $G(\cdot)$ . Its parameters at the  $t^{th}$  communication round are computed by
$$
\widetilde {\boldsymbol {w}} ^ {t} = \lambda \cdot \widetilde {\boldsymbol {w}} ^ {t - 1} + (1 - \lambda) \cdot \boldsymbol {w} ^ {t}, \tag {6}
$$