Resource-Aware Federated Self-Supervised Learning with Global Class Representations

Authors: mingyi  li
xiao  zhang
qi  wang
tengfei  liu
ruofan  wu
weiqiang  wang
fuzhen  zhuang
hui  xiong
dongxiao  yu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:21:44

Abstract: Due to the heterogeneous architectures and class skew, the global representation models training in resource-adaptive federated self-supervised learning face with tricky challenges: $\textit{deviated representation abilities}$ and $\textit{inconsistent representation spaces}$. 
In this work, we are the first to propose a multi-teacher knowledge distillation framework, namely $\textit{FedMKD}$, to learn global representations with whole class knowledge from heterogeneous clients even under extreme class skew. Firstly, the adaptive knowledge integration mechanism is designed to learn better representations from all heterogeneous models with deviated representation abilities. Then the weighted combination of the self-supervised loss and the distillation loss can support the global model to encode all classes from clients into a unified space. Besides, the global knowledge anchored alignment module can make the local representation spaces close to the global spaces, which further improves the representation abilities of local ones. Finally, extensive experiments conducted on two datasets demonstrate the effectiveness of $\textit{FedMKD}$ which outperforms state-of-the-art baselines 4.78\% under linear evaluation on average.

Main paper: 
 introduction: 
 The federated self-supervised learning (Fed-SSL) has emerged as a highly promising paradigm due to the extremely limited labeled data in real-world scenarios [4, 25, 30] . The Fed-SSL mechanism can learn common representations collaboratively across all the clients without labeled data [7, 8] , which could enable the aggregation of knowledge from diverse unlabeled data sources and overcome the limitations caused by the high cost and scarcity of labeled data [24, 28, 31] .Traditional Fed-SSL methods usually assume that each client should train the identical architecture model, such as FedU [33] , FedEMA [34] , FedCA [30] . But it would not be easy in resource-limited scenarios, especially for the existing arsing large-scale models [1, 26] . As shown in Fig. 1 , client A, client B and client C might train heterogeneous representation models due to the varying system resources. In addition, real-world data often exhibits skewed class distributions across clients.Therefore, how to learn global class representations under the heterogeneous architectures and class skew in resource-aware Fed-SSL paradigm is challenging, particularly comparing with existing F edU 2 [17] , FedX [5] with identical architectures.Although both Hetero-SSFL [20] and FedFoA [19] consider the heterogeneous client models, they can not learn a global representation model. In order to aggregate the knowledge from the clients to form global class representations, some tricky challenges arise. (1) Deviated representation abilities. Even for the same data samples, the different models might encode them into different latent spaces with deviated representation abilities. For example, client A and client B all have images with dog, cat, tiger, but client model A can encode cat, tiger well into different clusters, client model B can only learn better representations of dog. So how could global representation models take advantage of the best of both client models? (2) Inconsistent representation spaces. The skewed class distributions across clients lead to inconsistent representation spaces. For example in Fig. 1 , comparing with client A, client C has different kinds of images with dog, cat, airplane. Thus how to make global representation models encode the whole classes from all the clients well in a unified space? Therefore, different from the existing works, our goal is to break the gaps caused by the hybrid heterogeneity, which can learn the high-quality global representation model in federated self-supervised learning. Along this line, we propose FedMKD, a multiteacher knowledge distillation based resourceadaptive Fed-SSL framework, which can learn global representations over all classes from heterogeneous clients. First, an adaptive knowledge integration module is introduced to learn high-quality representations from all the heterogeneous models with deviated representation abilities. Then in order to encode all classes from clients in a unified space, the global model uses the weighted combination of selfsupervised loss and distillation loss to update. Besides, the global knowledge anchored alignment module is applied within the server to eliminate the inconsistency in representation spaces and reduce the burden on the clients. It uses global knowledge to additionally update the local models, which can not only make the local representation spaces close to the global space but also improve the representation capability of both local models and the global ones. Code is available at https://github. com/limee-sdu/FedMKD. The main contributions of this paper can be summarized as follows.• In resource-aware Fed-SSL, we are the first to delve into global class representation learning through revealing the deviated representation abilities and inconsistent representation spaces caused by the heterogeneous architectures and class skew.• We design a multi-teacher knowledge distillation framework, namely FedMKD, to adaptively aggregate positive knowledge from heterogeneous models with deviated representation abilities. Through combining the self-supervised loss and the distillation loss, FedMKD can encode skewed classes into a unified space.• Extensive experiments conducted on CIFAR-10 and CIFAR-100 show the representation abilities over all classes of the FedMKD perform better than state-of-the-art baselines. Our algorithm can improve 4.22% and 5.31% separately on the two chosen datasets. 
