MSfusion: Enabling Collaborative Training of Large Models over Resource-Constraint Participants

Thank you for your thoughtful and positive feedback. Here, we provide answers to your concerns point by point.

Weaknesses 1: Concern about the importance of the proposed contrastive learning objective and the dynamic overlapping method.

Our response: Thank you for your insightful assessment. The crux of our incorporation of the contrastive learning objective and dynamic overlapping method within MSfusion aimed to address the divergence between local participant models and the aggregated model, aiming for notable performance enhancements. To better substantiate the importance of these adaptations, we've extended our experiments to include the WikiText2 dataset, focusing on ablation studies. The results are shown in the following Table which is updated to Table 2 in the revision Section 5.1.

Notably, with a split model size of $\mu=18.75\%$ in the WikiText2 dataset experiments, the performance gain of MSfusion over the baseline variants are more significent than the CV. This main attributed to the larger relative size of local model parameters within the transformer architecture, where the model drift problem is further amplified.

These additional experiments underscore the tangible benefits and substantial impact of integrating these components into MSfusion, further solidifying their role in enhancing overall performance. We believe these supplementary findings reinforce the significance of the contrastive learning objective and dynamic overlapping method within MSfusion, lending stronger support to their effectiveness in improving performance.

Weaknesses 2: Need clarification for the association between the proposed collaborative training framework and the contrastive learning objective.

Our response: In our proposed decentralized collaborative training algorithm, we've highlighted two pivotal components: model shifting on a single server and model aggregation across servers. These aspects inherently introduce divergence within the local split model across rounds, as well as divergence between the local split model and the global full model. To explicitly address these divergence issues, our design integrates a contrastive learning loss that leverages representations from the local training model, local aggregated model, and the local model from the preceding round. This strategic utilization of contrastive learning aims to mitigate the challenges associated with model shifting within our MSfusion method. It's important to note that this contrastive learning objective is specifically tailored for our sub-model shifting algorithm, uniquely crafted for decentralized sub-model learning. This customized approach is a distinct facet of our decentralized framework, providing a targeted solution to the challenges of divergence encountered in collaborative training scenarios.

This contrastive learning method, while intricately woven into our collaborative framework, also holds potential applicability beyond our specific framework. Its adaptability to other distributed learning frameworks or its potential benefits to the backbone model itself can indeed be an area of further exploration. However, the unique synergy between the contrastive learning objective and our decentralized sub-model shifting algorithm remains a central and distinctive aspect of our proposed MSfusion method.

Question: Expectation for more discussion on the results of the ablation study.

Our response: Your feedback on the need for a more comprehensive discussion regarding the results of the ablation study and the proposed cross-sub-model contrastive learning is greatly appreciated. In response, we've taken careful consideration of these aspects and included an expanded discussion in our revised manuscript. Additional experiments are included in Table 2 and the following discussion on the result is added in Section 5.1 to the revision.

The results showcase MSfusion's superior performance over all baseline variants. Specifically, in NLP tasks, the performance gain of MSfusion is more significant. This is mainly attributed to the larger relative size of local model parameters within the transformer architecture, further amplifying the model drift problem.

Dear Reviewer PcAY,

We are approaching the end of the discussion phase, and we have unfortunately received no feedback from you on our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've added new experiments for the ablation studies to better show the performance gain of MSfusion. We have clarified the association between the proposed collaborative training framework and the contrastive learning objective. And added more discussion about the ablation studies to the revision.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer PcAY,

With less than 24 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've added new experiments for the ablation studies to better show the performance gain of MSfusion. We have clarified the association between the proposed collaborative training framework and the contrastive learning objective. And added more discussion about the ablation studies to the revision.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer PcAY,

With less than 3 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've added new experiments for the ablation studies to better show the performance gain of MSfusion. We have clarified the association between the proposed collaborative training framework and the contrastive learning objective. And added more discussion about the ablation studies to the revision.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Thank you for your review and thoughtful comments and suggestions. We have updated our paper based on your feedback and we will also address each point raised in the following separate comments.

Weaknesses 1: Novelty concern toward the previous splitting learning and FL works.

