Bridging the Gap Between Foundation Models and Heterogeneous Federated Learning

[Question 4] Further clarification on the key contributions in detail.

Thank you for requesting further details on the key contributions of our work. We have made revisions to our manuscript to provide clearer explanations, which are briefly summarized below:

1. Specialized FM compression algorithm: Our model compression algorithm is designed to transformer-based FMs and tailored to the unique constraints of edge clients, enabling efficient deployment of models in a heterogeneous environment. It dynamically scales FMs to suit the computational capabilities of each edge device, with more powerful devices receiving larger models and resource-constrained devices receiving smaller models. Hence, we refer to it as the specialized FM compression algorithm for edge-FL.

2. Enhance resource utilization: In edge-FL scenarios, local clients have varied resource capacities. Traditional approaches often lead to under-utilization of high-capacity devices due to the need to accommodate the limitations of the least capable devices. RaFFM addresses this inefficiency by deploying appropriately scaled models to each device, thereby optimizing the overall system efficiency.

3. Significant Reduction in communication overhead: RaFFM's strategy of using scaled-down models for less capable clients leads to a considerable reduction in communication costs. For instance, our experimental results (Table 4, Section 4.4, and Figure 2) show that with the FLAN-T5-base model, RaFFM achieved a reduction of 227.02GB in communication overhead to reach the target F1 score. This quantifies the significant reduction in communication overhead facilitated by our approach.

General response to Question on Introduction

Thank you for your constructive feedback on improving the introduction of our manuscript. We have made several modifications based on your suggestions, including the addition of citations to support our arguments, clarification of the points you raised in your review comments, and inclusion of further motivations, etc.

[Question 1] The problem statement “FMs are superior in few-shot transfer learning and FL” could be more explicit and clearer, and lack of reference support.

We greatly appreciate your feedback on making our problem statement more explicit and providing supporting references. In response, we have revised our manuscript to offer a clearer explanation and added relevant citations to substantiate the strengths of FMs in few-shot transfer learning, as well as their potential in non-IID FL.

Specifically, we have included citations [1] [2] to support our assertion regarding the efficacy of FMs in few-shot transfer learning. Additionally, we have referenced [3] to illustrate the potential advantages of FMs in non-IID FL environments.

Moreover, our experiments (see Table 3 on manuscripts) demonstrate the effective application of the FM model (ViT) in few-shot learning tasks within image classification. This empirical evidence further bolsters our claim regarding the suitability of FMs in these contexts.

Reference

[1] Brown, Tom, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D. Kaplan, Prafulla Dhariwal, Arvind Neelakantan et al. "Language models are few-shot learners." Advances in neural information processing systems 33 (2020): 1877-1901.

[2] Hu, Shell Xu, Da Li, Jan Stühmer, Minyoung Kim, and Timothy M. Hospedales. "Pushing the limits of simple pipelines for few-shot learning: External data and fine-tuning make a difference." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9068-9077. 2022.

[3] Zhuang, Weiming, Chen Chen, and Lingjuan Lyu. "When foundation model meets federated learning: Motivations, challenges, and future directions." arXiv preprint arXiv:2306.15546 (2023).

[Question 2] Why do fine-tuning FMs typically require significantly more resources compared to inference?

Thank you for highlighting the importance of clarity on why fine-tuning FMs typically requires significantly more resources compared to inference. We have expanded our manuscript to include a more detailed explanation and a relevant reference to support this point.

Here is a brief explanation: Fine-tuning FMs is more resource-intensive than inference because it necessitates maintaining the entire computational state of the model in memory for optimization purposes. This process involves not only storing the full model weights but also retaining gradients, activation states, and optimizer momentum states, all of which contribute to increased memory usage. This explanation is further detailed in the tutorial by Stern [1].

In our experimental setup, with a mini-batch size of 4, we observed that fine-tuning consumes approximately seven times more resources than inference. This disparity in resource consumption becomes even more significant with larger batch sizes.

Reference

[1] Stern, J. (2022, August 18). A comprehensive guide to memory usage in pytorch. Medium. https://medium.com/deep-learning-for-protein-design/a-comprehensive-guide-to-memory-usage-in-pytorch-b9b7c78031d3

[Question 3]: Lack of broader impact

Thank you for emphasizing the importance of addressing the broader impact of our work. Based on your suggestion, we have expanded the discussion on this aspect within the main body and the conclusion section of our manuscript.

Additionally, the broader impact of RaFFM can be summarized as follows:

1. Addressing Deployment Challenges of FMs in FL: RaFFM specifically tackles the challenges associated with deploying large-scale FMs in resource-heterogeneous edge-FL systems.