related work: 
 The federated self-supervised learning aims to learn high-quality representations from clients without large labeled datasets [11] . From the beginning, several works [12, 23] simply combine federated learning with self-supervised methods. Besides, FedU [33] designs a communication-efficient mech- ✓ = Client ✗ ✗ ✗ ✗ FedEMA [34] ✓ = Client ✗ ✗ ✗ ✗ L-DAWA [21] ✓ = Client ✗ ✗ ✗ ✗ FedX [5] ✓ = Client ✗ ✗ ✗ ✗ FedCA [30] ✓ = Client ✗ ✗ ✓ ✗ FLPD [29] ✓ = Client ✓ ✗ ✗ ✗ F edU 2 [17] ✓ = Client ✗ ✗ ✓ ✓ FedFoA [19] ✗ - ✓ ✗ ✓ ✗ Hetero-SSFL [20] ✗ - ✓ ✗ ✓ ✓ FedMKD(ours) ✓ ≥ Client ✓ ✓ ✓ ✓anism by only aggregating the online encoders under non-IID data. FedUTN [16] is proposed to use the aggregated online networks for the target network updating in the self-supervised framework. L-DAWA [21] proposes the layer-wise divergence aware weight aggregation to mitigate the influence of client bias. FedEMA [34] considers the divergence-aware moving average updating in clients, measuring the divergence between local models and global model. FedX [5] proposes a unsupervised federated learning framework to learn representations through a two-sided distillation method. However, all the above works intuitively gain the global model through parameters average due to the identical client models, which can not be applied in the heterogeneous clients setting directly. In addition, although FedCA [30] address the misaligned and inconsistent representation challenges by gathering features from clients, inducing potential privacy problems. FLPD [27, 29] introduces distillation method based similarity between prototypes from a labeled public dataset to update the local model. F edU 2 [17] focuses on mitigating representation collapse entanglement and obtaining unified representation spaces.Considering heterogeneous client models in federated self-supervised learning, Hetero-SSFL [6] introduces linear-CKA to align lower-dimensional representations between the local model and global model without architectural constraints. FedFoA [19] designs a factorization-based method to extract the cross-feature relation matrix from the local representations for aggregation. However, both Hetero-SSFL and FedFoA can not learn a global representation model considering the hybrid heterogeneity, which is the main focus of our work. The comparison details are shown in Table 1 . 
preliminaries: 
 The goal of federated unsupervised learning is to learn the generalized representation for some downstream tasks from several distributed unlabeled data sources. A federated learning setting consists of a central server and N clients. Each client k contains a local unlabeled dataset D k , and the server contains a public unlabeled dataset D pub . The local objective at k-th client isEQUATIONto minimize the expected local loss of client k on local dataset D k and ξ k is the unlabeled data.In traditional FL settings, it's assumed that {θ k } are identical and gain the global model using θ = N k=1 p k θ k , where p k is the weight of k-th client. But in real-world cross-device scenarios, each client might have a unique model and the architecture of the model might be different, which means that traditional aggregation methods are not available. We assume that θ k is not similar to others, and use {θ k } to collaborate in training the larger global model θ in server. Here we define the global update function isEQUATIONwhere L(θ, ξ) is global loss function in server and ξ is the unlabeled data sampled from global dataset.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: In this work, we provide the full set of theoretical assumptions and the complete process process of the proposed multi-teacher knowledge distillation framework FedMKD in the Appendix.B.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
