Paper: MSfusion: Enabling Collaborative Training of Large Models over Resource-Constraint Participants (ICLR-2024)

Reviewer comment (Evaluation):
The performance gap between the ablated models and the proposed model is not very large. This does not support the importance of the proposed contrastive learning objective and the dynamic overlapping method.

Author 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.

| Methods                | CIFAR10 ACC             | FLOPs       | WikiText2 PPL           | FLOPs       |
| :--------------------- | :---------------------: | :---------: | :---------------------: | :---------: |
| HeteroFL               | 37\.36 $\pm$ 0.6     | 35\.76M     | 579\.05 $\pm$ 8.4    | 1\.08B      |
| FedRolex               | 66\.41 $\pm$ 0.8     | 35\.76M     | 547\.32 $\pm$ 4.5    | 1\.08B      |
| MSfusion w/o Con       | 73\.16 $\pm$ 0.5     | 20\.32M     | 9\.57 $\pm$ 2.1      | 267\.2M     |
| MSfusion w/o Dyn       | 73\.31 $\pm$ 0.5     | 22\.33M     | 7\.35 $\pm$ 1.5      | 290\.1M     |
| MSfusion w/o Con & Dyn | 70\.17 $\pm$ 0.8     | **20\.32M** | 11\.24 $\pm$ 2.3     | **267\.2M** |
| MSfusion               | **75\.71 $\pm$ 0.5** | 22\.33M     | **5\.276 $\pm$ 0.4** | 290\.1M     |

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.

Reviewer comment (Experiments):
There is a gap between the experimental setup and the research background of the paper. From the title and abstract of the paper, I was looking forward to seeing collaborative training of large models, especially since the first sentence of the abstract mentioned models like GPT-3. However, I was surprised to find in the experimental setup that the visual tasks were performed using ResNet-18, and the model used for text-related tasks was not explicitly mentioned. But based on the calculated FLOPS in Table 1, I can infer that the parameter size of this model is far less than 1B. I am curious if a model with such parameter size can be considered a large model. If the authors decide to claim this framework is designed for Large Model, I highly recommend to conduct experiments on real large models such as models with at least billion-sized parameters.

Author response:
**Weaknesses 2:** Concern about limited model size and claim for the large model.

**Our response:** Thank you for pointing out this. Firstly, we note that many previous works claim to work for "large model", while the parameter sizes in their experiments do not reach billion level. In a well cited paper [4], experiments are conducted on 100 million size (GPT-2 117M); In an ICLR 2023 paper [5], they claim to work on "Large-scale linear models", while experiments are done on ResNet-18. 

Nevertheless, in order to better evaluate our proposed MSfusion for LLMs, we have now 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. The global full model parameter size for transformers applied to these datasets is 7.32M for PennTreebank, 19.3M for WikiText2 and 139.01M for WikiText103. The results are shown in the following table, which we have added to the Table 1 in revision.

|            | PennTreebank           |            |             | WikiText2              |             |             | WikiText103            |            |             |
| :--------- | :--------------------: | :--------: | :---------: | :--------------------: | :---------: | :---------: | :--------------------: | :--------: | :---------: |
| Methods    | Perplexity             | FLOPs      | $ \mu_n\$ | Perplexity             | FLOPs       | $ \mu_n\$ | Perplexity             | FLOPs      | $ \mu_n\$  |
| HeteroFL   | 55\.97 $\pm$  5.4    | 148\.6M    | 75%         | 579\.05 $\pm$  8.4   | 1\.08B      | 75%         | 784\.21 $\pm$  23    | 111\.9B    | 75%         |
| FedRolex   | 61\.52 $\pm$  6.8    | 148\.6M    | 75%         | 547\.32 $\pm$  45    | 1\.08B      | 75%         | 697\.42 $\pm$  33    | 111\.9B    | 75%         |
| Fed-ET     | N/A                    |            |             | N/A                    |             |             | N/A                    |            |             |
| MSfusion S | 9\.09 $\pm$  0.7     | **36\.2M** | 21\.875%    | 44\.33 $\pm$  3.6    | **198\.5M** | 18\.75%     |  13\.21 $\pm$ 1.1   | **32\.6B** | 21\.875%    |
| MSfusion M | 8\.02 $\pm$ 0.5     | 41\.8M     | 25%         | 5\.28 $\pm$  0.4     | 290\.1M     | 25%         | 8\.92 $\pm$  0.6     | 37\.3B     | 25%         |
| MSfusion L | **3\.11 $\pm$  0.2** | 91\.4M     | 50%         | **3\.59 $\pm$  0.2** | 633\.2M     | 50%         | **6\.31 $\pm$  0.4** | 74\.4B     | 50%         |

We note that these are the largest models we can experiment on given our computing resources. In addition to WikiText2, now the evaluations of MSfusion include 3 computer vision tasks and 3 NLP tasks. From the updated Table 1, it shows that MSfuion achieves high performance with low computation cost. 

From these additional experiment results, it is reasonable for us to infer that our propose method is suitable for even larger models.

[4] Tian, Yuanyishu, et al. "FedBERT: When federated learning meets pre-training." ACM Transactions on Intelligent Systems and Technology (TIST) 13.4 (2022): 1-26.

[5] Antorán, Javier, et al. "Sampling-based inference for large linear models, with application to linearised Laplace." ICLR 2023

Reviewer comment (Experiments):
Insufficient baselines. The authors claim that the proposed framework is compared with partial training-based approaches, however, there are some well-recognized approaches missed, e.g., [1][3].

Author response:
**Weaknesses 3:** Insufficient baselines.

