Probabilistic Federated Prompt-Tuning with Non-IID and Imbalanced Data

Fine-tuning pre-trained models is a popular approach in machine learning for solving complex tasks with moderate data. However, fine-tuning the entire pre-trained model is ineffective in federated data scenarios where local data distributions are diversely skewed. To address this, we explore integrating federated learning with a more effective prompt-tuning method, optimizing for a small set of input prefixes to reprogram the pre-trained model's behavior. Our approach transforms federated learning into a distributed set modeling task, aggregating diverse sets of prompts to globally fine-tune the pre-trained model. We benchmark various baselines based on direct adaptations of existing federated model aggregation techniques and introduce a new probabilistic prompt aggregation method that substantially outperforms these baselines. Our reported results on a variety of computer vision datasets confirm that the proposed method is most effective to combat extreme data heterogeneity in federated learning.

Paper

Similar papers

Peer review

Reviewer HV8G7/10 · confidence 4/52024-07-07

Summary

Mainly in the setting of imbalanced data, this paper proposes a probabilistic federated prompt-tuning pre-trained model from two aspects: local prompt generation and global prompt aggregation. In local prompt generation, each local set is assumed to be a random set of prompts distributed by a hierarchical generative model. In global prompt aggregation, non-parametric algorithm aligns local prompts of similar clients. On public computer vision datasets, experiment results demonstrate the effectiveness of the proposed method.

Strengths

1. This method is interpretable, and the local prompt is obtained based on Bernoulli distribution and Gaussian distribution sampling, and the EM algorithm is used to optimize the iterative process. 2. The quantitative results of imbalanced data and the prompts embedding results of t-SNE can fully show the effectiveness of the proposed method.

Weaknesses

1.The overall framework of the method is not clear enough. Although there is an overall framework diagram in the appendix, it omits a lot of details. It is confusing that the two DNNS are on the server side, and are the two DNNS on different clients consistent? 2.The quantitative comparison results of this method show GMM-PT, but the advantages of the proposed method and the theoretical distinction can be further explained.

Questions

1. The specific implementation details of the model, such as the DNN employed, the specific configuration of the pre-trained model seem to be unclear? 2. The lack of discussion of federated prompt-tuning pre-trained models in related work seems to be very relevant to this paper.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The lack of a discussion of limitations and deficiencies is important.

Reviewer h6E35/10 · confidence 5/52024-07-09

Summary

This paper studies the problem of prompt tuning in FL to address the data imbalance problem. The topic is interesting and broad enough for the community. The motivation and problem setting are good and promising. Experiments are sufficient to support the effectiveness of the proposed method.

Strengths

1. The problem of heterogeneous and imbalanced FL is interesting. 2. The overall writing and structure of this paper is good. 3. The experiments are sufficient to support the effectiveness of the proposed method.

Weaknesses

1. The local imbalance setting is very extreme, as described in section 4.1. There’s no real-world application to support this setting. 2. This paper lacks discussion of existing FL methods that address the class imbalance problem like: [1] Xinyi Shang, Yang Lu, Gang Huang, and Hanzi Wang, “Federated Learning on Heterogeneous and Long-Tailed Data via Classifier Re-Training with Federated Features,” in IJCAI, 2022. [2] Xinyi Shang, Yang Lu, Yiu-ming Cheung, and Hanzi Wang, “FEDIC: Federated Learning on Non-IID and Long-Tailed Data via Calibrated Distillation,” in ICME, 2022. [3] Wenke Huang, Yuxia Liu, Mang Ye, Jun Chen, Bo Du, “Federated Learning with Long-Tailed Data via Representation Unification and Classifier Rectification” in TIFS, 2024. Although the addressing problem is different in that the class imbalance is global, there are still connections in terms of problem settings. 3. The proposed prompt aggregation method seems not to directly address the problem of data imbalance. It is a more general solution to address the problem of heterogeneous FL. From the experiments, it can also be observed that PFPT achieves better performance on both non-iid and imbalance scenarios. 4. In the experiment, there should be more imbalance settings to support the robustness of PFPT. Currently, there is only one setting.

Questions

1. What are the real-world applications to support the scenario of extreme local imbalance in FL? 2. How does PFPT explicitly solve the problem of data imbalance? 3. Global imbalance seems more common in real-world applications. Can PFPT be effective in addressing the global imbalance problem?

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

Authorsrebuttal2024-08-07

Re: Additional detail regarding the extra experiments we run in response to the reviewer's Q3