designed fedmkd method: 
 We propose a multi-teacher knowledge distillation based federated self-supervised learning framework FedMKD, which is shown in Fig. 2 . In FedMKD, besides the local self-supervised learning (Sec. Each client performs self-supervised contrastive learning using an asymmetric Siamese network, inspired by BYOL [3] . The model M comprises an online encoder θ and a target encoder ϕ, both sharing the same architecture, with the online network incorporating an additional predictor p. That is M = {p(ϕ(•)), θ(•)}. Given an unlabeled image x, we can obtain two augmented views, v and v ′ , serving as inputs to online and target networks, respectively. The loss function is defined as follows:EQUATIONwhere r = θ(v) and r ′ = ϕ(v ′ ). This loss encourages the online network to produce representation p(r) that is similar to the positive sample generated by the target network r ′ . We then exchange the views, feeding v ′ to the online network and v to the target network, to compute L ′ self . At each training step, we use stochastic gradient descent to minimizeEQUATIONThe target network helps to provide regression targets to train the online network. Choosing α ∈ [0, 1] as the target decay rate, we employ the exponential moving average (EMA) of the online network to update ϕ:EQUATIONUsing this self-supervised training method, the model learns intricate representations from unlabeled data, capturing high-level features and patterns inherent in the dataset.In contrast to homogeneous federated learning, the presence of model heterogeneity poses a challenge: direct aggregation of local models into a global model is not feasible. To overcome this, we design a multi-teacher knowledge distillation mechanism to transfer local knowledge to the server.Given a batch of data B, the representation from the teacher model is denoted as r t and that from the student model as r s , the knowledge distillation loss is defined as follows:EQUATION)where τ is the temperature parameter controlling entropy and sim(•) is the similarity function between two representations.Then we extend this knowledge distillation learning method to multi-teacher. Although the data is heterogeneous, the knowledge of each local model is valuable, each local model captures the unique characteristics of local data. Our goal is to integrate the positive knowledge of all clients to guide the global model in learning a general representation of unlabeled data. We design a multi-teacher adaptive knowledge integration distillation that can adaptively weigh the representations from clients.Given a sample x i , representation from the n-th local model is R n,i ∈ R d where d is the dimension of the representation. Following [3] , a fully connected layer is employed to project the representation into a lower-dimensional space, enhancing the discriminate power of the learned representations. So, we map the representation from the global model R s,i into the same lower latent space, obtainingEQUATIONwhere r n,i , r s,i ∈ R k , k is the dimension of the new latent space and g(•) is the projector.In addition, we introduce an adapter module to learn instance-level teacher importance weights for knowledge integration. After getting r n,i , an attention block is used to generate the weighted sum of them. In this context, representation from global model r s,i is treated as the query, while those from local models R = [r 1,i , r EQUATIONwhere ri means the aggregated representation, and Attn(•) denotes the attention block.Returning to the knowledge distillation for unlabeled data proposed earlier, we treat the aggregate representation as the positive sample, and the remaining samples in the batch as the negative sample.The adaptive weight multi-teacher knowledge distillation loss function is expressed as follows:EQUATIONAbove all, the weighted combined loss for the global model is presented as:EQUATION)where γ is a hyper-parameter controlling the weight of the distillation process.As we mentioned, the representations from different models are inconsistent and the representation abilities of models are also deviated. So we introduce the global knowledge anchored alignment mechanism that each local model uses the global model as an anchor. It ensures that the local representation spaces are closer to the global ones.Unlike methods such as FedX [5] and MOON [15] , which align local models to the global model locally, our approach aims to train a better global encoder tailored for resource-constrained federated learning scenarios. EQUATIONandM = {θ ′ n , ϕ ′ n }.The training loss is updated toEQUATIONAccording to the idea of contrastive learning, the representations learned by the online network become more consistent with the knowledge captured by the target network. This global knowledge anchored contrastive learning suggests that the global model's knowledge is used as a positive example for the local model to train itself, thus making it more consistent with the target global network. Aligning local models with the global view representation helps create a comprehensive understanding of the overall data distribution. The process refines the knowledge acquired locally, ensuring that it contributes meaningfully to the overall federated learning process. Then, we use stochastic gradient descent to minimize L align to update theEQUATIONOnce the global model anchored alignment is finished, the server will send the online network of the local twin network θ ′ to the corresponding client to update the local model. The local online network θ n is replaced by the server-updated network θ ′ n . The target network is not replaced to retain more local knowledge and stabilize model training:EQUATIONso that, the local model can benefit from the alignment process and align to the representation under global view, which can further use the local data to train the model.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We propose a multi-teacher knowledge distillation framework FedMKD and describe this framework in detail in Sec.4. Furthermore, this paper disclose all the information needed to reproduce the main experimental results of it in Sec.5 and Appendix. C. 
