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

We thank the reviewer hPCk for the time and valuable feedback! We would try our best to address the comments one by one.

Response to Weakness:

Thank you very much for the insightful comments.

• First, we survey some federated transfer learning (FTL) methods, and add the following discussion of them. "Federated transfer learning aims to share the knowledge and insights derived from training machine learning models across various entities, all while keeping the raw data decentralized. It has been used widely in various regions, such as cross-domain recommendation [1,2], medical image classification [3,4] and financial service [5,6]" About attention mechanisms, we would add the following discussion as follows. "Attention mechanism is a component used in neural networks that dynamically focuses on the most relevant parts of the input data, enhancing the model's performance. By assigning varying levels of importance to different pieces of information, it allows the network to prioritize and efficiently process significant elements, has proven to be highly successful in various tasks, such as image classification [7,8], object detection [9,10] and image generation [11,12]." Both of these discussions will be added in Related Work in the final version.

• By the way, we would add more detailed description for readers who are unfamiliar with the specific datasets and techniques to understand our method as follows. "We use CIFAR-10 and CIFAR-100 to evaluate the performance. The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images. CIFAR-100 dataset is just like the CIFAR-10, except it has 100 classes containing 600 images each. There are 500 training images and 100 testing images per class."

• Finally, we present our contributions in Lines 65-76 of the Introduction. We emphasize that in resource-adaptive Fed-SSL, we are the first to delve into global class representations learning by addressing the deviated representation abilities and inconsistent representations caused by heterogeneous architectures and class skew. We designed FedMKD to tackle these challenges, and the experimental results demonstrate the superiority of our proposed FedMKD.

Response to Question:

