Mixture of Experts Made Personalized: Federated Prompt Learning for Vision-Language Models

We appreciate the thoughtful feedback provided by the reviewers. Below we address each concern in detail:

Reply to W1. Our approach incorporates several design choices that make this process both efficient and effective:

• Our KNN-based expert selection method (Section 3.2) provides an automatic, data-driven approach to expert inclusion that adapts to the distribution of client data. This removes the need for manual intervention in expert selection.

• The pool maintains natural diversity through our dynamic refresh mechanism (Eq. (7)), which ensures the pool stays current with the latest client updates while preserving historical knowledge through non-participating clients' prompts.

• The lightweight nature of prompts (only 8,192 parameters per prompt as noted in Sec. 1) makes maintaining even a large pool computationally feasible, allowing us to keep more prompts in the pool rather than aggressively removing them.

While finding the theoretically optimal selection method is impractical and beyond the scope of our work, our extensive experiments across 9 datasets demonstrate that the proposed approach consistently outperforms state-of-the-art alternatives, validating the effectiveness of our pool management strategy.

Reply to W2. Our design tries to specifically address the optimization challenge under data heterogeneity concern by the following strategies:

• The parameter-efficient design (Section 3.3) reduces the network from 4.2M to just 66.0K parameters through strategic dimension reduction, making optimization more tractable.

• Our multi-head attention mechanism learns relationships between image and text features rather than direct expert weights ($G(\cdot)_i$), providing greater robustness to heterogeneous data (as explained in Section 3.3) by leveraging the feature alignment ability of pre-trained CLIP-like VLMs.

• The weighted averaging approach (Equation 10) ensures stable training by maintaining a balance between local expertise and collective knowledge.

These design choices result in stable convergence, as evidenced by our experimental results across diverse federated settings.

Reply to W3. We agree with the reviewer that the proposed method increases the complexity of server maintenance complexity, and thus might not be a perfect fit for extremely large number of clients (e.g. 1 million clients). This being said, due to the compact size of prompts (8,192 parameters), both communication costs and server-side pool management remain efficient. Our approach maintains this efficiency even with large numbers of clients for several reasons:

• The extremely small size of prompts (8,192 parameters) means communication overhead remains minimal even with many clients. This is orders of magnitude more efficient than traditional federated learning approaches that share full models.

• Our KNN-based expert selection automatically sets a limit to the maximum transmitted prompts and the communication overhead, ensuring communication complexity scales linearly rather than quadratically with the number of clients.

• The server-side pool management is computationally efficient, requiring only simple matrix operations for updates and KNN calculations.

Furthermore, we would like to mention that the benefits of improved personalization and performance outweigh the modest additional coordination overhead, as demonstrated by our comprehensive experimental results.

Please feel free to let us know if you have any more concerns, and we will try our best to address them. Thank you for your feedback.

We would like to thank the reviewer for your time and patience in reviewing our manuscript. Below, we reply to every concern in weaknesses and questions.

Reply to W1. We conduct additional user-level differential privacy experiments (line 858-887 & Tab. 5). Below is a preview.

Table 5: Performance under ($\epsilon$, $\delta$)-differential privacy on CLIP datasets under pathological non-IID setting.

And we would also like to mention that in pFedMoAP, clients remain unaware of the origin of non-local prompts, including which client they come from or the training round they were last updated in. Besides, the non-local prompts are downloaded after a selection process on the server. These factors collectively make it extremely difficult for a client to infer the gradient associated with a prompt or the data it has been trained with, which largely mitigates the privacy risk.

Reply to W2. We have modified the description of the previous works in line 783 to avoid misunderstanding. Essentially, what we wanted to express is that the previous existing prompt-based FL works habitually adhere to the paradigm where locally learned prompts are not allowed for inter-client sharing before aggregation. This is because these works are inherent in the traditional FL framework for full-size models (as opposed to the prompts). Sharing full-size model before aggregation will cause unaffordable communication overhead, and thus, traditional FL does not consider inter-client sharing. However, since prompt is much smaller than the full-size model, this inheritance of the paradigm becomes unnecessary.

We would like to thank the reviewer for your time and patience in reviewing our manuscript. Below, we reply to every concern in weaknesses and questions.

Reply to W1. We sincerely apologize for the typos and lack of spaces between words. We have corrected them in the latest version.

Reply to W2a. We agree that pFedMoAP is under the umbrella of prompt-based FL with MoE (the synergy of which has not been explored by the community). Meanwhile, we also believe that this fact does not cover up our novelty and contributions, because:

• to the best of our knowledge, we are the first to propose a Mixture of Adaptive Prompt architecture for CLIP-like VLMs under federated settings. The proposed pipeline (Fig. 2) efficiently brings non-local knowledge in the adaptive prompts to boost the client's performance.
• we circumvent the paradigm of traditional FL with the pre-aggregated weights are not allowed to share. Existing prompt-based FL works habitually adhere to this paradigm, failing to fully leverage the lightweight nature of prompt.
• most importantly, from pure MoE perspective, our proposed attention-based gating network for VLMs is novel and favorable compared to the the linear projection-based gating network in traditional MoE. We detailed the reasons of it under federated settings in line: line 316-353. Briefly, attention-based gating network 1) is more robust to adaptive experts; 2) serves as linear probing with more capacity; 3) leverages CLIP's feature alignment with attention mechanism; and 4) is agnostic to experts' order. We conduct additional experiments to compare these two types of gating network to show the superiority of the proposed attention-based gating network. (line 889-909 & Tab. 6)

Reply to W2b. The intuition behind updating the non-local experts is that as the training progresses, the trained prompts become more and more like experts as they converge. The training can be considered as a process that keeps embedding various client-specific information in the prompts. When downloaded to a client, the more developed and thoroughly trained prompts offer more non-local information and help the local training as more qualified experts.