experiments: 
 In this section, we evaluate the representations learned from our proposed global model FedMKD on CIFAR-10 and CIFAR-100. We first describe the experimental setup and baselines, and then analyze the performance in comparison to other methods. Due to the space limitation, further hyperparameter analysis and communication cost analysis are represented in Appendix D.We use CIFAR-10 and CIFAR-100 [13] datasets to train all the models. Both of them contain 50,000 training images and 10,000 testing images. To construct the public dataset, we sample 4000 data samples from the training set, then divide the remaining data into N partitions to simulate N clients. To assess the validity of the public dataset, we use two sampling methods to construct it. First, we use a random sampling method over all classes to generate public dataset 'IID'. And for the public dataset 'Partial', data is selected randomly from 40% classes in two datasets.We utilize three settings to simulate heterogeneous data distributions among all the clients. For the IID setting, each client contains the same number of samples from all classes. For the class setting, each client only has 10/N and 100/N classes on two datasets and the classes between clients have no overlap. For the non-IID setting, data heterogeneity levels are described by the Dirichlet distribution Dir(β) [10] , where smaller β represents stronger heterogeneity levels, here we choose β = 0.5.Regarding the self-supervised learning framework design within each client, we use ResNet18 [9] and VGG9 [22] as the encoder network and Multi-Layer Perception (MLP) as the predictor. In order to construct the model heterogeneous setting, 2 clients train the Resnet18 encoder while 3 clients use the VGG9. And for the global representation model, Resnet18 is selected as the encoder in server.Firstly, we select several federated knowledge distillation frameworks FedMD [14] , FedDF [18] ,FedET [2] , MOON [15] , MOON-KL that use unlabeled public dataset for distillation. We then replaced the local model with a self-supervised model to evaluate the process of knowledge distillation in our method. And several federated self-supervised learning frameworks FedU [33] , FedEMA [34] , Hetero-SSFL [20] are also chosen as baselines. To verify the client's knowledge can improve the global model, we also train the global model separately on the public dataset, denoted as Std. ResNet18.Following FedEMA [30] , we evaluate the performance of learned representations using linear and semi-supervised evaluation. Due to limited space, please refer to Appendix C for more details.Table 2 and 3 shows the linear evaluation results and semi-supervised evaluation results of FedMKD compared with all the baselines on CIFAR-10 and CIFAR-100. We can gain the following observation.On the whole, our FedMKD outperforms all baselines under different public dataset settings and different data heterogeneity level settings on both two datasets. Compared to the second-best results, FedMKD achieves significant improvement. On average, our model improves CIFAR-10 by 4.22% and CIFAR-100 by 5.31% under linear evaluation and gain 3.66% and 2.07% improvement on two dataset under semi-supervised evaluation.The effectiveness of our multi-teacher adaptive knowledge integration distillation can be approved when compared with FedMD, FedDF, MOON-KL and MOON. Although these methods all designed new federated knowledge distillation frameworks based on unlabeled public dataset, since the original local model is supervised, they prefer using the class information from logits to distill. When observing the result of FedET, we find that although it also designs a larger global model in server which improves the model a lot, the final result is not satisfied. This is also because it designs a distillation method based on knowledge of probability distribution over classes. Next, compared with the federated self-supervised method, our FedMKD also achieves better performance. FedU and FedEMA are not applicable in model heterogeneity setting, so we cannot evaluate their effectiveness. Hetero-SSFL gain the best performance among all baselines but is worse than ours. That's because it aims to train personalised client models. Only the alignment module cannot hold the inconsistent representation space perfectly. But our global model can directly generate representation from the global model, it avoids using the representation from inconsistent clients.Apart from the client data heterogeneity, we consider the influence of the public dataset distribution.Here, we construct two public datasets, one is 'IID' to the whole data distribution and the other only has partial classes. In both two settings, our FedMKD also gets the best performance. The overall performance of 'IID' public dataset is better than 'Partial' setting. That's because the global model adapts the self-supervised learning on public dataset, and the diversity of the sample is important, which can help the model explore a wide range of features and patterns present in the data. And it's observed that when the public dataset is 'IID', the performance increases with the decrease of the data heterogeneity level for CIFAR-10, but it doesn't apply to the CIFAR-100. It's because there are too many classes in CIFAR-100 and the number of samples in each class is not efficient.In order to evaluate the effectiveness of FedMKD, we use the dimensionality reduction algorithm t-sne to visualize the representation on the test dataset of CIFAR-10 from different encoders. As shown in Fig. 3 (b)(c)(d), the global models in FedMKD trained on both 'IID' and 'Partial' public datasets both achieve better clustering results than Standalone training and MOON. These results further verify our model can gain better generalized representations although the representation spaces of clients are inconsistent. There's also an averaged global model in MOON, but it cannot tackle the problem of inconsistent spaces well using the average method, so it only gains a poor clustering performance. Additionally, the performance on 'IID' public dataset is better than 'paritial' ones from the observation of cluster performance. This suggests that the number of classes seen by the global model also affects how well the global model can encode all classes in a unified space. Comparing the class distributions in Fig. 3 (c ) and (d), we can find that although these two global models are trained on different public datasets, the final cluster layout is similar, which can further prove that our global model can encode all classes from clients even if it never sees some classes during the training. ClientA ClientB GlobalIn order to investigate the effectiveness of different parts of FedMKD, we design these comparison experiments:• Standalone training: The global model is trained alone using the public dataset, without the knowledge from client models.• FedMKD KL : The global distillation function is replaced by the KL-divergence function to measure the similarity of the aggregated representation r and global representation r s .• FedMKD w/o adaptive : The adaptive knowledge integration module is removed, each local representation has the same weight to generate the aggregated representation. Above experiment is conducted on the both CIFAR-10 and CIFAR-100 dataset, using the IID public dataset. The results of the ablation study experiment are shown in Table 4 . When comparing FedMKD with FedMKD KL , significant performance drop can be observed when the knowledge distillation function is replaced by the KL-divergence function. Thus we can conclude that for self-supervised learning, we prefer performing knowledge distillation based on representation, but the KL-divergence cannot capture the distribution characteristics from them. Therefore, the appropriate distillation method is critically important in self-supervised learning due to the lack of labels. And a worse performance on both two datasets can be observed when we use the equal weight instead of the adaptive weight. Because the client models are heterogeneous, their representation capabilities are different and the representation spaces are also inconsistent, so intuitively average representations may reduce the information contained in the representation. Finally, when we remove the alignment module, the performance under each setting all decreases, which demonstrates that the alignment is not only beneficial to the local models, but also improves the whole training process.In order to explore the scalability of our proposed algorithm FedMKD, we add the experiment that the number of clients is 5, 10, 30 on CIFAR-10, and 40% clients use the VGG model and 60% use ResNet18. And we repartition the data for each client under Dir(β = 0.5) and set the public dataset distribution as 'IID'. The results are shown in Table 5 . We can find that as the number of clients increasing, the performance decreases. The reason is that the total number of data is fixed, if the number of clients increases, the number of data in each client will decrease, which further affect the performance of the local model. 
