SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiaqi Wang $^{1}$  Xingyi Yang $^{2}$  Suhan Cui $^{1}$  Liwei Che $^{1}$
Lingjuan Lyu $^{3}$  Dongkuan Xu $^{4}$  Fenglong Ma $^{1*}$
$^{1}$ The Pennsylvania State University  $^{2}$ National University of Singapore
$^{3}$ Sony AI  $^{4}$ North Carolina State University
{jqwang,sxc6192,lfc5481,fenglong}@psu.edu,
xyang@u.nus.edu, lingjuan.lv@sony.com, dzu27@ncsu.edu
This paper focuses on addressing the practical yet challenging problem of model heterogeneity in federated learning, where clients possess models with different network structures. To track this problem, we propose a novel framework called pFedHR, which leverages heterogeneous model reassembly to achieve personalized federated learning. In particular, we approach the problem of heterogeneous model personalization as a model-matching optimization task on the server side. Moreover, pFedHR automatically and dynamically generates informative and diverse personalized candidates with minimal human intervention. Furthermore, our proposed heterogeneous model reassembly technique mitigates the adverse impact introduced by using public data with different distributions from the client data to a certain extent. Experimental results demonstrate that pFedHR outperforms baselines on three datasets under both IID and Non-IID settings. Additionally, pFedHR effectively reduces the adverse impact of using different public data and dynamically generates diverse personalized models in an automated manner<sup>2</sup>.
Federated learning (FL) aims to enable collaborative machine learning without the need to share clients' data with others, thereby upholding data privacy [1, 2, 3]. However, traditional federated learning approaches [2, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] typically enforce the use of an identical model structure for all clients during training. This constraint poses challenges in achieving personalized learning within the FL framework. In real-world scenarios, clients such as data centers, institutes, or companies often possess their own distinct models, which may have varying structures. Training on top of their original models should be a better solution than deploying new ones for collaborative purposes. Therefore, a practical solution lies in fostering heterogeneous model cooperation within FL, while preserving individual model structures. Only a few studies have attempted to address the challenging problem of heterogeneous model cooperation in FL [16, 17, 18, 19, 20], and most of them incorporate the use of a public dataset to facilitate both cooperation and personalization [17, 18, 19, 20]. However, these approaches still face several key issues:
- Undermining personalization through consensus: Existing methods often generate consensual side information, such as class information [17], logits [18, 21], and label-wise representations [22], using public data. This information is then exchanged and used to conduct average operations on the server, resulting in a consensus representation. However, this approach poses privacy and
Source code can be found in the link https://github.com/JackqqWang/pfedHR
Figure 1: Performance changes when using different public data. pFedHR is our proposed model.
security concerns due to the exchange of side information [23]. Furthermore, the averaging process significantly diminishes the unique characteristics of individual local models, thereby hampering model personalization. Consequently, there is a need to explore approaches that can achieve local model personalization without relying on consensus-based techniques.
- Excessive reliance on prior knowledge for distillation-based approaches: Distillation-based techniques, such as knowledge distillation (KD), are commonly employed for heterogeneous model aggregation in FL [19, 20, 24]. However, these techniques necessitate the predefinition of a shared model structure based on prior knowledge [20]. This shared model is then downloaded to clients to guide their training process. Consequently, handcrafted models can heavily influence local model personalization. Additionally, a fixed shared model structure may be insufficient for effectively guiding personalized learning when dealing with a large number of clients with non-IID data. Thus, it is crucial to explore methods that can automatically and dynamically generate client-specific personalized models as guidance.
- Sensitivity to the choice of public datasets: Most existing approaches use public data to obtain guidance information, such as logits [18, 21] or a shared model [20], for local model personalization. The design of these approaches makes public data and model personalization tightly bound together. Thus, they usually choose the public data with the same distribution as the client data. Therefore, using public data with different distributions from client data will cause a significant performance drop in existing models. Figure 1 illustrates the performance variations of different models trained on the SVHN dataset with different public datasets (detailed experimental information can be found in Section 4.4). The figure demonstrates a significant performance drop when using alternative public datasets. Consequently, mitigating the adverse impact of employing diverse public data remains a critical yet practical research challenge in FL.
Motivation & Challenges. In fact, both consensus-based and distillation-based approaches aim to learn aggregated and shared information used as guidance in personalized local model training, which is not an optimal way to achieve personalization. An ideal solution is to generate a personalized model for the corresponding client, which is significantly challenging since the assessable information on the server side can only include the uploaded client models and the public data. To avoid the issue of public data sensitivity, only client models can be used. These constraints motivate us to employ the model reassembly technique [25] to generate models first and then select the most matchable personalized model for a specific client from the generations.
To this end, we will face several new challenges. (C1) Applying the model reassembly technique will result in many candidates. Thus, the first challenge is how to get the optimal candidates. (C2) The layers of the generated candidates are usually from different client models, and the output dimension size of the first layer may not align with the input dimension size of the second layer, which leads to the necessity of network layer stitching. However, the parameters of the stitched layers are unknown. Therefore, the second challenge is how to learn those unknown parameters in the stitched models. (C3) Even with well-trained stitched models, digging out the best match between a client model and a stitched model remains a big challenge.
Our Approach. To simultaneously tackle all the aforementioned challenges, we present a novel framework called pFedHR, which aims to achieve personalized federated learning and address the issue of heterogeneous model cooperation (as depicted in Figure 2). The pFedHR framework comprises two key updates: the server update and the client update. In particular, to tackle C3, we approach the issue of heterogeneous model personalization from a model-matching optimization perspective on the server side (see Section 3.1.1). To solve this problem, we introduce a novel heterogeneous model reassembly technique in Section 3.1.2 to assemble models uploaded from clients, i.e.,  $\{\mathbf{w}_t^1,\dots ,\mathbf{w}_t^B\}$ , where  $B$  is the number of active clients in the  $t$ -the communication round. This technique involves the dynamic grouping of model layers based on their functions, i.e., layer-wise decomposition and function-driven layer grouping in Figure 2. To handle C1, a heuristic
Figure 2: Overview of the proposed pFedHR.  $K$  is the number of clusters.
rule-based search strategy is proposed in reassembly candidate generation to assemble informative and diverse model candidates using clustering results. Importantly, all layers in each candidate are derived from the uploaded client models.
Prior to matching the client model with candidates, we perform network layer stitching while maximizing the retention of information from the original client models (Section 3.1.3). To tackle C2, we introduce public data  $\mathcal{D}_p$  to help the finetuning of the stitched candidates, i.e.,  $\{\tilde{\mathbf{c}}_t^1,\dots ,\tilde{\mathbf{c}}_t^M\}$ , where  $M$  is the number of generated candidates. Specifically, we employ labeled OR unlabeled public data to fine-tune the stitched and client models and then calculate similarities based on model outputs. Intuitively, if two models are highly related to each other, their outputs should also be similar. Therefore, we select the candidate with the highest similarity as the personalized model of the corresponding client, which results in matched pairs  $\{\{\mathbf{w}_t^1,\tilde{\mathbf{c}}_t^i\} ,\dots ,\{\mathbf{w}_t^B,\tilde{\mathbf{c}}_t^m\} \}$  in Figure 2. In the client update (Section 3.2), we treat the matched personalized model as a guidance mechanism for client parameter learning using knowledge distillation<sup>3</sup>.
It is worth noting that we minimally use public data during our model learning to reduce their adverse impact. In our model design, the public data are used for clustering layers, fine-tuning the stitched candidates, and guiding model matching. Clustering and matching stages use public data to obtain the feedforward outputs as guidance and do not involve model parameter updates. Only in the fine-tuning stage, the stitched models' parameters will be updated based on public data. To reduce its impact as much as possible, we limit the number of finetuning epochs during the model implementation. Although we cannot thoroughly break the tie between model training and public data, such a design at least greatly alleviates the problem of public data sensitivity in FL.
Contributions. Our work makes the following key contributions: (1) We introduce the first personalized federated learning framework based on model reassembly, specifically designed to address the challenges of heterogeneous model cooperation. (2) The proposed pFedHR framework demonstrates the ability to automatically and dynamically generate personalized candidates that are both informative and diverse, requiring minimal human intervention. (3) We present a novel heterogeneous model reassembly technique, which effectively mitigates the adverse impact caused by using public data with distributions different from client data. (4) Experimental results show that the pFedHR framework achieves state-of-the-art performance on three datasets, exhibiting superior performance under both IID and Non-IID settings when compared to baselines employing labeled and unlabeled public datasets.
Model Heterogeneity in Federated Learning. Although many federated learning models, such as FedAvg [4], FedProx [2], Per-FedAvg [26], PFedMe [27], and PFedBayes[28], have been proposed recently, the focus on heterogeneous model cooperation, where clients possess models with diverse structures, remains limited. It is worth noting that in this context, the client models are originally distinct and not derived from a shared, large global model through the distillation of subnetworks [29, 30]. Existing studies on heterogeneous model cooperation can be broadly categorized based on whether they utilize public data for model training. FedKD [16] aims to achieve personalized models
Note that the network architecture of both local model and personalized model are known, and thus, there is no human intervention in the client update.
Table 1: A comparison between existing heterogeneous model cooperation works and our pFedHR.
Approach	Public Dataset	Model Characteristics
W. Label	W.o. Label	Upload and Download	Aggregation	Personalization
FedDF[19]	X	✓	parameters	ensemble distillation	X
FedKEMF[20]	X	✓	parameters	mutual learning	✓
FCCL [18]	X	✓	logits	average	✓
FedMD[17]	✓	X	class scores	average	✓
FedGH [22]	✓	X	label-wise representations	average	✓
pFedHR	✓	✓	parameters	model reassembly	✓
for each client without employing public data by simultaneously maintaining a large heterogeneous model and a small homogeneous model on each client, which incurs high computational costs.
Most existing approaches leverage public data to facilitate model training. Among them, FedDF [19], FedKEMF [20], and FCCL [18] employ unlabeled public data. However, FedDF trains a global model with different settings compared to our approach. FedKEMF performs mutual knowledge distillation learning on the server side to achieve model personalization, requiring predefined model structures. FCCL averages the logits provided by each client and utilizes a consensus logit as guidance during local model training. It is worth noting that the use of logits raises concerns regarding privacy and security [23, 31]. FedMD [17] and FedGH [22] employ labeled public data. These approaches exchange class information or representations between the server and clients and perform aggregation to address the model heterogeneity issue. However, similar to FCCL, these methods also introduce privacy leakage concerns. In contrast to existing work, we propose a general framework capable of utilizing either labeled OR unlabeled public data to learn personalized models through heterogeneous model reassembly. We summarize the distinctions between existing approaches and our framework in Table 1.
Neural Network Reassembly and Stitching. As illustrated in Table 1, conventional methods are primarily employed in existing federated learning approaches to obtain personalized client models, with limited exploration of model reassembly. Additionally, there is a lack of research investigating neural network reassembly and stitching [25, 32, 33, 34, 35, 36] within the context of federated learning. For instance, the work presented in [33] proposes three algorithms to merge two models within the weight space, but it is limited to handling only two models as input. In our setting, multiple models need to be incorporated into the model reassembly or aggregation process. Furthermore, both [25] and [34] focus on pre-trained models, which differ from our specific scenario.
Our model pFedHR incorporates two key updates: the server update and the local update, as depicted in Figure 2. Next, we provide the details of our model design starting with the server update.
3.1 Server Update
During each communication round  $t$ , the server will receive  $B$  heterogeneous client models with parameters denoted as  $\{\mathbf{w}_t^1, \mathbf{w}_t^2, \dots, \mathbf{w}_t^B\}$ . As we discussed in Section 1, traditional approaches have limitations when applied in this context. To overcome these limitations and learn a personalized model  $\hat{\mathbf{w}}_t^n$  that can be distributed to the corresponding  $n$ -th client, we propose a novel approach that leverages the publicly available data  $\mathcal{D}_p$  stored on the server to find the most similar aggregated models learned from  $\{\mathbf{w}_t^1, \mathbf{w}_t^2, \dots, \mathbf{w}_t^B\}$  for  $\mathbf{w}_t^n$ .
3.1.1 Similarity-based Model Matching
Let  $g(\cdot, \cdot)$  denote the model aggregation function, which can automatically and dynamically obtain  $M$  aggregated model candidates as follows:
$$
\left\{\mathbf {c} _ {t} ^ {1}, \dots , \mathbf {c} _ {t} ^ {M} \right\} = g \left(\left\{\mathbf {w} _ {t} ^ {1}, \mathbf {w} _ {t} ^ {2}, \dots , \mathbf {w} _ {t} ^ {B} \right\}, \mathcal {D} _ {p}\right), \tag {1}
$$
where  $g(\cdot, \cdot)$  will be detailed in Section 3.1.2, and  $\mathbf{c}_k^m$  is the  $m$ -th generated model candidate learned by  $g(\cdot, \cdot)$ . Note that  $\mathbf{c}_k^m$  denotes the model before network stitching.  $M$  is the total number of candidates, which is not a fixed number and is estimated by  $g(\cdot, \cdot)$ . In such a way, our goal is to optimize the following function:
$$
\mathbf {c} _ {t} ^ {*} = \underset {\mathbf {c} _ {t} ^ {m}; m \in [ 1, M ]} {\arg \max } \left\{\operatorname {s i m} \left(\mathbf {w} _ {t} ^ {n}, \mathbf {c} _ {t} ^ {1}; \mathcal {D} _ {p}\right), \dots , \operatorname {s i m} \left(\mathbf {w} _ {t} ^ {n}, \mathbf {c} _ {t} ^ {M}; \mathcal {D} _ {p}\right) \right\}, \forall n \in [ 1, B ], \tag {2}
$$
Algorithm 1: Reassembly Candidate Search
input:Layer clusters  $\{\mathcal{G}_t^1,\mathcal{G}_t^2,\dots ,\mathcal{G}_t^K \}$  , operation type set  $\mathcal{O}$  , rule set  $\mathcal{R}$  output:  $\{\mathbf{c}_t^1,\dots ,\mathbf{c}_t^M\}$    
1 Initialize  $\mathcal{C}_t = \emptyset$  .   
for  $k\gets 1,\dots ,K$  do   
//  $Q_{k}$  is the number of operation-layer pairs in  $\mathcal{G}_t^k$    
for  $q\gets 1,\dots ,Q_k$  do Initialize an empty candidate  $\mathbf{c}_q = []$  , operation type set  $\mathcal{O}_q = \emptyset$  group id set  $\kappa_q = \emptyset$  Select the q-th layer-operation pair  $(\mathcal{L}_{t,i}^{n},O_{i}^{n})$  from group  $\mathcal{G}_k$  Add  $(\mathcal{L}_{t,i}^{n},O_{i}^{n})$  to  $\mathbf{c}_q$  , add  $O_{i}^{n}$  to  $\mathcal{O}_q$  , add  $k$  to  $\kappa_q$  . for  $k^{\prime}\gets 1,\dots ,K$  do Check a pair from  $\mathcal{G}_{k^{\prime}}$  whether it satisfies: R1 (layer order): the layer index should be larger than that of the last layer added to  $\mathbf{c}_q$  and R2 (operation order): the operation type should be followed by the previous type in  $\mathbf{c}_q$  ; if True then Add the pair to  $\mathbf{c}_q$  , add its operation type to  $\mathcal{O}_q$  , add  $k^{\prime}$  to  $\kappa_q$  Move to the next pair;   
Check  $\mathcal{O}_q$  and  $\kappa_q$  with:   
R3 (complete operation): the size of  $\mathcal{O}_q$  should be equal to that of O, and   
R4 (diverse group): the size of  $\kappa_q$  should be equal to  $K$  . if True then Add the candidate  $\mathbf{c}_q$  to  $\mathcal{C}_t$  .   
return :  $C_t = \{\mathbf{c}_t^1,\dots ,\mathbf{c}_t^M\}$
where  $\mathbf{c}_t^*$  is the best matched model for  $\mathbf{w}_t^n$ , which is also denoted as  $\hat{\mathbf{w}}_t^n = \mathbf{c}_t^*$ .  $\sin (\cdot ,\cdot)$  is the similarity function between two models, which will be detailed in Section 3.1.3.
3.1.2 Heterogeneous Model Reassembly  $-g(\cdot, \cdot)$
To optimize Eq. (2), we need to obtain  $M$  candidates using the heterogenous model aggregation function  $g(\cdot)$  in Eq. (1). To avoid the issue of predefined model architectures in the knowledge distillation approaches, we aim to automatically and dynamically learn the candidate architectures via a newly designed function  $g(\cdot, \cdot)$ . In particular, we propose a decomposition-grouping-reassembly method as  $g(\cdot, \cdot)$ , including layer-wise decomposition, function-driven layer grouping, and reassembly candidate generation.
Layer-wise Decomposition. Assume that each uploaded client model  $\mathbf{w}_t^n$  contains  $H$  layers, i.e.,  $\mathbf{w}_t^n = [(\mathbf{L}_{t,1}^n,O_1^n),\dots ,(\mathbf{L}_{t,H}^n,O_E^n)]$ , where each layer  $\mathbf{L}_{t,h}^n$  is associated with an operation type  $O_{e}^{n}$ . For example, a plain convolutional neural network (CNN) usually has three operations: convolution, pooling, and fully connected layers. For different client models,  $H$  may be different. The decomposition step aims to obtain these layers and their corresponding operation types.
Function-driven Layer Grouping. After decomposing layers of client models, we group these layers based on their functional similarities. Due to the model structure heterogeneity in our setting, the dimension size of the output representations from layers by feeding the public data  $\mathcal{D}_p$  to different models will be different. Thus, measuring the similarity between a pair of layers is challenging, which can be resolved by applying the commonly used centered kernel alignment (CKA) technique [37]. In particular, we define the distance metric between any pair of layers as follows:
$$
\operatorname {d i s} \left(\mathbf {L} _ {t, i} ^ {n}, \mathbf {L} _ {t, j} ^ {b}\right) = \left(\operatorname {C K A} \left(\mathbf {X} _ {t, i} ^ {n}, \mathbf {X} _ {t, j} ^ {b}\right) + \operatorname {C K A} \left(\mathbf {L} _ {t, i} ^ {n} \left(\mathbf {X} _ {t, i} ^ {n}\right), \mathbf {L} _ {t, i} ^ {b} \left(\mathbf {X} _ {t, j} ^ {b}\right)\right)\right) ^ {- 1}, \tag {3}
$$
where  $\mathbf{X}_{t,i}^n$  is the input data of  $\mathbf{L}_{t,i}^n$ , and  $\mathbf{L}_{t,i}^n(\mathbf{X}_{t,i}^n)$  denotes the output data from  $\mathbf{L}_{t,i}^n$ . This metric uses  $\mathrm{CKA}(\cdot, \cdot)$  to calculate the similarity between both input and output data of two layers.
Based on the defined distance metric, we conduct the K-means-style algorithm to group the layers of  $B$  models into  $K$  clusters. This optimization process aims to minimize the sum of distances between
all pairs of layers, denoted as  $\mathcal{L}_t$ . The procedure can be described as follows:
$$
\min  \mathcal {L} _ {t} = \min  _ {\delta_ {b, h} ^ {a} \in \{0, 1 \}} \sum_ {k = 1} ^ {K} \sum_ {b = 1} ^ {B} \sum_ {h = 1} ^ {H} \delta_ {b, h} ^ {k} \left(\operatorname {d i s} \left(\mathbf {L} _ {t} ^ {k}, \mathbf {L} _ {t, h} ^ {b}\right)\right), \tag {4}
$$
where  $\mathbf{L}_t^k$  is the center of the  $k$ -th cluster.  $\delta_{b,h}^{k}$  is the indicator. If the  $h$ -th layer of  $\mathbf{w}_t^b$  belongs to the  $k$ -th cluster, then  $\delta_{b,h}^{k} = 1$ . Otherwise,  $\delta_{b,h}^{k} = 0$ . After the grouping process, we obtain  $K$  layer clusters denoted as  $\{\mathcal{G}_t^1,\mathcal{G}_t^2,\dots ,\mathcal{G}_t^K\}$ . There are multiple layers in each group, which have similar functions. Besides, each layer is associated with an operation type.
Reassembly Candidate Generation. The last step for obtaining personalized candidates  $\{\mathbf{c}_t^1,\dots ,\mathbf{c}_t^M\}$  is to assemble the learned layer-wise groups  $\{\mathcal{G}_t^1,\mathcal{G}_t^2,\dots ,\mathcal{G}_t^K\}$  based on their functions. To this end, we design a heuristic rule-based search strategy as shown in Algorithm 1. Our goal is to automatically generate informative and diverse candidates.
Generally, an informative candidate needs to follow the design of handcrafted network structures. This is challenging since the candidates are automatically generated without human interventions and prior knowledge. To satisfy this condition, we require the layer orders to be guaranteed  $(\mathcal{R}_1$  in Line 10). For example, the  $i$ -th layer from the  $n$ -th model, i.e.,  $\mathbf{L}_{t,i}^n$ , in a candidate must be followed by a layer with an index  $j > i$  from other models or itself. Besides, the operation type also determines the quality of a model. For a CNN model, the fully connected layer is usually used after the convolution layer, which motivates us to design the  $\mathcal{R}_2$  operation order rule in Line 11.
Only taking the informativeness principle into consideration, we may generate a vast number of candidates with different sizes of layers. Some candidates may be a subset of others and even worse with low quality. Besides, the large number of candidates will increase the computational burden of the server. To avoid these issues and further obtain high-quality candidates, we use the diversity principle as the filtering rule. A diverse and informative model should contain all the operation types, i.e., the  $\mathcal{R}_3$  complete operation rule in Line 16. Besides, the groups  $\{\mathcal{G}_t^1,\dots ,\mathcal{G}_t^K\}$  are clustered based on their layer functions. The requirement that layers of candidates must be from different groups should significantly increase the diversity of model functions, which motivates us to design the  $\mathcal{R}_4$  diverse group rule in Line 17.
3.1.3 Similarity Learning with Layer Stitching -  $\mathbf{sim}(\cdot ,\cdot)$