**Our response:** The main reason why we have small number of baselines is that the problem setting in our paper is rarely investigated in the literature of FL. Hence there are limited number of methods we can fairly compare with.
   As detailed in the response to W1, there are fundamental differences between the settings and approaches between the listed papers [1] [2] and our proposed MSfusion. For FedRep in [3], we note that it cannot be considered as partial training-based method, since in FedRep the local training is still on full model parameters; and the model is only partially transmitted and aggregated at the central server to deal with the data heterogeneity problem. 

[3] Collins, Liam, et al. "Exploiting shared representations for personalized federated learning." International conference on machine learning. PMLR, 2021.

Reviewer comment (Theory):
Lacks of theoretical analysis. Considering the there is a proposed loss function in the framework, it is better to have a theoretical convergence analysis on the framework.

Author response:
**Weaknesses 4:** Lacks of theoretical analysis. 

**Our response:** We would like to remind the reviewer that we indeed had a convergence analysis of our proposed method, which was placed in Appendix A.1 due to page limit. We now add the following bold **Detail analysis for DSS are provided in Appendix A.1.** at Section 4.1 in our revision for easier identification. 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.

Reviewer comment (Writing):
There are some typos, e.g., " for TinyImageNet experiments for WikiText2 experiments is 800" in section A.2.2.

Author response:
**Weaknesses 5:** There are some typos.

**Our response:** Thanks for your careful reading. We have now proof read the draft and corrected the typos.

Reviewer comment (Novelty):
The technical contribution of this paper is limited. Except introducing a double shifting model splitting scheme, the proposed design has no significant difference from existing partial training work as reviewed in Section 2.3.

Author response:
**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.

Reviewer comment (Experiments):
The experiments cannot support the previous design and analysis sections. While the previous design sections are claimed to study the collaborative training of large language models (e.g., GPT-3 in Abstract), the evaluation diverges to computer vision tasks and only take one natural language dataset for evaluation.

Author response:
**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. 

|            | PennTreebank           |            |             | WikiText2              |             |             | WikiText103            |            |             |
| :--------- | :--------------------: | :--------: | :---------: | :--------------------: | :---------: | :---------: | :--------------------: | :--------: | :---------: |
| Methods    | Perplexity             | FLOPs      | $ \mu_n\$ | Perplexity             | FLOPs       | $ \mu_n\$ | Perplexity             | FLOPs      | $ \mu_n\$  |
| HeteroFL   | 55\.97 $\pm$  5.4    | 148\.6M    | 75%         | 579\.05 $\pm$  8.4   | 1\.08B      | 75%         | 784\.21 $\pm$  23    | 111\.9B    | 75%         |
| FedRolex   | 61\.52 $\pm$  6.8    | 148\.6M    | 75%         | 547\.32 $\pm$  45    | 1\.08B      | 75%         | 697\.42 $\pm$  33    | 111\.9B    | 75%         |
| Fed-ET     | N/A                    |            |             | N/A                    |             |             | N/A                    |            |             |
| MSfusion S | 9\.09 $\pm$  0.7     | **36\.2M** | 21\.875%    | 44\.33 $\pm$  3.6    | **198\.5M** | 18\.75%     |  13\.21 $\pm$ 1.1   | **32\.6B** | 21\.875%    |
| MSfusion M | 8\.02 $\pm$ 0.5     | 41\.8M     | 25%         | 5\.28 $\pm$  0.4     | 290\.1M     | 25%         | 8\.92 $\pm$  0.6     | 37\.3B     | 25%         |
| MSfusion L | **3\.11 $\pm$  0.2** | 91\.4M     | 50%         | **3\.59 $\pm$  0.2** | 633\.2M     | 50%         | **6\.31 $\pm$  0.4** | 74\.4B     | 50%         |

Reviewer comment (Reproducibility):
Some important evaluation details are missing, including the detailed parameter size and the number of layers of the transformer model, as well as the size of participants and the dataset partition for the experiment over WikiText2.

Author response:
**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.

|                                               | CIFAR10             | CIFAR100            | TinyImageNet        | PennTreebank         | WikiText2            | WikiText103          |
| :-------------------------------------------- | :-----------------: | :-----------------: | :-----------------: | :------------------: | :------------------: | :------------------: |
| Data/Token size                               | 50,000              | 50,000              | 100,000             | 929,500              | 2,088,600            | 103,227,000          |
| Local data/token size  (10 participants) | 5,000               | 5,000               | 10,000              | 92,950               | 208,860              | 10,322,700           |
| Local epoch                                   | 1                   | 1                   | 1                   | 1                    | 1                    | 1                    |
| Batch size                                    | 10                  | 10                  | 40                  | 100                  | 100                  | 300                  |
| Model applied                                 | ResNet18            | ResNet18            | ResNet18            | Transformer          | Transformer          | Transformer          |
| Hidden size                                   | [64, 128, 256, 512] | [64, 128, 256, 512] | [64, 128, 256, 512] | [512, 512, 512, 512] | [512, 512, 512, 512] | [256, 256, 256, 256] |
| Embedding Size                                | N/A                 |                     |                     | 256                  | 256                  | 256                  |
| Number of heads                               | N/A                 |                     |                     | 8                    | 8                    | 8                    |
| Dropout                                       | N/A                 |                     |                     | 0\.2                 | 0\.2                 | 0\.2                 |
| Sequence length                               | N/A                 |                     |                     | 64                   | 64                   | 64                   |
| Parameter size of   global full model    | 11\.2M              | 11\.2M              | 11\.3M              | 7\.32M               | 19\.3M               | 139\.01M             |