conclusion: 
 In this work, we focused on how to solve the deviated representation abilities and inconsistent representation spaces caused by the heterogeneous architectures and class skew in federated selfsupervised learning. We proposed a multi-teacher knowledge based federated self-supervised learning framework FedMKD to learn a global model. Firstly, the adaptive knowledge integration module could learn high-quality representation knowledge from heterogeneous models. And the combination of the self-supervised loss and the distillation loss enabled the global model to encode all classes from clients in a unified space. Then a global knowledge anchored alignment module improved the local representation models in server and fed it back to corresponding clients. The experiments conducted on two datasets demonstrated that our proposed FedMKD was state-of-the-art and outperformed existing methods.Algorithm 1: Algorithm of FedMKD Random initialize θ q,n,0 .5 else 6 {θ q,n,0 , ϕ q,n,0 } ← {θ ′ q-1,n , ϕ q-1,n,T }. l q,n,t = L self (θ q,n,t-1 ; ϕ q,n,t-1 ; D n ).10 θ q,n,t ← θ q,n,t-1 -η∇l q,n,t .11 ϕ q,n,t ← αϕ q,n,t-1 + (1 -α)θ q,n,t .12 end 13Upload θ q,n,t . for n = 1 to N do 20 r q,n = g(θ q,n,t (b)).21 end 22 R = [r q,1 , • • • , r q,N ] 23 r s = g(θ q,s (b)). 24 r = Attn(r s , R)25 l q,s = L self (θ q,s ; ϕ q,s ; b) + γL distill (θ q,s ; r s ; r).26 θ q,s ← θ q,s -η∇l q,s .27 ϕ q,s ← αϕ q-1,s + (1 -α)θ q,s . for n = 1 to N do 32 {θ ′ q,n , ϕ ′ q,n } ← {θ q,n,T , θ q,s } 33 l ′ q,n = L align (θ ′ q,n ; ϕ ′ q,n ; D pub ). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss the broader impacts of this work in the last paragraph of Appendix.E.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
