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

We would like to thank the reviewer for recognizing our contribution. Your questions are addressed below:

Q1. Workflow Diagram. We apologize for the confusion. Our workflow diagram illustrates two different phases of our algorithm as annotated in the caption of Fig. 4. These two phases are iterative and thematically separated by the purple arrow:

On the left, the diagram highlights the local phase where each client (1) pulls a subset of prompts from the global set of the previous iteration and then (2) fine-tunes them.

On the right, the diagram highlights the global phase where local clients send their prompt sets to the server which then aggregates them.

Q2. Implementation Detail. Our work is based on a pre-trained Vision Transformer (ViT_B_32 pre-trained by ImageNet from PyTorch [a]. Each local client will fine-tune the pre-trained model using a local set of learnable prompts. At the end of each iteration, the local prompt sets will be shared with the server which then aggregates them into a global prompt set. Different baselines use different aggregation algorithms (see Appendix G).

[a] https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_32.html.

Q3. Discussion of federated prompt-tuning pre-trained models. We have provided this discussion in the introduction at lines 42-50. As the existing literature on federated prompt-tuning is relatively few [21-23] with most existing work utilizing FedAvg to aggregate local prompt sets, ignoring the prompt alignment issue, we decided to merge this discussion with our positioning in the introduction. If this work is accepted, we can use the 10th page to expand this discussion.

We hope our response has addressed your questions. Please let us know if you have any follow-up questions for us.

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

We would like to thank the reviewer for the detailed questions which we have addressed below.

Q1. Real-world application to support this setting. We would like to emphasize that our work has first showcased its improved performance on standard non iid settings simulated with Dirichlet prior -- see the first 2 rows of Tables 1-4 -- on local class distribution which is the common practice in prior heterogeneous FL work (including the references [1-3] suggested by the reviewer). This already implies a certain degree of imbalance in local data distribution. The extreme imbalance setting is considered as a limit test in the worst case. Previous works have also studied such setting [a-f]. Extreme class imbalance often happens in practical scenarios, e.g., (a) rare disease detection with extremely low prevalence rate (less than 1/10000) [e]; and (b) mobile keyboard prediction where on-device datasets are often imbalanced, with a few classes (e.g., commonly typed words or phrases) dominating [f].

[a] FedIIC: Towards Robust FL for Class-Imbalanced Medical Image Classification (https://arxiv.org/abs/2206.13803)

[b] FL with Non-IID Data (https://arxiv.org/abs/1806.00582)

[c] Addressing Class Imbalance in FL (https://arxiv.org/abs/2008.06217)

[d] FL with Matched Averaging (https://www.arxiv.org/abs/2002.06440)

[e] Complementary Pattern Augmentation for Rare Disease Detection (https://arxiv.org/abs/1911.13232)

[f] FL for Mobile Keyboard Prediction (https://arxiv.org/abs/1811.03604)

Q2.

A. Discussion of existing FL methods addressing class imbalance

We thank the reviewer for suggesting these additional references [1-3]. We note that these works focus on full fine-tuning approaches in settings where the centralized dataset is imbalanced. Instead, we focus on light-weight prompt-tuning scenarios in which the local datasets are imbalanced but their centralization would be reasonably balanced. As such, our setting and the settings adopted in these additional references are complementary. We will cite them and include the above discussion in our revision.

B. Explain how PFPT solves the problem of (local) data imbalance.

Local data imbalance will cause diversity across local prompt sets (i.e., clients). We refer the reviewer to Fig. 2, which shows that the local prompts indeed diverge to capture different data patterns. This means the same prompt position across different clients might encode different context information of the fine-tuning task. A simple aggregation that simply combines prompts at the same position across different clients from different contexts might therefore collapse them into less informative prompts. To avoid this, we must learn a prompt alignment so that we only aggregate prompts which encode the same aspect of context information. PFPT models this prompt alignment as a parameter of our generative model of the local prompts. Maximizing the observation likelihood of the local prompts will allow us to learn the most plausible alignment, hence solving the local data imbalance issue.

Re: the concern that the proposed prompt aggregation method is a more general solution to address the problem of heterogeneous FL and achieves better performance on both non-iid and imbalance scenarios:

We are not sure if this is considered a weakness. However, we agree that the title of our paper might be a bit too specific on imbalance data settings. Our main point is that existing prompt aggregation performs worse when the local data become more diverse and imbalanced. It can be observed (from all table) that the performance improvement over baselines increases when the data becomes more imbalanced. In the synthetic 4-dataset benchmark, with more diverse data collected from 4 different datasets, the performance gain over the baselines is much more pronounced.

Q3. More imbalanced settings & how effective PFPT is in global imbalance

Our work focuses on the common setting in existing heterogeneous FL work where the centralized dataset would be reasonably balanced but each local dataset can be imbalanced due to the client heterogeneity. This is the main cause of the solution drift issue.

However, to our understanding, we believe the reviewer wants us to explore the settings under which the centralized dataset is imbalanced (i.e., global imbalance). This presents an orthogonal challenge to the solution drift challenge considered in existing heterogeneous FL literature (including ours) which is beyond the current scope and deserves a separate treatment.

Nonetheless, we have run additional experiments comparing with the methods in the references [1-3] that the reviewer suggested. The experiments also use the same long-tailed datasets & FL scenarios detailed in those references. The results show that our method achieves higher accuracy than prior works substantially across all global imbalance settings with different imbalance factors (IF):

We hope the reviewer will consider increasing the rating if our response has addressed the questions sufficiently. We will be happy to answer any follow-up questions that the reviewer might have for us.

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

We thank the reviewer for recognizing our contribution and for the constructive feedback, which are addressed below.

Q1. Discussing other fine-tuning approaches.

We agree with the reviewer that a more comprehensive discussion on PFPT will make our contribution position clearer. However, we want to note that our main contribution is about addressing a technical issue (i.e., prompt alignment) of an existing federated prompt-tuning approach, which is essential in settings with heterogeneous and/or imbalance local data:

--

Existing federated prompt-tuning approaches [21-23] have ignored the important issue of prompt alignment across different clients. Aggregating misaligned prompts could result in suboptimal performance under imbalanced and/or heterogeneous local data (see our experiments). Our work aims to fix this problem. Therefore, our focus is necessarily specific to prompt-tuning.

--

Generalizing our proposed solution to other fine-tuning techniques and finding which one is most effective is indeed interesting but orthogonal to our contribution. We will investigate this in a separate follow-up work.

Nonetheless, we have provided a comparative analysis with adapter tuning in Appendix F (Table 7), which shows that prompt-tuning outperforms adapter-tuning (FEDAVG-Adapter, FEDPROX-Adapter) in most FL scenarios on 4 benchmark datasets. The results in Table 7 are quoted below and also supplemented with additional comparison with two other variants of federated adapter-tuning (FEDOPT-Adapter, SCAFFOLD-Adapter). We note that these new baselines were created during the rebuttal week for a more thorough comparison but they have not been investigated in prior literature.

In our revision, we will also include a broader discussion of existing fine-tuning techniques in Section 2, which is summarized in an additional comment below (due to limited rebuttal space).

Q2. Initialization of summarizing prompts.

The initialization of prompts would not affect the performance as it follows the standard initialization in [a] which is commonly used in deep learning. The reported std in our experiments is also relatively small, indicating that prompt initialization has little effect on performance variation.

[a] Understanding the difficulty of training deep feedforward neural networks. AISTATS (2010).

Q3. Each prompt captures a specific pattern of concept.

We refer the reviewer to Fig. 2 in our paper, which shows the t-SNE plots of the (learned) summarizing prompts on CIFAR-100 over 120 communication iterations. Yellow triangles denote the centroids of the t-SNE embeddings of the prompts. The dashed red line visualizes their trajectories. The plots show that each prompt follows a different trajectory, suggesting that each prompt does capture a specific pattern or concept.

Q4. For a single client, at most only one local prompt would be associated with the i-th summarizing prompt? This is due to the generative design in Section 3.2.1 (line 183-186). Per client, each summarizing prompt will flip a Bernoulli coin with learnable probability to decide whether to sample one local prompt from its vicinity. If the decision is positive, the sampled local prompt is said to be associated with the summarizing prompt. Hence, per client, for each summarizing prompt, there is at most one associated local prompt. This setup is similar to that of the setting in The Indian Buffet Process: An Introduction and Review. JMLR (2011)

Q5. More prompts improve performance. We have run additional experiments showing the increasing performance of our model with an increasing number of prompts -- see the figure in the attached PDF in our summary response.

Q6. Method that each client uses to select a subset of summarizing prompts? We adopt the prompt selection mechanism in “Learning to Prompt for Continual Learning” (CVPR, 2022)

We hope the reviewer will consider increasing the rating if our response has addressed the questions sufficiently. We will be happy to answer any follow-up questions that the reviewer might have for us.

We would like to thank the reviewer for recognizing the strength of our paper and for helping us catch a format issue.

We thought the broader impact section is counted as part of the checklist content which is not counted towards the page limit. It is possible that we had misunderstood the policy regarding the broader impact statement. We will move this section to the appendix in the revision to make sure it is not considered part of the main text.

We hope the reviewer will consider increasing the rating of our paper as this minor format issue is the only concern and it does not impact the technical contribution of our work.

Thank you for your consideration.