Thank you very much for the insightful comments. We would like to analyze the complexity of client and server separately. First in the client, each client chooses the appropriate model according to their resource, so the computation complexity is different among them. Then in server, the training process has two phases. The complexity of knowledge distillation is $O(T' \times |D_{pub}| \times (N+1))$, where $T'$ is the epoch of knowledge distillation, $|D_n|$ is the size of the public dataset, and $N$ is the number of client local model. The complexity of the global knowledge anchored is $O(|D_{pub}| \times N)$. By the way, considering that each deep network includes multiple operations and computations, we commonly use the amount of computation (FLOPS) to analyze the time complexity and the amount of memory access (Bytes) to analyze the space complexity. The results of each process are shown as follows:

If there are any further confusions/questions, we are happy to clarify and try to address them. Thank you again and your recognition means a lot for our work.

[1] Ammad-Ud-Din, Muhammad, et al. "Federated collaborative filtering for privacy-preserving personalized recommendation system." arXiv preprint arXiv:1901.09888 (2019).

[2] Minto, Lorenzo, et al. "Stronger privacy for federated collaborative filtering with implicit feedback." Proceedings of the 15th ACM Conference on Recommender Systems. 2021.

[3] Gong, Xuan, et al. "Federated learning with privacy-preserving ensemble attention distillation." IEEE transactions on medical imaging 42.7 (2022): 2057-2067.

[4] Sui, Dianbo, et al. "Feded: Federated learning via ensemble distillation for medical relation extraction." Proceedings of the 2020 conference on empirical methods in natural language processing (EMNLP). 2020.

[5] Shaheen, Momina, Muhammad Shoaib Farooq, and Tariq Umer. "Reduction in data imbalance for client-side training in federated learning for the prediction of stock market prices." Journal of Sensor and Actuator Networks 13.1 (2023): 1.

[6] Pourroostaei Ardakani, Saeid, et al. "A federated learning-enabled predictive analysis to forecast stock market trends." Journal of Ambient Intelligence and Humanized Computing 14.4 (2023): 4529-4535.

[7] Hu, Jie, Li Shen, and Gang Sun. "Squeeze-and-excitation networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[8] Woo, Sanghyun, et al. "Cbam: Convolutional block attention module." Proceedings of the European conference on computer vision (ECCV). 2018.

[9] Dai, Jifeng, et al. "Deformable convolutional networks." Proceedings of the IEEE international conference on computer vision. 2017.

[10] Carion, Nicolas, et al. "End-to-end object detection with transformers." European conference on computer vision. Cham: Springer International Publishing, 2020.

[11] Gregor, Karol, et al. "Draw: A recurrent neural network for image generation." International conference on machine learning. PMLR, 2015.

[12] Zhang, Han, et al. "Self-attention generative adversarial networks." International conference on machine learning. PMLR, 2019.

We thank the reviewer sGaU for the time and valuable feedback! We would try our best to address the comments one by one.

Response to Weakness1 & Question1:

Thank you very much for your recognition. As you mentioned, we have two mechanisms to transfer knowledge, but they are not duplicated. First, the multi-teacher knowledge distillation allows the global model to encode all classes from clients in a unified space. And the global knowledge-anchored alignment module is applied to eliminate inconsistency in representation spaces of clients, which is benefit to the global model training, also. These two machenisms are critical for learning global class representation and are indispensable.

Response to Weakness2 & Question2:

Thank you for your insightful comments. The additional global anchored mechanism aims to ensure that the local representation spaces are closer to the global ones so that the following knowledge distillation process could gain more from clients. This module can further eliminate the effect of inconsistent representation spaces and improve local performance. The experiment result of the improvement in client is shown in Section 5.4 which verifies the effectiveness of the module. Compared to personalized FL, the optimization objective is not identical. We aim to train a global model that can encode all classes in a unified space, and the improvement of clients is a middle process. But for PFL, it aims to learn better local model without global model, it's quite different from ours.

Response to Weakness3:

As you suggested, we would add the instructions about the whole model training algorithm in Appendix A in our final version.

Response to Weakness4:

We are sorry for our unclear discussion. Eq. (10) shows the weighted combined loss for the global model, it includes two parts, the self-supervised loss and knowledge distillation loss. So that the global model can encode all classes from clients in a unified space and improve its representation capability. The self-supervised one is identical to the local training loss, that is Eq. (3).

If there are any further confusions/questions, we are happy to clarify and try to address them. Thank you again and your recognition means a lot for our work.

We thank the reviewer 24fB for the time and valuable feedback! We would try our best to address the comments one by one.

Response to Weakness 1:

Thank you very much for the insightful comments. In the related works section (Lines 94-99), we have surveyed several existing studies that address heterogeneous client models in federated self-supervised learning. We analyzed the advantages and disadvantages of Hetero-SSFL and FedFoA in detail. Additionally, we select Hetero-SSFL as one of the baselines in the experiment to compare and evaluate the performance of our proposed method. And as you mentioned, we will further add the description of Hetero-SSFL and FedFoA in the introduction in the final version.

Response to Weakness 2:

Thanks for your sincere comments. As you are concerned, this construction method of public dataset is the experimental simulation part only does not pose a serious risk of data leakage, in our FedMKD, no data need to be transferred. Here, we just want to address one practical problem: if the distribution of public dataset is different from the data distributions of clients, can the proposed FedMKD still work?

In practice we can use appropriate public dataset to distill knowledge according to the specific task instead of this random sampling method. To verify this, we conducted another experiment using CINIC as the public dataset on the server. CINIC is an extension of CIFAR-10, augmented with downsampled ImageNet images, and it shares the same size and classes as CIFAR-10. We still use CIFAR-10 as the train dataset for the clients. The public dataset is 'Partial' and the data distribution in the client is 'Class'. The other setting is identical to the draft. The linear probing results are shown as follows:

Here we can see that even when intuitively choosing an appropriate public dataset, our method still works effectively.

Response to Weakness 3:

Thank you very much for the insightful comments. We use ImageNet-100 and ImageNet-1k to evaluate the effectiveness of our proposed method. The results are presented in Table 11 in Global Response PDF. We can conclude that our proposed FedMKD outperforms all baselines.

Response to Weakness 4:

Thank you very much for the insightful comments. We evaluate another single-encoder framework, SimCLR, on the clients and assessed the computation time and resources required for local model training. Although this method is more computationally efficient, its performance is much inferior to ours. The results are shown as follows:

Response to Weakness 5:

Thank you very much for the insightful comments. We've explored the scalability of our proposed algorithm FedMKD in Appendix D.3. In these experiments, we test scenarios with 10 and 30 clients, where 40% of the clients use the VGG model and 60% use ResNet18. We repartition the data for each client under Dir($\beta=0.5$) and set the public dataset distribution as IID. To compare scalability performance, we conduct the same experiments on the second-best baseline, Hetero-SSFL. The linear probing results are shown as follows:

The results show that as the number of clients increases, our proposed method still outperforms Hetero-SSFL, demonstrating the scalability of our FedMKD.

If there are any further confusions/questions, we are happy to clarify and try to address them. Thank you again and your recognition means a lot for our work.

We thank the reviewer de59 for the time and valuable feedback! We would try our best to address the comments one by one.

Response to Weakness 1:

Thank you very much for the insightful comments. As you said, knowledge distillation (KD) has been widely used in Fed-SSL, such as FedX and other works. However, our contrastive learning-based multi-teacher KD introduces novel approaches to address the challenges of deviated representation abilities and inconsistent representation spaces:

• First, to mitigate the impact of heterogeneous models with deviated representation abilities, we introduce an adaptive knowledge integration module to learn high-quality representations from them.
• Second, to encode all classes from clients in a unified space, the global model updates using a weighted combination of self-supervised loss and distillation loss to update.
• Finally, the global knowledge anchored alignment module is applied on the server to eliminate the inconsistency in representation spaces and reduce the burden on the clients.

In our experiments, FedMD and FedDF, which are based on traditional KD methods, were adopted as comparative baselines. As shown in Table 2 (Page 7) and Table 3 (Page 8), our proposed FedMKD achieves the best performance.

Response to Weakness 2:

Thank you very much for you valuable comment. Adopting larger models can indeed demonstrate the scalability of our proposed FedMKD framework. As suggested, we set one client's local model to ResNet50 and re-conduct the experiment under the setting of the 'Partial' public dataset and 'Class' distribution among clients on CIFAR-10. All other settings remained identical to those in the original paper. The linear probing results are shown below:

The results indicate that the introduction of a larger model can improve overall performance, and our proposed FedMKD still outperform all baselines.

Response to Weakness 3:

Thank you very much for the insightful comments. We use Imagenet-100 and Imagenet-1k to evaluate the effectiveness of our proposed method. The results are presented in Table 11 in Global Response PDF. We can conclude that our proposed FedMKD outperforms all baselines.

If there are any further confusions/questions, we are happy to clarify and try to address them. Thank you again and your recognition means a lot for our work.

We thank the reviewer QsUu for the time and valuable feedback! We would try our best to address the comments one by one.

Response to Weakness 1:

Thank you very much for the insightful comments. Our proposed FedMKD is designed for resource-aware settings, allowing each client to choose an appropriate model to train, even if the model architectures differ among clients. In this context, traditional aggregation methods are ineffective. However, if all clients use the same architecture, the model with the minimal parameters should be chosen.

In the newly designed experiments, both the clients' models and the global model are set to VGG9. The linear probing results compared to FedAVG are as follows:

In our FedMKD approach, each client can adaptively choose the appropriate local model to fully utilize its computational resources. For example, a client with ample computational resources can select a larger model like ResNet18, while another client with limited resources can opt for a smaller model like VGG9. This flexibility enables FedMKD to outperform FedAVG on the VGG9 model.

Response to Weakness 2 & 3

First, we apologize for any confusion caused by our previous presentation. To clarify, in Line 167, the negative sample is derived from the individual projection layers of the global model. In knowledge distillation (KD), we have an aggregated anchor representation $\bar{r_i}$, a positive representation $r_{s,i}$ from global model, and multiple negative representations $r_{s,j}$. The purpose of knowledge distillation is to push the anchor representation $\bar{r_i}$ and the positive representation $r_{s,i}$ closer together in the feature space while pulling the anchor away from the negative representations. Thus, the global model could encode all classes from clients in a unified space. Formally, we define this KD loss function, which is presented in the original paper on Line 168:

$L_{distill}=-log\frac{ exp(sim(r _ {s,i}, \bar{r} _ i)/\tau)}{exp(sim(r _ {s,i},\bar{r} _ i)/\tau)+\sum\limits _ {j\neq i}exp(sim(r _ {s,i},r _ {s,j})/\tau)}.$

We also illustrate this process in the Global Response PDF (see Figure 8).

Response to Weakness 4:

Thank you very much for the insightful comments. We use ImageNet-100 and ImageNet-1k to evaluate the effectiveness of our proposed method. The results, presented in Table 11 of the Global Response PDF, demonstrate that our proposed FedMKD outperforms all baseline methods.

About Presentation and writing:

Thank you very much for the insightful comments. We will revise the caption of Table 1 to "Comparison of federated self-supervised learning methods. √ indicates that the proposed method focuses on this challenge and × indicates that it does not.". Additionally, the columns labeled "inconsistent representation space" and "deviated representation ability" will be replaced with the objective description "data heterogeneity." The modified table is presented as Table 12 in the Global Response PDF.

If there are any further confusions/questions, we are happy to clarify and try to address them. Thank you again and your recognition means a lot for our work.