Please note that not all methods in those references are tested on all those FL scenarios. Hence, for each scenario (reported in a separate table -- see Q3 in the main rebuttal), we only compare with the (best) reported performance of the methods which were tested on that scenario. -- All local data distributions and FL settings are configured to be the same for fair comparison: Batch size=32 for $\alpha=0.5$, 128 for $\alpha=0.1$ and 32 for ImageNet-LT. No. total clients=20. No. online clients in each communication round=8. No. epochs in local training=10. -- We used the following released implementation of those references: FEDIC: https://github.com/shangxinyi/FEDIC/blob/main/options.py CReFF-FL: https://github.com/shangxinyi/CReFF-FL RUCR: https://github.com/liuyuxia211/RUCR/blob/main/options.py -- [1] Xinyi Shang, Yang Lu, Gang Huang, and Hanzi Wang, “Federated Learning on Heterogeneous and Long-Tailed Data via Classifier Re-Training with Federated Features,” in IJCAI, 2022. [2] Xinyi Shang, Yang Lu, Yiu-ming Cheung, and Hanzi Wang, “FEDIC: Federated Learning on Non-IID and Long-Tailed Data via Calibrated Distillation,” in ICME, 2022. [3] Wenke Huang, Yuxia Liu, Mang Ye, Jun Chen, Bo Du, “Federated Learning with Long-Tailed Data via Representation Unification and Classifier Rectification” in TIFS, 2024

Reviewer tCnK6/10 · confidence 3/52024-07-13

Summary

This paper introduces a novel approach to prompt-tuning within the federated learning (FL) framework, focusing on enhancing the adaptability of pre-trained models across diverse clients using a probabilistic method. They proposed a hierarchical approach to model the generation and aggregation of local prompts. They develop a way to associate local (client side) and summarizing prompts (server side), utilizing a weighted bipartite matching task that linearly interacts with the model's loss function to optimize prompt association. Their experimental results show that the proposed method is effective in data imbalance settings.

Strengths

- The authors introduce an innovative probabilistic approach to prompt tuning within the context of Federated Learning. - The logic of this paper is clear and easy to follow. - The insights provided in the paper have the potential to inspire further research within the community, suggesting directions for integrating probabilistic approaches into federated tuning.

Weaknesses

- The paper primarily focuses on prompt tuning; a comparative analysis with other efficient tuning methods would strengthen its persuasiveness. - In Sec. 2, the related work mainly focuses on the solution drift issue, yet it lacks a comprehensive discussion on various efficient tuning methods under scenarios of data scarcity. The authors need to provide additional motivation for focusing on prompt tuning and expand on related work concerning prompt tuning. - The paper addresses data imbalance in Federated Learning but provides limited evidence supporting the proposed method. More substantial theoretical or experimental evidence is needed to demonstrate its effectiveness in imbalance settings.

Questions

1. In Sec. 3.1's remarks, an alternative tuning strategy, adapter networks, is mentioned. Could the authors clarify the advantages of their proposed method over these adapter network approaches? 2. In Sec. 3.2.3, I think the initialization of summarizing prompts will affect the performance. Could the authors specify how these prompts were initialized in the experiments? 3. In lines 178-181, it is hypothesized that each prompt captures a specific pattern or concept. Could the authors provide evidence or further explanation to support this assumption? This clarification would strengthen the theoretical foundation of the study. 4. In lines 223-225, could the authors provide more evidence for why within a single client, at most only one local prompt would be associated with the i-th summarizing prompt? 5. In Appendix G, "10 learnable prompts" are mentioned. Could the authors explore whether increasing this number could enhance the performance of the proposed method? More prompts could potentially allow for the learning of finer-grained concepts. 6. In line 201, could the authors provide more details on the method each client uses to select a subset of summarizing prompts?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

N/A

Authorsrebuttal2024-08-07

Extended discussion on fine-tuning approaches (supplement content to our main rebuttal)

