Prompt Estimation from Prototypes for Federated Prompt Tuning of Vision Transformers
Authors: M Yashwanth
Sharannya Ghosh
Aditay Tripathi
Anirban Chakraborty
Venue: Accepted by TMLR ()

Abstract: Visual Prompt Tuning (VPT) of pre-trained Vision Transformers (ViTs) has proven highly effective as a parameter-efficient fine-tuning technique for adapting large models to downstream tasks with limited data. Its parameter efficiency makes it particularly suitable for Federated Learning (FL), where both communication and computation budgets are often constrained. However, global prompt tuning struggles to generalize across heterogeneous clients, while personalized tuning overfits to local data and lacks generalization. We propose PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning), a unified framework designed to achieve both generalization and personalization in federated prompt tuning of ViTs. Within this framework, we introduce the novel Class-Contextualized Mixed Prompt (CCMP) — based on class-specific prompts maintained alongside a globally shared prompt. For each input, CCMP adaptively combines class-specific prompts using weights derived from global class prototypes and client class priors. This approach enables per-sample prompt personalization without storing client-dependent trainable parameters. The prompts are collaboratively optimized via traditional federated averaging technique on the same. Comprehensive evaluations on CIFAR-100, TinyImageNet, DomainNet, and iNaturalist datasets demonstrate that PEP-FedPT consistently surpasses the state-of-the-art baselines under diverse data heterogeneity scenarios, establishing a strong foundation for efficient and generalizable federated prompt tuning of Vision Transformers.