2. Introducing novel FM compression algorithms for edge-FL clients:

We have developed specialized FM compression algorithms tailored for edge-FL systems. These algorithms enable dynamic compression of FMs to meet the resource constraints of individual edge clients, thereby facilitating more efficient and effective model deployment.

3. Limitations and future research directions: We also acknowledge the limitations in the Conclusion section. Notably, certain foundation models may remain unsuitable for deployment on resource-constrained edge devices, even after compression. We also highlight a need for advancements in hardware and algorithmic strategies, pointing towards potential directions for future research.

General response to Question on Abstract

Thank you for your feedback on the abstract. In response, we have significantly revised the abstract to provide a clearer and more detailed overview of our work. The revised abstract now includes specific quantitative data supporting our claims, addresses the key points you raised, and integrates relevant examples and motivations to provide a comprehensive overview of our work.

[Question 1] What metrics are used to measure this superiority? How “significant” are the improvements?

We appreciate your advice on providing clarity regarding the measurement metrics. In response, we have updated the abstract with specific quantitative measures as follows to avoid unclarity:

"RaFFM demonstrates up to 2.3 times improvement in resource utilization efficiency and saving up to 225.02 GB in communication costs when optimizing the FLAN-T5 Base model.”

Additionally, here are explanations of the evaluation metrics (detailed evaluations are available in Experiment section) :

Communication Cost Measurement: The network traffic between edge clients and the server is quantified in Gigabytes (GB) and Megabytes (MB). Detailed data and analysis can be found in Section 4.4, Table 4, and Figures 2(c) and 2(d) of the Experiment section.

System Resource Requirements: We quantify the system peak computational memory (RAM) requirements, in Gigabytes (GB), necessary to deploy and optimize a given FM across the entire edge-FL system.

Resource Utilization: This metric evaluates the average proportion of memory utilized during both the deployment and training phases, compared to the total computational memory available on edge devices.

[Question 2] Unclarity statements of performance are on par with traditional FL methods applied to full-sized FMs.

Thank you for highlighting the need for greater clarity regarding our comparative performance statements. We have taken your suggestion into account and have revised our abstract accordingly to provide a more quantifiable presentation.

In the revised abstract, we now include the following statement to offer specific metrics:

“In our experiments with the Vision Transformer (ViT) model, RaFFM's resource-aware deployment achieves a training acceleration of approximately 1.76 times compared to full-size FM deployments, with a minimal accuracy loss of 0.49%.”

On our original statements, the key information we want to convey is that RaFFM markedly enhances the training speed of FMs in edge-FL scenarios while incurring minimal performance loss compared to deployments using full-sized models. We followed your advice and added the quantitative example, which further underscores and clarifies our statements.

[Question 3] Clarification on the statement effective “across tasks in both natural language processing and computer vision domains."

We appreciate you for highlighting the need for specificity in our claims. Accordingly, we have revised our manuscript, we now detailed the scope and limitations of our framework’s effectiveness in NLP and CV domains. In the revised abstract, we now clarify:

“RaFFM demonstrates its effectiveness in a range of tasks within natural language processing, including question-answering and sentiment analysis, and in computer vision, particularly in image classification tasks.”

Furthermore, the Experiment section of our manuscript provides an in-depth analysis of these tasks. For NLP, we have thoroughly tested RaFFM across a spectrum of tasks like question answering, sentence classification, and sentiment analysis. In the CV domain, our focus has been on image classification tasks, where we have analyzed RaFFM’s strengths and weaknesses compared to baseline methods.

General response for reviewer U93Y

Dear reviewer U93Y,

Thank you for your valuable and constructive feedback. In response to your comments and suggestions, we have carefully revised the manuscripts to enhance clarity and address potential miscommunication. We have incorporated citations to support our arguments, clarified points raised in your comments, and added further motivations in the latest version of our paper.

Additionally, inspired by suggestions from other reviewers, we have conducted supplementary experiments to enrich our research. These include:

• Experiment on the Necessity of Specialized Salient Parameter Prioritization (Link 2 in supplementary material):

• Experiment Comparing with Additional Baselines (Link 3 in supplementary material)

• Post-training High-performance Resource-aware Model Deployment (Link 4 in supplementary material)

If there are any additional questions, please do not hesitate to let us know. We are committed to providing any necessary supplemental support material and explanations.

Dear Reviewer U93Y 😊,

I hope this message finds you well. We are writing to follow up on the revisions we submitted in response to your valuable feedback. We have made concerted efforts to address each of your concerns and suggestions comprehensively.