Reply to W3. The main reason why the gating network is not aggregated is due to the communication overhead. We provide the sizes of different gating network (Tab. 8) based on different (pooled) dimensions of the feature output by CLIP's image and text encoders. We can see that the size of the gating network in general is significantly larger than a prompt with size $16\times512\approx 8k$. Uploading and downloading the gating network will increase the communication overhead by a magnitude or more. Besides, we conduct additional experiments (line 889-909 & Tab. 6) and empirically show that aggregating the gating network would not much benefit the personalized performance. This shows that the global aggregating with a short personalization (i.e. the selected round) of the gating network can hardly work better than the completely localized gating network.

Thank you for your feedback and suggestions on the visualization, for which we added a new set of experiments (line 969-1006 and Fig.6). In these experiments, we visualize the contribution of the local and non-local experts towards the MoE on 6 datasets. In addition, we also conduct new experiments (Fig. 5, line 946) and plot the trend of the performance with different numbers of experts on CIFAR10 (Dir($\alpha$=0.5) with 100 clients). Both experiments suggest that the non-local experts are indispensable to the performance gain. Meanwhile, as the incorporation of the non-local experts has been empirically proven in these experiments, we would also like to politely mention that the other reviewers have all given credit to our novelty, and we briefly summarize our three major contributions/novelties as 1) the algorithm that personalizes prompt-based federated learning for VLMs with MoE architecture; 2) the attention-based gating network for VLM (not necessarily for FL) which outperforms the traditional projection-based gating network by large margins; and 3) the paradigm shift from existing FL approaches that limit pre-aggregated models from inter-client sharing.

The reason why the performance of aggregated gating is worse than non-aggregated gating is that, while aggregating the gating network embeds global knowledge to the parameter weights, as the deepest layer of the entire model that is parameterized, it should be fully "personalized" as opposed to "globalized" to achieve higher performance [1, 2]. The process of aggregating the gating network limits the extent of personalization to only one round of personalization (the round where the client is selected). This is much less than fixing the gating network locally throughout the entire federated training, especially in the case of few-shot learning (Tab. 6's results are on CLIP datasets, which are used for few-shot learning as mentioned in line 373) where one round of fine-tuning involves few images for the gating network to be trained with, hence the worse performance. With this being said, the design choice of not aggregating the attention-based gating network is also due to the large size of the gating network (Tab. 8). The attention-based network can be up to around 3 magnitudes larger than the prompts (gating with $d_{\text{feature}}=1024$: 4.2M parameters, prompt:16$\times$512 parameters), which would largely increase the communication overhead.

Please feel free to let us know if you have any more concerns. And we appreciate your consideration for raising the score. Thank you.

[1] Federated Learning with Personalization Layers [AISTATS 2020]

[2] Exploiting Shared Representations for Personalized Federated Learning [ICML 2021]

We would like to thank the reviewer for your time and patience in reviewing our manuscript. Below, we reply to every concern in weaknesses and questions.

Reply to W1. In pFedMoAP, the proposed mixture of adaptive prompt is different from classic prompt learning. Specifically, for Vision-Language Models (VLMs), classic prompt learning (e.g. CoOp [1]) does not involve mixture of expert prompt. On the other hand, in pFedMoAP

• we proposed to treat the adaptive prompts as the experts
• the prompts, trained from different distributions, are integrated by the proposed attention-based gating network, largely increasing the chance of better alignment for the output text feature with the image feature from CLIP-like VLMs (e.g. Table 1, pFedMoAP vs. CoOp). While CoOp only trains a local text feature ($T_L$), in pFedMoAP, we have $T_L$, $T_{NL_1}$, $T_{NL_2}$, ..., $T_{NL_K}$ with the proposed attention-based gating network to optimize for the alignment.
• compared to traditional linear projection-based gating network, the novel attention-based gating network for CLIP-like VLMs is favorable due to 4 reasons detailed in Sec. 3.3, line 316-348. Briefly, attention-based gating network 1) is more robust to adaptive experts; 2) serves as linear probing with more capacity; 3) leverages CLIP's feature alignment with attention mechanism; and 4) is agnostic to experts' order. We conduct additional experiments to compare these two types of gating network to show the superiority of the proposed attention-based gating network (line 889-909).

Besides, we would also like to mention that ''access prompts from other clients'' itself is a breakthrough from traditional prompt-based federated learning (FL) where existing algorithms only allow the aggregated prompt, instead of the raw pre-aggregated prompt, to be transmitted. This traditional paradigm unnecessarily embeds the global knowledge into the compressed aggregated prompt with loss of various local information, while pFedMoAP fully leverages the lightweight nature of the prompt and transmits the global knowledge in its raw form as the non-local adaptive prompt experts.

Reply to W2. Our experiments follow most recent prompt-based FL for VLMs works' standard [2,3] and used all 9 datasets they have used. The table below lists the datasets we used and the partition method for each dataset. We regretfully do not have enough resources (time/GPU) to finish extra experiments on large-scale datasets within the discussion period and leave this as future work.

[1] Learning to Prompt for Vision-Language Models [IJCV 2022]

[2] Global and Local Prompts Cooperation via Optimal Transport for Federated Learning [CVPR 2024]

[3] Harmonizing Generalization and Personalization in Federated Prompt Learning [ICML 2024]

Please feel free to let us know if you have any more concerns, and we will try our best to address them. Thank you for your feedback.

For easy tracking of our major modifications (apart from small modifications such as typo fixing) in the updated pdf, we mark the sentences with major modifications in blue (one appearance, line 788-793). We also mark blue and highlight the titles of the added sections (Sec. C (line 854) and its subsections). Thank you for your time.