Our response: Please note while we use the word "model splitting" in our approach, there is fundamental difference between model splitting approach in this paper and the "split learning" in Splitfed framework [1]. In our decentralized collaborative learning setting, there is no centralized server, and the model are split across different participants. In Splitfed, or more generally split learning, the model is split across a central server and clients, who work together to train the entire model through communicating intermediate data embeddings and gradients. Due to splitting of the model, the clients can not train its own part independently due to missing computations of the submodel at the server side. In sharp contrast, for the setting considered in this paper and the proposed MSfusion scheme, each participant is assigned a (dynamically changing) subset of model parameters, and the training of these parameters are done independently at each participant. We believe that this difference can be well illustrated by Figure 1 and Figure 2 in our paper, and Figure 1 in Splitfed paper.

While contrastive learning was utilized in Moon [2] to deal with data heterogeneity in federated learning, we use constrastive learning with a drastically distinct purpose of addressing the divergence within the local split model across rounds, as well as divergence between the local split model and the global full model, arising from the double shifting model splitting scheme introduced by MSfusion. To explicitly address these divergence issues, we design a novel contrastive learning loss in (7) that incorporates representations from the local training model, local aggregated model, and the local model from the preceding round. It's important to note that this contrastive learning objective is specifically tailored for our sub-model shifting algorithm, uniquely crafted for decentralized sub-model learning. This customized approach is a distinct facet of our decentralized framework, providing a targeted solution to the challenges of divergence encountered in collaborative training scenarios.

To summarize, our designs of model splitting and contrastive loss have fundamental differences with Splitfed [1] and Moon [2] respectively, in terms of design goals and specific implementations.

[1] Thapa, Chandra, et al. "Splitfed: When federated learning meets split learning." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 8. 2022.

[2] Li, Qinbin, Bingsheng He, and Dawn Song. "Model-contrastive federated learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

Dear Reviewer wET1,

We are approaching the end of the discussion phase, and we have unfortunately received no feedback from you on our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified the distinctiveness of MSfusion in comparison to splitting methods and prior FL works, expanded our experiments by integrating larger transformer models for a more comprehensive evaluation, explained reasons why baselines are limited, emphasized our theoretical analysis in the paper and corrected typos in the revised version based on your suggestions.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer wET1,

With less than 24 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified the distinctiveness of MSfusion in comparison to splitting methods and prior FL works, expanded our experiments by integrating larger transformer models for a more comprehensive evaluation, explained reasons why baselines are limited, emphasized our theoretical analysis in the paper and corrected typos in the revised version based on your suggestions.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer wET1,

With less than 3 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified the distinctiveness of MSfusion in comparison to splitting methods and prior FL works, expanded our experiments by integrating larger transformer models for a more comprehensive evaluation, explained reasons why baselines are limited, emphasized our theoretical analysis in the paper and corrected typos in the revised version based on your suggestions.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Thank you for your thorough feedback. We offer the following responses to address your concerns and questions:

Weaknesses 1: Contribution concern over existing partial training work.

Our response: We would like to emphasize that the introduced double shifting model splitting scheme (DSS) requires careful design of splitting scheme across participants and aggregation rounds (including initial placement of submodels, and the dynamic adjustment of model overlaps). Together with another important novel design of contrastive loss function, the proposed DSS enables each participant to operate with equally minimum computation and communication loads, achieving good performance on large models that cannot be trained in entirety on a single participant.

Note that the related works in Section 2.3, while also utilizing partial training, exhibit quite distinct goals and designs. For instance in [1], the goal of model splitting is to accommodate client heterogeneity, and achieving good model performance relies on existence of powerful clients who train the entire model locally. Such clients do not exist in our considered setting (all participants are resource constraint who cannot train the entire model locally), and designing model splitting scheme poses major challenges that are addressed in this paper.

[1] Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. FedRolex: Model-heterogeneous federated learning with rolling sub-model extraction, NeurIPS 2022.

Weaknesses 2: The convex assumption over the loss function (i.e., Assumption 1) for convergence analysis is quite strong.

Our response: We agree that the smoothness and strong convexity assumptions over the loss function are strong, but it is worth noting they are also widely adopted in the literature of distributed learning frameworks, like [2][3][4]. In these works, smoothness and strong convexity assumptions are also made to facilitate the convergence proofs. Our convergence analysis can also be viewed as an improvement over existing partial training works since no convergence analysis is provided in previous works like HeteroFL and FedRolex.

[2] Aleksandr Beznosikov, et al. "On Biased Compression for Distributed Learning." arXiv:2002.12410, 2020.

[3] Tian Li, et al. "Ditto: Fair and Robust Federated Learning Through Personalization" ICML 2021