end 36: 
 Send θ ′ q,n to client n. 37 end Lemma 2 For random variable x 1 , . . . , x n we haveEQUATIONLemma 3 For independent random variables x 1 , . . . , x n whose mean is 0, we haveEQUATIONBased on the above assumptions, we present the theoretical results for the non-convex problem.Lemma 4 (Deviation bound of the optimization function) In each communication round, the function value in server reduce after T ′ epochs and is bounded as:EQUATIONLemma 4 indicates the deviation bound of the optimization function of the server in each global round.Theorem 1 (Non-convex divergence for FedMKD) Let Assumption 1 to 3 hold and ∆ = F (θ 0 ) -F (θ T ′ ), given any δ > 0, suppose that the learning rates satisfy 0 ≤ η ≤ 2/L, afterEQUATIONcommunication round, we haveEQUATIONthe convergence can be guaranteed.Proof. Let's start the proof from L-Lipschitz condition: where (a) is from Assumption 1. Taking expectation of both sides, we obtainEQUATIONE[F (θq,t+1)] ≤ E[F (θq,t)] -ηE[∥∇F (θq,t)∥ 2 ] + Lη 2 2 E[∥F (θq,t, ξq,t)∥ 2 ] = E[F (θq,t)] -ηE[∥∇F (θq,t)∥ 2 ] + Lη 2 2 E[∥F (θq,t∥ 2 + σ 2 ] (b) ≤ E[F (θq,t)] -(η - Lη 2 2 )E[∥∇F (θq,t)∥ 2 ] + Lη 2 2 σ 2 ,where (b) follows from Assumption 2. Let's set the learning step at the start of training to T ′ ,EQUATION2 )EQUATIONProof of Theorem 1. According to Lemma 4,1 QT ′ (η - η 2 L 2 ) Q-1 q=0 T ′ -1 t=0 E[∇Lq,t)] ≤ 1 QT ′ Q-1 q=0 E[F (θ q,T ′ )] - 1 QT ′ Q-1 q=0 E[F (θq,0)] + LT ′ η 2 2 σ 2 ≤ δ(η - η 2 L 2 )Therefore,∆ Q ≤ δ(η - η 2 L 2 ) - LT ′ η 2 2 σ 2 ,which is equal toQ = 2∆ T ′ δ(2η -Lη 2 ) -T ′ (Lη 2 σ 2 ). 
 
 Appendix: θ ′ q,n ← θ ′ q,n -η∇l ′ q,n .We use CIFAR-10 and CIFAR-100 [13] datasets to train all the models. Both of them contain 50,000 training images and 10,000 testing images. To construct the public dataset, we sample 4000 data To assess the validity of the public dataset, we use two sampling methods to construct it. First, we use a random sampling method over all classes to generate public dataset 'IID'. And for the public dataset 'Partial', data is selected randomly from 40% classes in two datasets.We utilise three settings to simulate heterogeneous data distributions among all the clients. For the IID setting, each client contains the same number of samples from all classes. For the class setting, each client only has 10/N and 100/N classes on two datasets and the classes between clients have no overlap. For the Non-IID setting, data heterogeneity levels are described by the Dirichlet distribution Dir(β) [10] , where smaller β represents stronger heterogeneity levels. A value of β = 0.1 is chosen to simulate a high degree of heterogeneity, and β = 0.5 for a lower level. Fig. 6 shows the data distribution among clients on CIFAR-10 dataset. The x-axis represents 10 classes and the y-axis is the total of 5 clients. The size of each circle denotes the number of samples for the specific class in the respective client.Regarding the self-supervised learning framework design within each client, we use ResNet18 [9] and VGG9 [22] as the encoder network and Multi-Layer Perception (MLP) as the predictor. To compare with other baselines, we set the number of clients N = 5, and conduct the experiments for R = 100 rounds. In order to construct the model heterogeneous setting, 2 clients train the Resnet18 encoder while 3 clients use the VGG9. And for the global representation model, Resnet18 is selected as the encoder in server. The hyper-parameter γ in the loss of global model is set to 0.9. During the training process, each client trains locally for T = 5 epochs while the server also distills for T ′ = 5 epochs. Finally, we set the target decay rate α = 0.99, with a batch size of B = 128, and utilize SGD for optimization with a learning rate of η = 0.032. We implement all the methods in Python using EasyFL [32] based on PyTorch.Firstly, we select several federated knowledge distillation frameworks that use unlabeled public dataset for distillation. We then replaced the local model with a self-supervised model to evaluate the process of knowledge distillation in our method.• FedMD [14] : Each client trains convergence on the public dataset and then on local data.In each round, clients upload the embedding of the public dataset to the server. The server averages the embeddings and sends averaged embedding to clients. Clients first use average embedding to update local models on public dataset, then train on local data for a few epochs. There's no global model, the local models are evaluated to make comparison.• FedDF [18] : Clients in FedDF train locally and upload the model parameter to server. The server uses all client models to compute the embedding of the public dataset and uses the average embedding to train the global model. But it has no global model, so we only test the performance of local models.• FedET [2] : There is a global large model and several alternative small models on the server. Each client selects the appropriate small model for training locally and then uploads the parameter to the server. The server uses all the small models to compute representations of the public dataset and updates the global model with the averaged representation. Then the server utilizes the representation from the global model to update all the small models. • MOON [15] : Moon introduces a model-contrastive learning approach, treating the representation from the global model as positive knowledge to update the local models. And several federated self-supervised learning frameworks are also chosen as baselines.• FedU [33] : Clients in FedU upload the local online encoder and predictor for server aggregation. Then they update the local online encoder using the aggregate ones, and dynamically update the predictor based on the divergence. But it cannot applied to the model heterogeneity setting. • FedEMA [34] : FedEMA uses the aggregated local online network as a global model and measures the divergence between global and local networks. The divergence is then applied to the local exponential moving average update. It's also not applicable when models are heterogeneous. • Hetero-SSFL [20] : Each client trains locally and uploads the kernel matrix over the public dataset. The server aggregates them and sends to clients. Then client uses linear-CKA normalization term to align to the averaged kernel matrix. There's no global model, so we just test local models' performance.To verify that the client's knowledge can improve the global model, we also train the global model separately on the public dataset, denoted as Std. ResNet18.Following [33] and [30] , we evaluate the performance of learned representations using linear and semi-supervised evaluation. For linear evaluation, all the models are trained without any supervised labels. Subsequently, the model encoder is then frozen and the representations are utilized for training a new classifier over 200 epochs. For semi-supervised evaluation, we consider the scenario that only a small subset of the data has a label, here only 1% data are labeled. Then, different from freezing the encoder, we fine-tune the whole model with a new classifier using the labeled data for 100 epochs. For FedMD, FedDF and Hetero-SSFL, only local models can be evaluated. So we use the averaged representation from local models to do the linear evaluation. But in semi-supervised evaluation, we need to fine-tune the whole model, so we evaluate each client's performance alone and use the average result as the final result.Table 6 , Table 7 and Table 8 , Table 9 shows the linear evaluation results and semi-supervised evaluation results of FedMKD compared with all the baselines on CIFAR-10 and CIFAR-100. Comparing to the result in main text, we extra add the result on Dir(β = 0.1) to prove the efficiency of our algorithm totally.Considering the experimental setting, the communication and storage capabilities are limited, so we analyze the communication and computation efficiency of FedMKD compared to Hetero-SSFL and all the experiments maintain the original settings. Firstly, we calculate that the memory of ResNet18 and VGG9 is 42.63MB and 16.38MB separately. We use S(θ) to denote the storage cost of model θ.About the communication cost, the client in FedMKD only needs to upload the online encoder to the server, and download the updated encoder, so the cost per communication is equal to the ResNet18 or About local storage for each client, we compute the summary of the model and data. For FedMKD and Hetero-SSFL, each client runs the siamese network, so the model storage is double of the model memory, that is 2 * S(θ n ). But for the MOON, each client processes the local contrastive learning and the model contrastive learning, so they must store two local models and one global model, that is 2 * S(θ n ) + S(θ s ). Here, we neglect the storage of the predictor which is usually a 2-layer MLP. We use S(|D|) to denote the storage of the dataset. And each client in FedMKD only needs to store the To explore the influence of different hyperparameters, we conduct experiments on several key parameters of FedMKD.Impact of public dataset size. As mentioned in [20] , the size and construct method of the public dataset are both critically important, especially since our global model needs training based on it. About the latter one we've given an analysis in the last section, so we explore the influence of the size on the model performance. Considering the construct method of the public dataset, we choose part of the dataset as a public dataset and set it aside to be used. The performance change with respect to its size is shown in Fig. 7 (a)(d). We can observe that in both two datasets, when the public dataset is small, the performance of FedMKD is quite worse, and as the size of the public dataset increases, the performance gets better. But when the size is larger than 4000, the gain of performance is small, and because of the design of the model, the time cost of the server computation is larger, so we choose 4000 as our experiment setting.Impact of the number of global rounds. In order to investigate the impact of the training rounds to FedMKD, we fix the other hyperparameters and train the model for 200 rounds. The results showed in Fig. 7 (b)(e) demonstrate that while the total training rounds increase, the performance of the FedMKD gets better. For both two kinds clients data distribution in CIFAR-10, the global model converges at nearly 150 rounds. For CIFAR-100, the global model converges at 200 rounds. There's no denying that the increase of the global rounds will result in better performance, especially when the global round is small, the improvement of the performance is significant.Impact of the number of server epochs. Since our global model is trained on the public dataset, through standalone training and knowledge distillation. The lower server epoch means the global model may under-fitting not only the distribution of the public dataset but also the client knowledge.On the other hand, large server epoch will lead large computation cost, and over-fit the distribution of the public dataset, which is harmful to the generalization of the model, so the epoch in server is important. We choose 5 different settings T ′ = {1, 3, 5, 7, 10} to validate the influence of the server epochs. Based on the experiment results shown in Fig. 7 (c)(f), we can see that the performance of the global model improves when the server epoch increases with the distribution of the public is IID. But the change of improvement gets smaller when the server epoch is larger than 5.In this section, we discuss the limitations and broader impacts of the work.Limitations. Although we provide detailed explanations of the proposed algorithm and extensive experiments analysis, the theoretical proof of FedMKD is not rigorous enough. The characteristic of multi-teacher distillation of the global model has not been sufficiently theoretically justified and we only combined it with the self-training loss as one whole loss to analyze.Broader Impacts. FedMKD offers significant societal and technological benefits, which is crucial in kinds of domains like healthcare and finance. It promotes inclusivity by leveraging diverse data from various sources, thereby reducing biases and improving model generalization. Technologically, FedMKD lowers the dependency on labeled data, making federated self-supervised learning more efficient and scalable, and drives innovation in heterogeneous resource-limited devices. By carefully navigating these challenges: deviated representation abilities and inconsistent representation spaces, FedMKD can lead to responsible and equitable advancements in distributed AI technology.judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. Justification: We have reviewed the NeurIPS Code of Ethics. And we ensure that our work does not involve any inference of personal information and that it will only be used for academic purposes. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: The paper proposes a novel approach, which is a multi-teacher knowledge distillation framework for federated learning. In general, the reviewers are satisfied with the manuscript and the author responses. They have appreciated the proposed framework, writing and experiments. Though one reviewer response is still negative, I feel that the author's response adequately addresses the concerns regarding using other models and additional datasets. The recommendation is to accept the paper. The authors are strongly encouraged to address any remaining concerns in the camera-ready version.