Due to the limited rebuttal space, we defer this editorial discussion content to this separate comment: -- Existing fine-tuning approaches either use prompts to adapt the input [a] or adapter networks [b] to adapt the pre-trained weights. Both of which will help modify the pre-trained model to fit the context of a downstream task. Prompt-tuning methods focus on engineering cues, such as extra tokens which are appended as prefixes to the sequence of input embeddings to a multi-head self-attention unit. Such tokens or prompts provide beneficial context to performing the computational task, similar to how hints can be provided to assist puzzle solving. Adapter tuning is a parameter-efficient method for fine-tuning large foundation models. Instead of updating all the parameters of a model, adapter tuning involves adding small, trainable modules (adapters) to the model. During training, only these adapters are updated, while the rest of the model’s parameters remain fixed. The fundamental difference between these two approaches is that prompt-tuning only modifies the contextual information in the query, whereas adapter tuning instead alters how the pre-trained model behaves. In centralized learning, prompt-tuning is more memory efficient as the sizes of the prompts at the input embedding level only depend on the input sizes while the sizes of the update still depend on the model sizes. Thus, in the context of FL with resource-limited devices (e.g., wearable devices for health monitoring), prompt-tuning is more affordable in terms of memory usage. On the other hand, adapter networks can perform intricate updates to the model, and thus are suitable for more complex tasks. Both fine-tuning approaches are less investigated in the federated settings with heterogeneous and imbalance data. A few recent works [21-23] (as cited in our main text) have investigated a potential integration between FedAvg and prompt-tuning but have not addressed the prompt alignment issue (as elaborated in lines 51-59 in our main text) and have also not considered the heterogeneous and imbalance data setting. [a] "The power of scale for parameter-efficient prompt tuning." EMNLP (2021). [b] “Learning multiple visual domains with residual adapters”. NeurIPS (2017).

Reviewer DxxX5/10 · confidence 3/52024-07-13

Summary

This paper address the challenges of prompt-tuning pre-trained models in federated learning scenarios with diverse local data distributions. Specifically, it formulates the prompt summarizing procedure as a probabilistic set modeling task, treating each local set as an independent sample of a random point process and aligning similar prompts across different sets as part of the modeling parameterization. The research compares the proposed method's performance against various federated prompt-tuning baselines, demonstrating its effectiveness in combating data imbalance in extremely heterogeneous scenarios through a series of experiments and evaluations.

Strengths

1. The paper introduces a novel probabilistic set modeling approach for prompt summarization, which enhances the understanding and processing of local sets as independent samples of a random point process. 2. The method is good in addressing data imbalance in extreme heterogeneous scenarios, which is a significant challenge in federated learning environments. 3. The research includes comprehensive experiments and comparisons against existing federated prompt-tuning baselines, providing robust evidence of the method’s effectiveness. 4. The use of classical weighted bipartite matching within the generative model’s framework adds a layer of theoretical rigor to the research, grounding the practical contributions in solid mathematical foundations.

Weaknesses

1. “Papers to be submitted to NeurIPS 2024 must be prepared according to the instructions presented here. Papers may only be up to nine pages long, including figures. Additional pages containing only acknowledgments and references are allowed.” There is a minor formatting violation.

Questions

-

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

-

Reviewer h6E32024-08-08

Thanks for the authors' response. Most of my concerns are solved. I'm willing to increase my score to 5. However, I still suggest the authors reconsider the title because "data imbalance settings" are too broad to be covered by the proposed method.

Authorsrebuttal2024-08-08

Thank you for increasing the rating

Dear Reviewer h6E3, Thank you for the fast response. We appreciate the rating increase! We will update the title in our revision. Best regards, Authors

Reviewer HV8G2024-08-08

Thanks for the response. The authors have addressed my concerns, and I will raise my score to accept.

Authorsrebuttal2024-08-08

Thank you for raising the score to accept

Dear Reviewer HV8G, Thank you very much for agreeing to raise the score from weak accept to accept. We find this very encouraging! We would appreciate it if you could also update the score in the original review. Best regards, Authors

Reviewer DxxX2024-08-08

Rating

Given the authors’ response and other reviews' comments, I will maintain my rating.

Authorsrebuttal2024-08-09

Follow-up

Dear Reviewer tCnK, Thank you again for the detailed feedback. We hope our responses have addressed your questions sufficiently. We are happy to discuss further if you have follow-up questions for us. Best regards, Authors

Reviewer tCnK2024-08-13

Thanks for the nice rebuttal

Thanks for the author’s responses and additional experiments. Most of the concerns are addressed. And I will increase the score to 6.

Authorsrebuttal2024-08-13

Re: Thank you for increasing the rating

Dear Reviewer tCnK, Thank you for increasing the rating! We are glad our response has addressed your concerns. Best regards, Authors

Authorsrebuttal2024-08-14

Post-rebuttal summary

Dear AC and Reviewers, We would like to express our gratitude to the Area Chair for coordinating the review process of our paper. We are very glad that all reviewers found our rebuttal satisfactory and have given acceptance scores. We would also like to thank Reviewer DxxX for maintaining the original acceptance rating and Reviewers tCnK, h6E3, and HV8G for increasing their ratings of our work. We really appreciate your timely feedback and will revise the paper accordingly to incorporate our post-rebuttal discussion. Best regards, Authors

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC