Probabilistic Federated Prompt-Tuning with Non-IID and Imbalanced Data
Authors: Pei-Yau Weng
Minh Hoang
Lam M. Nguyen
My T. Thai
Tsui-Wei Weng
Trong Nghia Hoang
Venue: NeurIPS 2024 poster (2024)
Keywords: Probabilistic Learning
Primary area: probabilistic_methods
TL;DR: We propose a new probabilistic prompt aggregation method for fine-tuning on decentralized and imbalance data settings.

Abstract: 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.

Probabilistic Federated Prompt-Tuning with Non-III and Imbalanced Data
Pei-Yau Weng
Minh Hoang
Lam M. Nguyen IBM Research lamnguyen wawmmgyyn.mltd@iiee m1ld@iimm..... com
Washington State University pei-yau. weng@wsu. edu
Princeton University
minhboaug@princeton..co edu
My T. Thai
Tsui-War Weng
Trong Nghia Hoang" Washington State University trongnghia. woongughi..hoong@@sss... edu
University of Florida mythai@cti seciiiccieeeee edu
University of California San Diego IwengQucsd. edu
Abstract
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 more effective prompt-tuning method, optimizing for small set of input prefixes to reprogram the pre-trained model's behavior. Our approach transforms federated learning into distributed set modeling task, aggregating diverse sets of prompts to globally fine-tune the pre-trained model. We benchmark various basclines based on direct adaptations of existing federated model aggregation techniques and introduce 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.
Introduction
The proliferation of personal devices has transformed our data landscape into numerous federated systems with distinct resource constraints, data representations, distributions, and ownership. This has motivated the development of new machine learning (MI.) paradigms that enable collaborative learning across different systems while respecting their data privacy. A prominent framework to substantiate this collaborative scheme is federated learning (FL), which allows multiple systems to train a common model without sharing their private data [1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12].
Existing FL methods often assume that learning begins from scratch and does not build on prior expertise. On the other hand, fine-tuning pre-trained or foundation models [13] an emerging paradigm for efficient generation of ML solutions. Almost all state-of state-of-the-art models in natural language processing (NLP) are now fine-tuned versions of foundation models, such as BERT [14], BART [15], RoBERTa [16], and T5 [17]. Likewise, many top performing vision models have also benefited from the generalization capability of foundation models, such as the vision transformer model 8), which was trained on large-scale and generic datasets such ImageNet [19]. Nonetheless, this sune-tuning practice has only recently been investigated in federated learning [20, 21J.
Current research in this direction demonstrated various benefits of integrating fine-tuning into FL. For example, fine-tuning can utilize information better in decentralized data scenarios and thus improves the performance in various FL scenarios, mostly in NLP [22, 23, 24, 25]. Interestingly, has been pointed out that even simple initialization of local clients with a foundation model can prevent solution drift to some extent in large-scale scenarios with heterogeneous local data distributions [20, 21, 23). This is also consistently observed in the results of our case studies in Fig. 1.
"Corresponding authors: Pei-Yau Weng. Minh Hoang, Trong Nghia Hoang.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

However, in FL environments that depend on frequent communication and synchronization of model updates across multiple devices, fine-tuning the entire pre-trained model is often infeasible due to limited local storage and communication bandwidth. Our study also reveals that full-model fine- tuning approaches fall short when the federated data are not only heterogeneous but also imbalanced. Our case studies in Fig. 1 in particular show that when the local data are both heterogeneous and imbalance, federated learning with full-model adaptation suffers a huge performance drop, highlighting its instability in extreme data settings.
To address the resource bottleneck, several works in the recent federated learning literature have turned to a new class of parameter-efficiee tuning method called prompt tuning. This method focuses on engineering cues, such as extra tokens which are appended to the input embedding architecture. Such tokens or prompts provide beneficial context to performing the computational task, similar to how hints can be provided to assist puzzle solving. Local sets of prompts can then be aggregated or personalized using existing federated learning strategies. For example, the prior works of [26] and [27] use FEDAVG [2] to aggregate the local prompts while [28] integrates prompt-tuning into personalized federated learning, which helps generate client-specific sets of prompts that are well-customized to the corresponding local data distribution.
transformer
Although federated prompt-tuning approaches eliminate the need to update and communicate hun- dreds of millions of network weights, there is a substantial gap between the performance of these approaches in highly heterogencous, data- imbalanced settings, and the upper- upper-bound performance of fine-tuning in the centralized data setting, as shown in Fig. This gap can be attributed to the fact that locally learned prompt sets are learned in arbitrary orders, which are generally unaligned across clients. That is, the same prompt position across different clients might encode different contextual information about the local data. A simple aggregation that disregards prompt alignment might attempt to combine prompts from different contexts and collapse into less informative prompts. This has not been explicitly modeled in prior federated prompt-tuning work.
Intercgennto.u Dettbbu DIrO
Imbalence Distribution
Prompt Turing
Full Modill Taring
Prampi Raning
F.B-Mtdel Tuning
Acturayy (S0
100
05
05.2396.20 29
94.0594.71
66 0177696.100 1n
100
95. 25..2.66.99
90.21
91.08
60
78:75
20
Centralized
FEDAAG
PLUPROX
Centralized
FEDAYG
FEDPHOX
Test Accuracy (%) achieved on the CIFAR- dataset by solving Eq. (3) via centralizing data, using FEDAVG, and using FEDPROX (orange) full-model (FM) and (blue) prompt-tuning (PT) setups. The evaluation is performed under (left) a standard (non-extreme) heterogeneous data partition; and (right) an extremely inbalanced data partitioning scheme (see Section 4).
Figure 1:
To address this issue, we adopt a hierarchical probabilistic modeling approach to characterize both the generation and alignment of local prompts in this paper. We cast the server aggregation step as an alignment of local prompt sets. This alignment discovers and aggregates local prompts encoding similar contextual information into summarizing prompts. On the client side, we view each local prompt set as an independent sample drawn from a generative model parameterized by the global summarizing prompts. In this view, each local prompt can be seen probabilistic exploration initialized by randomly selected summarizing prompt. The alignment between prompts can be characterized in terms of their association with the summarizing prompts, which can be inferred via learning the parameters of the above generative model. We summarize our main contributions below:
C1. We formulate the prompt summarizing procedure each local set is assumed to be an independent sample of a random point process. The alignment of similar prompts across differrnt local sets can be set as part of the modeling parameterization, whose optimization can be interleaved with the optimization of the point process's parameters (Section 3.2).
probabilistic set modeling task in which
C2. We develop an algorithm to find the most probable association between the local and summarizing prompts. This association is viewed as a latent variable in the above generative model. Specifically, we cast its inference as a classical weighted bipartite matching task via an interesting observation that the association between the summarizing prompt and each local set of prompts participate linearly to the loss function of our proposed generative model (Section 3.3).

---

C3. We compare the performance of our method against various federated crompt-tuning baselines based on existing heterogeneous federated learning techniques to demonstrate its effectiveness. Our reported results on variety of experiments and baselines demonstrate consistently that our method is most effective to combat data imbalance in extreme heterogeneous scenarios (Section 4).
2 Related Work
Federated learning (FL) [1, 2] is a collaborative framework that allows multiple parties to collaborate on training a common model without sharing their private data. In FL, the training data are distributed across m clients. The t-th client owns a local, private dataset D; {Xem 9tm m=1 comprising ne data points. The goal is to fit model w. using all private datasets without centralizing them. That is,
W. arg min L(w) A
T4e
(w())) where L((w)
&
1 ThE
e(xmm: Yems w)
(I)
m=1
and n n2 + n2 + ... Pmy denotes the total number of data points while {(Xim: Yomi W denotes some loss function of choice. Clients collaborate via sharing their local models instead of data. Each client can run multiple local updates before sharing its local model for aggregation. This helps reduce the number of communication rounds while still preserving the convergence guarantee. [2] names this the FEDAVG algorithm, which iterates between local optimization and global aggregation:
= U USe,
w w. )) Vt E [m], W.
f4
W
(2)
The local update routine -( .) is typically standard gradient updates such as SGD or Adam. At the beginning of an iteration, the local weight is set to be the global estimate from the previous communication round. In practice, local data distributions tend to diverge which consequently causes the local updates in Eq. (2) to have different convergence points across clients. This is known the solution drift phenomenon, which often decreases the performance of the federated model. Numerous approaches had been proposed mitigate this phenomenon, which includes the following directions:
Client Regularization. These methods focus on augmenting the local update strategies to prevent clients from drifting apart. For example. FEDPROX [29] introduces an fg-rgullarizatoon term to penalize updates that diverge the local model from the global model. FEDDYN [30] uses a FL-adapted version of [31] distributed optimization the regularizer. SCAFFOLD [32] attempts to correct the direction of local gradient updates at each client using their sum of gradients. [33] utilizes the similarity between model representations correct local training via contrastive learning.
Server Regularization. Another approach to prevent the drifting effect caused by heterogeneous local distributions is to replace the model average in Eq. (2) with a different mechanism of model aggregation. For example, 17, 6) decomposes client models into bags of neurons and performs a non-parametric clustering of neurons. Cluster centroids are used synthesize aggregated model for the next communication round. This approach bypasses the drifting phenomenon as the number of cluster is non-parametric and can be adjusted to accommodate new information patterns. In a similar vein, [34] adopts a probabilistic perspective of model aggregation, sampling higher-quality global models and combining them via Bayesian model ensemble, leading
more robust aggregation. [35] presents a data-free knowledge distillation methods for FL that help to train generators without compromising clients' data. [36] redistributes each client's shared model to others for b consecutive communication iterations, exposing it to various heterogeneous data sources before performing aggregation, curving their solution divergence.
Personalization. Instead of learning a single, universal model, [29, 37, 38, 39, 40] seek to Iearn personalized models for all clients. [3. 39] formulates this problem musti-task learning, whereas [37] and [38] respectively ground it in meta-learning and regularization frameworks. Several recent works impose shared data representation across clients [40, 41, 42, 43, 44], which is optimized using existing FL techniques, and locally tne-tune a small classification head to achieve personalization. However, personalized models trained in this manner tend to only perform well on its corresponding local test set, and not on a comprehensive test set combining all local test data (see Section 4).
Finally, most previous work assume that each FL client has sufficient data to adequately train its local model. This often contradicts real-world situations where data are scarce [45, 36, 46, 47, 48). While data scarcity is not new challenge of machine learning, it has not been thoroughly investigated in the
3

---

context of FL. [36] points out that with limited, scarce data, local models often have bad qualities, and aggregating such models tend result poor global performance. Although fine-tuning pre-trained large models is an increasingly popular technique to combat data shortage, most existing FL works (including [36] which aims to address data scarcity) have not tried to leverage this resource. This motivates us to investigate prompt-tuning as new FL paradigm.
3 Probabilistic Federated Prompt-Tuning
First, Section 3.1 provides a concise background on the standard prompt tuning technique in FL setting. Motivated by the result of our case study (see Fig. 1). Section 3.2 further introduces probabilistic federated prompt-tuning framework that aims to close up the performance gap between federated prompt-tuning and centralized full-model fine-tuning in data imbalance settings. Our framework characterizes each local model as random of prompts distributed by hierarchical generative model. An effective optimization algorithm to learn these is then detailed in Section 3.3.
An overall diagram featuring a bird view of our framework also provided in Fig. 2.
Gibbl Samwer
Global Server Probabllbttcc Prompe
Pooling Local homes Sero
Pronge Siledie
Locd Poosttall
Pronss
Cooaliatonn Hell
Ciien:1
Clent
Olert N
Clent
Local Phase
Global Phase
Figure 2: Workflow of Probabilistic Federated Prompt Aggregation: (left) each client selects a subset of prompts from the global set of summarizing prompts using the prompt-selection mechanism adapted from [49], and fine-tune them using local data; and (right) the server collects all local prompt sets and updates the global summarizing prompts that aggregate similar local prompts. This is achieved by our proposed probabilistic federated prompt aggregation (PFPT) algorithm.
3.1 Prompt-Tuning with Pre-Trained Model
We use pre-trained Vision Transformer [18] in all our experiments. The pre-trained model is composition Fz Fa F. where Fg is the classification head, Fa is the stack of attention blocks, and F. is the input embedding block. Keeping both F, and F. frozen, the local solution for each client t can be represented as Fo, F2 o (Fa Uw) which applies a personalized prediction head on the output of the frozen attention block Fa whose (set) inputs are the union of F. and set w of reprogramming prompts [49]. w is in turn union of & individual prompts w = wi W2 U... wk. Each local solution is characterized by a tuple of (Fe,,, w) comprising personalized head Fe,, and common prompt set w. The FL framework in Eq. (1) is now re-cast as
W. arr min
T2 {
m
L((w)
with Le(w) A max
( The Σ
Xtms Nomi Fe,es w) (3)
w)
Fee
The
where E(Xem Momi Pe,t, w) f(xim= ytmm Fe,, o Fa o (F. Uw)) which makes explicit the leverage of the existing (pre-trained) expertises (Fa, Pe). The optimal solution w. of Eq. (3) can be approximated using any of the existing federated learning algorithms, such FEDAVG [2] and FEDPROX [29].
There is an alternative approach to enabling light-weight federated fine-tuning. Instead of using prompts, approaches in this direction use (lcarnable) adapter network that adapts the output of an intermediate (frozen) neural network segment before passing to the rest of the (frozen) neural network. Local adapters across client then be aggregated using FEDAVG [21]. This direction is however orthogonal to federated prompt-tuning, which is our main focus here. Further investigation into probabilistic methods for adapter aggregation would be potential follow-up of our current work.
Remark.
3.2 Probabilistic Prompt Model
Intuitively, our federated prompt tuning framework is an analog to traditional FI. on the space of prompts. Prior to each communication round, we suppose that each local client f, via prompt-tuning

---

has obtained a set of n, prompts W, {wi (11::::2..... W2,............................. W.t....,,tm As the participating clients upload their local prompt sets, the server aggregates the combined set and subsequently returns a of n summary prompts, denoted as
{0:)--1- We will elaborate on this aggregation step in Section 3.2.2. Similar to standard FL., this aggregated set is also distributed to all clients at the beginning of the next communication round. Each client then selects the most relevant prompts from this set for further fine-tuning. This sampling step follows generative process described in Section 3.2.1.
3.2.1 Generative Model
At the beginning of each communication round, each client constructs a set of n, prompt initializations given the server-broaccast summary Φb We model this construction by random generative process. First, each local prompt initialization, Web, modeled as sample drawn from Gaussian with learnable neural parameterization,
N (we diag (a(bu:)))
Wek
where the diagonal covariance matrix are parameterized as the output a neural net a (with weight 7) on the mean parameter V The set of mean parameters w, & + is turn modeled as a random subset of the server-broadcast summarizing prompts Φ, following Beroulli point process prior with a finite mixture as the base measure:
BeP ((o -u( 0.))
(+bb)
(5)
Here, a denotes the sigmoid activation function, and & is another leep-parameterized function parameterized by weight w. By definition of the Bernoulli process, the sampled measure is given by:
Q+(b) = Σ C (v.=
(6)
where c E {0,1} 1 is the outcome of a Bernoulli trial with bias (g(0;; w)). that is:
(c,) = - (olo::)) (1-o (1-((((o::)) 1-ce
In layman term, this process simply summarizing prompt Φ, If this coin lands on head (c: = 1), the summarizing prompt Φ will be included in the set of mean parameters. Vice versa, if this coin lands tail, will skip Φ; Because of this, we remark that n, could be different from the number of prompts in the previous round.
will observe the result of coin toss & for every
3.2.2 Prompt Aggregation
Given the above generative story, prompts { which maximize the likelihood of observing local prompts {wr] The resulting optimal set of summarizing prompts then be used to re-program the pre-trained model, which correspond to the global tne-tuned solution. Assuming that each prompt captures a particular fine- tuning pattern concept, summarizing the prompt sets characterizing the local solutions will allow local models to be aggregated on the concept level, naturally mitigating the solution drift effect caused by compounding impact of data heterogeneity and pre-trained weight interference. This is substantiated in the overall computation process below.
describe algorithm to find the set of summarizing
Let 2/e E 10 denotes whether the local prompt Wis sampled from Gaussian centered at the i-th summarizing prompt Φ, That is, the assignment variable ztx = if and only if there exists k E [ne] such that whe Φ, Each client t can now be represented as (w,, z,) where z, [z])
zi (2)1, 2/2 ... 2,m,) ) and w, (W W 12. Wt..)
Now, using the generative process specified in the previous section and all its defining parameters
40 (w, y, ( the log likelihood of each client be derived as follow.
log P (w), (p) = log P (w) ze: sp) log IP (z, (4)
(8)
where each of the summand on the right-hand side is computed below. First, given (2+, 4),
Wi Ze. (
= II
N we whe diag a (prxe 7)) with 2he = (zta Φ1 +... + 2/1 Φn (9)
1=1
5

---

Algorithm/ Probabilistic Federated Prompt Tuning (PFPT)
input: pre-trained model F, of iterations, no. of sampled clients per iteration
output: optimized set of prompts +
1: initialize global summarizing prompts +
2: for
do 3: sample m clients with private datasets {D:)
4: for = to m do
5: W. f- client(F, F, Dr. 3) // local coampt-tuning-opimiiini - (w(w) in Eq. (3)
6: end for
7:
4. - server( (erver([w:)?)) // solving Eq. (11)
8:
Φ- + Φ, E Φ
Σ >x 2tk > 0 // remove inactive prompts
9:
end for
10:
return the set o of optimal prompts
Then, given and let Cri = 2/11 + 2/2 +... tamm
P zt (p) = II 2=1
(o(o:::)))
x Il (1-((o(o:))) ==1
(1-c1)
(10)
Eq. (9) follows from the Gaussian likelihood of local prompt in Eq. (4) with the additional identifica- tion form for u which computable if ze is given. Eq. (10) on the other hand is the consequence of the Bernoulli process in Eq. (6), which might not be immediately straight-forward. Its detailed
derivation will be given in Appendix B.
3.2.3 Overall Workflow
Using Eq. (9) and (10), we have the following prompt aggregation formulation,
+" = argmax
max W. W.
log
Wi Zt w, 7, ②
(11)
where + - ( z = [2 ) 1- and 7 is previously defined in Eq. (4). The probability term on the right-hand side can be further expanded into computable terms using Eq. (8), Eq. (9), and Eq. (10). Solving Eq. (11) is therefore the key step in our federated fine-tuning framework, whose overall flow is summarized in Alg. 1. Our algorithm proceeds in multiple iterations. At each iteration, a set of m clients is sampled. Each client selects the most relevant subset of summarizing prompts from the server using mechanism adapted from continual visual prompt-tuning [49]. The selected subset of summarizing prompts are then fine-tuned using local data (see lines 3-4). The sets of fine-tuned prompts are subsequently uploaded to the server, which performs the aggregation solving Eq. (11) (see line 6). The aggregation step has two subtlcties: (1) the total number of global prompts, n = n2 + ... + n,, is proportional to the total number of prompts across clients, making the entire algorithm solving Eq. (11), global prompts that were not assigned any client will be removed (see lines 7 and 8). Intuitively, if the sampled clients are similar, their prompt sets will substantially overlap, which reduces the total number of distinct prompts selected by all clients. Otherwise, the sampled clients are dissimilar, there will be less overlapping and the union set of selected global prompts will expand, increasing the complexity of the resulting model to match the increased data heterogeneity.
parametric; and (2) once the optimal assignment parameters z are found via
3.3 Optimization
Solving Eq. (11) above is however not trivial due to its mixed set of discrete/continuous variables. To sidestep this intractability, we instead formulate it as bi-level optimization that alternates between two sub-tasks: (1) optimizing given y and Φ; and (2) optimizing C, w, Y given z. Among which, solving sub-task (2) is straight-forward since only requires being able to differentiate expression in Eq. (9) and Eq. (10), which is trivial if we use the current estimation of the (discrete) assignment parameters z. On the other hand, solving sub-task (1) seemingly prohibitive expensive because it involves optimizing the (discrete) assignment parameters Fortunately, this can be

---

mitigated by casting both Eq. (9) and Eq. (10) into a linear form with respect to the assignment parameter z. Intuitively, might be confusing why such linear form can be achieved at all given that z is input to non-linear log probability function in Eq. (9) and Eq. (10). However, we must recognize that for each client I, there is at most one local prompt from which can be associated with the th summarizing prompt. This implies z è (2f1 242..... zim,)) ) either zero or one-hot vector. This is an important observation because linear function of such zero one-hot vector will remain linear regardless of any (non-linean post-processio post-prcesssing transformation.
shown Lemma 3.1 below.
Lemma 3.1. For any scalar function g(r) and binary vector C
: Sol such that [St. .2....,
G E {0, } and Ç has at most one non-zero component, we have
9
Σ S r.
Σ
Sa Sa- ((rs))))
(12)
i=1
with respect to any set (ri)e=1 of valid inputs to g(r).
The proof of Lemma 3.1 is straight-forward. First, there no non-zero component, both sides of Eq. (12) evaluate to g()). Otherwise, suppose the only non-zero component appears at position both sides of Eq. (12) will evaluate to g(rw). In both cases, Eq. (12) holds. Using Lemma 3.1, we can establish the desired linear forms for Eq. (10) and Eq. (11) which are its immediate consequences.
Lemma 3.2. Let IP(w, () defined as in Eq. (9). Let Li(z) = E log IP(w: Zz 9))
considering (Wis (p) as constants. We have
Li(z) ΣΣΣ
Zuu llgg (Weill,, diag (a (a())))))
(13)
which is linear in terms of the assignment purameter z. Lemma 3.3. Let P( (p) defined as in Eq. (10). Let L2(z) Ei log P(zz ()), considering 4p as constants. Then, we have
" m n1
m
L2(z) = Σ ΣΣΣ
log (1- (s(0;:)))
2kk log
(14)
(1-
i=1
which is linear in terms of the assignment parameter z.
The detailed proof for Lemma 3.2 and Lemma 3.3 are deferred to Appendix C. Using these results, we can put together the overall optimization task for z while fixing the rest of the parameterization
argmax Lz(z) + L2(z) }
log P (w. Zu 1
(0)
argmax
(15)
which is a weighted linear optimization task. The second equality follows from Eq. 8 and the results of Lemma 3.2 and Lemma 3.3. Now, we further choose optimize z, iteratively while fixing in addition z-c. Eq. (15) reduces to a weighted bipartite matching task, which can be solved effectively in ((max²] n;) n,)*) processing time using the Hungarian algorithm [50]. detailed pseudo-code implementing the above scheme is deferred to Appendix D.
4 Empirical Results
This section presents our empirical studies on a variety of computer vision datasets, including CIFAR- 10 and CIFAR-100 [5]], TinyImageNet [52] and synthetic, diverse dataset created by pooling together the MNIST-M [53]. Fashion-MNIST [54], CINIC-10 [55] and MMAFEDB (available on Kaggle) datasets, which is referred as the 4-dataset.
Our experiments are conducted on two data settings: (a) Dirichlet- based heterogeneous partition following previous setup in [6]; and (b) manual imbalanced data partition, as detailed below.
Heterogeneous Partition. We partition the training data into m subsets (for m clients). Each client has observations of all classes but the distributions of classes across clients are different. We simulate this using Dirichlet(a- 1,) distribution over an s-dimensional simplex where is the number of classes and is the concentration parameter.
7

---

Table 1: Accuracy (%) achieved the CIFAR-10 dataset by PFPT and other baselines.
FEDAVG-PI
PT
SCAFFOLD-PT
FEDOPT-PT
PFEDPG
GMM-PT
PHPT (ours)
0.5
94,06 + 0.45
94.26 0.34
92.26
90.70 0.54
58.00 0.22
92.63
94.39 0.51
a= 0.1
93.05 + 0.03
93.05 = 0.23
91.82 0.31
88,20 + 0.39
17.25 = 0.23
92,68 : 0.28
93.39 +
imbalance
91.21 + 0.16
91.08 + 0.25
87.06 1.27
79.32 + 1.73
43.16 0.27
90.50 + 0.17
91.45 + 0.08
Table 2: Accuracy (%) achieved on the CIFAR-100 dataset by PFPT and other baselines.
FEDAVG-PI
SCAFFOLD-PT FEDPROX-PT
FEDOPT-PT
PFEDPG-PT
GMM-PT
PFPT (ours)
0.5
0.27
79.21 = 0.31
0.18
43.47 0.39
77.50 0.12
80.24 + 0.224
a= 0.1
71.50 0.46
73.60 = 0.38
09.94
00.31 = 0.39
29.16 0.80
70.19 + 0.46
75.08 0.51
imbalance
70.07 + 0.35
70.67 + 0.25
69.35 + 0.96
58.86 + 0.77
11.57 = 0.25
66.95 + 0.46
72.05 0.93
Table 3: Accuracy (%) achieved on the TinylmageNer dataset by PFPT and other baselines.
FEDAVG-PT
SCAFFOD-PPT FEDPROX-PT
PEDOPr-PT
PFEDPG
GMM-PT
PFPT (ours)
0.5
86,38 0.18
0.33
01255:0.22
40.220.22
81.73 0.20
86.91 0.11
a= 0.1
78.58 0.57
79.19 0.28
78.02 + 0.33
63.81 + 0.79
3433=00
76,90 + 0.28
82.31 0.26
imbslance
75.21 4: 0.73
75.49 1 0.38
76.88 0.45
64.33 1: 0.61
11.92 = 0.13
73.85 4 0.36
78.21 1: 1.25
Our experiments are conducted with a 0.1 and a 0.5. For each client r, we drawn sample Pr ~ -Dirichet(a: 1.) where Pr. E (0, specifies the percentage of examples in class e to be assigned to client r. For the CIFAR-10, CIFAR-100 100,, and TinyImageNet datasets, we set m 100. For the 4-dataset, we simulate 20 partitions for each of the 4 sub-dataset using Dirichlet(a- - 1,) with s 10, 10, 10, and for MNIST-M, to total of 80 clients with diverse and heterogeneous data distribution.
ASSHION-MNISS CINIC-10, 10, and MMAFEDB. This amounts
Imbalance Partition. For the CIFAR-10 (10-class), CIFAR-100 (100-class) and Tinylm (200-class) datasets, the training data of each client is set to be dominated by a particular subset of classes, which amounts to 10% of the total number of classes. 99% of the local dataset of each client is set to belong to a certain 10% of the total number of classes. For our synthetic 4-dataset, we partition each of the sub-dataset into 20 subsets. Each subset is set so that 99% of its data points belong to single class. The remaining 1% data of each class are then pooled together and evenly distributed among all clients. This amounts to 80 clients with extremely imbalance local datasets.
The above schemes are applied only on the train partition of each dataset. We use the default train/test partition for CIFAR-10, CIFAR-100, and TinylmageNet. For the 4-dataset, we sample 30K data points from the default train partition of each of its 4 sub-datasets. We also sample 2.5K data points from the default test partition of each sub-dataset. This amounts synthetic dataset with 120K data points in the train partition and 10K data points in the test partition.
For each data setup on each dataset, we compare the performance of our probabilistic federated prompt-tuning (PFPT) algorithm with those of a representative set of state-of-the federated learning algorithms adapted to the prompt-tuming setting in Eq. (3), which include FEDAVG [2], FEDPROX [29], SCAFFOLD [32], FEDOPT [56] and PFEDPG [28]. PFEDPG is a personalization method that is originally measured based on how well its personalized models perform on their corresponding local test sets. In this context, the performance target is however set on a global test set we use PFEDPG's common model for evaluation. We will refer these adapted baselines as FEDAVG-PT PT, FEDPROX-PP PT, SCAFFOLD- compare PFPT against simple prompt clustering baseline, which replaces the prompt averaging of FEDAVG by Gaussian Mixture Model (GMM) clustering in which the cluster centroids are returned as the aggregated prompts. We refer to this baseline GMM-PT. All results are averaged over 5 independent runs and reported in Tables 1 to 4 below. The result of each run is evaluated on a global test set, which comprises the entire test partition.
FEDOPT-PT, and PFRDPG. In addition, we also
4.1 NoD-IID Data with Locally Skewed Class Distributions
Table reports the performance of PFPT and various federated pomppttunnng baselines on the CIFAR-10 10 dataset. In all three settings (i.e.. heterogeneous partitions with a 0.5 and 0.1, as well as the imbalance partition), PFPT consistently achieves the best performance. On average, the classification accuracy of our method improves by 0.3% over the closest competitors, FEDAVG-PT and FEDPROX-PT. While this gain seems modest, we note that the prompt tuning performance tends to fall off more significantly on several FL frameworks that were devised to counteract the effect of imbalanced data, such as SCAFFOLD-PP PT and PT.. In fact, the accuracy our method
8

---

Table 4: Accuracy (%) achieved on the synthetic 4-dataset, which combines data from MNIST-M [53], Pashion-MNIST [54], CINIC-I0 [55], and MMAFEDB by PFPT and other baselines.
FEDPROX-PT FLDAVG-PT
SCAFFOL--PT
FEDOPT-PT
PFEDPG
GMM-PT
PEPT (ours)
0.5
59.48 0.59
59.50 0.70
11.5
40.69
30.54
57.05 1.04
76.89 0.191
0.1
48.74 0.03
48.71 = 2.13
16.65 3.17
43.51 = 0.78
25.40 0.17
16.82 + 2.80
70.29 0.332
imbalance
39.28 + 0.55
39.65 + 1.89
23.79 1.24
44.43 + 1.78
22.54 0.32
41.37 + 1.12
62.23 1.02
05
0-011
Inbolance
Centralzed Loarning
200
10
10
08
04
02
62
62
02
Figure t-SNE plots of the (learned) summarizing prompts of PFPT on CIFAR- over 120 communication iterations with different heterogeneity settings. Yellow triangles denote the centroids of the t-SNE embeddings of the prompts. The dashed red line visualizes their trajectories.
0.5
a 0.1
Imbalance
40
Pool Size
30
10
0
40
80
120
0
40
80
120
0
40
120
Communication Rounds
Communication Rounds
Communication Rounds
Figure 4: Variations in CIFAR- global prompt pool size across 120 communication rounds under different heterogeneity settings.
respectively improves by 2.7% and 7.0% over these baselines. This result clearly suggests that prompt tuning performance is more vulnerable to the heterogeneity challenge in FL, and thus requires more specialized treatment. Finally, we observe that PFEDPG achieves the worst performance among all baselines (23.6% worse than PFPT average). This performance gap is especially pronounced on the imbalance setting. While PFEDPG claimed to deal with both heterogeneity and prompt-tuning this result is not surprising since it assumes that the local test sets partitioned similarly to the local training sets, whereas our results are obtained on a holistic global test set.
We further repeat this experiment on the CIFAR-100 and TinylmageNet datasets, which are both more challenging than CIFAR-10 10 due to the increased number of classes (i.e., 100 and 200 classes respectively). The results of these experiments are reported in Table 2 and 3. As expected, our method still consistently outperforms other baselines. However, the performance gap between PFPT and other methods tends to widen proportionately to the increased difficulty of the dataset. For example, in the CIFAR- 100 experiment, our method 1.2% better than the second best baseline (FEDAVG-PT) and 47.6% better than the worst baseline (PFEDPG).
In the TinylmageNet experiment, our method is 2.4% better than FEDAVG-PT and 50.1% better than PFEDPG. These results are therefore consistent with our findings in Table 1. Finally, Table 4 records the prompt tuning performance of the same baselines on the most challenging scenario, which combines 4 datasets from completely different vision domains. We observe a significant performance gap between PFPT and baselines that have been keeping up with its performance in earlier experiments, such as FEDAVG-PT (now 20.6% worse) and FEDPROX-PT (now 20.5% worse). More interestingly. GMM-PT, which has performed on par with our method in previous experiments, observes significant drop in performance (now 21.4% worse). These extremely wide performance gaps are due to the fact that local prompts are more diverse (to account for highly heterogencous tasks), which require more refined alignment technique to facilitate effective aggregation.
4.2 NID Data with Globally Skewed Class Distributions
In addition to the above setting which features non-III data and locally skewed class distribution, we also evaluated the performance of our framework in a more adverse setting where the class

---

distribution will remain skewed even if all local datasets are aggregated. That is, local datasets in this setting are non-IID draws from long-tailed data distribution parameterized by an imbalance factor IF = (max, nte) mee//min, Re)) characterizing the ration between sample sizes (n₂) of the most frequent and least frequent classes [57, 58]. Such (global) long-tailed datasets are synthetically created from the original CIFAR- -1FAR-100 and ImageNet datasets following the data simulation in [59]. Following previous experiment setup in previous (non- VIT and no prompt-tuning) baseline methods FEDIC [57] and CReFF [58] 2, the resulting long-tailed datasets SCIFAR-100-LE and ImageNet-LT) are configured to have IF = 10, 50, 100 for CIFAR- 100 and IF = 1280/5 5 for ImageNet. The reported results in Table 5 show that method is also highly effective in this setting, achieving significantly higher accuracy than the current SOTA [57, 58] across all datasets and imbalance setups.
Table 5: Accuracy (%) achieved on long-tailed datasets by PFPT, CREFF [58] and FEDIC [57].
CIFAR-100-LT
Non-IID (a- = 0.1)
ImageNet-LI LT (IF = 1280/5)
IF
100
IF
50
IF
FEDIC
33.67
34.74
41.93
28.93
CREFF
26.19
28.32
35.49
26.31
PFPT (ours)
60.74
65.54
71.66
75.54
4.3
Prompt Convergence and Diversity
To generate insights regarding the learned prompts, we plot the 2-dimensional I-SNE embeddings of all CIFAR-I00 100 global prompts discovered by our method across 120 communication rounds. This is repeated for the two heterogeneity and one data imbalance FL settings as well as the centralized data setting (Fig. 3). The yellow triangles in each plot mark the corresponding centroids of the prompt embeddings (one per 10 communication iterations). We also plot best-fit spline curve (dashed and colored in red) through these centroids to visualize their update trajectory. In all settings, the distance between successive centroids consistently gets smaller as training progresses, which suggests that the learned prompts generally converge well. This is also corroborated by Fig. 4, which shows that the number of global prompts also converges over 120 communication iterations. It also observed that in the centralized learning scenario in which no heterogeneity is present, the convergence happens much faster. The prompts quickly converge within 10 communication iterations. We also observe a gradual increase in the spread of the prompts with respect to their centroid. This suggests that the learned prompts are optimized by our method to capture the data diversity across participating clients. More interestingly, we move from the standard heterogeneity settings (with 0.1 and 0.5) to the extremely heterogeneous setting with imbalance (local) data, the magnitude of this spread becomes larger, confirming the strong correlation between task heterogeneity and prompt diversity.
Conclusion
Our paper presents new and effective approach to address the challenges of prompt-tuning pre- trained models in federated data scenarios with diverse local data distributions. Our proposed approach is a hierarchical probabilistic framework that models server aggregation as non-parametric alignment of locally learned prompt sets into summarizing prompts. In each subsequent communication round, local clients sample and perform local updates relevant prompts from the previous set of summarizing prompts. In this manner, our approach bypasses the effect of solution drift and outperforms existing federated learning techniques (applied on prompt-tuning) vision datasets. The reported results emphasize the cost-efficiency and efficacy of our approach in combating data heterogeneity within extremely diverse federated scenarios.
various computer
Aednoweesgmments and Disclosure of Funding. This work used GPU compute resource at SDSC through allocation CIS230391 from the Advanced Cyberinfrastructure Coordination Ecosystem: Ser- vices and Support (ACCESS) program [60], which supported by U.S. National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296. My T. Thai acknowledges the support of National Seience Foundation grants SCH-2123809 and III24466006. T.-W. Weng is supported by National Science Foundation awards CCF-2107189, IIS 23S-231311 IIS IIS-2430539, the Hellman Fellowship, and Intel Rising Star Faculty Award.
"These federated mhe-tuning baselines use full fine- fine-tuning with ResNet as their backbone. Both baselines were not evaluated in main setting in Section 4.1.
10

---

References
[I] Jakub Konečny, H. Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated Opti- mization: Distributed Machine Learning for On-Device Intelligence. CoRR, abs/1610.02527, 2016.
[2] H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Effiiii Learning of Deep Networks from Decentralized Data. In Proc. AISTATS, pages 1273-1282, 2017.
[3] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated Multi- Task Learning. Advances Neural Information Processing Systems, 30, 2017.
[4] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D' D'Oliveira, Salim EI Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaïd Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and Open Problems in Federated Learning. CoRR, abs/1912.04477,, 2019.
[5] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic Federated Learning. CoRR, abs/1902.00146,, 2019.
[6] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian Nonparametric Federated Learning of Neural Networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 7252-7261. PMLR, 09-15 Jun 2019.
[7] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, and Nghia Hoang. Statistical Model Aggregation via Parameter Matching. In H. Wallach, H. Larochelle, A. Beygelzimer, d' d'Alché-Bue, E. Fox, and R. Garnet,, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[8] Tengfei Ma, Trong Nghia Hoang, and Jie Chen. Federated learning of models pre-trained on different features with consensus graphs. In Uncertainty in Artificial Intelligence, 2023.
[9] Ziwei Fan, Hao Ding, Anoop Deoras, and Trong Nghia Hoang. Personalized federated domain adaptation for item-to-item recommendation. In Uncertainty Artificial Intelligence, 2023.
[10] T. N. Hoang, Q. M. Hoang, K. H. Low, and How. Collective online learning of Gaussian processes in massive multi-agent systems. In Proc. AAAI, 2019.
[11] Minh Hoang and Trong Nghia Hoang. Few-shot learning via repurposing ensemble of black-box models. In AAAI Conference on Artificial Intelligence, 2024.
[12] Trong Nghia Hoang. Effective knowledge representation and utilization for sustainable collabo- rative learning across heterogeneous systems. Al Magazine, 45(3):004440, 2024.
[13] Rishi Bommasani, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, S. Buch, Dallas Card, Rodrigo Castellon, Niladri S. Chatterji, Annie Chen, Kathleen A. Creel, Jared Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah D. Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Gcoff Keeling,
11

---

Fereshte Khani, O. Khattab, Pang Wei Koh, Mark S. Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovee, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning, Suvir P. Mirchandani, Eric Mitchell, Zanele Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Benjamin Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, J. F. Nyarko, Giray Ogut, Laurel OTT, Isabel Papadimitriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Robert Reich, Hongyu Ren, Frieda Rong. Yusuf H. Roohani, Camilo Ruiz, Jack Ryan, Christopher R e' Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishna Parasuram Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose E. Wang. William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie. Michihiro Yasunaga, Jiaxuan You, Matei A. Zaharia, Michael Zhang, Tianyi Zhang. Xikun Zhang, Yuhui Zhang. Lucia Zheng, Kaitlyn Zhou, and Percy Liang. On the Opportunities and Risks of Foundation Models. ArXiv, 2021.
[14] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACL, pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[15] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, AbdeIrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising Sequence-to-Seqqueec Pre-training for Natural Language Generation, Translation, and Comprehension. CoRR, 2019.
[16] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: Robustly Optimized BERT Pretraining Approach. CoRR, abs/1907. 16992, 2019.
[17] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the Limits of Transfer Learning with Unified Text-to-Text Transformer. JMLR, 21(1):5485-555 5551, 2020.
[18) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv preprint arXiv:2010. 11999, 2020.
[19] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A Large-Scale Hicrarchical Image Database. In CVPR, pages 248 255. IEEE, 2009.
[20] John Nguyen, Jianyu Wang, Kshitiz Malik, Maziar Sanjabi, and Michael Rabbat. Where to Begin? On the Impact of Pre- Tre-Trainng and Initialization in Federated Learning, 2023.
[21] Wang Lu, Xixu Hu, Jindong Wang, and Xing Xie. FedCLIP: Fast Generalization and Personal- ization for CLIP in Federated Learning, 2023.
[22] Yuanyishu Tian, Yao Wan, Lingjuan Lyu, Dezhong Yao, Hai Jin, and Lichao Sun. FedBERT: When Federated Learning Meets Pre- Training. ACM Transaction of Intelligent Systems and Technology, 13(4), aug 2022.
[23] Tuo Zhang, Tiantian Feng, Samiul Alam, Dimitrios Dimitriadis, Mi Zhang, Shrikanth S. Narayanan, and Salman Avestimehr. GPT-FL: Generative Pre-trained Model-Assisted Federated Learning, 2023.
[24] Yeachan Kim, Junho Kim, Wing- Lam Mok, Jun oun-Hyung Park, and SangKeun Lee. Client- Customized Adaptation for Parameter-Efiiee Federated Learning. In Anna Rogers, Jordan Boyd-Graber and Naoaki Okazaki, editors, Findings of the Association for Computational Linguistics: ACI. 2023, pages 1159-1172, Toronto, Canada, July 2023. Association for Compu- tational Linguistics.
[25] Ankur Agarwal, Mehdi Rezagholizadeh, and prasanna Parthasarathi. Practical Takes on Federated Learning with Pretrained Language Models. In Andreas Vlachos and Isabelle Augenstein, editors, EACI. pages 454471, Dubrovnik, Croatia, May 2023. Association for Computational Linguistics.
12

---

[26] Tao Guo, Song Guo, Junxiao Wang, Xueyang Tang, and Wenchao Xu. PromptFL: Let Federated Participants Cooperatively Learn Prompts Instead of Federst-d Learning in Age of Foundation Model. IEEF Transactions on Mobile Computing, 2023.
[27] Haodong Zhao, Wei Du, Fangqi Li, Peixuan Li, and Gongshen Liu. FedPrompt: Communication- Efficient and Privacy Preserving Prompt Tuning in Federated Learning. In Proc. ICASSP, 2023,
[28] Fu-En Yang, Chien-Ye Yi Wang, and Yu-Chiang Frank Wang. Efficient Model Personalization in Federated Learning via Client-Specific Prompt Generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 19159-19166 October 2023.
[29] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and Robust Federated Learning Through Personalization. In Marina Meila and Tong Zhang, editors, International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6357-6368. PMLR, 18-24 Jul 2021.
[30] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated Learning based on Dynamic Regularization. In International Conference on Leurning Representations, 2021.
[31] Ohad Shamir, Nati Srebro, and Tong Zhang. Communication--ffiiiee Distributed Optimization using an Approximate Newton-type Method. In International Conference on Machine Learning, pages 1000-1008. PMLR, 2014.
[32] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5132-5143, PMLR, 13-18 Jul 2020.
[33] Qinbin Li, Bingsheng He, and Dawn Song. Coneel-oottrattive Federated Learning. In CVPR, 2021.
[34] Hong- You Chen and Wei-Lun Chao. FedDistill: Making Bayesian Model Ensemble Applicable to Federated Learning. In ICLR, 2021.
[35] Lin Zhang, Li Shen, Liang Ding, Dacheng Tao, and Ling- Yu Duan. Fine-Tuning Global Model via Data-Free Knowledge Distillation for Non-IID Federated Learning. In CVPR, pages 10174-10183, June 2022.
[36] Michael Kamp, Jonas Fischer, and Jilles Vreeken. Federated Learning from Small Datasets. arXiv preprini arXiv:2 10.00469, 2021.
[37] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3557-3568. Curran Associates, Inc., 2020.
[38] Canh T. Dinh, Nguyen Hoang Tran, and Tuan Dung Nguyen. Personalized Federated Learning with Moreau Envelopes. In NeurlPS, 2020.
[39] Qiying Yu, Yang Liu, Yimu Wang, Ke Xu, and Jingjing Liu. Multimodal Federated Learning via Contrastive Representation Ensemble. In ICLR, 2023.
[40] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting Shared Representations for Personalized Federated Learning. In Marina Meila and Tong Zhang, editors, ICML, volume 139 of P Proceedings of Machine Learning Research, pages 2089-2099. PMLR, 18-24 Jul 2021.
[41] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive Personalized Federated Learning. CoRR, abs/2003.1344 2020.
[42] Paul Pu Liang, Terrance Liu, Ziyin Liu, Ruslan Salakhutdinov, and Louis- Morency. Think Locally, Act Globally: Federated Learning with Local and Global Representations. CoRR, abs/2001.01523, 2020.
13

---

[43] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated Learning with Personalization Layers. CoRR. abs/1912.008 2019.
[44] Filip Hanzely and Peter Richtárik. Federated Learning of Mixture of Global and Local Models. CoRR, abs/2002.055 06561,, 2020.
[45] Dianwen Ng. Xiang Lan, Melissa Min-Szu Yao, Wing P Chan, and Mengling Feng. Federated Learning: A Collaborative Effort to Achieve Better Medical Imaging Models for Individual Sites that Have Small Labelled Datasets. Quantitative Imaging in Medicine and Surgery, 11(2):852, 2021.
[46] Ying Zhao, Junjun Chen, Di Wu, Jian Teng, and Shui Yu. Multi-Task Network Anomaly Detection using Federated Learning. In Proceedings of the 10th International Symposium on Information and Communication Technology, pages 273-279 2019.
[47] Yu Zhang, Guoming Tang. Qianyi Huang, Yi Wang, Kui Wu, Keping Yu, and Xun Shao. FedNILM: Applying Federated Learning to NILM Applications at the Edge. IEEE Transactions on Green Communications and Networking, 2022.
[48] Ying Zhao, Junjun Chen, Qianling Guo, Jian Teng, and Di Wu. Network Anomaly Detection using Federated Learning and Transfer Learning. In Security and Privacy in Digital Economy: First International Conference, SPDE 2020, Quzhou, China, October 30-November 2020, Proceedings, pages 119-231. Springer, 2020.
[49] Zifeng Wang, Zizhao Zhang, Chen- Yu Lee, Han Zhang, Ruoxi Sun, Xiaoqi Ren, Guolong Su, Vincent Perot, Jennifer Dy, and Tomas Pfister. Learning to Prompt for Continual Learning. In CVPR, pages 139-149, 2022.
[50] Harold W Kuhn. The Hungarian Method for the Assignment Problem. Naval Research Logistics Quarterly, 2(1-2):83 97, 1955.
[51] Alex Krizhevsky, Geoffrey Hinton, et al. Learning Multiple Layers of Features from Tiny Images, 2009.
[52] Ya Le and Xuan Yang. Tiny Imagenet Visual Recognition Challenge. CS CS 23/N, 7(7):3, 2015.
[53] Seunghun Lee, Sunghyun Cho, and Sunghoo Im. DRANet: Disentangling Representation and Adaptation Networks for Unsupervised Cross-Domain Adaptation. In CVPR, pages 15252- 15261, June 2021.
[54] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: A Novel Image Dataset for Benchmarking Machine Learning Algorithms. CoRR, abs/1 708.0774777 2017.
[55] Luke N. Darlow, Elliot J. Crowley, Antreas Antoniou, and Amos J. Storkey. CINIC-10 is not ImageNet or CIFAR-1O, 2018.
[56) Muhammad Asad, Ahmed Moustafa, and Takayuki Ito. FedOpt: Towards Communication Efficiency and Privacy Preservation in Federated Learning. Applied Sciences, 10(8):2864, 2020.
[57] Xinyi Shang, Yang Lu, Yiu-Ming Cheung, and Hanzi Wang. FEDIC: Federated Learning on Non-IID and Long-Tailes Tailed Data via Calibrated Distillation In 2022 IEEE International Conference on Multimedia and Expo (ICME), pages 1-6, Los Alamitos, CA, USA, July 2022.
IEEE Computer Society.
[58] Xinyi Shang, Yang Lu, Gang Huang. and Hanzi Wang. Federated learning on heteroge- neous and long-tailed data via classifier re-training with federated features. arXiv preprint arXiv:2204.1 13399, 2022.
[59] Kaidi Cao. Colin Wei, Adrien Gaidon, Nikos Aréchiga, and Tengyu Ma. Learning imbalanced datasets with label-distriuutin--aaare margin loss. In Neural Information Processing Systems, 2019.
[60] Timothy J. Boerner, Stephen Deems, Thomas R. Furlani, Shelley L. Knuth, and John Towns. Access: Advancing innovation: Nsf's advanced cyberinfrastructure coordination ecosystem: Services & support. In Practice and Experience in Advunced Research Computing 2023: Computing for the Common Good, PEARC '23, page 173-176, New York, NY, USA, 2023. Association for Computing Machinery.
14

---

Code Release. Our experimental code is released and maintained at http::///gggtuuu..cceee/aaaaaaee
A Broader Statement of Impact
This research focuses on developing an effective prompt-tuning and summarizing algorithm for feder- ated learning scenarios featuring set of clients with diversely skewed local data. The mathematical approaches and insights developed in this paper will help bridge the gap between large, pre-trained models and their applications in private data settings with extremely skewed data distribution. While applications of our work to real data could result in ethical considerations, this is an indirect (and unpredictable) side-effect of our work. Our experimental work uses publicly available datasets to evaluate the performance of our algorithms; no ethical considerations are raised.
B Derivation of Eq. (10)
Note that 2, {z which is set of independent random vector given sp Hence,
n
n2
IP (z: (p) = T =
(ztk (p)
(16)
1= ==k=1
where the last equality follows from the fact that z = (=/1 = =f2...., ,mm)) which is turn a set of independent Bernoulli variables z E 1) indicating whether the local prompt Wis was sampled from Gaussian centered the i-th summarizing prompt di. Furthermore, following the
Bernoulli process definition in Eq. (6), the i-th summarizing prompt Φ; is selected independently with probability (g(ó;; w)) to sample the k-th local prompt wer via setting th., de in Eq. (4). Hence,
PP (2tx
(p)
(o(0)
(1-
(17)
Plugging Eq. (17) into Eq. (16) and using Cri 2/1 + =/2 + ... + zim, result in Eq. (10).
C Proof of Lemmas 3.2 and 3.3
Lemma 3.2. To derive the result of Lemma 3.2, note that Eq. (9) implies
log P (w;) Zz, sp) =
log N
Wik
Σ
z/R
Φ
diag Σ at 23
A=
1=1
Ml
zix Φ2
(18)
i=1
where we define
log N Wik r, diag (1:7))
g (r)
(19)
In addition, since =/e = with 2/p E 10, 1}. Lemma 3.1 implies
Σ
Ekke Φ
= Σ
z ge ge()))
(20)
i=1
which can be plugged into Eq. (18) arrive at
n
log P (w;)
Zr, sp)
gk
2 $
=
2
(0,)
(21)
k=
2/s log N
Wekk
Φ diag (0::7)
(22)
15

---

where the last equality follows from the definition of in Eq. 20. Finally, taking summation over
t 1, ..,,...,7 on both sides of Eq. (22) and using the definition of 21, we arrive Lemma 3.2.
Lemma 3.3. To derive the result of Lemma 3.3, we start with taking log on both sides of Eq. (10), which results in
log P (ze) (p)
Σ cu log o
(g(0; w))
Σ
(-cu) log (1 o (1-((p((:)))
(23)
o ( ( w)
Σ
Cti log
+
log (1- (1-00(0;;))) Σ ==1
(24)
1 (0;; w)
1=1
( (0,5 w)
R
n +
Σ log 1 (1-o(9(0:w)) (25)
Atk
log
1 ($:: w))
where the last equality is due to the fact that #/1 + #12 +............... "tm." Hence, taking the summation overt = 1,2, m on both sides of Eq. (25) and using the definition of L2, we arrive at Lemma 3.3.
D Optimizing Eq. (15) via Iterative Weighted Bipartite Matching
For any given index 1, Eq. (15) can be rewritten
F (1; (1: Σ
n
(og
Wi*
Φ., (inig((((;;:))) + log
(0, w)
1 - (0,; w)
i=1
12
nl
+ G(z-1) = ΣΣ
z
+ G(z-c)
(26)
i=1k=
where G(z-r) denotes the summation over the remaining terms in the definition of Ln and L2 that do not depend on 22, and
(s(())))
Ci A
log Wex Φi diag (a (a(o;:))) + log 11-0g((0;10
(27)
Thus, treating all but 2, as constants, we can solve for the optimal value of z via solving the corresponding weighted bipartite matching. As we iterate over optimizing for z, one time, this results in an iterative weighted bipartite matching approach, which is described below:
Algorithm 2 Iterative Weighted Bipartite Matching
input: generative parameters 7, w, w and Φ output: optimized set of values for initialize z randomly
2: for to do
3: compute the cost matrix Cj using Eq. (27) and fixing z-r as constant
4: Z: <- argmax,, F(z; 2-)) // solve the weighted bipartite matching algorithm in Eq. (27)
5: end for
6: return the optimal set of alignment
E Additional Ablation Studies
Fig. 5 shows the learning trajectories of the summarizing prompts over 120 communication iterations of our method PFPT and GMM-PT, which the vanilla baseline for modeling prompt alignment via clustering. Overall, the visual plots show that the learning trajectories of GMM-PT appear less
16

---

0.1
Imbalance
120
100
0.8
0.8
0.8
Iterations
0.4
4-
0.4
0.4
40
20
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
1.0
8
0.1
Imbalance
120
1.0
1.0
1.0
100
0.8
0.8
0.8
-00
Iterations
0.6
60
vig
0.4
0.4
0.4
40
0.2
0.2
0.2
20
0.0
0.0
0.2
0.4
0.8
1.0
0.0
0.6
0.8
0.2
0.4
0.6
1.0
Figure 5: t-SNE plots of the (learned) summarizing prompts of (top) our method PFPT and (bottom) GMM-PT over 120 communication iterations with different heterogeneity settings. Yellow triangles denote the centroids of the t-SNE embeddings of the prompts. The dashed red lines visualize the trajectories of the centroids. The figure best viewed with color.
diverse than those of PFPT, hinting that the better performance of PFPT can be attributed to its more diverse exploration over the prompt space, allowing it to contextualize information better into its prompt set. This appears consistent with the reported performance gap between GMM-PT and PFPT across a variety of datasets (see Table 1-Table the main text).
Additional Experiment Results
In addition to the results reported in the main text, we have also conducted extra experiments comparing the performance of PFPT (ours) with a more challenging set basclines. These are the improved variants of the original baselines now improved with a client-specific prompt selection mechanism [49], which enables each client to selectively contextualize its local data into subset of most relevant prompts. This prevents different information contexts from being collapsed into the same prompt. This modification improves the performance of the baselines substantially the diverse 4-dataset (Table 6) while preserving their performance on the Tiny ImageNet dataset (Table 7). However, it is observed that even in this more challenging setting, PFPT still outperforms all baselines both datasets substantially. featuring an accuracy improvement up to 4.5 4.52%. Similar to our observations in the main text, the improvement most pronouncing when client data becomes more diverse, which the case with the 4-dataset.
used in the main text experiment. Each baseline is
Remark. We also consider this improvement to be a part this work's contribution (although is not the main contribution) since the incorporation of the prompt selection mechanism in [49] has also never been investigated the existing literature of federated fine- tne-tunnng. The comparison results in this section also helps enrich ablation studies of the key component of PFPT.
also run experiments to compare the performance of our PFPT algorithm against another set of federated fine-tuning baselines that incorporate adapter-tuning approach into the FedAvg and FedProx backbones. The results are reported Table 8, which shows that performance achieved by incorporating the adapter approach to fine-tuning in FedAvg and FedProx is only comparable
17

---

Table 6: Accuracy (%) achieved synthetic 4-dataset by our proposed PFPT algorithm and other baselines improved with a similar client-specific prompt selection mechanism. The symbol + indicates the improved variant of the original baseline.
FEDAVG-PT+
FEDPROX-PTS
FEDOPT-PT+ GMM-PT+
PFPT
0.5
75.92
76.85
57.37
74.03
76.89
a = 0.1
66.21
67.77
61.97
66.47
70.29
imbalance 61.00
61.75
57.15
61.66
62.23
Table Accuracy (%) achieved on the TinylmageNe dataset by our proposed PFPT algorithm and other baselines improved with similar client-specific prompt selection mechanism. The symbol indicates the improved variant the original baseline.
FEDAVG-PT+
FEDPROX -P++
FEDOPT-PT+
GMM-PT+
PFPT
a= 0.5
86.30
85. 78
60.52
84.71
86.91
0.1 77.79
75.46
64.60
76.35
82.31
imbalance 74.25
74.15
64.24
74.17
78.21
that of PFPT on CIFARIO. Since FedAvg and FedProx were consistently the best performing baselines (after PFPT) across all datasets, we believe that the comparison of FedAvg and FedProx configured with adapter-tuning is sufficient to demonstrate the advantage PFPT over potential adapter- based FL approaches.
Table 8: Accuracy (%) achieved on the all dataset by FEDAVG with Adapter-Tuning, FEDPROX with Adapter- and our proposed PFPT algorithm.
Method
Setting
CIFARTO
CIFARTO0
TinylmageNet
synthetic 4-datasets
0.5
93.8640.17
75.95+0.40
78.88+0.23
55.55+0.79
FEDeVG-Adapter
0.1
92.66*0.26
65.0420.68
57.6240.80
30.58+4.66
imbalance
92.3340.26
49.820.79
40.90+ .34
10.86*8. 994
0.5
93.69+0.20
75.75+0. 16
79.0140.56
58.39+1.1
FEDPROX-Adapter
0.1
93.04+0.33
64.59=0.82
58.6240.56
32.92+ .34
imbalance
92. 13
50.75
37.6542.119
13.1947.99
0.5
94.39+0.51
80.2420.24
86.9140.11
76.89=0.17
PFPT
0.1
93.39+0.22
75.08=0.51
82.31=0.26
70.29=0.32
imbalance
91,45#0.08
72.05+0.93
78.21+1.25
62.23+1.02
G
Hyperparameter Settings
All experimented baselines use batch size of 16 and are fine-tuned with 10 learnable prompts. experiments are performed on a V100 GPU with 32GB GPU RAM. We employ same hyperparameter setting in all baselines (except for SCAFFOLD-PT), The hyperparameter settings are all presented in Table 9. SCAFFOLDS. F--TT requires a larger learning rate since SCAFFOLD'ss variate control mechanism is designed to work with SGD. Hence, unlike other baselines, SCAFFOLD-PP PT is configured with SGD instead Adam. Empirically, we find that a larger learning rate required for SCAFFOLD-PT with SGD reach its best performance in our setting.
H
Limitations
Despite the clear advantage of PFPT in terms of communication bandwidth, opposed to full-model FT and other forms of FL that typically have communicate the entire architecture weights per round of communication, we foresee that PFPT will face some challenges when the pre-trained and downstream task domains differ (which goes even beyond the distribution shift setting that we investigate in this paper). Intuitively, a significant domain shift will likely require using more trainable prompts. Can we quantify this shift and bound the minimum number of prompts required to enable accurate adaptation? Although this question is beyond the scope of this paper, will investigate
18

---

Table 9: Hyperparameter setting for all baselines and our PFPT
Method
Secting
Commuemicatin round
Eps, local training
Optimizer Icarning rale
Sampled
Sizc
clients
elients
CIEARIO: 100
U.5
16
120
betam((.9. (0.9. 9,0.9 0.98), epemte-6 Ir: 5c-4
CIFARIO0: 100
TinylmageNet: 100
10
Synthetic 4-dataset:
PEDAYG-PT FEDPROX-FT -P FEDAVG-Adige -ddapeer HNDRROX-Adiite
CIFARIO:
Adam:
CIFARI00: 100
0.1
16
120
betu-(0.9, 0.98), eps-le.6
TiuyImageNet:
Synthetic 4-dataset:
FEDOPT-PT
CIFARIO: 100
PEEDRG-PT GMM-PT PEPT
Adam:
100
imbolance
16
120
beta-(0.99 0.98), epe-1e-6
TinyImageNet:
In: Ic-4
Synthetic 4-dataset: 80
0.5
16
200
Ir: 5e-4
SGD
TinylmageNet: 100
CIFARI00: 100
10
Synthetic
CIFRRIO 100
SCAFFOLD-PT
SGD
CIFARI00: 100
TlnylmageNer:
Synthetic d-dataset: 80
CIEARIO: 100
SGD
CIFARIO0:
imbolance
16 200
Ir. le-4
TinyImageNet: 100
10
Synihetic 4-dataset:
and address it in future works, especially in the context of heterogeneous FL where each client operates in different domain. We have also raised the issue of prompt tuning having compute all intermediate gradients since the prompt tokes are appended at the top level of the pre-trained model. Going forward, can explore new prompt placement strategies to reduce this memory consumption.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Our contributions can be found in Sections 3 and 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please refer to Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

Justification: Our results mostly establish useful probabilistic equalities which do not require any assumptions. Complete proofs can be found in Appendices B, and D.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Our experiment protocols are described in Section 4. The detailed hyperpa- rameter settings are fully specified in Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
21

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: All the used datasets are publicly available. Upon acceptance, we will release our code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
sssisssssooosssssssi oneooonoooooooooooooo olycy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/ public/guidee SCoeeS/mmmmSe Sanasuuumssssss wwsssssollll.o..oo onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Such details can be found in Section 4 and Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We have provided error bars for all reported results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The detail regarding our compute resources are provided in Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that NeurIPS Code of Ethics https://neut httpss//////iis...eeee cc/ppblic/ ctt.i///eehhicuieellnlllnlc
didn't I make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: We have read the NeurIPS Code of Ethics and not find our work violate any aspects of the code
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We provide statement of impact in Appendix A.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Our work do not create new dataset. We only use existing, publicly available datasets and their combination. We also do not create any new pre-trained NLP or vision models. We only used existing, publicly available pre-trained models.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
We cite the source of all datasets and pre-trained models used in our experi- ments.
Justification:
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
24

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasttss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does not release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our work does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
25

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "4CqlMhOe1l", "reviewer": "Reviewer_DxxX", "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.", "questions": "-", "limitations": "-", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. The method is good in addressing data imbalance in extreme heterogeneous scenarios, which is a significant challenge in federated learning environments.\n3. The research includes comprehensive experiments and comparisons against existing federated prompt-tuning baselines, providing robust evidence of the method’s effectiveness.\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for recognizing the strength of our paper and for helping us catch a format issue.\n\nWe 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.\n\nWe 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.\n\nThank you for your consideration.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4BJ6VpCkzr", "reviewer": "Reviewer_tCnK", "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.", "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?\n  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? \n  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.\n  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? \n  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.\n  6. In line 201, could the authors provide more details on the method each client uses to select a subset of summarizing prompts?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The authors introduce an innovative probabilistic approach to prompt tuning within the context of Federated Learning. \n  - The logic of this paper is clear and easy to follow.\n  - 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.\n  - 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.\n  - 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing our contribution and for the constructive feedback, which are addressed below.\n\n**Q1. Discussing other fine-tuning approaches.**\n\nWe 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:\n\n--\n\nExisting 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.\n\n--\n\nGeneralizing 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. \n\nNonetheless, 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.\n\n| $\\alpha=0.5$ | CIFAR10 | CIFAR100 | TinyImageNet | synthetic 4-dataset |\n|------------------|----------------|----------------|----------------|---------------------|\n| FEDAVG-Adpater | 93.86±0.17 | 75.95±0.40 | 78.88±0.23 | 55.55±0.79 |\n| FEDPROX-Adapter | 93.69±0.20 | 75.75±0.16 | 79.01±0.56 | 58.39±1.15 |\n| FEDOPT-Adapter | 89.34±0.77 | 35.96±1.32 | 23.51±0.65 | 31.85±0.85 |\n| SCAFFOLD-Adapter | 78.10±0.39 | 18.64±1.54 | 22.41±0.65 | 29.72±0.71 |\n| PFPT (Ours) | **94.39±0.51** | **80.24±0.24** | **86.91±0.14** | **76.89±0.17** |\n\n| $\\alpha=0.1$ | CIFAR10 | CIFAR100 | TinyImageNet | synthetic 4-dataset |\n|------------------|----------------|----------------|----------------|---------------------|\n| FEDAVG-Adpater | 92.66±0.26 | 65.04±0.68 | 57.62±0.80 | 30.58±4.67 |\n| FEDPROX-Adapter | 93.04±0.33 | 64.59±0.82 | 58.62±0.56 | 32.92±1.34 |\n| FEDOPT-Adapter | 85.32±1.40 | 22.07±3.03 | 14.76±1.06 | 20.01±1.82 |\n| SCAFFOLD-Adapter | 79.06±0.81 | 12.27±1.76 | 17.87±1.19 | 18.68±2.44 |\n| PFPT (Ours) | **93.39±0.22** | **75.08±0.51** | **82.31±0.26** | **70.29±0.32** |\n\n| Imbalance | CIFAR10 | CIFAR100 | TinyImageNet | synthetic 4-dataset |\n|------------------|----------------|----------------|----------------|---------------------|\n| FEDAVG-Adpater | **92.33±0.26** | 49.8±0.79 | 40.90±1.34 | 10.86±8.94 |\n| FEDPROX-Adapter | 92.13±0.13 | 50.75±1.71 | 37.65±2.19 | 13.19±7.92 |\n| FEDOPT-Adapter | 74.22±1.11 | 11.75±3.64 | 10.16±1.00 | 21.13±5.20 |\n| SCAFFOLD-Adapter | 80.56±1.08 | 20.56±2.12 | 22.49±1.79 | 4.33±1.91 |\n| PFPT (Ours) | 91.45±0.08 | **72.05±0.93** | **78.21±1.25** | **62.23±1.02** |\n\nIn 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).\n\n**Q2. Initialization of summarizing prompts.** \n\nThe 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.\n\n[a] Understanding the difficulty of training deep feedforward neural networks. AISTATS (2010).\n\n**Q3. Each prompt captures a specific pattern of concept.** \n\nWe 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.\n\n**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)\n\n**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.\n\n**Q6. Method that each client uses to select a subset of summarizing prompts?** \nWe adopt the prompt selection mechanism in “Learning to Prompt for Continual Learning” (CVPR, 2022)\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5gSC13TDdk", "reviewer": "Reviewer_h6E3", "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.", "questions": "1.\tWhat are the real-world applications to support the scenario of extreme local imbalance in FL?\n2.\tHow does PFPT explicitly solve the problem of data imbalance?\n3.\tGlobal imbalance seems more common in real-world applications. Can PFPT be effective in addressing the global imbalance problem?", "limitations": "N/A", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1.\tThe problem of heterogeneous and imbalanced FL is interesting.\n2.\tThe overall writing and structure of this paper is good.\n3.\tThe experiments are sufficient to support the effectiveness of the proposed method.", "weaknesses": "1.\tThe local imbalance setting is very extreme, as described in section 4.1. There’s no real-world application to support this setting.\n2.\tThis paper lacks discussion of existing FL methods that address the class imbalance problem like:\n[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.\n[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.\n[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.\nAlthough the addressing problem is different in that the class imbalance is global, there are still connections in terms of problem settings. \n3.\tThe 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. \n4.\tIn the experiment, there should be more imbalance settings to support the robustness of PFPT. Currently, there is only one setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed questions which we have addressed below.\n\n**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.,\n(a) rare disease detection with extremely low prevalence rate (less than 1/10000) [e]; and \n(b) mobile keyboard prediction where on-device datasets are often imbalanced, with a few classes (e.g., commonly typed words or phrases) dominating [f]. \n\n[a] FedIIC: Towards Robust FL for Class-Imbalanced Medical Image Classification (https://arxiv.org/abs/2206.13803)\n\n[b] FL with Non-IID Data (https://arxiv.org/abs/1806.00582)\n\n[c] Addressing Class Imbalance in FL (https://arxiv.org/abs/2008.06217)\n\n[d] FL with Matched Averaging (https://www.arxiv.org/abs/2002.06440)\n\n[e] Complementary Pattern Augmentation for Rare Disease Detection (https://arxiv.org/abs/1911.13232)\n\n[f] FL for Mobile Keyboard Prediction (https://arxiv.org/abs/1811.03604)\n\n\n**Q2.** \n\n**A. Discussion of existing FL methods addressing class imbalance**\n\nWe 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.\n\n**B. Explain how PFPT solves the problem of (local) data imbalance.**\n\nLocal 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.\nTo 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.\n\n**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:**\n\nWe 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.\n\n**Q3. More imbalanced settings & how effective PFPT is in global imbalance**\n\nOur 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. \n\nHowever, 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. \n\nNonetheless, 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):\n\n| CIFAR100-LT ($\\alpha=0.1$) | IF = 100 | IF = 50 | IF = 10 |\n|-------------------------|----------|---------|---------|\n| FEDIC [2]  | 33.67    | 34.74   | 41.93   |\n| PFPT (ours) | 60.74    | 65.54   | 71.66   |\n\n| CIFAR100-LT ($\\alpha=0.5$) | IF = 100       | IF = 50        | IF = 10        |\n|-------------------------|----------------|----------------|----------------|\n| CReFF [1]  | 34.67 | 37.64 | 47.08 |\n| RUCR [3]                 | 36.83          | 40.80          | 50.90          |\n| PFPT (ours)             | 60.69          | 65.41          | 73.68          |\n\n| ImageNet-LT ($\\alpha=0.1$) | |\n|-----------------------------------|:---------------:|\n| CReFF [1]                          |      26.31      |\n| FEDIC  [2]                         |      28.93      |\n| PFPT (ours)                       |      75.54      |\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SaRTUOM8wi", "reviewer": "Reviewer_HV8G", "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.", "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?\n2. The lack of discussion of federated prompt-tuning pre-trained models in related work seems to be very relevant to this paper.", "limitations": "The lack of a discussion of limitations and deficiencies is important.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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?\n2.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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for recognizing our contribution. \nYour questions are addressed below:\n\n**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:\n\nOn 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. \n\nOn the right, the diagram highlights the global phase where local clients send their prompt sets to the server which then aggregates them.\n\n**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).\n\n[a] https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_32.html. \n\n**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.\n\nWe hope our response has addressed your questions. Please let us know if you have any follow-up questions for us.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their constructive comments. We summarize below our responses to the reviewers’ questions and concerns, as well as additional results to support our method.

**Reviewer HV8G** requested several clarifications and minor adjustments of our manuscript, which we have thoroughly addressed. 

**Reviewers tCnK and h6E3** requested additional discussion regarding other fine-tuning techniques and FL techniques that deal with class imbalance, which we have provided in the respective responses. We have also conducted extra experiments to compare our method with these techniques. We refer the reviewers to our response to **Q1 of reviewer tCnK** for an empirical comparison with adapter tuning techniques, and to **Q3 of reviewer h6E3** for another empirical comparison with FL techniques handling data imbalance (in the global settings suggested by the references provided by **reviewer h6E3**). These additional results show that our method performs robustly against other baselines across different settings.

**Reviewer tCnK** requested additional experiment showing that the model performance will increase with a larger no. of prompts (Q5). We have provided that experiments in the attached PDF.

**Reviewer h6E3** also requested for evidence that our experiment setting is realistic. We would like to highlight that all of our Dirichlet partitioning scenarios are standard in many previous FL work. Our extreme imbalance partitioning scenario has also been investigated in several well-cited works, which we have provided references in our response to **Q1 of Reviewer h6E3**.

**Reviewer DxxX** raised a format issue with the broader impact statement which will be fixed in our revised draft.

We thank all reviewers again for the constructive reviews and we hope our response has sufficiently addressed all questions. We will be happy to answer any follow-up questions that the reviewers might have during the Reviewer-Author discussion week.

Meta-review: The proposed method is novel and practical. All reviewers are positive to accept the paper. It could be further enhanced by justifying the relation between the proposed method and the proposed scenario.
Decision: Accept (poster)