Reviews: review: paper_summary: paper_summary: - The paper proposes a framework called FedMKD which is a multi-teacher knowledge distillation framework for federated learning. 
- To allow for different clients to have different resources, a resource adaptive approach is designed. The approach handles class skew and different architectures. 
- The knowledge distillation framework allows for learning representations from the different client representations. An SSL loss is used in addition on the common representation space. 
- Different from prior works, a global representation model is learnt despite differences in clients. 
- Approach is shown to outperform prior works on CIFAR-10 and 100 datasets.
strength_weakness: strength_weakness: Weaknesses/Concerns & Questions:
- How does the approach compare to the traditional FL settings when aggregation like L107 is done (assuming all client architectures are the same) ? 
- Did the authors experiment with applying the knowledge distillation to individual projected features instead of the aggregate ? 
- L167: Are the negatives also after the aggregation step ? or from individual projection layers ?   
- Additional larger-scale datasets will be helpful to understand whether the approach will scale. 

Presentation and writing:
- Table 1 needs more details to be helpful. A more detailed caption should help. More subjective columns like "inconsistent representation space", "deviated representation ability" are not apt for a table like this.
questions: questions: Please refer to the weaknesses section.
limitations: limitations: The paper has adequately discussed limitations and potential negative impacts.
paper_summary: paper_summary: This paper studies the problem of federated self-supervised learning. A multi-teacher knowledge distillation framework is proposed to address the two challenges: deviated representation abilities and inconsistent representation spaces. Specifically, the adaptive knowledge integration mechanism is designed to learn better representations from all heterogeneous models, and a global knowledge-anchored alignment module is used to make the local representation spaces close to the global spaces. Experiments on two datasets demonstrate the effectiveness of the proposed method.
strength_weakness: strength_weakness: * Knowledge distillation is widely used in federated learning and federated self-supervised learning. The contribution of the proposed method (i.e., multi-teacher knowledge distillation framework) is incremental.
* This paper uses small models such as resnet18 and vgg9.
* Only CIFAR datasets are used to evaluate the proposed method. Experiments on ImageNet-1k are important to evaluate the effectiveness.
questions: questions: Weakness.
limitations: limitations: The authors adequately addressed the limitations in the appendix.
paper_summary: paper_summary: This paper proposes a multi-teacher knowledge distillation framework named FedMKD for resource-adaptive federated self-supervised learning (Fed-SSL). The method aims to address the challenges of global representation learning in Fed-SSL caused by heterogeneous architectures and imbalanced class distributions. Experimental results demonstrate that FedMKD is significantly effective in addressing heterogeneity and class distribution imbalance issues in federated self-supervised learning.
strength_weakness: strength_weakness: 1. This paper combines federated self-supervised learning and resource-aware federated learning. However, in introduction and related works, the authors lack a discussion about some works about resource-aware federated learning with heterogeneous clients.
.
2. The authors assume to build a global dataset to perform knowledge distillation and anchored alignment. However, random sampling makes the global dataset and client dataset have the similar data distribution, which violates the original intention of data privacy in FL. Please provide some explanation for this.