As the open discussion period is approaching its deadline on November 22nd, we understand that you have many commitments and appreciate the time and effort you dedicate to the review process. If any further questions or clarifications are needed, please do not hesitate to let us know. We are more than willing to provide any additional information that can assist in your review.

Thank you once again for your insightful comments and valuable guidance. Your feedback is crucial in helping us refine our work, and we await your thoughts on our response.

Best regards 😄,

Anonymous authors

[Question 1] Motivation and contribution of salient parameter prioritization

Thank you for your insightful comments. We value this opportunity to clarify the novelty and specific contributions of our work. We have included the references you suggested ([1], [2], [3]) in our revised manuscript to acknowledge the established work in this field.

Firstly, we wish to clarify that our use of the L1 norm as a saliency metric is part of a broader framework. We have developed a specialized salient parameter prioritization (SPP) algorithm specifically for transformer-based Foundation Models (FMs), as detailed in Section 3. This approach is integrated with the other components of RaFFM and addresses the specific challenges of applying FMs in FL environments.

While we recognize that identifying salient parameters is an established method in the domain of network pruning. However, our proposed SPP is more than just measuring saliency, prioritization implies weight reordering, and the reordering benefits of the scaled FMs always retain the most salient parameters. Our novel contribution of SPP lies in retaining the inherent attention saliency of transformer-based FMs after model scaling, particularly in heterogeneous resource FL contexts.

Because transformers, particularly multi-head attention layers, are designed to capture positional and sequential information. An example is Vision Transformer (ViT), where adding positional embeddings significantly improve the performance. Traditional pruning-based parameter prioritization methods, if applied directly, would destroy the inherent attention characteristic in FMs. To circumvent this, we propose our specialized SPP to preserve the essential capability of multi-head attention layers.

Summary of our novel contribution in SPP:

• Maintain the integrity of pre-trained knowledge in FMs . (Demonstrate by the supplementary materials link 2 and Theorem 1 in Appendix, details refer Question 2 response)

• Specialization for Transformers: Tailoring the salient parameter prioritization specifically for multi-head attention-based transformer models.

• Prototype for Heterogeneous FL model aggregation: Strategically placing salient parameters at the forefront for efficient extraction and aggregation in heterogeneous model FL.

General response to Reviewer Y3SH

Dear reviewer Y3SH,

We appreciate your insightful comments. We have taken your feedback into careful consideration and have made the following enhancements: first, we emphasized the novelty of our proposed transformer-based salient parameter prioritization (SPP) algorithm. Secondly, we have conducted additional experiments to further validate the effectiveness of our SPP algorithm compared to existing rank-based pruning algorithms in transformer models. Furthermore, suggested by the reviewer, we have conducted additional experiments comparison with PruneFL and PriSM. Additionally, we have also emphasized the distinctions between RaFFM and network pruning, referenced the suggested papers, and provided a detailed experimental setup.

The follow-up experiments are available in the supplementary materials, where we provide comprehensive information, step-by-step instructions, and Jupyter notebook tutorials to verify our results:

• Experiment on the Necessity of Specialized Salient Parameter Prioritization (Link 2 in supplementary material):

• Experiment Comparing with Additional Baselines (Link 3 in supplementary material)

• Post-training High-performance Resource-aware Model Deployment (Link 4 in supplementary material)

If there are any additional questions, please do not hesitate to let us know. We are committed to providing any necessary supplemental support material and explanations.

[Question 4] Difference between Network Pruning

Network pruning: Network pruning typically involves reducing the complexity of a pre-trained model by eliminating less significant neurons or connections. This process often occurs post-training and requires subsequent fine-tuning to restore or enhance the model’s performance. A key limitation, especially relevant in the context of federated learning at the edge, is the computational intensity of this post-pruning fine-tuning process. When dealing with edge clients with limited computational resources and local data, this approach can be particularly challenging.

Salient Parameter Prioritization in RaFFM: RaFFM’s Salient Parameter Prioritization (SPP) extends beyond mere measurement of saliency. It encompasses strategic weight reordering, ensuring that the scaled Foundation Models (FMs) retain the most crucial parameters. Significantly, SPP maintains the attention mechanism's saliency within scaled FMs. SPP enables RaFFM to collaboratively and dynamically optimize scaled FMs at the edge of resource constraints, based on their local resources. The outcome is a scalable FM capable of generating multiple high-performance subnetworks that meet these constraints. A notable advantage of RaFFM is its ability to scale an FM post-federated learning without necessitating additional fine-tuning.