[4] Yifan Shi, et al. "Improving the Model Consistency of Decentralized Federated Learning" ICML 2023

Weaknesses 3: Limited NLP evaluation.

Our response: Thank you so much for pointing out this. In order to better evaluation our proposed MSfusion for LLMs, we have added two new wildly adapted NLP datasets PennTreebank and WikiText103. For PennTreebank, it contains about a million tokens with a vocab size of 10,000, and for WikiText103, it contains about 103 million tokens with a vocab size of 267,735. Together with WikiText2, our evaluation now covers 3 computer vision tasks and 3 NLP tasks, with different global full model parameter sizes, ranging from 7.32M for PennTreebank to 139.01M for WikiText103. The results are shown in the following table, which we have added to the Table 1 in the revised draft. The table demonstrates that our proposed MSfusion can successfully finish the collaborative training process with good performance, while greatly reducing the computation cost for each participant. We hope that this additional experiment can address the reviewer’s concern and strengthen our contribution.

Weaknesses 4: Experiment details not clear.

Our response: We agree that our previous evaluation details were not clear enough, and we have added a Table 4 to the revision in Appendix A.2.2 to present the parameter size, data partition, and model architecture for each experiment. The participant size for each experiment is emphasized in the paper. We hope that this table can provide a clearer understanding of our experimental settings and results.

Question: From Table 1, why the baselines of HeteroFL and FedRolex perform so badly over WikiText2?

Our response: We thank the reviewer for the question. The poor performance in Table 1 regarding the baselines of HeteroFL and FedRolex over the WikiText2 dataset primarily stems from their inherent inability to address the challenge of model drift. As clarified in our paper, our study focuses on collaborative training of large models, where participants possess comparable yet limited split model sizes. In the experiments outlined in Table 1, the split model sizes for all participants are fixed. Particularly concerning transformers, the local split model parameters are relatively larger compared to ResNet in computer vision tasks. This discrepancy results in a situation where HeteroFL leaves a significant proportion of global full model parameters untrained. Similarly, FedRolex exacerbates the model drift problem by not training a considerable number of global full model parameters between communication rounds. Given that these methodologies primarily concentrate on the splitting scheme and do not address the model drift issue, their performance is notably compromised when collaboratively training larger models. To better elucidate these observations, we have included the following explanation in revision Section 5.1.

In NLP tasks, MSfusion significantly outperforms HeteroFL and FedRolex. The suboptimal performance of HeteroFL and FedRolex in collaborative training larger models is due to untrained global model parameters and exacerbated model drift, issues unaddressed in their splitting-focused methodologies. This emphasizes the advantage of MSfusion in handling larger models more effectively.

Dear Reviewer wTYd,

We are approaching the end of the discussion phase, and we have unfortunately received no feedback from you on our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified our novel contributions in problem setting, split scheme design, dynamic overlapping, and contrastive loss design compared to previous partial training works. Additionally, we've addressed the utilization of smoothness and strong convexity assumptions in distributed learning, expanded our NLP experiments for a more comprehensive evaluation, introduced an experiment details table as suggested, and provided an explanation for the performance drop of HeteroFL and FedRolex in Table 1 within the revised version.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer wTYd,

With less than 24 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified our novel contributions in problem setting, split scheme design, dynamic overlapping, and contrastive loss design compared to previous partial training works. Additionally, we've addressed the utilization of smoothness and strong convexity assumptions in distributed learning, expanded our NLP experiments for a more comprehensive evaluation, introduced an experiment details table as suggested, and provided an explanation for the performance drop of HeteroFL and FedRolex in Table 1 within the revised version.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors

Dear Reviewer wTYd,

With less than 3 hours remaining in the discussion phase, we regret to note that we haven't received any feedback from you regarding our rebuttal.

Please can we kindly ask you to take a look at our responses, and let us know whether we have clarified your questions and addressed your concerns?

Specifically, we've clarified our novel contributions in problem setting, split scheme design, dynamic overlapping, and contrastive loss design compared to previous partial training works. Additionally, we've addressed the utilization of smoothness and strong convexity assumptions in distributed learning, expanded our NLP experiments for a more comprehensive evaluation, introduced an experiment details table as suggested, and provided an explanation for the performance drop of HeteroFL and FedRolex in Table 1 within the revised version.

Thank you very much again for the time you spent reviewing.

Paper3368 Authors