3. CIFAR-10 and CIFAR-100 are both image datasets with relatively small sizes and low image resolution, which cannot fully reflect the complexity of real-world applications. It is recommended to add experiments on large-scale image datasets such as Tiny-ImageNet and ImageNet-100.

4. The framework of dual encoders increases the training load in clients. It is recommended to evaluate the computation time and resources required and provide analysis for local model training.

5. To evaluate the robustness and effectiveness of this framework, it is suggested to explore experiment settings with more clients, such as [1].

[1] ScaleFL: Resource-Adaptive Federated Learning with Heterogeneous Clients
questions: questions: See weakness.
limitations: limitations: The characteristic of multi-teacher distillation of the global model has not been sufficiently theoretically justified and the authors only combined it with the self-training loss as one whole loss to analyze.
paper_summary: paper_summary: The authors propose a multi-teacher knowledge transfer framework, FedMKD, based on two challenges in resource-adaptive federated self-supervised learning: deviated representations abilities and inconsistent representations. This framework uses an adaptive knowledge integration mechanism and a weighted combination of different loss functions to ensure that global representations with class-wide knowledge can be learned from heterogeneous clients even in the presence of extreme class skew. Extensive experiments on two datasets demonstrate the effectiveness of FedMKD, which outperforms the state-of-the-art baseline by an average of 4.78% under linear evaluation.
strength_weakness: strength_weakness: 1.	The mechanism design seems to be confusing. As shown in the FedMKD, multi-instructor knowledge distillation is allowing global models to learn local models, while global knowledge anchored alignment is allowing local models to learn global models, knowledge transfer happens alternately.

2.	FedMKD aims to learn global class representations according to a global model. But an additional global-anchored mechanism improves the local performance. What’s the difference between personalized federated learning and the proposed FedMKD?

3.	It’s suggested that the authors could add the instructions about the whole model training algorithm in Appendix A, which can help reviewers understand the whole process.

4.	The presentation needs to be improved, especially for some equations. The self-supervised loss value in Eq. (10) was not discussed earlier. If it is equal to the loss value in Eq.(3), but with a different notation.
questions: questions: 1.	As shown in Weakness 1, it seems to that multi-instructor knowledge distillation is allowing global models to learn local models, while global knowledge anchored alignment is allowing local models to learn global models, so does the two mechanisms result in duplication of knowledge transfer?

2.	In the problem statement, the main objective is the global loss. Since recent studies on cross-device scenarios also consider personalized FL (PFL), how about the personalization? Can your method be extended to PFL?
limitations: limitations: The authors pointed out the limitations of the paper that is the theoretical proof of FedMKD is not rigorous enough, the characteristic of multi-teacher distillation of the global model has not been sufficiently theoretically justified.
paper_summary: paper_summary: This paper introduces a novel approach for enhancing global representation models in resource-adaptive federated self-supervised learning through a multi-teacher knowledge distillation framework, named FedMKD. The proposed method addresses the challenges posed by heterogeneous architectures and extreme class skew, demonstrating significant improvements in representation abilities across diverse clients. The authors provide detailed experimental results to support their claims.
strength_weakness: strength_weakness: The Related Work section could be expanded to include a more thorough discussion of existing approaches in both transfer learning and attention mechanisms.
The paper may have a limited audience due to its specialized focus on image classification tasks. It could benefit from providing more context for readers unfamiliar with the specific datasets and techniques discussed.
The contributions of the work could be more explicitly stated in the Introduction to enhance clarity for readers.
questions: questions: It would be beneficial for the authors to include a complexity analysis of the FedMKD algorithm to provide insights into its computational efficiency. If a complexity analysis is not feasible, the authors should provide a rationale for its omission.
limitations: limitations: The limitations of the FedMKD framework are well-addressed in the paper. These constraints appear to be inherent to the design choices made in the method and are necessary trade-offs for its functionality.

Comments: title: Response to Reviewer sGaU
Response to Reviewer QsUu
Thanks for the responses
comment: Thank you very much for taking the time to read and respond to our rebuttal. We sincerely appreciate your recognition of our work.
After reviewing the responses and other reviewers' comments, I appreciate the detailed clarifications and new experiments provided. The explanations effectively addressed my concerns, leading me to update my rating to weak accept.
We appreciate you taking time to read and respond to our rebuttal. And thank you very much again for recognizing our work.
We appreciate you taking time to read and respond to our rebuttal. And thank you very much again for recognizing our work.
I have read the responses and other reviews. Thanks for the detailed detailed clarifications and new experiments. I have updated my rating for the paper.
Thank you for your response ! I appreciate the thoughtful consideration given to my comments. The explanations and revisions provided are clear and effectively address my concerns, and I strongly recommend the acceptance of this paper.