To practically demonstrate RaFFM's efficacy, we provide a demonstration in the supplementary materials (link 4). This demo illustrates RaFFM's capability to generate multiple high-performance subnetworks post-training without the need for extra fine-tuning. It exemplifies RaFFM's efficiency in adapting FMs for diverse and resource-constrained FL environments while preserving their performance and integrity, particularly with transformer architectures.

Dear Reviewer Y3SH 😊,

I hope this message finds you well. We are writing to follow up on the revisions we submitted in response to your valuable feedback. We have made concerted efforts to address each of your concerns and suggestions comprehensively.

As the open discussion period is approaching its deadline on November 22nd, we understand that you have many commitments and appreciate the time and effort you dedicate to the review process. If any further questions or clarifications are needed, please do not hesitate to let us know. We are more than willing to provide any additional information that can assist in your review.

Thank you once again for your insightful comments and valuable guidance. Your feedback is crucial in helping us refine our work, and we await your thoughts on our response.

Best regards 😄,

Anonymous authors

General response to Reviewer erSL

Dear reviewer erSL,

We are grateful for your supportive and insightful comments.

In line with your comments, we have revised our original manuscript. The motivation has been clarified, particularly highlighting the broader impacts from a security perspective of integrating Foundation Models (FMs) into Federated Learning (FL). Additionally, we have augmented our experimental results with quantitative measurements that illustrate the performance of our algorithms in end-user FL scenarios.

Additionally, Following constructive suggestions from other reviewers, we have also conducted three additional experiments:

• Experiment on the Necessity of Specialized Salient Parameter Prioritization (Link 2 in supplementary material):
• Experiment Comparing with Additional Baselines (Link 3 in supplementary material)
• Post-training High-performance Resource-aware Model Deployment (Link 4 in supplementary material)

If there are any additional questions, please do not hesitate to let us know. We are committed to providing any necessary supplemental support material and explanations.

[Question 1] Clarification on Motivation to Train Larger Foundation Models with Federated Learning

Thank you for your remarks regarding the motivation behind using FL to train larger FMs. We have revised our manuscript to provide a clearer motivation on the abstract, introduction, and conclusion section.

Indeed, our primary focus is on the fine-tuning of FMs within FL scenarios, specifically addressing the unique challenges that arise when fine-tuning these models in a collaborative edge-FL environment.

Given the substantial size of FMs and their rich pre-trained knowledge, our proposed RaFFM method prioritizes resource efficiency and the preservation of this pre-trained knowledge. Additionally, RaFFM is designed to support post-FL model deployment, enabling the generation of a significant number of scaled models that cater to various resource constraints encountered at the edge.

[Question 2] Broader Impact on Privacy and Security

Thank you for highlighting the importance of discussing the broader impact of our solution on privacy and security. In response, we revised our manuscripts correspondingly and added related discussion on the Conclusion section.

While our primary focus has been on addressing the challenges of resource heterogeneity in applying FMs to FL, we recognize the implications for privacy and security. Incorporating FMs into FL can indeed have a positive impact on FL's security. Due to their pre-training on large corpora of public datasets, FMs have the potential to significantly mitigate threats such as inference attacks. As highlighted in the work [1], amidst the growing concern for statistical notions of privacy, the application of FMs may contribute to achieving perfect secrecy for certain sensitive tasks, given their rich pre-trained knowledge.

Reference

[1] Arora, Simran, and Christopher Ré. "Can Foundation Models Help Us Achieve Perfect Secrecy?." arXiv preprint arXiv:2205.13722 (2022).,

[Question 3] Quantitative Measurement of RaFFM’s Speed Up on End-User Devices

Thank you for your valuable suggestions regarding the need for quantitative measurements of RaFFM’s efficiency on end-user devices. In response, we have conducted additional experiments and provided specific data to demonstrate RaFFM’s performance in this context, as shown in Table 1.

We focused on optimizing the BERT-base model for the SST-2 dataset in an FL scenario involving 100 clients, with a participation rate of 10% among local clients. Recognizing that actual training time can vary due to several uncontrollable factors such as device temperatures and electricity voltage, we chose to measure energy consumption as a more stable indicator of efficiency.

Our experiments were conducted across five different edge system settings, categorized as Xsmall, Small, Medium, Large, and Xlarge (detailed settings information can be find in our supplementary materials link 5). were designed to represent a range of end-user devices, from resource-constrained to more powerful systems. As shown in Table 1, by tracking the peak energy usage for each setting, we provide a more accurate and reliable measure of RaFFM's efficiency and its adaptability to diverse hardware environments.

Table 1. Peak Energy Usage at Edge System