Published in Transactions on Machine Learning Research (03/2026)
Prompt Estimation from Prototypes for Federated Prompt Tuning of Vision Transformers
M Yashwanth* Deportment of Computational and Data Sciences, Indian Institute of Science.
yoshwanthm wwi.c.c...ii............
Sharannya Ghosh
Accenbure, Japan
Aditay Tripathi Google, India.
aditaytr@gmail.om com
Anirban Chakraborty
anirban Wiscc oc in
Department of Computational and Data Sciences, Indian Institute of Science.
Reviewed on OpenReview: https: // openrevee........ net/ forun? id- 901CpPRj6A
Abstract
Visual Prompt Tuning (VPT) of pre-trained Vision Transformers (ViTs) has proven highly effective as a parameter-efficiet fine-tuning technique for adapting large models to down- stream tasks with limited data. Its parameter efficiency makes particularly suitable for Federated Learning (FT.), where both communication and computation budgets are often constrained. However, global prompt tuning struggles to generalize across heterogeneous clients, while personalized luning overfits to local data and lacks generalization. We propose PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning), unified framework designed to achieve both generalization and personalization in federated prompt luning of ViTs. Within this framework, we introduce the novel Class-Cone Contextualized Mixed Prompt (CCMP) - based on class-specific prompts maintained alongside globally shared prompt. For each input, CCMP adaptively combines class-specific prompts using weights derived from global class prototypes and client class priors. This approach enables per-sample prompt personalization without storing client-dependent trainable parameters. The prompts are collaboratively optimized via traditional federated averaging technique the same. Comprehensive evaluations CIFAR-100, Tiny Net Net,in and iNaturalist datasets demonstrate that PEP-FedPT consistently surpasses the state-of-the-an baselines under diverse data heterogeneity scenarios, establishing a strong foundation for efficient and generalizable federated prompt tuning of Vision Transformers.
1 Introduction
Federated learning (FI.) (McMahan et al., 2017 is a collaborative machine learning approach which a central server coordinates multiple clients to jointly train a global model, while preserving the privacy of the client's data by keeping the local data decentralived. A major challenge in FL is data heterogeneity: datasets each client can differ significantly in distribution, resulting in non-identically distributed (non-iid) data across the network. These discrepancies often cause client models to converge to different local minima, a phenomenon known as "client drift" (Karimireddy et al. 2020), which in turn degrades the generalization performance of the global model. To address this issue, personalized FI. methods (Chen & Chao Ma et al. 2022 Shamsian et al. 2021) have been proposed, aiming to better accommodate diverse local data distributions. However, a key drawback personalized methods is their roduced generalization performance new unseen clients (Deng el. al. 2024
"Equal contribution.
Work done during internship Indian Institute of Science. Providddd valuable insights that helped this work.

---

Published in Transactions on Machine Learning Research (03/2026)
Inspired by their success in centralized learning, large foundation models (FMs) are increasingly adopted in FI. to mitigate data heterogeneity (Bommasani et al.) 2021 Dosovitskiy et al. 2021 Radford et al. 2021 p. FMs demonstrate enhanced robustness in non-iid data settings (Qu et al. 2022). Yet, their substantial computational demands for tuning on resource-constrained edge devices and high communication costs present major hurdles. Parameter-efficint tuning methods, such as Visual Prompt Tuning (VPT) (Jia el al. 2022), offer a promising solution for efficient FM adaptation within FL, significantly reducing communication overhead while harnessing the power of large models.
Prompt tuning has recently gained attention in Federated Learning (FL), with methods like FedPR (Feng ct al. 2023) and PFedPG (Yang et al.) 2023) exploring its potential. However, cach comes with notable limitations. FedPR employs global prompts and is primarily suited for cross-silo FL. This design struggles in highly heterogeneous client settings, where shared global prompts fail to adapt to diverse local data distributions (Li et al. 2020, Deng et al. 2024p. pFodPG, on the other hand, generates client-specific prompts at the server and sends them to clients for local fine-tuning. Although personalizod, this approach implicitly assumes full client participation in every round an impractical assumption in many real-world FL systems. Moreover, such personalized strategies risk fitting to local data (Wu et al. Deng et al. 2024), limiting their generalization to unseen or pon-participating clients (as seen in Table
To address these issucs, SGPT (Deng et al. 2024) employs shared and group-specific prompts to enhance generalization. However, its two-stage training process and non-differentiable mechanism add optimization complexity and computational overhead. Furthermore, when data heterogeneity is high, it still struggles to generalize across diverse client distributions (Tables [ and 2). These limitations highlight key trade-off in FL prompt tuning: global prompts generalize well but lack expressiveness, while personalized prompts offer local adaptability but suffer from poor generalization and scalability. This raises the central question:
Con we achieve effective personalization while relying solely on globally shared prompts?
We answer this by proposing a novel prompt-tuning strategy tailored for fine-tuning of vision transformer in FL. Our method introduces class-specific prompts that are jointly optimized with shared prompts to address data heterogeneity across clients. To induce personalization without local prompt storage, we propose Prompt Estimation from Prototypes for Federated Prompt Tuning (PEP-FedPTT. It generates a Class Contextualized Mixed Prompt (CCMP) by combining global class-specific prompts. The combination weights determined by per-class membership scores, computed using global cls-token protolypes and the client's 's local class priors. The global cls prototype aggregates class centroids across clients, where each centroid is computed by averaging the cls token representations of data points within a given class. For a given input, we estimate class membership scores refined by each client's class priors. These scores act as soft weights to combine class-specific prompts into a single, differentiable mixed prompt (CCMP). Owing to strong semantic structure learned by pre-trained vision transformers, samples from the same class tend to form compact clusters in the represcntation space. This property allows the global class centroids to serve as reliable anchors for estimating class membership under domain heterogeneity. Incorporating class priors accounts for label heterogeneity across clients: clients may observe the same label set but with markedly different class frequencies. Local priors reflect the client's underlying data distribution. Together, the similarity scores and class priors provide a robust estimate of class membership that captures both domain and label heterogeneity. Each client can dynamically personalize prompts using only shared global information, without the need for local prompt storage or specialized server-side generation. Our approach integrates scamlessly into standard FL pipclines and delivers strong empirical performance across heterogeneous datasets, as demonstrated in Tables D and 2 Despite relying solely on global prompts, it effectively utilizes clients' data distribution, achieving scalable generalization. Since carly layers lack abstraction, while late layers limit prompt influence, CCMP is injected at intermediate layers, where representations are sufficiently semantic yet allow adequate depth for effective optimization. Our key contributions are as follows:
1. We introduce a unified framework PEP-FedPT that jointly optimizes class-specific and shared prompts to address data heterogeneity in federated learning of ViTs. We exploit achieve personalization by using only global prompts. Thus our proposed strategy aims to strike an effective balance between generalizability and personalization.
the clients' distribution to
2. We design novel prompt-mixing strategy that generates Class Contexualized Mixod Prompts (CCMP) as a function of class-specific prompts shared globally across clients. We empirically demonstrate its
2

---

Published in Transactions on Machine Learning Research (03/2026)
superiority over existing methods through extensive experiments on datasets exhibiting feature and label imbalance.
3. We provide theoretical insights into our design by showing that CCMP minimizes a quadratic upper bound and is optimal in the Minimum Mean Squared Error (MMSE) sense.
2 Related Work
2.1 Federated Learning (FL):
FT. is machine learning paradigm that emphasizes data, privacy by enabling collaborative model training under the coordination of a central server, without requiring direct data sharing. FedAvg (McMahan et al.)
2017) is the most commonly used technique for aggregating local models. This has led to its broad adoption
across various domains, such as Internet of Things and mobile devices (Mills et al. 2019
Nguyen et al.,
2021 Hard ct al., 2018 Ramaswamy et al. 2019 healthcare (Ricke et al., 2020; Xu et al., 2021 Nguyen et al., 2021: Brisimi et al. 2018 Feng et al. 2023), person re-identification (Zhuang et al. 2020 and
face recognition (Liu et al. 2022a Under hata-eeterogeneitty, training the models using FedAvg leads client-drift. Addressing this, regularization techniques (Acar et al.) Li et al.
2020: Gao et al. 2022 and
variance reduction methods (Karimireddy et al. 2020), and several studies on improving the genera lizatinnnzatin
performance in FL. by inducing flatness during the local training (Sun et al. 2023 Caldarola el. al.
2022 have
come to light. In case of extreme data heterogeneity across the clients, training a single model for all the clients will be difficult. As an alternative, personalizod FL approaches have been proposed (Tan et al. 2023 Chen &
Chao Ma et al.] 2022 Shamsian et al. 2021). These frameworks share some model parameters with the server while keeping others client-specific, enabling adaptation to local data distributions. Federated Domain Generalization aims learn models robust to client-level domain shifts; representative approaches include FedSR Nguyen et al. 22022 GA Zhang et al. 2023). and FedGaLA Pourpanah et al. (2025), which address this via representation regularization, variance-aware aggregation, and gradient alignment, respectively.
2.2 Prompt Tuning and Federated Learning:
As the fine-tuning of the foundational models became ubiquitous for downstream tasks in centralized learning (Bommasani et al. 2021 Dosovitskly et al.) 2021 Radford et al. 2021 prompt tuning techniques were originally proposed in the NLP community (Li & Liang. 2021 Liu et al. 2022b). Recently, Visual Prompt Tuning (VPT) was proposed for prompt tuning in ViT models, demonstrating its efficiency. VIT based FL (Qu el al. 2022) shows robustness to heterogeneity. However, due to heavy communication costs, prompt-tuning on pre-trained VIT models gained attention. In the recent literature, prompt tuning for FL has been proposed in methods like SGPT (Deng et al., 2024), FedPR (Feng et al.) 2023), pFedPG (Yang) et al. 2023) and also in the context of Vision and Language models FedOTP Li et al. 2024). SGPT relies on group-specific prompts and requires alternate training due to a non- differentiable selection mechanism. PFedPG depends on full client participation, while FedOTP assumes access to both text and image encoders. We address these limitations by introducing class-specific prompts and a novel per-sample prompt-mixing strategy based on class priors and cls-token prototypes. [Jnn et al. (2024) explore prompt-based domain adaptation by tuning prompts in pre-trained vision language models for cross-domain transfer.
3 Preliminary
In this paper, we use boldface letters to denote matrices and vectors. The operator '.' refers to element-wise multiplication and * refers to the standard matrix multiplication. We define TL, as the 8th transformer layer, and E denotes the expectation operator. The term & represents the probability of observing samples from class at client k. Additionally, sim(p.q) denotes the cosine similarity, while I represents the indicator function. [M] denotes the set [M)
3.1 Visual Prompt Tuning (VPT)
VPT is a parameter-efficient fine-tuning method for the pre-trained ViT (Jia el al. 2022 It an efficient. alternative to fine-tuning the full model. JJia al. (2022) propose VPT, where prompts are inserted at the input of the ViT. with d dimensional trainable prompt P. E Rdx1 follows:
cls, P, = TL,([cls,,,, P, E,-I),
(1)
The detailed notations and definitions are in Sec. A.1 of Appendix
3

---

Published in Transactions on Machine Learning Research (03/2026)
H * clsM-
(2)
Ej E [Rdxar denotes the image tokens at layer i, ny denotes the number of image tokens, M denotes the number of layers in the transformer. The final layer's cls token i.c, cls is used for classification. In the above model, the classification head H and Po are trainable.
3.2 Federated Learning (FL)
In FL, the server orchestrates the training with n clients with the goal of minimizing the following training objective:
min f(0) = 1 Σ (0),
(3)
0
Su. denotes the tthh client local objective function. O denotes the model parameters shared across the clients. In general, it can be written as fO) E (0; (x, 8)). D3 denotes the data distribution of the client k
(x.y)>>z
and (0; (x,y) denotes the task-specific loss function. For a classification task, x denotes the imput and y is the ground truth. In FL training, at each round I, the server broadcasts the global model 0' to a randomly selected subset of clients S.. Each client St performs several steps of local training starting from 0 and then sends its updated model of back to the server. The server aggregates these updates using federated averaging:
0
1
Σ %
|5
kES.
The updated model g++1 is then broadcast to clients the next round. This procodure describes the basic FedAvg algorithm McMahan al. (2017).
4 Proposed Method: Prompt Estimation from Prototypes- Federated Prompt Tuning (PEP-FedPT)
We consider a federated learning setup with clients coordinated by central server, where each client's 's data is drawn from a distinct distribution Dr. Following VPT (Jia et al. 2022 assume each client uses a pre-trained VřT-B/10 as its local model architecture. We introduce Shared Prompts and ClassCContxxtullizee Mixed Prompts (CCMP). Also, by utilizing the information in the local class priors and the global class prototypes, we softly combine the class-specific prompts, leading to per-client customization while sharing the global class-specific prompts. A highlevel overview of our method is shown in Algorithm
4.1 Prompt Design
Upon inscrtion of prompt Ps.,, into layer of a ViT, the input and output for that layer can be written as:
clsi, P E, (cls E,
(4)
where clsy denotes the cls-token representation at output layer i' of ViT, and Eg denotes the combined representation of the remaining tokens.
We have two sets of trainable parameters: Shared Prompts Ps and Class-Specific Prompts Pc, both of which are common across all clients. While we insert the shared prompts directly, the Class-Specific Prompts are used to compute Class Contexualized Mixed Prompts (CCMP) (denoted by m ) which are then inserted in subsequent layers in the VIT. We now describe each of these in detail.
Shared Prompts (Ps): Inspired by Jia el al. 2022)) we added shared prompts at the very first. layer of the ViT model. These are shared across the clients and are given by Ps [Ps, Psg Pais..Piss)) where [S] is the number of shared prompts inserted and are processed as follows:
cls,, s E₁ TT.((cllo, Ps, Eol).
(5)
4

---

Published in Transactions on Machine Learning Research (03/2026)
Communication
(b) Client Updates
Class Prior
Takens
Client 1
Insert shared prompt
Transformer Layer
Transformer Layer
Server
Probabiiities
Insert COMP
Transformer
Logits
MLP Head
Centroids
Class Prompts
cis Token
Trainable
Figure 1: The left panel (a) illustrates server-client communication during federated training. In each communication round, clients (right panel (b)) insert shared prompts at the input of the transformer and class-contextualized prompts- derived by mixing class prompls using probabilities computed from local class priors, cls-tokens and centroids at intermediate layer(s).
Algorithm 1: Pseudocode for PEP-FedPT
Compute global class prototypes before training: for each communication round do Server sends model parameters to clients; for each client in parallel do for local epochs do
5
6
Insert shared prompts at layer 0;
7
for layers in which CCMP is inserted do
Compute class scores from priors and CLS-prototype similarity;
Mix scores with class prompts to form CCMP; Insert CCMP;
10
Train shared and class prompts;
Send prompt updates and class prototypes to server;
Server aggregates prompt updates;
Periodically, server aggregates class prototypes;
As discussed in Ostapeuko et. al. (2022) and Deng el al. (2024). the early layers capture the low-level representations and can be shared across the classes. This allows the model to have better generalization. It shown that the representations in the initial layer of pro-trained ViT are uniformly distributed on the manifold, indicating that the information shared aeross the classes. This true even when the data distribution different across the clients shown See A.5.4 of the appendix.

---

Published in Transactions on Machine Learning Research (03/2026)
Class Contextualized Mixed Prompts (CCMP) (m (n))): This prompt is obtained by softly combining the class-specific prompts given by Pc = Pe, P.....PPPPp Po with scores driven by client-specific data distribution, where |C| is equal the total number of classes and Pc [dx[el, These class-specific prompts Pc are shared across the clients, but the scores that act as weights to combine these class-specific prompts are local to each client. These soft weights for client lt at the input of layer I on the i-th training input denoted by 1, E [0, 1||C|x1 are designed as the function of imput data point X cls token prototypes and class priors. Finally, the CCMP mi-1 is added at the input of layer I, and it's given below
m (k) = Pc S
(6)
The overall input and output after adding the CCMP my-1 at the input of layer I is shown below:
clst, mi mi)k, Ps.
Er (TL:([csss,,,,
Ps Ex-1))
The soft weights are explained Sec. 4.2 In Table 13 we show that inserting CCMP prompts too carly in the ViT is not advantageous, as the token representations the initial layers are not sufficiently informative for reliable estimation of CCMP scores. Conversely, adding prompts later layers also suboptimal: although the cls token representations are stronger at these stages, the prompts not inserted deeply enough to learn meaningful adaptations. Therefore, the most effective placement is within the mid-to-late layers (namely 5, 6 and 7 for all our experimenls), where the representations are both informative and feasible for effective prompt learning.
If the ViT has M layers, the final logits are given by:
H clsM-
H denotes the classification layer. Finally, we aim to solve the following federated optimization problem involving the shared and class-specific prompts.
min H
Σ
fu pre: Ps. Pc,H..
(9)
kml
Here Wpre denotes the pre-trained VIT parameters and fe denotes the loss for the client k.
4.2 Estimation of Soft Weights for CCMP
We present the design of the soft weights, which are aimed to provide class-specific information to the model. We exploit the information present in the class prototypes of cls;-1 token at a layer I. Our empirical observations show that
Top-5 Accuracy predicted across layers prototypes
Prompt Accuracy (top-5) (%6)
90
80
cls tokens of the pro-trained ViT model carry significant information regarding the downstrean ream lask. In the Figure 2 we observe the Top-5 zero-shot accuracy of the CIFAR- 100 datasct computed at each layer. The accuracy layer is computed by taking the minimum distance between the cls token corresponding to the test input and the class prototypes of the cls token at the input of layer I.
70
60
40
0
3
5
8
10
11
We now describe the soft weights computed for a layer i.e, 81-1- Let us denote the cls token at the input of layer cor- responding data point Xi, for client & in communication round t as cls;- The cls token's class prototype for the class at communication round t is denoted by 1 and it is computed as in Eq. 10
Layer Number
Figure 2: The Top-5 accuracy computed based the minimum distance between the cls token correspouding to the input and the cls proto- types. This shows that the cls representations in the middle layers have coarse information of the task.
Mi clsi- Lyaaos Tke > O,
Pi
(10)
0,
mk.c 0.

---

Published in Transactions on Machine Learning Research (03/2026)
Here 7i.cc = =E data points of client k, and = denotes indicator function. It takes value if the data point i of client k belongs to class otherwise it is 0. After every fixed update period R, the server aggregates the prototypes from the clients to compute the aggregated prototype f at the r-th period as Eq. Let. the set of communication rounds within this update period be A {rR, r 1.....,rr (r+ 1)R-
u. where N& denotes the number of
D. EiEr Exes, [ng 10) 1i-1K De >
i 11-1
(11)
(o.
De 0.
here D. D EnES, - #0): St S [n] denotes the subset clients sampled by the server at a round t. As the training progresses, the server uses the momentum to update the aggregated prototypes Pi-1, to form the global class prototype 1f-1, as in Eq. 12 which is then communicated to the clients. The parameter p denotes the momentum. The updated prototypes f are sent to the clients.
-.- 16-1-11-11 + (1 p) f8-11
(12)
If D. is we set p 1.
We now define the un-normalized score function & 1, assigned by the cls token cls:-1,i, corresponding to input X, to class-specific prompt Pe at the client k. Here we drop the index of the communication round and the update period r for better readability.
sim (cls) 1,i,ks -1,,)- mi-1) 0%.
811-1, = exp
(13)
T
We define sim(p. q) g_g [
T is the hyper-parameter and of is the prior probability of the class e at the client k. This can be obtained by computing the empirical label distribution. Availability of such class prior information a common assumption in the works such as Lee et al. (2022)
(14)
N
Di.k =c
t=1
The final scores 11,1 are obtained
gC 11,1-188
(15)
111-1AA
E 8T
4,1-1
The scores 1., can be interpreted as the probability assigned to the class-specific prompt Pe, given the cls token clsi-1,i,A- All the probabilities across the classes form the desired weight vector S1,-1, and the final CCMP m;-1 is computed using Eq. 6 CCMP achieves personalization even when client priors are uniform, due to its ability to capture the domain gap within the scores. This is direct result of the cls;-1,i,k tokens, as they will have different representations for different domains.
Privacy considerations: We acknowledge that PEP-FedPT requires the transmission of class prototypes to facilitate adaptive prompt mixing. However, these transmitted statistics represent aggregated, intermediate layers', low-dimensional summaries of the local data rather than raw data. Moreover, without formal privacy protection measures, strict privacy guarantees the Sec A.6 of (Xu et a it is usually hard for the server to extract sensitive data from category-level feature statistics (in our case cls-prototypes) alone without having access to the feature extractor layers. Prototype- based learning in FL has been used previously in (Dai et al. 2023 Tan et al. 2022 Xu et al.). Moreover in practical scenarios, even gradient-sharing algorithums like FedAvg would require a privacy preserving technique for the system to be completely reliable. In A.4.I we have discussed the impact of adding Laplace DP noise to the class prototypes on the overall performance of our method. We have observed that it has had minimal impact the final accuracy of our proposed method. We also show that (c, 0) privacy be attained.
challenging in virtually any FL framework. As discussed in
7

---

Published in Transactions on Machine Learning Research (03/2026)
4.3 Theoretical Underpinnings of CCMP
4.3.1 CCMP minimizes the Quadratic Upper bound on the Loss around the class prompts
CCMP constitutes the core component of our methodology and by design, it specific to each client. We show that CCMP minimizes the quadratic upper bound on the loss. We denote the estimate of the class prompt for class i in any round Pe; and the class prompts will be denoted by Pc = [Pa, > Pey .... Pacjl- Let m(k denote the CCMP prompt. used for client P and let the total number of clients be n, and of denote the empirical probability that a data point client k belongs to class i. Let be the set of all possible prompts across all the clients, such that m(k) P, vk (1,2.....,h) n.. We denote the average loss corresponding to data point whose true label y= as l.. We rewrite this loss in terms of prompt p and class i as I'(p (p). Then the global loss across all clients can be computed as = = - We now state the following assumptions:
o% 1 (m(())
A 1. P is compact subset of IR.", where d is the token dimension. This means that all the possible values for the prompts lie within bounded region and they do not drift off to infinily.
A 2. i is Lipschitz smooth with parameter B, Vi E [[[]] E [m].
This implies that the gradients do not change abruptly.
A 3. (P) achieves its' minimum value for porp=p- p
Intuitively this assumption implics that clients class-specific loss function has same minimum. This is true for label imbalance and is not true for feature imbalance setling.
Proposition 1. Iʃ the above assumplions hold, we show that J can be upper bounded as S S = Smax [|m(())))))))))){)) - Pe,, + +C and it is minimized at m(k) = E| &epe, Vk E [n].
I Ir Ex 5% (Px( 2 which is equivalent to the (CCMP) described in 4.2 as T >>> 1. Smax = maxielIc] 8,, Ö is a constant which depends on P. Under the assumplion. 3 which correspond to label heterogeneity setting this vanishes when Pa. = P. Vi E [[C']] which makes L a tight upper bound of f.
A detailed proof is provided in Section A.6.1 of the appendix. The proof sketch proceeds by constructing upper bound on the class-wise loss using smoothness assumptions. Specifically, the first-order term is bounded via compactness and smoothness properties. The upper bound is then minimized with respect to each m(k), The key insight from this proposition is that each client's 's prompt differs due to client-specific class priors, even though the underlying class prompts Pa, are shared globally. This mechanism enables the prompts to adapt to each client's 's data distribution. The mixing scores in Eq. 15 depend both on the data instance and the class priors; this is further discussed in Sec AA.6.4 of appendix.
4.3.2 CCMP is Optimal in Minimum Mean Squared Estimate (MMSE) Sense
We denote p& (p/cls as the posterior probability of the prompt after observing the cls token cls;-1 at the input layer | It should be noted that this discrete probability measure over the class-specific prompts (Pay Peg ... }. If we assume that the density over the cls tokens follows pa (clss-1)) and the posterior over the class given the cls token is modeled as (based on Eq. 15 i.e.,
Ph (p Pel cll -
(16)
This induces the joint probability density over the cls tokens observed and the class-specific prompts (P Pa ..., P...p. J. We denote this by pk pc(sss-1. p) and is given below
P& (alss-_11, p) Ph (p)csssss))))) (cls
(17)
If the cls tokens and the class-specific prompts at imput of layer I has the joint density given by px(cls(-1.P) as in Eq. 17 then the CCMP prompt for a client k, m)-1(k) obtained in Eq. 6 is Minimmm Mean Squared Estimator (MMSE) of the true class prompl.
Proposition 2.
2or notation convenience, we drop the layer index from (k)-)
In clay- 1 we omit the subscripts of client. k data point i and round for simplifying notation.
8

---

Published in Transactions on Machine Learning Research (03/2026)
The proposition 2 says that the CCMP obtained in Eq. G is optimal in MMSE sense. The detailed proof is given in Sec. A.6.2 of the appendix. This is done by showing MMSE optimality of the estimator R[p|cls/-1]- The discussion regarding the convergence is provided in Sec. A.6.3 of the appendix.
5 Experiments
Datascts: We conducted extensive experiments on four popular datasets (1) CIFAR-100 (Krizhevsky &
Hinton 2009) dataset consists of 50,000 training images and 10, 000 test images distributed across 100 classes.
(2) TinylmageNet (Le & Yang 2015) contains 100K images of 200 classes, with each class containing 500
training images and 50 test images. (3) DomainNett (Peng et al., 2019) comprises 0,6 million images of 345
classes distributed across six domains: Clipart, Infograph, Painting, Quickdraw, Real, and Sketch; however, following the protocol of (Yang al. (2023), we use the top ten most. frequent, classes for our experiments. (4)
(Naturalist originally introduced in Van Horn et al. (2018) is large-scale fine-grained visual classification dataset comprised of images of natural species. In this work use federated version of this dataset.
Setup: We split CIFAR-100 and Tiny-ImageNet datasets among the clients with two different settings of
data heterogeneity: pathological splitting (Li et al., 2023 Oh et al.) Deng et al. 2024), where each client observes only 10 classes, and Dirichlet- based splitting denoted by Dir(E) (Acar et al.), where each client has a non-identical label distribution. A lower & value indicates higher heterogeneity, and we set & 0.3. For CIFAR-100 we consider 100 clients and for TinyImageNet we consider 200 clients. Only randomly chosen clients participate in every round. For the DomainNet dataset consider the setting as Deng et al. (022
el al. p where each domain is allocated to 10 clients among 60 clients, this considers the scenario of feature imbalance setting. 6 randomly sampled clients participate in each communication round. Finally, in the iNaturalist dataset (Hsu et al.) 2020), we make sure each client gets at least 16 training samples. This will have around 100k training samples distribuled among the 1018 clients and 1203 classes. The partition is performed to mimic the cross-device (Kairouz et al., (021) non-iid setting. Tn our method (PEP-FedPT) the CCMP is insterted the layers 5, 5,6,7, 7. In all the above experimental setups, the partition of the dataset across clients is completely disjoint, i.c., two clients contain the same data sample in all cases. The visualization of data heterogeneity is provided in Sec. A.3.1 of appendix. The detailed hyperparameter settings is provided in A.3.2 and the evolution of cls representations are provided in the Sec. A.5.4 of the Appendix.
Table 1: Quantitative comparisons on CIFAR-100, Tiny-ImageNel datasets using VIT-B/16. We report the accuracy under two non-iid data partitioning setups 1) pathological: Each client observes only 10 classes 2) Dirichlet: Label distribution of each client is drawn from Dirichlet distribution.
Datasets
CIFAR- 100 (%) f
Tiny- Net (%)
Pathological
(0r(0.3)
Pathological
(ir(
Method
Mean Ace
Worst Ace
Mean Acc
Worst Acc
Mean Acc
Worst, Aue
Mean Acc
Worst Acc
Hcad-Tuming
27.85
50.87
79.56-0.25
66.66+2.79
68.39 20.76
44.09
70.73
45.63
FedVPT
23.6220.02
70.1940.11
84.91200.0
74.6440.74
14.20
54.00-2.46
76.57 10.34
50.341251
FedVPT-D
85. 2.77
70. 70.12
88, 60.
79.17=n..6
79.60 79.60+0.42
59.83_1..6
83.30+0.16
60.33
FedPR
81.77 10.30
08.9910.44
82.27 L0.22
73.29 11.38
68.8610.17
47.50_1.63 1.63
68, 68.93
47.37 11.44
SCPT
24.14
70.7940.50
85.90=0.21
76.73gg1.0
75.65 .81
55. 35.66_3..2
78.844
53.87
PFedPG
92.9611.33
84.5811.1
77.27 10.77
62.34
82.93
50.21 1.00
55.91 10.65
49.31 11.05
P-PT
75.37
55. 14.144
80.10.m.25
68.33 08.33+40.0
61.68
38. 02. 08.09
62. 62.78
40.30
PET-FedPT(Ouus)
95.46
84.74+3.12
88.75
81.00=a.00
91.52++6.1
77.33
83.444...2
61.00+0.a1
pre-trained on the ImageNet- 21K dataset (Feng et al.) 2023 Yang et al. 2023) as our base model. ViT-B-16 was originally trained on images with à resolution of 224x224 pixels, utilizing a patch size of 16. To maintain compatibility, we resize our input images to 224x224 pixels. In the prompt-tuning stage, we specifically focus on optimizing shared and class prompts, as well as the classifier head. The papem-parreeter details are in Sec. A.3.2 of the appendix.
Model Details: We use the Vision Transformer (VIT-B/16) (Neil & Dirk 2020)
Baselines: We compared our method against several global and personalized Federated Learning (FL)
methods that use the Prompt Tuning including Head-Tuning (Sui et al. 2022) FedPR (Feng et al. 2023
9

---

Published in Transactions on Machine Learning Research (03/2026)
FedVPT and FedVPT-D (Jia et al., 2022 PFedPG Yang et al. 2023) and SGPT DDng et al 2024p. To thoroughly assess our method's performance, we introduce new baseline called P-PT which personalizes the prompts, giving insights into how the personalization of prompts impacts performance.
Evaluation Methodology: We use two key metrics to assess the performance of both the baselines and our proposed method (Deng et al. 2024]: (1) Mean Accuracy calculates the average accuracy across individual clients' test data, reflecting adaptation to diverse client data distributions. (2) Worst Local Accuracy reflects the performance of the worst-performing client, indicating adaptation to the most challenging local data. Furthermore, utilize a heldout evaluation strategy (Yuan et al. 2021), where 90% of clients participate in training and 10% are reserved for testing. All aforementioned metrics are reported scparately for both participating and heldout clients. This setting demonstrates the model's effectiveness in onboarding new clients and adapting to previously unseen data without sharing updates with the central server. All experiments are done over 3 different runs, and the mean and standard deviations are reported as (mean
5.1
Results and Discussion
5.1.1 Label Heterogeneity Results
The class hetcrogencity results are presented in Table- 0 PEP-FedPT outperforms all the basclines, c.g., when compared against pFedPG with CIFAR-100 pathological setting we observe an improvement of 2.5% in mean accuracy. For Tinylmagenet the improvement is 8.59% (mean accuracy) over pFedPG and 17.5% (worst accuracy) over FedVPT-D. This shows that our prompt-mixing mechanism effectively addresses data heterogeneity independently and performs even better in scenarios of higher heterogeneity or lower class overlap between clients. Similar improvements are seen in other settings. Additional experiments are provided in Sec. A.4 of the appendix. The visualization of the accuracy vs communication rounds is shown in Sec. A.55I
Table 2: Experimental results on DomainNet and iNaturalist. For DomainNet we consider each domain belonging to a client and we report the accuracy attained by each client and the average accuracy. On the iNaturalist dataset we report the average test accuracy of all the clients and the 15th percentile worst accuracy. Our method significantly outperforms all the baselines on this challenging datasct.
Datasels
Donsiiaet(%))
(Naturalia(()) f
Method
Clipart
Infograph
Painting
Quickiraw
Roal
Sketch
Moon Ace
Wonst Aee
Meam Ave.
Worst Acc (13%)
Head. Tuning
01
27.45g1.07
91.39
24
96.68_o
26.4661142
83.7
38. 38.8
49.41
20.50..0.7
FddVPT
10784_1a7
58.56
92.25 12.25
77.81 10.30
96.78 0.
89.24 1. sm
10.23
37.02 19.41
52. 22 12.22
23:21 10:50
FedVPT-D
04
60.29
43.4
84.58.1.5
96.96 06
91.58 m
47.31
57.96.1.12 42
30.00.0
FedPR
91.62 m1
M6.00
91.1611.17
73.72 10.0
96. ..33
86.41 11
82.96
35. 18.18
41. 21.255
DS. 62, 30:30
SGPT
02.64g005
60.622a:2
91.541o.78
83.55
96.55_0117
29.93
85.5640.00
37.34g1441
55. 58.78
27. 39
PRodPG
82.80_naa
63.561as
92.2771..0
87.333
97.16 _s_
80.34 10:30
17.40
52.05 10.32
52.422a..0
12.54_.0
P-PT
40.11g1
56.73
90.25g1..o
74.81.0.a1
18
85.26
82.30
35.67
45.60g00.
16.603030
PEP-PeIPT(Ous))
95.46...1
21.6821.4
93:00go.55
86.80
97.6730.55
91
89.16 00
59.79.2.5
63.48g1..0
41.10.g0.s
5.1.2 Feature Heterogeneity Results
Feature-hettergeenityy results presented in Table FedVPT-D serves strong baseline due to its inclusion of prompts at each layer. We present the results for our method. Tn this setting, personalized method PFedPG also proves highly beneficial due to the pronounced feature imbalance among clients. Our method PEP-FedPT improves on average by 5.52% over the best performing baseline FedVPT-D on the iNaturalist. dataset. For the iNaturalist dataset, under a low client participation rate (1%), many clients are sampled infrequently, and a large number of clients (over 150) have extremely small test sets (fewer than 10 samples). Thus, misclassification of only a few test samples leads to a worst-client accuracy of zero for all methods. Consequently, the standard worst-client metric becomes uninformative. To address the evaluation difficulty in this setting, we report the client accuracy at the 15th percentile a redefined worst-client metric, which provides a more informative assessment of performance on underperforming clients.
5.1.3 Heldout Evaluation
In the Table 3 we present the results for the heldout evaluation, where we report the accuracies of the participating clients and the new clients. Participating clients are the ones who participate in the foderated training and the new clients do not participate in the FL training. In this setup, the testing accurcay implies the zero-shot predictions on the held-out clients' test data. We consider the CIFAR-100 dataset.
10

---

Published in Transactions on Machine Learning Research (03/2026)
with pathological partitioning where each client observes only 10 classes. Most baseline methods achieve competitive accuracy on participating clients but show reduced performance on unseen clients. Personalized methods like pFedPG achieve very high participating accuracy but fail in held-out testing resulting in poor performance, highlighting poor generalization. Since pFedPG is not explicitly designed for this ovaluation protocol, its performance under this setting should not be interpreted as indicative of its ellectiveness. We include pFedPG to illustrate the behavior of personalization and therefore mark the corresponding results as Not Applicable (NA). In contrast, our method consistently achieves the best results across both datasets, with 95.66% vs. 93.71% on CIFAR-100 and 92.53% vs. 90.60% on Tiny-ImageNet, showing strong generalization to unseen clients. The results on DomainNet and iNaturalist. are provided in Sec. A.4.7 of the appendix.
methods a strict zero-shot setting,
Table 3: Quantitative comparisons on CIFAR-100 and Tiny-ImageNet with held out evaluation: We report the accuracy with pathological partitioning where each client observes classes. can be observed that the personalized methods like pFedPG perform the worst in the held-out evaluation Our method performs well on the clients participating in the FL training and also on the unseen clients.
CIFAR-100 (t)
Inngeeee (1)
Method
Participating Aoe
Testing Are
Participating Ace
Testing
Head
77.81
77.10g0.1
67.97 1000
97 m
PsdVPT
83.62.0.44
82.39+132
44.15
74.15 19
FedVPT-D
85.06
81.87
77.38 13.5
76. 76.89
FedPR
81.62
80.61
69.37 11.4
23 so
SGPT
g1.90
83.632n.64
76.38
78.10.1 50
PRedPG
03.3220.85
N'A
86.09+1.42
N.A
P-PT
75.97
72.19
60.76.0.a0
60.41
95.66
93.71
92.53 1o.35
90.60_..m1
5.2
Analysis of PEP-FedPT
5.2.1 Ablations
We conducted ablation studies shared and class-specific prompts evaluating their impact by varying the number of shared prompts and the influence of class priors the prompt mixing strategy. Table g reports the effect of combining shared prompts with CCMP under both Pathological and Dirichlet splits. For the Pathological split, the average accuracy improves from 83.62% with only shared prompts to 95.46% with shared + CCMP. Similarly, for the Dirichlet split, the performance increases from 84.91% to 88.75%. These results highlight the consistent benefit of incorporating CCMP across different data partitioning strategies.
Table 4: Shared and CCMP ablation on the CIFAR-100 dataset with Dirichlet and Pathological Partitions. We report the Accuracy(%).
Prompt Strategy Pathological Split (f) Dirichlet Split (?)
Only Shared Shared CCMP
83.62.0.02
84.91 88.75
95.46.0.10
The Table 5 presents the effect of incorporating class priors into the prompt design on CIFAR-100 under both Pathological and Dirichlet splits. The results show that using Shared + CCMP with Class Priors (CP) consistently improves performance over the baseline without CP. In particular, the Pathological split benefits, with accuracy increasing from 84.01% to 95.46%, while the Dirichlet split. also shows notable gain from 86.12% to 88%7.. Similar analysis for other datasets is given in section A.4.3 and A.4.2 of the appendix. Further additional experiments can be found in Section A.4 of the appendix.
Table 5: Impact of Class Priors on CIFAR-100 dataset with Dirichlet and Pathological Partitions.
Prompt Strategy
Pathological Split (f) Dirichlet Split (f)
Shared + CCMP Without CP
84.01 95.4020016
86, 1240.1
Shared CCMP With CP
88.75g035
5.2.2 Computation and Communication
We denote that d and dn are token and attention head dimensions, C' and denote the number of classes and layers respectively, and T are the tokens. The minimum computations required by ViT forward is given as:
11

---

Published in Transactions on Machine Learning Research (03/2026)
Accuracy vs Communication Rounds
Table 6: Comparison of computation and communi- cation. Resources required to achieve 83% accuracy on CIFAR-100.
Accureey 180
Method
Training Time (see)
Params Communlcated
Rounds Required
RddVPT
M
100
RedVPT-D
4760
7.67 M
90
SGPT
PdPPR
55380
58.4 M
>00
PEP FedPT (Ours)
1153
4.6 M
12
Communication Rounds
Figure 3: Convergence CIFAR-100 with pathological partitioning.
The Query (Q), Key (K) and Value (V) requires Tdd, multiplications each. The inner product matrix QKT requires 7h, multiplications. The feedforward computations requirement is d'. If H heads are present and there are C' classes for the classification then we need (3Tddd,, + T²d,)= + LT& + Cd multiplications. For CIFAR-100 on ViT-B/ //16/ 16 the CCMP computation takes only 0.008% of total computations, which is very negligible implying the efficiency of the CCMP computation. Table G compares the computational and communication complexity of our proposed method, PEP-FedPT against the different baselines. For a fair comparison, we compare the methods that use global prompts. We analyze the resources required to achieve 83% accuracy, which is the highest accuracy reported for FedVPT. Our results show that PRP-FedPT achieves this accuracy in just 12 rounds, roquiring lowest training time and significantly roducing communication overhead (4.6M) compared to SGPT (13.0M), where M denotes million. The claim of 12 rounds can be verified in the Figure 3 FedPR only attains 81.66% in 100 rounds so we report this as (> 100). The training times reported are measured on an Nvidia RTX-A6000 GPU. The detailed computation of why 4.6 M is : Head requires (100 x 768), shared prompt (1 x 768) class prompts (100 x 768), prototypes (100 x 768 3) scaled by 3 because of three layers with CCMP prompts. Total rounds 12 and in total, yields 4.6M parameters. In this communication analysis, since our method shares the global prompts, we compared only the methods that are not personalized.
6 Limitations and Scope for Future Work
Our approach relies empirical estimates of class priors and cls token centroids, which require access to labeled data on the client side and introduce several limitations. In semi-supervised or unsupervised settings, where labeled data is scarce or unavailable, these estimates may become unreliable, potentially degrading prompt construction and making adaptation to such scenarios non-trivial. Similarly, in long-tailed data distributions, the limited presence of certain classes can negatively affect centroid quality and, consequently, model performance. Finally, the method is primarily designed for non-IID data distributions, leveraging heterogeneity to enable personalization via global parameters; under IID) settings, this advantage diminishes and the method effectively reduces to standard Fed-VPT. Our theoretical result in Proposition. is relatively tighter for label heterogeneity setting in comparison to feature drift owing to our Assumption 3 Exploring extensions to address these limitations presents promising direction for future research.
7 Conclusion
We propose a novel prompt-tuning methodology for Vision Transformers (ViTs) by introducing class-specific prompts alongside shared prompts. Our approach leverages the cls-token representations in pretrained ViT layers to extract prototypes, which are then combined with each client's prior label distribution to compute soft scores that guide the mixing of class-speeific prompts into a unified, optimized prompt (CCMP). This dynamic mixing allows (CCMP) lo achieve personalization while using global prompts only. This combined prompt (CCMP) is subsequently embedded within the VIT layer. Our method PEP-FedPT, achieves State of the Art performance, surpassing previous methods across the benchmark datasets!!
1 'Acknowledgment: This work is partially supported by the P3DX project, scod-funded by the Ministry of Electronics and Information Technalogy (MeitY). The authors would also like to acknowledge compute supports received from Kotak-IISc AT- Centre (KTAC), TTSc and the PMRF fellowship.
12

---

Published in Transactions on Machine Learning Research (03/2026)
References
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 Learning Represenlations.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint (rXi::2188.07256 2021.
Thcodora S. Brisimi, Ruidi Chen, Theofanic Mcla, Alex Olshovsky, Ioannis Ch Paschalidis, and Wei Shi. Federated learning of predictive models from federated electronic health records. International journal of medical informatics, 112:59 67, 2018.
Debora Caldarola, Barbara Caputo, and Marco Ciccone. Improving generalization in federated learning by seeking flat minima. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXIII, pp. 654-672. Springer, 2022.
Hong-Yon Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In International Conference Learning Representations.
Yutong Dai, Zeyuan Chen, Junnan Li, Shelby Heinecke, Lichao Sun, and Ran Xu. Tackling data heterogeneity in federated learning with class prototypes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7314-7322, 2023.
Wenlong Deng, Christos Thrampoulidis, and Xiaoxiao Li. Unlocking the potential of prompt-tuning in bridging generalized and personalized federated learning. In Proceedings of the IEEE/CVF Conference on Compuler Vision and Pallern Recognition, pp. 8087-6092 2024.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https::/ wtw.pppwwwweeeeeeeeeeeoerrrreeea
Cynthia Dwork and Aaron Roth. The algorithmic foundations of differential privacy. Foundations and trends@ in theoretical computer science, 9(3-4):211-44) 2014.
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of cryptography conference, pp. 265-284. Springer, 2006.
Chun-Mei Feng, Bangjun Li, Xinxing Xu, Yong Liu, Huazhu Fu, and Wangmeng Zuo. Learning fedcrated visual prompt in null space for mri reconstruction. In Proceedings of the IEEE/CVI Conference Computer Vision and Pattern Recognition, pp. 8073, 2023.
Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng-Zhone Xu. Fedde: Federated learning with non-iid data via local drift decoupling and correction. In Proceedings of the IEEE/CVP Conference Computer Vision and Pottern Recognition (CVPR), pp. 10112-10121, June 2022.
Andrew Hard, Chloé M Kiddon, Daniel Ramage, Francoise Beaufays, Hubert Eichner, Kanishka Rao, Rajiv Mathews, and Sean Augenstein. Federated learning for mobile keyboard prediction, 2018. URL
1811. 03604
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real-world data distribution. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part X 16, pp. 76 92. Springer, 2020.
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim. Visual prompt tuning. In European Conference on Computer Vision, pp. 709-727. Springer, 2022.
13

---

Published in Transactions on Machine Learning Research (03/2026)
Cheng Jin, HaiTing Zheng, and Hang Yu. How design learn prompt for domain adaptation? In 2024 International Joinl. Conference on Newral Networks (IJCNN), pp. 1-8, 2024. doi: N.1109/1JCNN60899 2024.10651258.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trendsl in Machine Learning, 14(1 2):1 210, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scalfold: Stochastic controlled averaging for federated learning. In Inter- national Conference on Machine Learning, pp. 5132-5143 PMLR, 2020.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Canadian Institute for Advanced Research, 2009.
Ya Le and Xuan Yang. Tiny imagenel visual recognition challenge. CS 231N, 7(7):3, 2015.
Gihun Lee, Minchan Jeong, Yongjin Shin, Sangmin Bae, and Se- Se-Young Yun. Preservation of the global knowledge by not-true distillation in federated learning. Advances in Neural Information Processing Systems, 35:38461-3844 2022.
Hongxia Li, Zhongyi Cai, Jingya Wang, Jiangnan Tang, Weiping Ding, Chin-Teng Lin, and Ye Shi. Fedtp: Federated learning by transformer personalization. IEEE transactions on neural networks and learning systems, 2023.
Hongxia Li, Wei Huang, Jingya Wang, and Ye Shi. Global and local prompts cooperation via optimal transport for federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Poltern Recognition, pp. 12151-12161, 2024.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet. Talwalkar, and Virginia. Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 450, 2020.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 4582 4597, Online, August 2021. Association for Computational Linguistics. doi:
URL
353
Xiaoxiao Li, Meirui JIANG, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learning on non-IID features via local batch normalization. In International Conference Learning Representations, 2021. URL
Chih-Ting Liu, Chien-Y' Yi Wang, Shao-Y: Yi Chien, and Shang- Hoon Lai. Fedfr: Joint optimization federated framework for generic and personalized face recognition. In Proceedings of the AAAT Conference Artificial Intelligence, volume 36, pp. 1656-1666000 1656-166,,226 2022a.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Jic Tang. P-tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks. In Proceedings of the 60th Annual Meeling of the Association for Computational Línguistics (Volume 2: Short Papers), pp. 68, 2022b.
Xiaosong Ma, Jie Zhang, Song Guo, and Wenchao Xu. Layer-wisee wised model aggregation for personalized federated learning. In 2022 TEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10082 10091, 2022. doi:
SNVPR52688.0022.00088
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient. learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp- PMLR, 2017.
Jed Mills, Jia Hu, and Geyong Min. Communicttion-fffcien foderated learning for wireless edge intelligence in iot. IEFF Internet of Things Journal, 7(7):5986 5994, 2019.
14

---

Published in Transactions on Machine Learning Research (03/2026)
Houlsby Neil and Weissenborn Dirk. Transformers for image recognition at scale. Online: https://ai. googleblog. com/2020/ / /2
hlml, 2020.
A Tuan Nguyen, Philip Torr, and Ser Nam Lim. Fedsr: A simple and effective domain generalization method for federated learning. Advances in Neural Information Processing Systems, 35:38831 188413,38333111 2022.
Dinh C Nguyen, Ming Ding, Pubudu N Pathirana, Aruna Seneviratne, Jun Li, and H Vincent Poor. Federated learning for internet of things: A comprehensive survey. IEEE Communications Surveys & Tutorials, 23 (3):1622-1658, 2021.
Jachoon Oh, SongMook Kim, and Se- So-Young Yun. Fedbabu: Toward enhanced representation for federated image classification. In International Conference on Learning Represenlations.
Oleksiy Ostapenko, Timothee Lesort, Pau Rodriguez, Md Rifat Arefin, Arthur Douillard, Trina Rish, and Laurent Charlin. Continual learning with foundation models: An empirical study of latent replay. In Sarath Chandar, Razvan Pascamu, and Doina Precup (eds.), Proceedings of The 1st Conference on Lifelong Learning Agents, volume 199 of Proceedings of Machine Learning Research, pp. 60-91. PMLR, 22 Aug 2022.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1406 1415, 2019.
Farhad Pourpanah, Mahdiyar Molahasani, Milad Soltany, Michael Greenspan, and Ali Etemad. Federated unsupervised domain generalization using global and local alignment of gradients. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 39, pp. 19948-19958, 2025.
Liangqiong Qu, Yuyin Zhou, Paul Pu Liang, Yingda Xia, Feifei Wang, Ehsan Adeli, Li Fei-Fei, and Daniel Rubin. Rethinking architecture design for tackling data heterogeneity in federated lcarning. In Proceedings of the IEEE/CVF conference on compuler vision and pallern recognition, pp. 10061-1007,, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Adilya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 9748-8763 PMLR, 2021.
Swaroop Ramaswamy, Rajiv Mathews, Kanishka Rao, and Françoise Beaufays. Federated learning for emoji prodiction a mobile kcyboard. arXiv preprint arXiv:1906. 04329, 2019.
Nicola Ricke, Jonny Hancox, Wenqi Li, Fausto Milletari, Holger R Roth, Shadi Albarqouni, Spyridon Bakas, Mathieu N Galtier, Bemnett A Landman, Klaus Maier-Hein, et al. The future of digital health with federated learning. NPJ digital medicine, 3(1):1-7, 2020.
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernet- works. In International Conference Machine Learning, pp. 9489- 9502. PMLR, 2021.
Guangyu Sun, Matias Mendieta, Taojiannan Yang, and Chen Chen. Exploring parameter-efficiee fine-tuning for improving communication efficiency in federated learning. 2022.
Yan Sun, Li Shen, Tiansheng Huang, Liang Ding, and Dacheng Tao. Fedspeod: Larger local interval, less communication round, and higher generalization accuracy. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://ppenrevie..c
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 34(12):9587-99) 9603, 2023, doi: 10. IN.LLS. 2022.3160699
Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogencous clients. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pp. 8432 8440, 2022.
15

---

Published in Transactions on Machine Learning Research (03/2026)
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pictro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8769 8778, 2018.
Shanshan WU, Tian Li, Zachary Charles, Yu Xiao, Ken Liu, Zheng Xu, and Virginia Smith. Motley: Benchmarking heterogeneity and personalization in federated learning. In Workshop Federated Learning: Recent Advances and New Challenges (in Conjunction with NewrIPS 2022).
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference Learning Representations.
Jie Xu, Benjamin Glicksberg, Chang Su, Peter Walker, Jiang Bian, and Fei Wang. Federated learning for healthcare informatics. Journal of Healtheare Informatics Research, 5:1-19, 2021.
Fu-En Yang, Chien-Yi Wang, and Yu-Chiang Frank Wang. Elficient model personalization in federated learning via client-specific prompt generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 19159 19168, 2023.
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? arXio preprint arXiv:21100 11221100.1111. 2021.
Ruipeng Zhang, Qinwei Xu, Jiangchao Yao, Ya Zhang, Qi Tian, and Yanfeng Wang. Federated domain generalization with generalization adjustment. In Proceedings of the IEEE/CVP Conference Compuler Vision and Pattern Recognition, pp. 3954-3963 2023.
Weiming Zhuang, Yonggang Wen, Xuesen Zhang, Xin Gan, Daiying Yin, Dongzhan Zhou, Shuai Zhang, and Shuai Yi. Performance optimization of federated person re-identification via benchmark analysis. In Proceedings of the 28th ACM International Conference on Multimedia, pp. 955 963, 2020.
A Appendix
A.1 Overview of Notation and Definitions
We give the overview of the notations and definitions used in the paper
Boldface letters to denote matrices and vectors.
refers to dlement-wise multiplication.
refers to the standard matrix multiplication.
We define TL; the tt transformer layer.
E denotes the expectation operator.
The term & represents the prior probability of class e occurring at client
Additionally, sim(p.g) q) denotes the cosine similarity, while represents the indicator function.
Mf-LA, denotes the client k computing the CLS token prototypes in communication round t and the
input of layer
H denotes the classification head.
16

---

Published in Transactions on Machine Learning Research (03/2026)
II-II denotes Euclidean norm or 2-norm.
M denotes the set .M)
A (rR,rR+ (r ( 1)R 1} denotes the set of communication rounds in r-th update period of length R.
7kee denotes the number of samples corresponding to class in client k.
Ps and Pa denote shared prompts and class-specific prompts respectively.
m denotes the CCMP.
Ng denotes the total number of datapoints at client and denotes the number of clients.
A.2 Method Details: Algorithm
We briefly go over the prototype update and the CCMP computation equations. Client level prototype aggregation at communication round is given in the below Eq.
18
Σn clsi-1, Likk
Ak.c > O,
Mi.c
1i-1,k
(18)
0,
Ak.cc 0.
Server prototype aggregation during the warm-up phase is given by Eq. 19
1
1
(19)
Sao
RESD
Server aggregation the prototypes the end of r-th update period is Eq. 20
De EIEn ExEs, I(G Pi-IKD
De >
(20)
(0.
D.
Sever updating the prototypes based on the momentum given Eq. 21 De is 0, we set p =
p = 1i-1,-11--
+ (1- p) ff-1.r
(21)
The class prior for class at client I is computed in Eq.
22
N₂ 0
Iy.
(22)
No
The soft scores are computed based on similarity between the class prototypes and cls representations as in
Eq. 23
sim (, pf_ 86
80
111-1.k
eap
(23)
17

---

Published in Transactions on Machine Learning Research (03/2026)
The scores are converted to probabilities using Eq.
24
8 1.1
1
(24)
Σ
8
1,1-1,
The probabilities serve as weights class-specific prompts which produce the Class Contexualized Mixed Prompts (CCMP) as Eq. 25
m = Pc * 81
(25)
1, is the vector containing 1,1-1. for different values of c.
Algorithm 2: PEP-FedPT
Input: H, Pc μ,, Pretrained Vision Transformer w Wpre: Training data (æ, y) ~ D, Set of class labels C, Learning rate y, Number of local epochs E, Update period R, Total number of communcatiion rounds T, Total number of clients n, CCMP layer index
Output: 0 (H,Ps. Pc}. 14-11
Server samples the subset So [n]
pi-1 W WarmStartUpp(w pree So) Ve E
for round E [T] do
Server samples participating clients S; C. [n] for client &o E [St] do
L o 1i-1, LocalTrain (0', 11-1, k,,,,, Vc E
O++1 +- FeddAveraging((0 E S.) McMahan et al. 2017
if 1 mod R = then
r+ H
10
for CE C do
Pi-1r + (enregattCentrroodss(e--aaaaae A3) Eq. 20
11
12
M + U'plateCCentroidss(f--- Ma 1 [Eq 21
M pi-1.
14
Return 0, 11-1
15
Function WarmStartUp(w (t,
16
for client I in Sin do
17
obtain 1 [Eq.
10
18
Return 1L Vc E C
19
Server obtains 19-10 Eq. 19
20
Return 15-10
Function LocalTrain 11-1 1.,,,, D):
22
compute 1, Vc C' Eq. 10
I 0¹
24
for c E' do
25
m + PromptMizzngg (H, Ps, Pc, W.G.WWn Wpres 11)1 [Eq.
22
23
24
25
26
Define loss 1 = I(H, m,,,y)
27
02 02 17- Vle.
Return 0z. 16
18

---

Published in Transactions on Machine Learning Research (03/2026)
A.3 Experimental Setup
A.3.1 Details on Heterogeneity
We consider two different kinds of heterogeneily label imbalance and feature imbalance. In the label imbalance we again consider different settings, pathological and the Dirichlet based non-iid settings as shown in Figure
For pathological settings select few classes of data points for each client and allocate the data among those labels. For Dirichlet we allocate the data by drawing a sample from the Dirichlet distribution. We consider these settings using the CIFAR-100 and Tiny-lmageNet Datasets by distributing the data among the 100 and 200 clients respectively and sampling only 5 clients in each communication round. For Dirichlert settings the degree of non-iid is controlled by the parameter å and its denoted by Dir(6). The lower delta implies higher heterogeneity and higher value implies the lower heterogeneity. Throughout the work we consider the value of 0 to be 0.3.
Dom Domain-A
Domain-A
0
Client ID
Client ID
(a) Non-IID of Label Shift: Pathological
(b) Non-IID of Label Shift: Dirichlet
Figure 4: Comparison of Non-IID Label Shift due to Pathological setting and the Dirichlet setting
By feature imbalance, we mean clients are distributed with different domains. Tt be seen in the Figure. The DomainNet dataset can be viewed analogous to the one described in the Figure 5a In the Figure 5b the split shows the mix of feature and the label imbalance.
i
!
:
Domain-A
Domain-A Demain-B Domaii Domain-O
:
Domain-a -B Domsin-C Domain-D Domain-E
Clinn ID
Client ID
(a) Feature Imbalance
(b) Feature Imbalance along with label shift
Figure 5: Comparison Non-lID) Feature Shift
A.3.2 HyperParameter Details
We follow stochastic Gradient Descent with momentum (Deng et al. 2024) as the default optimizer with learning rate 0.1 with exponential decay and the momentum 0.9. For all the experiments we consider number of shared prompts (ns) to be 1, unless explicitly mentioned. We add the class specific prompts at the layers
19

---

Published in Transactions on Machine Learning Research (03/2026)
5, 6 and 7. We also set the gradient clipping to 10 following Acar et al.) For all our experiment consider number of shared prompts to 1 except the Tiny-ImageNel Dirichlet where we set it to 5. The CCMP is inserted at the layers 5, 6 and 7. We set the the temperature parameter т to 0.05 for all our experiments. We show the datasct-specific hyperparameters in the table.
Table Sppaset-specific hyperparameter settings
Datases
# Classes
# Clients
Classes / Client
Comm. Roands
Participation Rate
Local Epochs
Centrold Update Interval
CIFAR-100
100
100
(Pathological / Dir(0.3)
100
5%
Tiny- EnageNe
200
200
10 (Pathological / Dir(0.3))
100
2.5%
DomsinNet
10
60
30
10%
10
INaturalist
1203
500
A.4
Additional Experiments
A.4.1
Class-Level Differential Privacy via Laplace Mechanism
Table 8: Impact of class-level DP noise (e 0.2) mean accuracy across datasets.
Method
CIFAR-100 (Path)
CIFAR- 100 (Dir-0.3)
Tiny ImageNet (Path) Tiny Imageet (Dir-0.3)
With DP Noise
93. 03.23_0.07
86, 86.92=0.07
91.164a.1
82.924a.1
Without DP Noise
95.46=00116
88.75=0.25
91.52+a.11
83. 44.0.2
To prevent leakage of individual client privacy a federated learning setting through sharing of cls-tokens, we employ the most common Laplace mechanism as described in Dwork al. (2006) for class prototype everytime the client. shares it with the server. After client aggregates els-tokens and forms its corresponding class prototypes, estimate the sensitivity of each class for client k based on the maximum L1 deviation of its CLS- token representation from the corresponding class prototype, normalized by the number of samples Ne.:
2 max, ccs (c)
Mell1
Sc,
No
where c denotes the CLS token of the i-th sample belonging to class client k, .N.kk k is the total number of samples belonging to class for client k and Me is the prototype reprosenting class in the embedding space across all clients. To enforce differential privacy, Laplace noise is added to each class prototype 0 based on its sensitivity Se, and predefined privacy budget e:
0c.k + 0g + Laplace(O, Sc,,
This class-aware noise injection is performed during every client-server communication, for all the CCMP layers. As a result, individual class-level contributions are obfuscated, thereby enhancing privacy while preserving model performance under nou-IID data distributions. In the Table g we have shown the impact of dp noise on our overall accuracy. We have used = 0.2 for our experiment.
Thcoretical guarantees for differential privacy: The sensitivity function defined in A.4.1 is an upper bound of the true maximum possible difference of average calculated from neighboring datsets.
Consider two datasets D and D2, each containing an equal number N of [cls] token representations, that differ at exactly one data point: clsi E D, and cls₂ E D2. The sensitivity of the dataset average is given by
Af max cls; cls2l11
N
For any J
Af = max//clsı - 11 + 41 clsa|l1
marl/elsı
ull1 + 1| cls2|11
N < maxa; ||cls, #4/12
N <
N
20

---

Published in Transactions on Machine Learning Research (03/2026)
By taking a be the empirical average of all cls-token representations. We upper bound the sensitivity by Se,k- Following this, it is straightforward to establish privacy guaranties. We refer to Dwork & Roth (2014) for the formal proof. In particular, Theorem 3.6 shows that adding Laplace noise to the class prototype ensures (c, didffferrel privacy.
A.4.2 Impact of Class Priors
Table 9: Ablation on class priors for iNaturalist, DomainNet and ImageNe datasets. We report the Mean Accuracy (%)
Prompt
iNaturalist
DomainNet
TTiy--mmaeeN
Shared + CCMP Without CP
54.3840.06
86 86.3440.52
81 $1.08
Shared CCMP With CP
63.4841.10
89.1540.70
83.4440.02
The ablation results in the Table 9 highlight the effect of incorporating class priors into the Shared+ CCMP strategy. For both iNaturalist and DomainNet, adding class priors consistently improves mean accuracy compared to using - without priors, with gains of nearly 9% on iNaturalist and about 3% on DomainNet.
A.4.3 Impact of CCMP and Shared Prompts
Table 10: Ablation on prompts for iNaturalist, DomainNet and Tiny-mmageNet datasets.
Prompt
iNaturalist
DomainNet
Tiny- Fing--mageNet
Ouly Shared
52.22 +0.50
84.23++.r2 89.1540.70
79.02 83.4440.02
Shared + CCMP
63.4841.10
The ablation results in the Table T0 compare the effect of using only shared prompts versus combining them with CCMP. On iNaturalist, the mean accuracy improves from 52.22% to 63.48%, while on DomainNet, the performance rises from 84.23% to 89.15% when CCMP addod.
A.4.4 Impact of increasing shared prompts
Table 11: Impact of Accuracy on increasing the number of shared prompts with non-iid partitioning of (Or(0.3). Incrcasing ns results in minor improvements for CIFAR-100 and DomainNet
Dataset
ng =
ng =3 5
ng 10
CIFAR 100
88.75-0.25
89.65
90.53+0.59
DomainNet
89.15
89.29
90.22+0.33
In the Table 11 show the impact of varying the number of shared prompts. It can be observed that the impact is quite minimal.
A.4.5 On the Gain of CCMP
Introducing class-specific prompts increases the total parameter space compared to using a single global prompt. However, the performance gain achieved by our method is not solely due to this increased parameter count. To validate this, augment the FedVPT baseline by adding 50 and 100 prompts (matching the scale of our class prompts). The mean accuracy improves initially but quickly saturates, with only marginal gains between 50 and 100 prompts. We can see this in Table 12 we have used the CIFAR-100 datasel with pathological data partitioning for this experiment. This indicates that merely increasing the number of prompt tokens is not sufficient achieve better performance. Instead, our method's distinct soft mixing of class-specific prompts using global class prototypes and local client priors plays key role in boosting accuracy, demonstrating the effectiveness of our proposed personalized prompt tuning mechanism.
21

---

Published in Transactions on Machine Learning Research (03/2026)
Table 12:
Effect of increasing number of prompts in FedVPT baseline. The accuracy saturates despite increasing parameter space, indicating that gains from our method are not due higher parameter count.
Number of Prompts
Mcan Accuracy (%)
83.62
50
87.15
100
87.45
A.4.6 Varying the Location of CCMP Injection
We perform the analysis of our method PEP-FedPT. It can be that adding the CCMP prompts too early in the VIT is not beneficial as the cls token representations at the very early layers do not have better representations. Adding the prompts at later layers is also not beneficial, even tough the cls tokens have better representations, since the prompts inserted are not deep enough to learn useful ropresentations. The choice of using the three prompts is to be efficient and, at the same time, to provide a fair comparison with methods like SGPT (Deng et al.] 2024).
In the Table 13
Table 13: Impact of adding the proposed CCMP prompts at different layers of ViT on CIFAR-100.
Position of CCMP Mean Accuracy
1, 2, 3
90.05 95.4640.16
5, 6, 7
8, 9, 10
93.55+0.14
A.4.7 Heldout Evaluation DomainNet and iNaturalist
The comparison shows that methods like Fed-VPT-D and SGPT provide competitive results, especially on DomainNet. However, our method achieves the best overall performance, with 62.41% participating and 54.16% 16% lesting accuracy on iNaturalist, and 90.32% participating and 88.73% testing accuracy on DomainNet. This highlights its robustness across both datasets and evaluation settings. For iNaturalist about 916 clients participated in the training while 102 clients were held out. For DomainNet, clients, per domain, were held out, and 54 clients, 9 from each domain, participated in the training.
Table 14: Comparison of methods iNaturalist and DomainNet datasets with the held-out setting
INaturalist (t)
DomainNet (f)
Method
Partieipating Ace
Testing Ace
Participating Ace
Testing Acc
Head
48.8740.41
45.27
82.31
83.1942.00
Fed-VPT
51.6940.41
48 48.05
82.92
13.68
Fod-VPT-D
57. 57.13
53.20
87, 2..0
87. 17.54+1.52
P-PT
43.87 40.94
41.20
82.89 82.89
83.05
FedPR
38 88.6220.16
36.03
83.59
22.62211.7
SGPT
55.8210.12
53.81
86.55
87.27 69
pFedPG
55.61
NA
88.34
NA
(oP-(eo))((Orr)
62.41
54.1640.39
90.32
88.73
A.4.8 Alternative view of Worst Client Accuracy
Table. 15 reports worst-client accuracy on iNaturalist, measured as the lower-tail (5%, 10%, and 15%) percentiles of per-olient lest. accuracy, which reflects the performance of the most disadvantaged clients under data heterogeneity. Most baseline methods achieve near-zero accuracy at the 5% percentile, indicating limited robustnoss. In contrast, our method consistently attains the highest worst-client accuracy across all percentiles, with substantial margins over competing approaches. The gains the 5% percentile demonstrate clear improvement for the worst-performing clients, while the consistent advantages 10% and 15%
22

---

Published in Transactions on Machine Learning Research (03/2026)
Table 15: k% percentile accuracy on iNaturalist.
Method
5%
10%
15%
Head
12.50-0.32
20.5640.47
VPT
14.28
23.21+0.55
VPT-D
10.50
20.00+0.23
30.00+0.74
P-PT
10.93
16.6644.30
SGPT
00.30
18.91 18.41
27 27.27
FedPR
03.14
08. 08.6220.339
pFedPG
12.54400.41
FEP-FedPT(Ours)
20.00
33.00
41.10
percentiles indicate more equitable performance across the federation. Overall, the results highlight that the proposed approach improves robustness to client heterogeneity beyond average accuracy gains.
A.5 Visualization and Further Analysis
A.5.1 Accuracy Vs Communication Rounds
Accuracy Communication Rounds
300
Accuracy 1%6)
40
SGPT
FedVPT
FedVPT-D
FedPR
PEP-FedPT(Ours)
60
100
Communication Rounds
Figure Comparison the convergence different methods across the Communication rounds on the Tiny-Iman ImageNet dataset with pathological non-iid partitioning where each client only observes 10 classes.
The Figure G shows how accuracy is improving across the FL communication rounds across the various algorithms. is clearly evident that our proposed PEP-FedPT algorithm attains the best accuracy in fewer communication rounds compared to the other algorithms, thus minimizing the computation and communication costs.
A.5.2 t-SNE visualization of
we show the t-sne visualization of the trained class prompts on CIFAR-100 pathological 10-class setting and observe that each class prompt learns its own representation, which beneficial to making the final classification decision.
In the Figure 7
A.5.3 Visualization of the soft weights for CCMP
In the Figure g we plot the soft weights averaged across all the test examples belonging class and class across all the clients. It can be observed that on an average the soft scores gives high score for the relevant class prompts.
23

---

Published in Transactions on Machine Learning Research (03/2026)
t-SNE Plot
15
17.5
50
t-SNE component
25-
10.0
25
-35
-15
50
25
25
50
t-SNE component
Figure 7: (-SNE visualization of the learned class prompts, it be seen that each prompt learns its own representation implying no collapse of dimensions.
0.49
0.00
0.173
0.373
Salection Probabiltty
0.350
Salection Probabiiity
0.150
0.22
0.325
0.200
0.300
0.000
0.000
165
class irdex
class index
(a) class:0
(b) class:
Figure 8: soft weights Averaged over all the data points that belong to class 0 and 1. It shows that on Average the soft weights give more importance to the prompt corresponding to the true class.
24

---

Published in Transactions on Machine Learning Research (03/2026)
A.5.4 Visualization of Representations at different layers
In the Figures 9d and 9h the manifold post-training, which suggests that the utility of shared prompts is distinct from that of CCMP.
observe that in initial layers the represenlations are uniformly distributed across
SNE Vesaiiation of Claes ans
SUG Vieus tistton of Clvvc Repreeso the
1.90 Wesulization of Class Reprosentic
Oimonaton
Dimoraton
Dimonaton
(a) Layer-0
(b) Layer-7
(c) Layer-11
(d) t-SNE representations of cls tokens for different layers using DomainNet dataset. indicates that the initial layer representations are distributed uniformly over the manifold. The representation gets better once CCMP is incorporated the later layers.
Vecaiiiation Cluse Regreceebiln ann
1:SUP Visu Regireen
1:5NF Vecalination Closs Regureaeeeennn Regueeaenetrno
Drnercken
1
Dimumaion's
Dinumioo
(e) Layer-0
Layer-7
(g) Layer-11
(h) t-SNE representations of cls tokens for different layers using CIFAR-100 dataset, denoting that shared prompts at layer learn only common class representations, unlike CCMP introduced later in the model. The representations are obtained using the fine-tuned ViT-B/16.
Figure 9: Comparison of (-SNE representations across layers for DomainNet and CIFAR-100 datasets.
A.5.5 Robustness on varying the Dirichlet Concentration
Table 16: CIFAR results under different Dirichlet partitions.
Method
Dir(0.1)
Dir(O.5)
Mean Acc
Worst Acc
Mcan Acc
Worst Acc
Head
79.21
65.484n.02
79.95
72.24
VPT
84.13
71.80+0.0 01
84.97+0.00
75.63=0.01
VPT-D
87. 87.08 0.02
74.00
87.92
79.00.o.12
P-PT
79.16-0.24
67 67.2441.24
78.80
64.54-0.68
SGPT
85.36-0.01
73.00+0.16
86.04+0.02
73.49=0.01
FedPR
81.64-0.32
65.0841.80
82. 82.11+0.59
71.42=0.84
PFedPG
84.14=0.49
73.62
73.90
60.00
PEP-FedPT(Ours)
90.8540.13
84.28 11.1
87.75
79.5 10.16
25

---

Published in Transactions on Machine Learning Research (03/2026)
This table presents the performance of different methods the CIFAR dataset under two levels of label heterogeneity, modeled by Dirichlet partitions (Dir(0.1) and Dir(0.5) ). For each method, report the average (avg) and worst-case (worst) accuracy across clients, along with the standard deviation. Our proposed method consistently achieves the highest or comparable average and best Worst Acc accuracy, indicating superior robustness and elfectiveness compared baseline methods.It can be seen that performance gain is higher when data beterogeneity more. Our method becomes FedVPT the iid setup as the scores used mix class prompts will become identical.
A.5.6 Sensitivity to Temperature
CPA1:200
Accuracy (%)
62
10
10
20
10
15-0
15"
101
Temperature (T)
Figure 10: Sensitivity of the Average Accuracy to the
Figure. [0 shows the effect of the temperature parameter, т, on model accuracy for two datasets: C1FAR-100 and
Accuracy is measured for different values of ranging from 10-5 to 100. For both datasets, increasing initially improves accuracy, reaching peak at the same 7), after which further increases in lead to drop in performance. This indicates that moderate temperature helps optimize model performance, while very small or very large temperatures can reduce accuracy.
A.5.7 Impact of temperature on optimal location for CCMP injection
Table 17: Impact of temperature on optimal location for CCMP injection: Accuracies shown in CIFAR- (0.3).3) setting
Layers
т 0.0001
== 0.05
T 100
1, 2, 3
86.18
86.70+0.04
86.32_0.16
5, 6, 7
86.2440.14
88.75
87.90
9, 10,
87.7740.01
85.68
87.35=0.18
In table 17 we show how the accuracy varies on CIFAR-100 under dirichlet setting. At low temperature setting (T 0.0001), the similarity scores computed from the cls token representations receive significantly higher relative weight. In this regime, performance is primarily driven by the quality of these scores rather than the depth at which CCMP is injected. Since cls token representations at later layers (c.g., layers 9-11) are more expressive, they provide more accurate similarity estimates, resulting in improved accuracy. As the temperature increases (7 = 0.05 and - = 100), the influence of the similarity scores is reduced, and the depth of prompt insertion becomes a critical factor. At higher lemperatures, sufficient insertion depth is required enable the prompts to learn meaningful representations, and shallow insertion no longer adequate achieve strong performance.
A.5.8 Personaization and generalization Trade-off
Figure. [ visualizes the tradeoff between personalized performance and generalized performance across four datasets: CIFAR-100, Tiny-ema iNaturalist, and DomainNet. Fach point corresponds to a method, with
26

---

Published in Transactions on Machine Learning Research (03/2026)
CIFAR: 100
Tiny-Imagetiet
INaturalist
DomainNet
Generalizen Accuracy (N)
giorio 08
& SERIT g-
453
641
6007
715 D3. 81 n.o. 175
20.5 Personalized Accuracy (56)
Personalized Accuracy (1)
Personalized Accuracy (%)
Personalized Accuracy 196
Figure 11: Personalization and Generalization Trade-off of different methods
the x-axis representing personalized accuracy (participating accuracy on clients) and the y-axis representing generalized accuracy (tosting accuracy on held-out data). Methods closer to the top-right corner achiove a better balance between personalization and generalization. Across all datasets, PEP-FedPT consistently occupies dominant region, achieving simultancously higher personalized and generalized accuracy compared to prior methods. In contrast, several baselines improve personalization the expense of generalization vice versa, highlighting an inherent tension between the two objectives. The results demonstrate that PEP-FedPT bridges this tradeoff.
A.5.9 Evolution of Class Prompts across rounds
round 0
round 10
round 20
100
tsne, component
tame cowponent 2
Lsne comporoooooooooo component ornent 2
-300
component_1
component
component
Figure 12: Evolution of Class Prompts across Rounds on CIFAR-100 Datasct
The Figure. 12 illustrates the evolution of class prompts across training rounds on the CIFAR-100 dataset. the initial stage (round 0), all class prompts are tightly clustered, indicating that they start from a nearly idaltical or uninformative initialization. out, reflecting the model's ability to differentiate between classes. By later rounds (round 20), the prompts form well-soparated representations, suggesting that each class prompt has adapted to capture class-spocific semantic information.
training progresses (round 10), the prompts gradually spread
A.6
Theoretical Details
A.6.1 CCMP minimizer of quadratic upper bound around class prompts
For clarity and completeness, we restate the relevant proposition from the main paper.
Proposition 3. If the assumptions to 3 hold, we show that f can be upper bounded as i - 5% (I.(P... Panz [[m(k) Pes +C and it is minimized at m(k) = -! &eper Wk [n]. which is equivalent to the (CCMP) described in sec.4.2 as >> 1. Bmax max4ell01 Bi, C is constant which depends on P. This vanishes when Pe. = = Vi [[C|] which makes L a tight upper bound of
Proof. We begin by applying the smoothness assumption on the loss function f for each class i. By Assumption 2, & is 3-smooth, which implies that for prompts m(k) P and Pa; E P. for k [n] and
27

---

Published in Transactions on Machine Learning Research (03/2026)
E [[C]] we have
(m(())) S (Pa.) + ( - Pa.) + Br 2
[|m(k)) Pe:l??
(26)
We define
fmax
max S rE[[C]]
(27)
which gives
(m())) & (P(a) + (Pa,) (m((k) - Pe,)) + maax |m()) Po.)
(28)
2
Now we know that P is compact, let the diameter be D = sup ||x yll which gives us 2,yEP
Pi P2ll D for all Pr. P2 E P
(29)
[lmmk)) Peill S D.
(30)
Since by Assumption (p)) is smooth, we have for x. E P
[V
vely)ll S S yll
(31)
I 8max yll
From equation 27
(32)
If 0294 ||x yll 8, for 6,max have
1V - v < (x)x - (v) S e, E [d]
(33)
Vl(x), is a continuous mapping of compact metric space p into metric space IR Let Bi, be the diameter of (P()),, j E [d]. then
(x)) is compact Vj E [d].
d Σ Biy
(p)())) S B
(34)
Using Cauchy-Schward inequality and from 34 30 we have
(p. (m(())
(m(k) - Pc.) $ DB;
(35)
From 28
(m(k)) < (p..) C* + Bunx
[m(s) - Pe..l.
(36)
2
where Ö DB, The global loss of the clients given by
G
81
& (m(k))
(37)
< i
Σ 8/
(p..) +
3max
||(kk) Pes Pctl13) +C,
(38)
2
28

---

Published in Transactions on Machine Learning Research (03/2026)
where Č= n -=-
1
Cks
(39)
which proves the first part of our main proposition 1 in the paper.
If Pc; P we have tight upper bound L = n E E ! 8% (Po..
-
|m(k) Pc. pa.11)] because
V4 (Pa.) vanishes, according to Assumption 3
We interested in finding the optimal client prompts m(k) for each client k.
IC1 max (m(k)
OLx @m(k)
Σ
B
Pa)
(40)
n
|0
Bmax N
m(k)
since
8
(41)
aLx
[C|
Setting
O, we have m(A)
8 Pe.
(42)
dm(k
which gives the second part of our proposition [ and completes our proof.
A.6.2 CCMP as MMSE estimator of the true class prompt
Here we clarify the details of the distribution ps (cls;-1, p) pc especially ps(cls;-1)- We define as the density induced by the deterministic transformation of the data distribution through the preceding network layers.
Let T X - Rd represent the composite non- linear mapping performed by the first (- layers of the ViT, such that for any input x, the representation is given by cls = (x). We can equip a probabitty space on the input as (X,F,)) P), where FF is a o-algebra (typically the Borel o-algebra) and the probability measure on X. We also equip the Measurable space of the Rdd as (RR)).
If B mess-reee subset of Rdd, Pr(B) given by the pushforward measure of on B by Ti-1 which is (B)). We can always do this the map Ti-1 is continous and hence mcasurable.
Consequently, P& (cls:-1) is the probability density induced by this distribution derived via P(x) under the mapping Ti-1- The posterior probability Pa (p Pc, cls; only implies that once we observe cls; the probability that it belongs to a class i.
This how we model the joint distribution ps (clss-1,,))
We assume joint data distribution P(x, y) over the input space X and the set of class labels y with marginal P(x). To formalize the notions, define the input space probability space
If the cls lokens and the dlass-ppccffic prompls inpul of loger has the joint densily given by pk (clss-11.)) as in Eq.. 17 then the CCMP prompt for client k, ml)((k) obtained in Eqq. 6 Minimum Mean Squared Estimator (MMSE) of the true class prompt.
Proposition 2
Proof. Consider the following mean- mean-squared error
J(p) Ellp $1
(43)
29

---

Published in Transactions on Machine Learning Research (03/2026)
where the expectation is taken across the joint distribution of (k clst-1). The that's minimizes the J(p) is the MMSE estimator, and p is our true class prompt. We have the following
(()) Elpp F1
Epp [| -
Elpp E fbl
+ 2P Elpiclat-1) p)
pp E|plcls:-11| E|ipppps 611
The equality is obtained as the cross term is zero i. we have (p Elpless p)] = 0. It follows by using the iterated expectation as shown below.
Ep - cls; P) = [p - - Eippcls- -
(44)
[ (p E/plcls.- ptppp///llllss..ppp.......s Els/css - -)]]
(45)
0.
(46)
We now have
J(P) = Elpp - E[p|cls/-]]] B -p1
(47)
From the above Eq. 47 it can be readily seen that (()) is minimized by setting the value of E[p]cls
c
E Elplcls- = Σ p(p Pe.. |cls--) Po...
(48)
m=1
From Eq. 16 we can rewrite the above equation
Ep
iil-1 Pem
Pc S
From Eq. 25 we conclude that E/p|cls--1] = m--1-
A.6.3 Convergence
We assume the following assumptions on the loss functions based on (Karimireddy et al. 2020 Acar et al.
A 4. The loss functions fü Lipschiltz smooth, i.e., (0,) fx2)))) S <3110 02
A 5. Nhd E [v)1 s G + ( where f(0) = E (0). This is referred to bounded
gradient dissimilarity assumption,
A 6. let E| VI(e, (z.v)) (e)]l < o², for all lo and 0. Here I(O, (x, 3)) is loss evaluated on the sample (x, y) and (0) is expectation across the samples drawn from Dj. This is a bounded variance assumption.
In the above assumptions, the parameter O denotes the trainable, shared, and class-specific prompts along with the classification head parameters.
The entire computation of the soft scores S,1-1,8 for the client k based on cls, can viewed as part of the model architocture itself (F and encapsulated inside the client's loss function.
We then have the following proposition.
30

---

Published in Transactions on Machine Learning Research (03/2026)
Proposition 3. Theorem V of Karimireddy et ol. (020)) in Appendix D.2: let 0 arg min f(0), the 0
global step-size be og and the local slep-size be 04. When the update period R. is very large or T >> 1, the
PEP-FedPT algorithm will have contracting gradients. If Initial model is 0", = f(0)) f(0)) and for
constant M, then in T rounds, the model 0T satisfies [ - O (BMVE (FG)2)))
VTLS (T+1)973 + BBB
The above proposition states that the PEP-FedPT algorithm requires O ( communication rounds to make
the average gradients of the global model smaller, i.e,, /[IlV](0")]" we only employ global prompts and parameters for the training.
The result plug and play because
Training Lass vs. Communication Rounds
RoJPA
P-PT
SCPT
169
Head TeduPT
Cross-Entropy Loss (log scale)
Communication Rounds
Figure 13: Comparison of training loss of various algorithms CIFAR-100 dataset
13 illustrates the training loss (cross-entropy, log scale) versus communication rounds for various baselines. We observe that P-PT struggles to converge, while FedPR and Head show limited improvements with early plateauing. Methods such as SGPT and FedVPT achieve more stable convergence, and FedVPT-D further reduces the loss by incorporating additional regularization. In contrast, our proposed PEPFedPT consistently outperforms all baselines, achieving both faster convergence and the lowest final loss. Our theory, which minimizes the quadratic upper bound at convergence is expoected to have lower training loss. This empirical trend aligns with our predictions, thereby ensuring improved stability and convergence in practice.
Figure
A.6.4 Analysis of CCMP when the scores the function of data
We show that CCMP minimizes the quadratic upper bound on the loss even when the scores are functions of both the data and class-priors. We show an upper bound on the loss for a given input x. This helps us remove assumptions related to the temperature parameter and thus gives more general theoretical analysis. Our notations are restated accordingly. We denote the estimate of class prompt for class any round be Let m(k, x) denote the prompt used at client & for data, point x. 1 and let the total number of clients be n, and 8% denote the empirical probability that a data point at client & belongs to class We assume that the joint density of the data in client p& (x, y) is modeled pk (x, y)
Pe. We denote the class prompts by Pa = [Pe, Pea ...,, Pejcy P.jc
(x))px(y)x)) the posterior (x(r)x) assumed to be given by the scores in Eq. 15 which we denote by s4.4xx and model Ps(y = ixx) by defining pæ(y i|x) == Sk,ix- Let P be the set of all possible prompts across all the clients, such that m(k,x) E P, Vk E 21. the client k is denoted by the (m4k, x, 3)). Note the expectation is over the pk (x, y). The goal is to estimate m (k; x) as a function of class prompts (Pey, .eg...,,eej The global loss across all clients can be computed as ʃ 1 Σ E112 (m(k, x), w)]]
Vx. The overall loss of
We now state the following assumptions:
for notation convenience, we drop the layer index from (k,
31

---

Published in Transactions on Machine Learning Research (03/2026)
A 7. P is compact subset of RR, where d is the token dimension.
A 8. (0(0, x, is 3 smooth in argument O with parameter B Vy E [|C'],,x, Vk E [n].
Proposition 4. If (2 (x, satisfies the above assumptions 7 to B we show that over-
all loss function f = - Blfx(( x), (x, v))] can be upper bounded as f I =
Wk E [n]. which is equivalent to the (CCMP) described in sec.4.2 C is constant which depends P. The E is over the distribution of the data x. Here we defined m(k) = m(k, x), ()) = (, = i) and s = Skiix
1 E sk ((()),)) Smax |lm(k) -P, Pe. + and it is minimized at m(k) = Ee
|C| c=1 sip
Proof. we expand the clients loss B[(x(m(kk x),x,)x,...))) x, v)] as below
E (m (k. x) x. y) [4 (m x). x. y)]lx
(49)
E
Lu (m (k, x), i)px(y i|x]]
(50)
E
Eg (m x), x, [)s
(51)
(k))))s
(52)
In the last step we use the definitions in the proposition i.e, m(k) = m(k, x)) (m(k)) = (,(m(k,x),,,11 y = and st Sk,ix
If the Lipschitz smooth and compactness: assumptions hold, then by following similar arguments from 26 till 38 we will have the global loss of clients given by,
C
$ t (m{)))
(53)
C
e (Pc.) + 8 |lm(k) - -pa Pc.11) +C,
<
(54)
T2
which proves the first part.
We are interested in finding the optimal client prompts m(k) for each client k and for each data point x. This is obtained by optimizing the argument inside the expectation which is Σ CC 5 (p.. + & ||lmm))))))))))))@>) Pe. -c..1²) ] with respect to mk).
2 CC
$ (et
[m(k) P.. po.
$
(m(k) (k) - Pc.)
(55)
dm(k)
i=1
|C|
3 N
m mkk)
S/Pr,
since
8
(56)
0Lx 0, we have Setting dm(k)
m(k)
8
P
(57)
im
which gives the second part of our proposition I and completes our proof.
32


Reviews:
[{"review_id": "P3tItsYGCb", "reviewer": "Reviewer_72NJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "## Summary\n\nThis paper studies **federated visual prompt tuning (VPT)** for pre-trained **Vision Transformers (ViTs)** under **client heterogeneity**, focusing on the tension between (i) *global prompt tuning* (better generalization but weaker personalization) and (ii) *personalized prompt tuning* (strong local fit but poor generalization / scalability).\n\nThe proposed method, **PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning)**, introduces:\n\n* **Shared prompts** ($P_S$) that are inserted at the input and are globally shared.\n* **Class-specific prompts** ($P_C = [p_{c_1},\\dots,p_{c_{|C|}}]$) that are also globally shared.\n* A **Class-Contextualized Mixed Prompt (CCMP)** (m) constructed **per client / per sample** by mixing the class prompts using **soft weights** derived from:\n\n  * similarity between a sample’s **cls token** representation at a chosen layer and **global class prototypes** aggregated on the server, and\n  * the client’s **local class prior**.\n\nThe paper provides:\n\n* A method to compute and aggregate **class prototypes** across clients over rounds (with a momentum-style update),\n* A **theory section** arguing CCMP is optimal under certain assumptions (quadratic upper bound minimization and an MMSE-style argument),\n* Extensive experiments on **CIFAR-100 / Tiny-ImageNet** (label heterogeneity) and **DomainNet / iNaturalist** (feature/domain heterogeneity), including a **held-out client evaluation**, ablations, and a computation/communication comparison.\n\n## Key strengths\n\n* **Clear and fairly intuitive problem setup** and notation in the preliminaries and method sections (overall, the formulation is easy to follow).\n* **Strong empirical performance** across multiple datasets and heterogeneity settings (Tables 1–3; Appendix held-out table), with consistent gains in “mean accuracy” metrics.\n* **Figure 2** is a genuinely interesting diagnostic: it highlights how predictive power varies across ViT layers, motivating where to insert CCMP.\n* Helpful ablations on:\n\n  * **Shared vs Shared+CCMP** (Table 4; Appendix Table 9),\n  * **Impact of class priors** (Table 5; Appendix Table 8),\n  * **Layer placement of CCMP** (Appendix Table 12).\n* Includes a communication/computation discussion (Table 6) and additional appendix analyses.\n\n## Key weaknesses / concerns\n\n* **Evaluation metric consistency issues** (notably for iNaturalist “Worst Acc”) and **unclear fairness/interpretation** of held-out evaluations for baselines like pFedPG.\n* **Theory–practice mismatch around the temperature ($\\tau$)**: theory relies on ($\\tau \\gg 1$) but experiments fix ($\\tau = 0.05$); there is no sensitivity analysis.\n* **Heterogeneity robustness is under-explored**: the Dirichlet label-imbalance setting is shown for essentially **one** concentration value (0.3), despite the paper positioning itself as addressing heterogeneity broadly.\n* **Clarity/notation problems**: the paper uses **($\\delta$)** for both class priors and (in the appendix) the Dirichlet parameter, and there are places where the prompt-injection description/pseudocode is hard to implement directly.", "claims_and_evidence_yn": "No", "claims_explanation": "The empirical evidence for “PEP-FedPT performs strongly on the reported benchmarks” is largely supported by the reported tables/plots. However, several claims and supporting evidence are **not yet consistent/clear enough** for me to answer “Yes” overall:\n\n1. **Metric definition inconsistency (high impact on evidence clarity).**\n   In the main paper, “Worst Local Accuracy” is defined as the worst-performing client (Evaluation Methodology). But in the appendix, the paper states that for **iNaturalist** the “Worst Acc” is instead **the minimum accuracy over the final 250 rounds**, and also notes that “worst client accuracy is zero for all methods” due to participation issues. This makes the meaning of the iNaturalist “Worst Acc” column unclear and potentially incomparable to other datasets/methods.\n\n2. **Held-out evaluation vs. baseline applicability (fairness/interpretation).**\n   Table 3 shows extremely low held-out “Testing Acc” for **pFedPG** (e.g., single-digit accuracy), which is consistent with the *expected behavior* of certain personalized methods under a “new clients, no adaptation” protocol—but then the comparison needs to be framed explicitly as such, and/or baselines should be given a fair adaptation strategy. Right now the evidence does not clearly separate “method is bad” from “method is not designed for this held-out protocol”.\n\n3. **Theory–experiment mismatch for ($\\tau$).**\n   The theoretical claims tie the CCMP form to regimes like ($\\tau \\gg 1$), while experiments fix ($\\tau = 0.05$). Without sensitivity analysis or an argument that the theoretical form is robust to very small ($\\tau$), the theoretical evidence is not convincingly connected to the empirical configuration.\n\n4. **Some quantitative statements are slightly inaccurate.**\n   Example: the text reports an improvement over pFedPG on CIFAR-100 pathological of 2.57%, but Table 1’s mean accuracies imply 95.46 − 92.96 = 2.50%. This is minor but reinforces the need for careful verification.\n\nGiven these issues, I believe the core idea is promising and the empirical results look strong, but the **evidence is not yet fully accurate/convincing/clear**.", "audience_yn": "Yes", "audience_explanation": "Federated adaptation of large foundation models (here, ViTs) using **parameter-efficient tuning** is an active and practically important topic. This submission addresses a real and well-known tension in FL (**generalization vs personalization under heterogeneity**) and proposes a mechanism (prototype-guided prompt mixing) that is simple enough to be broadly useful if clarified and validated. The combination of:\n\n* strong empirical results,\n* a prototype-based mixing strategy,\n* and a focus on prompt tuning in FL\n \n should interest readers working on FL, PEFT/prompting, and robust generalization.", "requested_changes": "## Major changes\n\n### 1) Clarify the held-out evaluation protocol and baseline fairness (especially pFedPG)\n\n* Clearly specify: for “Testing Acc” on held-out clients, **are methods allowed any adaptation** (e.g., local fine-tuning of head/prompts, few-shot tuning, or any client-side steps), or is it **zero-adaptation**?\n* If it is a **zero-adaptation** protocol, explicitly state that and discuss which baselines are applicable under this setting.\n* The extremely low pFedPG held-out accuracies suggest either (a) the baseline is evaluated in a setting it is not designed for, or (b) there is an implementation/config mismatch. Either way, this requires explicit clarification and discussion.\n\n###  2) Fix metric definitions and ensure consistent reporting (esp. iNaturalist “Worst Acc”)\n\n* Align the definition of “Worst Acc” across datasets, or rename metrics so they are unambiguous.\n* If iNaturalist “Worst Acc” is **min accuracy over the final 250 rounds** (as stated in the appendix), then:\n\n  * label it accordingly in the main paper and Table 2 caption, and\n  * ideally report an additional robustness metric that matches the “worst client” concept (e.g., worst over participating clients, or a lower quantile like 5th percentile).\n\n### 3) Provide a sensitivity analysis on the temperature parameter ($\\tau$)\n\n* The paper uses ($\\tau = 0.05$) for all experiments but theoretical statements relate CCMP properties to ($\\tau \\gg 1$).\n* Please add an ablation varying ($\\tau$) (e.g., log-scale sweep) and show:\n\n  * performance stability,\n  * behavior of the soft weights (entropy/sharpness),\n  * and whether the “best” layer choices change.\n* Alternatively (or additionally), revise the theoretical discussion to match the regime used in practice.\n\n### 4) Vary the Dirichlet heterogeneity parameter (not only 0.3)\n\n* The Dirichlet label-heterogeneity results use essentially one concentration value (0.3).\n* Since the paper positions itself as addressing heterogeneity, please include results for multiple concentration values (e.g., ($\\alpha \\in {0.1,0.3,0.5,1.0}$) or similar) to demonstrate robustness of performance and the personalization/generalization tradeoff.\n\n## Minor changes\n\n### 1) Add a 2D personalization–generalization tradeoff plot\n\nA 2D plot could provide a clearer picture of the tradeoff across methods. For example:\n\n* x-axis: personalization,\n* y-axis: generalization,\n* each method as a point (or Pareto front).\n  This would make the personalization/generalization behavior immediately interpretable and may further highlight the proposed method’s strengths.\n\n### 2) Make the early overview more intuitive / higher-level\n\n* The introduction currently has a text description, but it would benefit from a more “first-glance” explanation of the design choices and motivation:\n\n  1. why cls-token prototypes,\n  2. why combine with class priors,\n  3. why insert at intermediate layers (and how layer choice was decided).\n* Consider adding a simplified high-level summary or a short step-by-step algorithmic summary early (before deep notation), so readers can quickly form the mental model.\n\n### 3) Expand related work with Federated Domain Generalization\n\nSince the paper frames itself around generalization under distribution shift/heterogeneity, it would be beneficial to include a small subsection on **Federated Domain Generalization (FDG)** beyond foundation model prompt tuning, e.g.:\n\n* “FedSR: A simple and effective domain generalization method for federated learning” (NeurIPS 2021)\n* “Federated domain generalization with generalization adjustment” (CVPR 2023)\n* “Federated unsupervised domain generalization using global and local alignment of gradients” (AAAI 2025)\n\n  This would better position the contribution in the broader literature and make the paper more self-contained.", "broader_impact_concerns": "N/A", "additional_comments": "Please refer to requested changes.", "extra_scores": {}, "extra_text": {}}, {"review_id": "rrOUY1lOjC", "reviewer": "Reviewer_rvUd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes PEP-FedPT (Prompt Estimation from Prototypes for Federated Prompt Tuning), a novel federated learning framework designed for fine-tuning Vision Transformers (ViTs) in heterogeneous data settings. The key contributions are:\n1. Unified Prompt Framework: Introduces a combination of shared prompts (global) and class-contextualized mixed prompts (CCMP) to balance generalization and personalization without storing client-specific parameters.\n2. CCMP Mechanism: Dynamically mixes class-specific prompts using soft weights derived from global class prototypes and local class priors, enabling per-sample personalization using only globally shared information.\n3. Theoretical Grounding: Shows that CCMP minimizes a quadratic upper bound on loss and is optimal in the minimum mean squared error (MMSE) sense.\n4. Strong Empirical Results: Demonstrates superior performance across multiple datasets (CIFAR-100, TinyImageNet, DomainNet, iNaturalist) under label and feature heterogeneity, outperforming state-of-the-art baselines in both participating and held-out client evaluations.\n5. Efficiency: Achieves high accuracy with significantly reduced communication and computation costs compared to existing methods.\n\nKey Strengths:\n1. Effectively addresses the personalization-generalization trade-off in federated prompt tuning.\n2. Lightweight and scalable, with minimal added computational overhead.\n3. Well-supported by both theoretical analysis and extensive experiments.\n\nKey Weaknesses:\n1. Relies on labeled data for estimating class priors and prototypes, limiting applicability to semi-supervised or unsupervised FL settings.\n2. Privacy guarantees are not formally proven; relies on empirical measures like differential privacy.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Comprehensive experiments across four diverse datasets under varying heterogeneity settings (pathological, Dirichlet, feature imbalance).\nClear ablation studies that validate the contributions of shared prompts, CCMP, and class priors.\nTheoretical proofs provided in the appendix, linking CCMP to loss minimization and MMSE optimality.\nComparative analysis against strong baselines, including personalized and global FL methods, with consistent outperformance.\nEfficiency analysis showing reduced communication and computation costs.", "audience_yn": "Yes", "audience_explanation": "The paper sits at the intersection of several active research areas: federated learning, vision transformers, parameter-efficient tuning, and personalized learning. The proposed method addresses a critical challenge in FL—balancing personalization and generalization—while being computationally efficient. Researchers and practitioners in FL, computer vision, and efficient deep learning would find the approach novel and practically relevant.", "requested_changes": "1. Clarify Privacy Implications: While differential privacy (DP) is briefly discussed in Appendix A.4.1, the privacy-risk trade-off should be more explicitly addressed in the main paper, especially since prototypes are shared.\n2. Expand on Limitations: The limitation regarding reliance on labeled data should be elaborated, suggesting potential extensions to semi-supervised or self-supervised FL as future work.\n3. Visualization Enhancements: More figure could be better annotated to clearly show how class prompts and representations evolve across layers.\n4. Reproducibility: Provide clearer hyperparameter settings and potentially release code upon publication.\n5. The article mentions the usage of prompt in vision tasks, which can refer to the article \" How to Design or Learn Prompt for Domain Adaptation?\". This article proposes image encoder tuning, named Im-Tuning, for more separable image features.", "broader_impact_concerns": "No major ethical concerns are raised by this work. The method is designed for efficient and privacy-aware federated learning, which aligns with ethical ML practices. However, as with any FL method, there is a risk of unintended bias propagation if data heterogeneity reflects societal biases. The authors should consider adding a Broader Impact Statement discussing:\n1. Potential biases in class prototypes if client data distributions are skewed.\n2. The importance of ensuring fairness across clients in personalized FL systems.\n3. The environmental impact of reduced communication rounds as a positive contribution.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "22I7ym5i1E", "reviewer": "Reviewer_AyXx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an efficient training strategy for the ViT architecture in an FL setting with heterogeneous clients. Here, the authors adapt the visual prompt tuning method to the FL setting by adding class-specific prompts as training parameters. The authors mix class-specific prompts for each client based on their closeness to a global class prototype, yielding class contextualized mixed prompts. The authors show that their mixing strategy achieves the optimal error under the assumption on the type of client heterogeneity (Proposition 1). Additionally, under the assumption of a \"well-specified\" model (Eq. (16)), the authors show that CCMP prompts are optimal estimators of class prompts (Proposition 2). The authors show an advantage of their method in several class- and domain-heterogeneous settings.\n\n**Strengths**\n1. Motivation is clear.\n2. The authors provide a theoretical justification for some elements of the method.\n3. Strong performance in the proposed evaluation.\n\n**Weaknesses**\n1. I find Section 4.1 quite confusing.\n2. I think the authors do not sufficiently discuss the limitations of their approach.", "claims_and_evidence_yn": "Yes", "claims_explanation": "I think the evidence is sufficient. However, I believe the limitations and implicit assumptions of the proposed approach should be further discussed.\n\n1. The method implicitly assumes that it deals only with clients that differ only in the distribution of classes across them. It can be seen from the construction in Section 4.1, where the only way of model personalization is the mixing of class prompts. I believe this assumption should be made explicit and discussed as one of the method's limitations.\n2. The assumptions in Section 4.3 lack a semantical explanation. In particular, Assumption A 2 in Section 4.3.1 again seems to imply a specific form of clients' heterogeneity. Similarly, Eq. (16) looks quite specific and implicitly assumes a lot about the outcome of optimization. I believe these assumptions should be discussed more thoroughly.", "audience_yn": "Yes", "audience_explanation": "I think this paper could be interesting to people working on applying ML in the medical domain, since here FL seems like a promising technique for private model optimization.", "requested_changes": "1. Please make Section 4.1 clearer. In particular, explain for which layers we need to add CCMP and what happens before and after this addition.\n2. Please clarify the scope and limitations of your work regarding client heterogeneity.\n3. Please discuss the assumptions in Section 4.3.", "broader_impact_concerns": "I don't see a need for a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)