Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation
Authors: Xinghao Wu
Xuefeng Liu
Jianwei Niu
Guogang Zhu
Mingjia Shi
Shaojie Tang
Jing Yuan
Venue: NeurIPS 2025 poster ()
Keywords: Personalized Federated Learning
Feature-Classifier Mismatch
Prompt-Driven Feature Transformation
Primary area: general_machine_learning

Abstract: Federated Learning (FL) faces challenges due to data heterogeneity, which limits the global model’s performance across diverse client distributions. Personalized Federated Learning (PFL) addresses this by enabling each client to process an individual model adapted to its local distribution. Many existing methods assume that certain global model parameters are difficult to train effectively in a collaborative manner under heterogeneous data. Consequently, they localize or fine-tune these parameters to obtain personalized models. In this paper, we reveal that both the feature extractor and classifier of the global model are inherently strong, and the primary cause of its suboptimal performance is the mismatch between local features and the global classifier. Although existing methods alleviate this mismatch to some extent and improve performance, we find that they either (1) fail to fully resolve the mismatch while degrading the feature extractor, or (2) address the mismatch only post-training, allowing it to persist during training. This increases inter-client gradient divergence, hinders model aggregation, and ultimately leaves the feature extractor suboptimal for client data. To address this issue, we propose FedPFT, a novel framework that resolves the mismatch during training using personalized prompts. These prompts, along with local features, are processed by a shared self-attention-based transformation module, ensuring alignment with the global classifier. Additionally, this prompt-driven approach offers strong flexibility, enabling task-specific prompts to incorporate additional training objectives (\eg, contrastive learning) to further enhance the feature extractor. Extensive experiments show that FedPFT outperforms state-of-the-art methods by up to 5.07%, with further gains of up to 7.08% when collaborative contrastive learning is incorporated.

Tackling Feature-Classifir Mismatch in Federated Learning via Prompt-Driven Feature Transformation
Xinghao Wu'; Xuefeng Liu Jianwei Niu 1.2.5g Guogang Zhu',
Mingjia Shie, Shaojie Tang", Jing Yuan', State Key Laboratory Virtual Reality Technology and Systems, School of Computer Science and Engineering, Beihang University, Beijing, China
Hangzhou Innovation Institute of Beihang University, Zhejiang Key Laboratory of Industrial Big Data and Robot Intelligent Systems, Hangzhou, China
Center for AI Business Innovation, Department of Management Science and Systems, University at Buffalo, Buffalo, New York, USA
University of North Texas, Denton, Texas, USA Zhongguancun Laboratory, Beijing, China 6 Sichuan University, Sichuan, China
Abstract
Federated Learning (FL) faces challenges due to data heterogeneity, which limits the global model's performance across diverse client distributions. Personalized Federated Learning (PFL) addresses this by enabling each client possess individual model adapted to its local distribution. Many existing methods assume that certain global model parameters difficult train effectively in collabo- rative manner under heterogeneous data. Consequently, they localize or fine-tune these parameters obtain personalized models. In this paper, reveal that both the feature extractor and classifier the global model are inherently strong, and the primary cause of its suboptimal performance is the mismatch between local features and the global classifier. Although existing methods alleviate this mis- match to some extent and improve performance, find that they either (1) fail fully resolve the mismatch while degrading the feature extractor, or (2) address the mismatch only post-training allowing it to persist during training. This increases inter-client gradient divergence, hinders model aggregation, and ultimately leaves the feature extractor suboptimal for client data. To address this issue, we propose FedPFT, novel framework that resolves the mismatch during training using per- sonalized prompts. These prompts, along with local features, are processed by a shared self- lerf--attettonbbiie
transformation module, ensuring alignment with the global classifier. Additionally, this prompt-driven approach offers strong flexibility, enabling task-specific prompts to incorporate additional training objectives (e.g., contrastive learning) to further enhance the feature extractor. Extensive experi- ments show that FedPFT outperforms state-of-the-art methods by up to 5.07%, with further gains of up to 7.08% when collaborative contrastive learning is incorporated. The code is available at http:://ggiubb.omm//ggggoWWWuW
Introduction
Federated Learning (FL)
enables clients to collaboratively train global model without sharing their raw data. major challenge in FL is data heterogeneity, where data
34
65 24
21
23
(2)
"Email: wuxinghao@i buaa. edu.en "Coresponding author: niujianwei @ buan.edu.cn
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

CIFAR 10, a = 0.5
CIFAR 10,
1.0
Preteicee by Fl. algaribe
Retrained by client's data
Methods
Match Acc. Origin Ace. Probe Acc. Match Acc.
Origin Acc. Probe Acc.
FedAvg FedBABU
70.71%
60.05% 60.87%
71.36% 71.33%
67.34% 68.55%
60.90%% 62.96%
67.87% 68.27%
replace
71.86%
FedPer
70.23%
69.50%
70,20%%
66.42%
65.92%
66. 66%
FodBN
69.42%
66.19%
69.80%
15.34%
63.39%
65.61%
(a) Transfoom
Probe
FedCAC
70.26%
68.15%
70.52%
67.22%
65.77%
67.47%
Ondimt Ditt
RodAvg-FT
71.21%
70.63%
71.36%
66.84 86.84%
66.44%
67.87%
E 5
wio addressing mismatch
FeIBABU-FT
72.01%
71.69%
71.33%
68.97%
68.42%
68.27%
addressing mismatch
FedPIT 72.81%
72.66%
71.91%
69.26%
69. 00.0
69.38%
0,7
FedPFT+Con
76.71%
76.83%
76.80%
73.52%
73.
74,50%
100
200
300
400
500
Communication
Figure 1: Left: Match Acc. represents the accuracy after applying linear transformation to the features to adapt them to the classifier. Origin Acc. indicates the accuracy of the original model. Probe Acc. refers to the accuracy achieved by retraining the classifier with local data. All accuracies are obtained on the client testing data. The disparity between Origin Acc. and Match Acc. indicates the degree of mismatch (see Appendix A for formal definition). Hight-Top: (a), (b), and (c) are toy examples of the models used to calculate the three types of accuracy. Right-Bottom: Addressing mismatch (orange line) during training reduces inter-client gradient divergence compared to not addressing it (blue line).
distributions across clients are not independently and identically distributed (non-IIID. Under such conditions, a single global model often fails to generalize well across all clients, leading to suboptimal local performance.
To address this, Personalized Federated Learning (PFL) has been proposed, enabling each client to maintain an individual model tailored to its local data distribution. A predominant approach in PFL involves generating the personalized model by localizing or fine-tuning selected subset of the global model parameters, based on the premise that these parameters are inherently challenging to optimize effectively through collaborative training. For instance, FedPer [ personalizes the classifier, while FedBABU (37) fine-tunes it. FedBN [26] introduces personalized batch normalization (BN) layers, and methods like AlignFed 663 personalizes feature extractors. Additionally, advanced approaches have adopted more sophisticated techniques, including hypernetwork [3 reinforcement learning [43, or quantitative metrics [50]. to select parameters for personalization.
The above research in PFL reveals an intriguing phenomenon: regardless of which components of the global model are personalized (e.g.. fea- ture extractors, classifiers, or intermediate layers), various PFL methods tend to achieve comparable performance levels. This observation leads to a non-trivial hypothesis: each individual compo- nent within the global model may already be sufficiently optimized through conventional FL processes, whereas the suboptimal performance stems primarily from the mismatch between these well-trained modules.
Global Feature Extractor
Global Classifier Proxies
Local Class Features
Transformed Class Features
Personalized Linear Transformation Layers of Clients and 2
Clbeat
Client 2
(a) Mismatched
(b) Matched
Figure 2: A toy example illustrating the mis- match in FedAvg and how a linear transform ad- dresses it. (a) The local features from two clients have well-clustered structures, but due the ef- fects of non-lll data, their features mismatch with the global classifier proxies. (b) By ap- plying personalized transformations to the local features of each client, the features are aligned with the global classifier.
To validate this hypothesis, we conduct a linear transform experiment by inserting a simple lin- car layer between the FedAvg-trained global fea- ture extractor and its classifier, then retraining this layer on each client's 's local data. As reported in Fig. ("Macch Acc."), this adjustment greatly im- proves accuracy over the original global model ("Origin Acc."). This result reveals that FedAvg already learns strong feature extractor and classifier, but their direct combination leads to feature-classifier mismatch when applied to client-specific data. A toy example illustrating this mismatch is shown in Fig.

---

Notably, existing PFL methods that localize (lines 5-7) or fine-tune (lines 8-9) only part of the model parameters mitigate this mismatch, which explains their strong performance despite modifying small portion of the model. However, these personalization strategies come with important limitations. iindicate that localizing partial parameters degrades the quality
Our linear probe experiments [ of the global feature extractor (ie., lowering the "Probe Acc.", thereby lowering the model's upper performance bound.
In contrast, fine-tuning the classifier preserves the strong feature representations learned by the global feature extractor, leading higher "Match Acc.". However, it remains post-training adjustment. addressing mismatch only after the global model has been trained. Since mismatch persists throughout training, clients must further adjust the feature extractor align with the classifier in each round, exacerbating inter-client model divergence. This leads to detrimental interference during model aggregation, making the global feature extractor suboptimal for client data. As illustrated in Fig. (right-bottom), when the mismatch exists (blue line), the gradient divergence among client models is larger.
To resolve the mismatch problem throughout training, we introduce a novel PFL method, FedPFT. While inserting a personalized linear layer between the shared feature extractor and shared classifier is a straightforward approach, it can easily overfit the limited local training data due to the large number of personalized parameters, and insufficient for handling complex datasets (as demonstrated in the Appendix HH. To overcome these limitations. we draw inspiration from prompt technology [T] which uses prompts to guide model behavior. FedPFT incorporates personalized prompts with minimal trainable parameters and shared self ifff-attention-aasee feature transformation module (FTM). Both the prompts and local features are fed into the FTM, where features are transformed via the attention mechanism. In each round, FedPFT first trains the prompts to align local features with the global classifier. Subsequently, training the model parameters based on this alignment reduces the model divergence among clients and enhances the synergy between the feature extractor and classifier. The results Fig. demonstrate that FedPFT not only resolves the mismatch problem but also improves the feature extractor.
Another advantage of our designed prompt-driven FTM its strong flexibility across different tasks. It can leverage task-specific prompts to incorporate various tasks beneficial for client collaboration, such as contrastive learning 20 48 feature alignment [6] [4 etc. Taking contrastive learning as an example, as shown in Fig. FedPFT+COn further improves model performance by introducing collaborative contrastive learning through prompts.
Our main contributions can be summarized as follows:
We identify that the global model's ''s inadequate performance on client local data in non-llI) scenarios is primarily due to the mismatch between local features and the classifier. We show that the reason localizing or fine-tuning some parameters improves performance is that it indirectly alleviates this issue. This insight offers a new perspective for future PFL approaches to better address the non-lID problem.
We propose a new PFL framework that incorporates a prompt-driven feature transformation module align local features with the global classifier. This approach not only resolves the mismatch problem but also provides flexibility for incorporating various collaborative tasks to further enhance PFL performance.
Our experiments on multiple datasets and non-lID scenarios (including both label shift and feature shift) demonstrate the superiority of FedPFT, outperforming state-of-the-ard (SOTA) methods by up to 5.07%. When further incorporating contrastive learning task, this improvement can reach up to 7.08%.
2
Related Work
has emerged as an effective approach to addressing the challenges posed by non-IID data in FL. In recent years, numerous methodologies have been proposed in this domain 6 25 31
PFL
45
60
46 59
57
52 62
39 35 In this section, we focus on three categories of methods most relevant
our work.
Parameter decoupling methods aim to decouple a subset of parameters from the global model for personalization. Approaches such as FedPer [I. FedRep [4, and GPFL (58) focus personalizing
3

---

the classifier. In contrast, methods like LG-FedAvg 28 and AlignFed 64 advocate for the personal- ization of the feature extractor. Additionally. FedBN 26 and MTFL (36] propose personalizing batch normalization (BN) layers within the feature extractor. Techniques employing deep reinforcement learning 43 or hypernetworks [3 have been used to determine which specific layers to personal- ize. The recent FedCAC [50] and FedSelect 44 further advance this by enabling parameter-wise
selection. FedDecomp [49] extends this idea by decomposing each parameter into a personalized component and shared component, achieving finer-grained decoupling. These decoupling methods indirectly help alleviate the mismatch issue within the global model by allowing local parameter adjustments. However, these methods do not completely resolve the mismatch issue. Furthermore, excessive parameter personalization reduces information exchange across clients, which can degrade the overall quality of the feature extractor.
Classifier fine-tuning-based methods such as FedBABU 37 and FedETF 27 propose that when the global model is well-trained, fine-tuning the classifier alone can yield strong personalized performance. Consequently, these methods adopt a two-stage training strategy: in the first stage, they focus on training high-quality global model, and in the second stage, they fine-tune the classifier to obtain a personalized model. By improving the feature extractor in the first stage, these methods provide a stronger foundation for addressing the mismatch through classifier fine-tuning, leading to better personalized performance. However, since the mismatch persists throughout the global model training process, clients must continuously adjust the model to align the feature extractor and classifier in each round. This exacerbates inter-client model divergence, causing interference during aggregation and ultimately making the feature extractor suboptimal for client data.
Feature alignment methods have shown that non-Ill data can lead to discrepancies in feature distributions across clients, which in turn cause client drift and hinder model training. To mitigate this issue, existing approaches align features across clients, for example, through global prototype alignment [54] 61.53 or classifier calibration [32 7 56 37 thereby reducing inter-client feature divergence. However, these methods focus solely on inter-client alignment and overlook the fea- ture-classifier mismatch that occurs within each client, which we identify as a key bottleneck. In contrast, our approach directly addresses this mismatch by aligning local features with the global classifier, which also implicitly promotes inter-client feature alignment through shared semantic space.
Prompt-tuning-base based methods. Prompt tuning has been widely adopted in vision and language tasks as lightweight way to adapt pre-trained models [1] 29 16 30). Recent works such as PFedPG [55], SGPT [5), FedOTP (19) FedAPT (42), and PFedPT 18 extend this idea to PFL, where prompts are optimized to adapt a shared, well-trained backbone to local data. These works primarily focus on how to adapt a well-trained global model to fit local client distributions. In contrast, our approach addresses a fundamentally different problem: how to mitigate feature-classifer mismatch during the training of the backbone itself. Rather than applying prompts fine-tuning tool, we incorporate them as a mechanism to improve alignment and reduce gradient divergence throughout training. Our experiments (see Section 4.4b further demonstrate that post-training prompt tuning is insufficient to match the performance of our training-time alignment strategy.
3 Methodology
3.1 FedPFT Framework:
Architecture and Workflow Overview
As illustrated in Fig. Bia, the core of FedPFT is the introduction of a prompt-driven feature transfor- mation module (FTM) т; between the feature extractor d and sseeifici heads (e.g., classifier). This module transforms local features during training to better align them with their correspond- ing tasks. As shown in Fig. 3(),, prompts p and image features f are fed into the FTM, where self-attention operations transform, f to f', which is then used by task heads. Specifically, FedPFT employs classification prompts Pw,; to transform f, ensuring alignment with the classifier hrra-
Each training round of each client in FedPFT includes five key steps:
1. Client downloads the global models, including feature extractor Φ, FTM r, and classifier hig-
2. Client freezes the global feature extractor & and updates т; and prompts Px,i using the cross-entropy loss LcB to align local features with the frozen global classifier h...

---

Trainable modules
ggrregate
Models client
Frozen modules
Server
Models of client
Shared modules
Personalized modules
①
Download Φ, hx
A Upload $i Ti. Ax.i
Pact
Rx
Plac
Rygi
Add
Lce
LOE
Self-attention
Features different cluss
Layer
Global classifier proxies
Clignt
(a) The overview FedPFT in one communication round
(b) The feature transformation module (T)
Figure 3: Overview of FedPFT. (a) The training process of each client in one communication round.
(b) The feature transformation module in FedPFT.
3. Based on the alignment between the features and the classifier, the client freezes the prompts Px,i and updates Φi Ts, biç with LeE learn client local knowledge.
4. Client uploads {oi, Tr, he,) the server while retaining (px.s) locally.
5. The server aggregates the models uploaded by the clients.
3.2 Problem Formulation
In PFI., N clients train their personalized models wi, i E [N] under the coordination server, aiming for each w; to perform well on client data distribution D2 This objective can be formalized as N min(w::eiN NV E> (w,; D:), where L; represents the loss function of the i-th client.
In this work, we aim to enhance personalized models by addressing the mismatch problem between local features and the classifier in the global model. The training objective FedPFT thus formulated as:
min min ExLLi T, h,,, Pe.s; di) = Ea, LLCE(Ø, he, d:)]] @. (PRa)ee(N)
where & and he represent the global feature extractor and global classifier, respectively. is the newly introduced global FTM. This module, along with the personalized classification prompt Px.s. transforms local features to align with the global classifier. Lce denotes the cross-entropy loss for classification tasks. d, represents the dataset the client.
3.3 Feature Transformation Module (FTM)
In FedPFT, given sample zj d, processed by the feature extractor é, the extracted feature is fa Rm where is the feature dimension. A collection of n prompts is denoted {pr R' e N, N,15 S n). The FTM operates as follows:
[f P(f
(2)
where [ signifies stacking and concatenation along the sequence length dimension, yielding [f). p'] E ()+ Here, the f/ E R" represents the transformed feature, which subsequently used for downstream tasks. An illustration of the FTM is provided in Fig. 3). In FedPFT, denote n2 as the number of prompts contained in Pr.i-
The FTM dynamically adapts local features to downstream tasks, offering strong flexibility. By em- ploying task-specific prompts p. it can introduce auxiliary objectives beneficial for client collaboration.
We explore this further in Section 3.6
3.4 Classification Task with Personalized Prompts
FedPFT employs a set of personalized prompts Px, as inputs to the FTM, enabling each client i to transform its local features for better alignment with the global classifier. Specifically, the
5

---

classification loss for client i is defined as:
LCB(Ø, T, Pu, bxi
log Me log(or,e), where æ, ~ dj.
is the number of classes, and Softmax 7((((2), Pr,,i)) represents the predicted proba- bilities, with oi,e being the ones class
c.
3.5 Dual-Phase Alternating Optimization Strategy
To effectively resolve the mismatch problem and coordinate the training of different modules in FedPFT, we propose an alternating optimization strategy, which partitions each local training round into two phases: the feature transformation phase and the model training phase.
Feature transformation phase. In this phase, the objective to align local features with the classifier by optimizing the classification prompts px,i and the FTM 7; while keeping the global feature extractor & and global classifier hg frozen. The training objective is formulated as:
min (Te,Te, Pa, Φ, hig ds). Pu.
Model training phase. Following the above phase, this phase focuses on learning client-specific knowledge by updating the feature extractor, FTM, and classifier while keeping the prompts fixed. The training objective is given by:
min Len(d Ty. Ne, Prc.s dz). oi
(5)
Let R. represent the total number of local epochs per training round. We divide it into Ry epochs for the feature transformation phase and R, epochs for the model training phase, satisfying + = R. It is crucial that Rg is always larger than R. ensure that the mismatch between local features and the classifier resolved before training the model parameters.
After completing local training, the updated parameters di T,, and bw are aggregated at the server to facilitate client collaboration, while the personalized prompts Ps.i remain local. We simply adopt the aggregation method used in FedAvg. The pseudo-code of FedPFT is summarized in the Appendix D
3.6 FedPFT with Additional Tasks: An Example of Contrastive Learning
As discussed in Section 3.3 our FTM provides strong flexibility. Benefiting from this, the FedPFT framework can seamlessly integrate additional tasks beneficial for PFL, such as contrastive leam- ing (20 48), feature alignment [61] [54], multi- task learning, etc., by simply incorporating task- specific prompts. In this section, we utilize con- trastive learning as an example.
Client
Px.i
hgg,
LCE
Φ
Loon
Pol
hpp
Figure 4: FedPFT with contrastive learning. The orange solid line and the green solid line repre- sent the forward propagation of the classification task and the contrastive learning task, respec- tively.
As depicted in Fig. 4 we introduce another set of personalized prompts Pp.s. which are fed into to transform features for the contrastive learning task with a global projection head hp.s The goal is to enable all clients to collaboratively optimize the contrastive learning loss, thereby improving the performance of the feature extractor. The training objective can be formulated as
E, (L O. (eIN))
min
min
Nn, i dz) = E,g, [L T, hic Pn, d.) + (;; (6)
where Teon represents the contrastive learning loss function. The detailed optimization process and the definition T.co are provided in the Appendix F
In Section 4 we refer to this extended method
as FedPFT+Con, validating the effectiveness of combining FedPFT with contrastive learning.

---

Table 1: Test accuracy (%) of different methods under Dirichlet non-III scenarios CIFAR-10, CIFAR- 100 and Tiny ImageNet. The best results are highlighted first
second and third.
CIFAR-10
CIFAR-100
Tiny ImageNet
Methods
0.1
0.5
1.0
0.1
0.5
1.0
0.1
0.5
1.0
FodAvg
56.81=0.07
60.734:188
61.53+0995
34.470085
31.9640.70
03.20-+0.2
19.15-1.15
18.38:0.88
17.73++073
FedBABU
55.87=1168
61,4409922
63.57+1.36
36.37-0.76
33.02++1115
33.61=00355
20.05=1.29
19.07=114
17.77+0037
FedETF
56.27:20.4
19.62::1.22
61.51+140
33.84=1005
11.00:1005
32.24:20224
17,96=0.54
05.50=00:00
15.18:00.33
FedFA
56.93=0.73
61.0241998
62.51=0.88
35.16=1.088
34.07=00788
34.69+:.55
15.86=0.95
17.49:0.61
16.75=0.60
Local
81.91:33009
60.15:1:.86
52.24:00.41
4.66
22.65::0.51
18.76:20.63
24.07=0.62
8.75::0.30
6.87:1:028
FodAvg-FT
85.55+22.21
12.00+22.05
67.78+0.67
60.08=0.90
40.72:+0.55
07
1..36=1.13
25.11=0.33
29.
FedBABU-FT
85.57=11.78
12.16:21.34
6..4311.20
60.10:0.54
40.95:10:92
38.05-0.899
40.62=1322
25.47:10004
21.37:20.51
RodETF-Por
83.49:22.31
67 67.90:22.58
63.51=0499
57.00=0.92
36.87=11977
35.1
34.04=0.66
19.98=11.68
15.44+0.28
FedRA-FT FT
80.70:1.2.98
67.68.12.26
66.50:10488
42.62-00.59
31.19:10.41
13.78-10.79
24.03=0.11
18.07
16.51:10.50
FodAMP
84.99::1.82
68, 08.26::0079
64.87=00.95
46.68=1006
24.74:20:58 58
18.22:0041
21.85=0.71
10.70:10.32
7.13=0.21
FedPer
84.43.10.47
64.92.10.66
0.94
2..
21.53
0.333-0.31
12.69.1.42
8.67:0..0
FedRep
84.59:21.58
67 06
60.52:+0.72
51.25
26.97-00.33
20.63=0042
30.83=
12.14:20.28
8.37=00.25
FedBN
83.55.12.32
66.79.11008
62.20
0.63
36.94.10.94
13.67
0..3
19.61:10.35
16.57.10.44
FedRoD
86.23:22.1.
1.77
1..4
00.17
39
1.
06.80-00.66
01.06=0.77
1.
22.32::1.13
pFedSD
86.34:12.61
71.97
1.,21
04.14
41.06.10.83
38,271.0.20
19.31
19.25 1.801118000
15.91
PFedGate
87.25-11.91
11
07.85
08.54
27.47:10.79
22.98.10.03
0..5
24.09:10.67
19.69.10.14
FedCAC
86.82::1118
69.83-10.46
65.39:10.51
57.22
08.64:10.63
32.59110322
40.19 1,20
23.70:10.28
18.58:10.22
PFcdPT
82.38:22.11
67.33:+1.33
64.37+1.22
43.21=1.66
35.23=+0.87
36,25=0.377
23.55=0068
22.35=:0.49
21.69:0..24
RdPPFT
87.23:2.69
74.10::195
0.723:0.76
60.98-0.39
44.87:10.76
41.83:0.37
0.14
28.61:0040
25.10:10.59
4 Experiments
4.1
Experimental Setup
Datasets. In this section, we mainly verify FedPFT in the label shift non-lID scenario, which is one of the most commonly used scenarios in FL research. Specifically, we examine two settings: Dirichlet non-lID and Pathological non-IID, In each setting, we employ three datasets: CIFAR-10 [1], CIFAR-100 [K], and Tiny ImageNet [15. Additionally, in Section4.66 we also verify FedPFT
in the feature shift non-IID scenario with PACS 17 and DomainNet (38] datasets.
In the label shift scenario, each client is assigned 500 training samples. For CIFAR-10 10 and CIFAR- 100 datasets, each client has 100 test samples; for the Tiny ImageNet dataset, each client has 200 test samples. Both training and test data have the same label distribution.
Baseline methods. We compare our method against 17 state sfate-of-thh-arr art (SOTA) methods: FedAvg
34], FedBABU (37). FedETF (27). FedFA (61, FedAvg-FT. FedBABU-FT, FedETF-Pert FedFA-FT,
FedAMP 10. FedPer (1) RedRep [4) FedBN (26), FedRoD 3. PFedSD [12. pFedGate (2].
FedCAC 500 and pFedPT [8] Here, FedAvg-FT, FedBABU-FT, and FedFA-FT indicate local
fine-tuning the classifier of the global model, as we find that fine-tuning the classifier consistently
outperforms full-model fine-tuning. FedETF-Per uses the official fine-tuning method to obtain personalized models.
Hyperparameter settings. For the general hyperparameters of FL, we set the number of clients N - 40, batch size = 100, and local update rounds R 5. The total rounds T are set to 1000 to ensure all methods reach full convergence. We select the highest average accuracy achieved by all clients across all rounds as the result. Each experiment repeated with three random seeds, and the mean and standard deviation are reported. We employ the ResNet [9] model architecture, specifically ResNet-8 for CIFAR-10 10 and ResNet-10 for CIFAR-100 and Tiny ImageNet. Please refer to the Appendix F for more details.
4.2 Comparison with SOTA Methods
We compare our proposed FedPFT against 17 SOTA methods across three datasets in Dirichlet non-IID scenarios. The experimental results on CIFAR-10, CIFAR-100, and Tiny ImageNet are presented in Table Results in Pathological non-III scenarios are provided in the Appendix G
Results in Dirichlet non-III scenario. In this setting, we evaluate the impact of varying non-IlI degrees by adjusting a. The results, as detailed in Table [ demonstrate that performance varies significantly depending on the underlying design principles of cach method. Among all methods,
7

---

FedRoD demonstrates robust performance across all datasets and non-IlD degrees. This is primarily due its clal-ccassifier design: personalized classifier for local feature alignment and global classifier for assistance from other clients to improve generalization. FedPFT explicitly addresses the mismatch problem and achieves superior results across all scenarios.
To further validate the benefits of incorporating also evaluate
Table 2: Test accuracy (%) of FedPFT+Con un- der Dirichlet non-IID on CIFAR- 100 and Tiny ImageNet.
additional tasks into FedPFT, TedPFT+Con. As shown in Table FedPFT+Con significantly outperforms existing SOTA methods by up to 7.08% by integrating contrastive learn- ing into FedPFT. This improvement highlights the strong flexibility of our proposed FTM in enhanc- ing PFI. performance.
Datasets
a 0.1
a = 0.5
a
CIFAR-10
88.60+2.19
7..54+1.88
74.81+0077
CIFAR-100 62.03+1141 47.98+0.78
44.29:t0.74
Tiny
43,42+1162 32,44+0.58 27.84+0.41
4.3 Ablation Study
In this section, validate the effectiveness of each component in FedPFT on the CIFAR-100 dataset.
The experimental results are presented in Table
Table 3: Experiments on the CIFAR-100 with
Setting represents FedAvg. Setting II incorpo- rates classification prompts P., to allow each client to adjust the global model individually to obtain a personalized model, resulting in a performance im- provement. Setting III (i.e., FedPFT) incorporates alternating training, where prompts are first up- dated to align local features with the global classi- fier to address the mismatch problem, followed by training model parameters. This approach effec- tively resolves the mismatch issue and enhances the synergy between the feature extractor and classifier, resulting in the greatest performance im- provement (20.01% %) compared to other modules.
a 0.1 to illustrate the effectiveness of different modules.
Settings Ps. Alter. Leon Pp Accuracy (%)
33.87+
40.97 1.28
III
v
60.98
IV
v
61.13
V
62.03=+1441
VI
03.76
Setting IV adds contrastive learning loss to Setting III, aiming to further enhance the feature extractor's performance through contrastive learning techniques. Setting (i.e, FedPFT+Con) incorporates specific prompts pp to better transform features for the contrastive learning task, reducing mutual interference between the two tasks during training.
Setting VI represents incorporating contrastive learning into PFL without addressing the mismatch issue. As shown, while contrastive learning can improve model accuracy by enhancing the quality of the feature extractor, its performance is far inferior to FedPFT (i.e., Setting III). This further highlights the critical role of addressing the mismatch problem in the PFL.
4.4 Disentangling Model Capacity and Training Strategy Effects
FedPFT introduces an FTM, which slightly in- creases model capacity compared to standard FL backbones. A natural concern is whether the per- formance gains arise simply from this added ca- pacity, or if they can be achieved by applying cxisting fine-tuning strategies to the same archi- tecture. To disentangle these factors, reuse the architecture of FedPFT but train it with different strategies: (i) standard FedAvg without personal- ization (FedPFT-Avg). (ii) full local fine-tuning after FedAvg (FedPFT-Avg-Full FT), and (iii) prompt-only fine-tuning after FedAvg FeddFTT-Av--Pomme
Table 4: Test accuracy (%) of FedAvg variants using the FedPFT backbone under Dirichlet non- IID setting on 100.
Methods
a == 0.1
0.5
a = 1.0
FedPFT-Avg
32.66-0.84
31.20 1.20
32.69:0.77
FedPFT-Avg
0.59
1..44.1.26
36. 16.10.0..8
FodPFT-Avg-Prompr FT
56.76=0.46
38.884 1.00
36.6220.18
FT). Table rreports the results on
CIFAR- 100. Compared with FedAvg and FedPFT in Table I none of these alternatives match the performance of FedPFT, despite using the same model capacity.
These results demonstrate two key insights: (1) Model capacity alone does not explain the improve- ments of FedPFT, in fact, larger models are harder to train under limited local data. (2) Simple
8

---

fine-tuning strategies, including prompr tuning. fail to achieve competitive results, highlighting the alignmenn necessity of our training-timg time alignment design.
4.5 Impact of Prompts on Feature Representation
In this section, we analyze how the introduced personalized prompts influence feature representation.
The impact on the alignment of features with downstream tasks. We visualize the features transformed by different prompts in FedPFT+Con using t-SNE. The experimental setup is consistent Larger markers in the figures represent feature centroids of corresponding classes for each client.
with Appendix The results are shown in Fig. [
is evident that the features transformed by dif- ferent prompts adapt to the specific requirements of their respective downstream tasks: features from classification prompts Px align more with client data distribution than image similarity; for instance, two classes within client may appear close even if the images are not visually similar. Additionally, these features exhibit well-defined decision boundaries between classes, as expected in classification tasks. In contrast, features from contrastive learning prompts Pp are more aligned with image similarity. For example, in Fig. (5)),, the feature centroids of "cat" and 'dog" are close, while those of 'airplane' and 'dog' are farther apart, consistent with contrastive learning principles.
(a) Features transformed by classification prompts
(b) Features transformed by contrastive learning
The Figure 5: The effect of different prompts on fea- ture space.
prompts Pe
The impact on the linear separability of features. We assess the linear separability of features three stages during the forward propagation process: 'None' represents the fea- tures extracted by o. "px" represents the features transformed by the FTM using classification prompts. 'Pp" represents the features transformed by the FTM using contrastive learning prompts. The results are detailed in Table55] Interestingly, the accuracies across different prompt conditions are generally similar, suggesting that the prompts do carry additional knowledge to enhance fea- ture separability.
Table The effect of prompts Px. and Pp on linear probe accuracy.
CIFAR-10
CIFAR-100
Prompt
0.1
0.5
These experiments confirm that prompts work by transforming features into the required format to align with downstream tasks, rather than improv- ing feature separability. This finding is consistent with the motivation of our paper.
None
87.69%
77.12%
73.93%
64.08%
46 16.50%
40.79%
P=
87.83%
77.25%
64.12%
46.43%
40.95%
87.82%
77.25%
74.02%
64.18%
46.40%
40.95%
4.6 Experiments in Feature Shift Scenarios
We conduct experiments on two feature shift datasets, PACS [17 and DomainNet 38 PACS and DomainNet have four and six domains, respectively, with each domain assigned to one client. The number of clients corresponds to the number of domains. For each client, we allocate 1000 training samples and 500 testing samples. The experimental results are shown in Table 6
Existence of the mismatch phenomenon. The results show that in both datasets, there are still noticeable gaps between Origin Acc. and Match Acc. across all methods, especially in FedAvg. This indicates that the mismatch problem persists in feature shift scenarios and is a major reason for the suboptimal performance of FedAvg.
Table 6: Comparison of different methods in the feature shift non-llI scenarios.
DomainNet
Methods Origin Ace.
Match Acc. Origin Acc. Match Acc.
FedAvg 71.48%
75.49%
63.53%
67.17%
FedPer
74.86%
75.31%
65.70%
65.67%
Superiority of our method. From the perspective of mismatch degree, the small gap between Origin Acc. and Match Ace. FedPFT demonstrates its effectiveness in addressing the mismatch problem in feature shift non-IID scenarios. In terms
FedBN
73.91%
74.71%
67.57%
68.57%
FodCAC
74.94%
75.94%
67.80%
68.53%
FedPFT
77.67%
77.64%
70.37%
70.57%

---

of Origin Acc., FedPFT significantly outperforms SOTA methods (e.g., by up to 2.73% on PACS), further highlighting the superiority FedPFT in addressing the feature shift non- -ID problem.
4.7 Robustness
In this section, we evaluate the robustness of FedPFT under partial client participa- tion. We conduct experiments on CIFAR- 10. CIFAR-IOO and Tiny ImageNet with = 0.5, considering scenarios where only random 50%. 70%, and 90% of clients participate in training each round. The ex- perimental results are presented in T
Table 7: Accuracy (%) of FedPFT with partial client participation. Values in 'O" indicate the performance change relative to 100% client participation, as pre- sented in Table T
Datasets
70%
CIFAR- 73.8 (.0.2)) 74.21 (+0.1 74.33.1 1,38 (+0.23)
CIFAR-100 45.74 (40.87) 45.46 14 (+0.59) 45.87=0..1 (+1.00) 28.53110.62 (.0.0)) 29.2410.16 (+0.63) 29.90.1..10 (+1.29)
Tiny
The results indicate that FedPFT maintains high accuracy even when only a subset of clients participate in training. Notably, on CIFAR- and Tiny ImageNet, performance may even improve under reduced client participation. This is likely because fewer participating clients per round can mitigate the impact of non-IID data distribution on the global model. These findings demonstrate the robustness of FedPFT in scenarios with partial client participation. For a detailed comparison with SOTA methods under different participation ratios, please refer to Appendix N
4.8 Training Efficiency
As illustrated in Table 8 FedAvg uploads the fea- ture extractor Φ; and classifier Nx. in each round. FedPFT adds FTM т3 increasing communication overhead by 20.47% for ResNet-8 and 21.17% for ResNet-10. 10. While FedPFT incurs additional com- munication costs, important to weigh against the performance enhancements and flcxibility of- fered by TІ- The improved accuracy and robustness to non-lID data might justify the additional costs in scenarios where model performance is critical.
Table 8:
The communication cost of each client in FedAvg and FedPFT in one round. The percentages in '0) represent the increase compared to FedAvg.
Models d: T
bus FodAvg
FedPFT
ResNet-8 1.24M 0.26M 25,70K 1.27M 1.53M (20.47%) ResNet-10 4.1M
1.05M
51.30K
4.96M
6.01M (2117%)
We also evaluate computation efficiency. FedPFT incurs slightly higher per-round computation time than FedAvg, but remains computationally more efficient than most SOTA methods. Please refer to the Appendix for details.
5 Conclusion
We identify that in non-IID scenarios, the fundamental cause of the global model's poor performance on client data is the mismatch between local features extracted by the feature extractor and the global classifier. This mismatch not only reduces model inference accuracy but also exacerbates inter-client gradient divergence during training, hindering model aggregation and ultimately rendering the feature extractor suboptimal for client data. To address this issue throughout training, we propose FedPFT, novel PFL method that incorporates prompt-driven feature transformation module. In each training round, FedPFT first trains the prompts to align local features with the classifier, followed by updating model parameters. Our experiments demonstrate that FedPFT not only effectively resolves the mismatch but also improves feature extractor quality, leading to significant performance gains over state-of-ther -hrt--rtttt-tttt methods.
Acknowledgments
This work was supported by the National Natural Science Foundation of China under Grants 62372028 and 62372027.
References
[1] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint arXiv:1912.008/1 2019.
10

---

[2] Daoyuan Chen, Liuyi Yao. Dawei Gao. Bolin Ding, and Yaliang Li. Efficient personalized federated learning via sparse model-adaptation. arXiv preprint arXiv:2305.02776. 2023.
[3] Hong- You Chen and Wei-Lun Chao. On bridging generie and personalized federated learning for image classification. In International Conference on Learning Representations, 2022.
[4] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. pages 2089-2099, 2021.
[5] Wenlong Deng, Christos Thrampoulidis, and Xiaoxiao Li. Unlocking the potential of prompt-tuning in bridging generalized and personalized federated learning, arXiv e-prints, pages arXiv-2310, 10, 2023.
[6] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3566 2020.
[7] Ziqing Fan, Jiangchao Yao, Bo Han, Ya Zhang, Yanfeng Wang, et al. Federated learning with bilateral curation for partially class-disjoint data, Advances in Neural Information Processing Systems, 36:32006- 32019, 2023.
[8] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729 9738, 2020.
[9] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pottern recognition, pages 770-778, 2016.
[10] Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangehuan Liu, Jian Pei, and Yong Zhang Personalized cross-silo federated learning on non-iid data. In Proceedings the AAAI Conference on Artificial Intelligence, 2021.
[11] 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, 2022.
[12] Hai Jin, Dongshan Bai, Dezhong Yao, Yutong Dai, Lin Gu, Chen Yu, and Lichao Sun. Personalized edge intelligence via federated self- knowledge distillation. IEEE Transactions on Parallel and Distributed Systems, 34(2):667-538) 580, 2022.
[13] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[14] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 10 (canadian institute for advanced research). URL http://www. cs. toronto. edu/krizcifarr himl, 5, 2010.
[15] Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
[16] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.0886 08691, 2021.
[17] Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain genoralization. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[18] Guanghao Li, Wansen Wu, Yan Sun, Li Shen, Baoyuan Wu, and Dacheng Tao. Visual prompt based personalized federated learning. Transactions on Machine Learning Research. 2023.
[19] Hongxia Li. Wei Huang, Jingya Wang, and Ye Shi. Global and local prompts cooperation via optimal transport for federated learning, arXiv preprint 2xXiv22433.0044 2024.
[20] Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastiv federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10713-10722, 2021.
[21] Shiwei Li, Yingyi Cheng, Haozhao Wang, Xing Tang, Shijie Xu, Weihong Luo, Yuhua Li, Dugang Liu, Xiuqiang He, and Ruixuan Masked random noise for communication-effiiii federated learning. In Proceedings of the 32nd ACM International Conference on Multimedia, pages 3686-3694, 2024.
[22] Shiwei Li, Qunwei Li, Haozhao Wang, Ruixuan Li, Jianbin Lin, and Wenliang Zhong Fedbif: Communication-efice efffcciet federated learning via bits freezing. IEEE Transactions on Parallel and Dis- tributed Systems, 2025.
11

---

[23] Shiwei Li. Xiandi Luo, Haozhao Wang, Xing Tang, Shijie Xu. Weihong Luo, Yuhua Li, Xiugiang He, and Ruixuan Li. The panaceas for improving low-rank decomposition in communicaaton---flccinnn federated learning. In Forty-second International Conference on Machine Learning, 2025.
[24] Shiwei Li, Wenchao Xu, Haozhao Wang, Xing Tang, Yining Qi, Shijie Xu, Weihong Luo, Yuhua Li, Xi- uqiang He, and Ruixuan Li. Fedbat: communication-effiike federated learning via lcarnable hinarization. In Proceedings of the 41st International Conference on Machine Learning, pages 29074-29095, 2024.
[25] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 5357-6368. PMLR. 2021.
[26] Xiaoxiao Li, Meirui JINGG, Xiaofei Zhang. Michael Kamp, and Qi Dou. FedBN: Federated learning on non-III features via local batch normalization. In International Conference on Learning Representations, 2021.
[27] Zexi Li, Xinyi Shang, Rui He, Tao Lin, and Chao Wu. No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5319-5329, 2023.
[28] Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Moreney. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01522) 2020.
[29] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 36, 2024.
[30] Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks, arXiv preprint arXiv: 1e..:22000.70222 2021.
[31] Jun Luo and Shandong Wu. Adapt to adaptation: Learning personalization for cross-silo federated learning. In Lud De Raedt, editor, Proceedings of the Thirty-Firs! International Joint Conference on Artificial Intelligence, IJCAI-22, pages 2166-2173. International Joint Conferences Artificial Intelligence Organization, 7 2022. Main Track.
[32] Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No fear of heterogeneity: Classifier calihration for federated learning with non-iid data, Advances in Neural Information Processing Systems, 34:5972-5984, 2021.
[33] Xiaosong Ma. Jie Zhang. Song Guo, and Wenchao Xu. Layer-wised model aggregation for personal- ized federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10092-10101, 2022.
[34] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficne learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
[35] Lei Meng, Zhuang Qi. Lei Wu, Xiaoyu Du, Zhaochuan Li, Lizhen Cui, and Xiangxu Meng. Improving global generalization and local personalization for federated learning. IEEE Transactions on Neurul Networks and Learning Sysiems, 36, 2024.
[36] Jed Mills, Jia Hu, and Geyong Min. Multi-task federated learning for personalised deep neural networks in edge computing. IEEF Transactions on Parallel and Distributed Systems, 33(3):630-641, 2021.
[37] Jachoon Oh, SangMook Kim, and Se-Young Yun. Fedbabu: Toward enhanced representation for federated image classification. In International Conference on Learning Representations.
[38] Xingchao Peng, Qiaxun 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 (ICCV), October 2019.
[39] Zhuang Qi, Lei Meng, Zitan Chen, Han Hu, Hui Lin, and Xiangxu Meng. Cross-silo prototypical calibration for federated learning with non-iid data. In Proceedings of the 31st ACM International Conference on Multimedia, pages 3099-3107, 2023.
[40] Zhuang Qi. Lei Meng, Zhaochuan Li. Han Hu, and Xiangxu Meng. Cross-silo feature space alignment for federated learning on clients with imbalanced data. In The 39th Annual AAAI Conference on Artificial Intelligence (AAAI-25), pages
2025.
12

---

[41] Mingjia Shi. Yuhao Zhou, Kai Wang, Huaizheng Zhang, Shudong Huang. Qing Ye, and Jiancheng Lv. Prior: Personalized prior for reactivating the information overlooked in federated learning. In A. Oh, 'T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 28378 283922. Curran Associates, Inc., 2023.
[42] Shangchao Su, Mingzhao Yang, Bin Li, and Xiangyang Xue. Federated adaptive prompt tuning for multi-domain collaborative learning. In Proceedings of the AAAI Conference Artificial Intelligence. volume 38, pages 15117-15125, 2024.
[43] Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai. Partialfed: Cross-domain personalized federated lcarning via partial initialization. Advances in Neural Information Processing Systems, 34:23309-23320, 2021.
[44] Rishub Tamirisa, Chulin Xie, Wenxuan Bao, Andy Zhou, Ron Arel, and Aviv Shamsian. Fedselect: Personalized federated learning with customized selection of parameters for fine-tuning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23985-23994, 2024.
[45] 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-966 2023.
[46] Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across hetcrogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8432-8442 8440, 2022.
[47] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using I-sne, Journal of machine learning research, 9111), 2008
[48] Lirui Wang, Kaiqing Zhang, Yunzhu Li, Yonglong Tian, and Russ Tedrake. Does learning from decentral- ized non-llD unlabeled data benefit from self supervision? In The Eleventh International Conference on Learning Reprexentations, 2023.
[49] Xinghao Wu, Xuefeng Liu, Jianwei Niu, Haolin Wang, Shaojie Tang, Guogang Zhu, and Hao Su. Decou- pling general and personalized knowledge federated learning via additive and low-rank decomposition. In Proceedings of the 32nd ACM International Conference on Multimedia, MM '24, page 7172-7188 New York, NY, USA, 2024. Association for Computing Machinery.
[50] Xinghao Wu, Xuefeng Liu, Jianwei Niu, Guogang Zhu, and Shaojie Tang. Bold but cautious: Unlocking the potential of personalized federated leamingg through cautiously aggressive collahoration. In Proceedings of the IEFF//VF International Conference on Computer Vision (ICCV), pages 19375-19384, October 2023.
[51] Xinghao Wu, Jianwei Niu, Xuefeng Liu, Tao Ren, Zhangmin Huang, and Zhetao Li. pfedgf: Enabling personalized federated learning via gradient fusion. In 2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pages 639-649. IEEE, 2022.
[52] Xinghao Wu, Jianwei Niu, Xuefeng Liu, Guogang Zhu, Shaojie Tang, Wanyu Lin, and Jiannong Cao. The diversity bonus: Learning from dissimilar clients in personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 36(10):18613 186111187777777 2025.
[53] Xinghao Wu, Jianwei Niu, Xuefeng Liu, Guogang Zhu, Jiayuan Zhang, and Shaojic Tang. Enhancing visual representation with textual semantics: Textual semantics powered prototypes for heterogencous fedcrated learning, 2025.
[54] Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference on Learning Representations, 2023.
[55] Fu-En Yang, Chien-Y Yi Wang, and Yu-Chiang Frank Wang. Efficient model personalization federated learning via client-specific prompt generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 19159-19168, 2023.
[56] Felix Yu, Ankit Singh Rawat, Aditya Menon, and Sanjiv Kumar. Federated learning with only positive labels. International conference on machine learning, pages 10946 10956. PMLR, 2020.
[57] Jianqing Zhang, Yang Hua, Jian Cao, Hao Wang. Tao Song. Zhengui Xue, Ruhui Ma, and Haibing Guan. Eliminating domain bias for federated learning in representation space. Advances in Neural Information Processing Systems, 36, 2024.
13

---

[58] Jianqing Zhang, Yang Hua, Hao Wang, Tao Song. Zhengui Xue, Ruhui Ma, Jian Cao, and Haibing Guan. Gpfl: Simultaneously learning global and personalized feature information for personalized federated learning. Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages October 2023.
[59] Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedala: Adaptive local aggregation for personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 1237-111244 2023.
[60] Jianqing Zhang, Xinghao Wu, Yanbing Zhou, Xiaoting Sun, Qiqi Cai, Yang Liu, Yang Hua, Zhenzhe Zheng, Jian Cao, and Qiang Yang. Hullibb comprehensive heterogeneous federated learning library and benchmark. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2025.
[61] Tailin Zhou, Jun Zhang, and Danny H. K. Tsang- Fedfa: Federated learning with feature anchors to align features and classifiers for heterogeneous data. IEEF Transactions Mobile Computing, 236))6631-6674, 2024.
[62] Guogang Zhu, Xucfeng Liu, Jianwei Niu, Shaojie Tang, Xinghao Wu, and Jiayuan Zhang. Dualfed: Enjoying both generalization and personalization in federated lcarning via hicrachical representations. In Proceedings of the 32nd ACM International Conference Multimedia, MM '24, page 11060- New York, NY, USA, 2024. Association for Computing Machinery.
[63] Guogang Zhu, Xuefeng Liu, Shaojie Tang, Jianwei Niu. Aligning before aggregating: Enabling cross- domain federated learning via consistent feature extraction. In 2022 IEEE 42nd International Conference on Distributed Computing Systems (ICDCS), pages 809-819. IERE, 2022.
[64] Guogang Zhu, Xuefeng Liu, Shaojie Tang, and Jianwei Niu. Aligning before aggregating: Enabling com- munication efficient cross- domain federated learning via consistent feature extraction. IEEE Transactions on Mobile Computing, 23(5):5880- 5896, 2024.
[65] Guogang Zhu, Xuefeng Liu, Xinghao Wu, Shaojie Tang, Chao Tang, Jianwei Niu, and Hao Su. Estimating before debiasing: A bayesian approach to detaching prior bias in federated supervise learning, 2024.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the core contributions: identifying feature-classifier mismatch as a key limitation in PFL, introducing FedPFT with a prompt- driven transformation module, and demonstrating improvements over SOTA methods. These claims fully supported by experimental evidence throughout the paper (see Sections and 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
14

---

Justification: The limitations are explicitly discussed in Appendix Q including the narrow focus on global model-based PFL frameworks and the need for deeper understanding of how parameter personalization affects the feature extractor.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiscless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
The paper includes a theoretical analysis (Appendix R clearly listing assump- tions, propositions, and proof' supporting convergence guarantees for the FedPFT method in non-convex settings.
Justification:
Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
15

---

The paper includes detailed descriptions of datasets, model architectures, non-IID scenarios, and hyperparameters in Sections g and Appendix The experiments are repeated with multiple seeds, and standard deviations are reported. The code is also provided the supplementary material. Guidelines:
Justification:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
The paper uses standard publicly available datasets (c.g., CIFAR-10, CIFAR- 100, Tiny ImageNet) with well-known benchmarks. Code provided in the supplementary material.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
public/guide htte///gids//eee/ee www.ddes//oeeuuuiisseeeeeeeeeeee
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips http://///iii/oooe////////a/aaaa
for more details.
16

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section and Appendix Frtenddii describe all training settings, including model type (ResNet), batch size, communication rounds, number of clients, optimizers, learning rates, and prompt configurations.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Results reported with means and standard deviations across three random seeds. Variance and error margins are provided in all key results tables.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
17

---

Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiment??
Answer: [Yes]
Justification:
Section F.4 reports compute resources in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: //te:://neuspspp wwtpriip..ccppuupeeee wttppublic/Bhhicccuuudellll.llll
Answer: [Yes]
Justification: The research adheres to the NeurIPS Code of Ethics. No personally identifiable data human subjects are involved; all datasets are publicly available image benchmarks.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is an algorithm research study and is not related to specific ap- plications. The model architectures and datasets used in our experiments are all publicly available.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
18

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper does not release models or datasets with high risk misuse. uses standard vision datasets and does not introduce sensitive content.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual- use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The paper uses standard datasets like CIFAR-10 and CIFAR-100 which are publicly available under clear terms. The ResNet architecture is also standard and open-source.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pattassts.paperswiii.
package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new datasets assets are introduced in this work.
Guidelines:
19

---

The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve any human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does involve human subjects, and no IRB approval is necessary.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: No LLMs were used the design or core methodology of the research. LLMs were used for writing assistance, that is not required to be declared.
Guidelines:
20

---

The answer NA means that the core method development in this research does not involve LLMs as any important. original, non non-standard components.
Please refer to our LLM policy (https://neurip:: /ttp://pp.ccc//offfeeeee//////// for what should or should not be described.
21

---

A Formal Definition of Feature-Classifer Mismatch
We formally define the feature-classifier mismatch
Mismatch(f,, W) = (...u))n, [C(f', (C(f,,W) C(f W)].
where D; denotes the data distribution of client i, represents the original local features extracted by the feature extractor, W is the classifier, f is the transformed feature generated by the feature transformation module (e.g., linear transformation layer), and L(.,-) is the loss function (e.g., cross-entropy).
This metric quantifies the improvement in performance (i.e., reduction in classification loss) when local features are aligned with the classifier. Since loss values are often less intuitive, we report the accuracy gap between "Origin Acc." and "Match Acc." (defined in the caption of Figure I as practical measure of mismatch throughout the paper. To complement this definition, we also explore additional metrics for quantifying feature-classifee mismatch, as discussed in Appendix
B Further Discussion on Feature-Classifir Mismatch
Beyond accuracy-base measures, we also compute the cosine similarity between local features and their corresponding classifier weights. As shown in Table 9 FedPFT achieves values comparable to strong baselines such as FedBABU-FT.
Table 9: Cosine similarity between local features and their corresponding classifier weights in different methods CIFAR-10.
Scenarios
FedAvg
FedAvg-FT
FedBABU
FedBABU-FT
FedPFT
a = 0.3
0.20
0.26
0.28
0.39
0.42
1.0
0.22
0.24
0.31
0.37
0.36
However, this similarity metric is not strongly correlated with accuracy or loss, as can be influenced by confounding factors such as feature norm and classifier bias terms. For this reason, we do not adopt as primary definition of mismatch in our paper. Instead, we define mismatch through the performance gap before and after feature transformation, which directly reflects model effectiveness. Our method therefore designed to learn a transformation module that minimizes classification loss, rather than enforcing explicit angular constraints.
C More Explanation of the Experiments Discussed in the Introduction
To facilitate understanding of the experiments mentioned in the introduction, we give a toy example to visualize the models used to calculate Origin Acc., Probe Acc., and Match Acc. in Figure
Pretrained by FL algorithm
Retrained by client's data
insert
replace
(a) Origin Model
(b) Linear Probe (c) Linear Transform
Figure 6: A toy example illustrating the model structures used calculate Origin Acc., Probe Acc., and Match Acc.
Figure aa) a) represents the model trained using the FL algorithm, where the accuracy measured on the client's 's local data is referred to as Origin Acc. Figure ((b) illustrates the model obtained from the linear probe experiment, where the classifier in the pre-trained model is replaced by linear classifier retrained the client data. The accuracy corresponding this model is referred to as
22

---

Probe Acc. Figure ( depicts the model obtained from the linear transform experiment, where linear transformation layer, retrained on client data, is inserted between the feature extractor and classifier of the FL. pre-trained model. The accuracy of this model is referred to as Match Acc.
D Pseudo-code of FedPFT
The pseudo-code of FedPFT is summarized in Algorithm
Algorithm FedPFT
Input:
Each client's initial personalized prompts P The initial global models (6) ,7()) h Client Number N; Total round T; Epochs of two learning phases R, and Ra.
Output: Personalized model (5(1)),(1) ,1("), ((2) each client.
for = to 7
Client-side:
for i = to N in parallel do
Initializing (Øt (t) (e)
An with (0(0),((1) (.1")
Updating {+. Px.i) ) by (4) for R, epochs to obtain (T(()) ((+1)))
Updating {0) (t) (( } by (5)(5) for Ra epochs obtain (0 2/,0/(10
(1+2) 3)
Sending (0) (t+1)) (e+1), ((+1) to the server.
Kii
end for Server-side:
Aggregating of global model (0(1-1), ((44)) (e044,,,+* 1. Sending (o(t+1) ((1+1), 1111) to each client i.
end for
E Details of Combining FedPFT with Contrastive Learning
Following Section of the main text, this section provides details of the combination of FedPFT with contrastive learning, including the definition of Leea and the optimization process of Eq. (6)
E.1
Definition of Lcon
We adopt the Momentum Contrast (MoCo) framework for contrastive learning. The associated contrastive loss function is defined as:
exp (q k5) where x ~ d,.
Le T, P.P..ss hp: x) = log
(8)
E exp (q k5/B)
In this formula, hp the projection head used for contrastive learning. v y= h, T((((²)) Pr.i)) represents the query vector, and h = ho o T([o(="), Po.i)) denotes the positive key vector. Here, x' and x are augmented versions of the sample x, and ho refer to the momentum-updatted encoder and projection head, respectively. is temperature hyperparameter, and K is the number of negative samples drawn from MoCo's ' queue, comprising the set {kk [k, J==0.
E.2 Optimization Process of FedPFT+Con
As discussed in Section 3.5 of the main text, FedPFT employs an alternating training strategy. In FedPFT+Con, we extend this approach by incorporating contrastive learning into the optimization process.
Feature transformation phase. In this phase, FedPFT+Con additionally utilizes Leva update ø.. and hp to enhance the feature quality. The objective can be formulated as:
min LCE(T;, Pr.ii 6s, hr, de) + Lool Tú, hosi Paás d.).
(9)
23

---

Model training phase. FedPFT+Con additionally updates Pp.s in this phase to align the features with the contrastive learning task, reducing interference from the classification task. Its training objective can be formulated as:
min E
Len(drs Thy Masii Pie,s, d:) + LCon(Past Ti; &i hp,s dz).
(10)
Figure 7 illustrates the training process of the contrastive learning and classification tasks in CedPFT+Con. The blue modules represent components from FedPFT, while the orange mod- ules represent additional components introduced in FedPFT+Con. Solid arrows indicate forward propagation and dashed arrows represent backpropagation.
Forward Propagation
Backward Propagation
Frozen Module
Trainable Module
Modules of FedPFT
Additional Modules of FedPFT+Con
Phase 1
Ppl
hp.
Phase
Ppd
hp.
Loo
Con
dj
d2
Ø
CF
LE
Pk.C
fgy.i
hg,
Client
Figure 7: Training process of FedPFT+Con in each client i.
F
Experiment Setup
F.1 Introduction non-IID Scenarios
Pathological non-IID. In this setting, each client is randomly assigned data from subset of classes with equal data volume per class. For the CIFAR-10, CIFAR-100, and Tiny ImageNet datasets, we assign 2, 20, and 40 classes of data each client, respectively.
16-
14
Client ID
Client ID
5
Class I ahels
Class Labels
Class Labels
(a) 0.1, 10-class
(b) 0.5, 10-class
(c) 1.0. 10-class
Cleat
20
Class Labels
Class Labels
(d) a = O. 50-class
(e) a = 0.5, 50-class
Figure 8: Visualization of data partitioning in Dirichlet non-lID scenarios with different a.

---

Dirichlet non-IID. This is commonly used setting in current FL research [51, 30 41 In this scenario, the data for each client is generated from a Dirichlet distribution denoted as Dir(a). As the value of a increases, the class imbalance within each client's 's dataset progressively decreases. This Dirichlet non-IID setting enables the evaluation of different methods across a broad spectrum of non-IID conditions, reflecting various degrees of data heterogeneity.
For clearer, more intuitive understanding, we involve 20 clients with 10-class and 50-class datasets to visualize the data distribution among clients with varying a values. As depicted in Figure [8] the horizontal axis labels the data class indices, while the vertical axis lists the client IDs. Each red dot indicates the class data assigned to client, with larger dots signifying higher volume of data in that class.
F.2 Introduction to Comparative Methods
FedBABU [37], FedETF [27], and FedFA [61] focus on enhancing the feature extractor of the global model. In FedBABU and FedETF, the classifier is frozen during the training of the global model. FedFA utilizes prototypes align client feature extractors and calibrate the classifier. FedAMP [TO] a weiightd--ggrregation-bbaee based method where clients with similar data distributions are given higher aggregation weights during model aggregation. Because it mainly encourages the collaboration of clients with similar data distribution, it is a method that pays more attention to the local data distribution of clients. FedPer M. FedRep 4 FedBN (6). FedRoD [3 and FedCAC [50] are methods, which personalize the global model by retaining certain parameters locally based on FedAvg. FedRoD additionally introduces balanced global classifier to obtain assistance from other clients, alleviating the overfitting issue caused by personalized classifiers alone, pFedSD [2] and pFedGate are fine-uning-bar -baed methods that adapt the global model to local data through fine-tuning, pFedSD directly fine-tunes the global model by distilling local models, while pFedGate trains an additional gating network and applies to the global model. pFedPT [88, prompt- pomptthased method, can also be viewed as fine-tuning approach, enhancing the global model's adaptation to local data distributions by adding prompts to images. FedAvg-FT FedBABU-FT, and Fe are methods where all clients collaboratively train global model during the first stage, followed by client-specific classifier fine-tuning of the global model to obtain personalized models. FedETF-Pec Per uses the official lne-tuning method to obtain personalized models.
phaameter-decupllngg
F3 Hyperparameter Settings in Different Methods
For the unique hyperparameters of each baseline method, we utilize the optimal parameter combina- tions reported in their respective papers. For learning rates, we adjust within {le-1, 1e-2, 1e-3).
In FedPFT, to simplify the hyperparameter tuning process and enhance usability, we provide a default set of hyperparameters: for all scenarios, we T4 10 and (R,, , R.) (4, 1). We use the SGD optimizer with a learning rate of 0.05 for the FTM and 0. for other components. In FedPFT+Con, for the Dirichlet non-IID scenario with a 0.1, we set (Ry, f Ra) (3. 2). while in other scenarios. we use (R,, RR,) = (4, 1). The learning rate for the FTM remains 0.01, with other hyperparameters consistent with FedPFT. Unless otherwise specified, the above hyperparameter settings are used in our experiments, though fine-tuning these parameters for specific scenarios may yield better performance.
F4 Compute Resources
Experiments are implemented using PyTorrh and conducted on 4x NVIDIA RTX 2080 GPUs. For the methods we compared, as well as FedPFT, single training session requires 12-36 hours. For FedPFT+Con, the training process takes longer due to the use of the MoCo algorithm, which requires data augmentation that can only be executed the CPU, Consequently, a single training session for FedPFT requires 48-72 hours.
G Comparison with State-of-the-art Methods
We present the comparative results of FedPFT and FedPFT+Con against established methods on CIFAR-10, CIFAR-100, and Tiny ImageNet datasets under Pathological non-III scenarios Table TO
25

---

Table 10: Test accuracy (%) of different methods under Pathological non-IID setting on CIFAR-10 CIFAR-100, and Tiny ImageNet.
Methods
CIFAR-10
CIFAR-100
Tiny ImageNet
FedAvg
54.33 + 3.03
34.27 + 0.44
18.05 + 0.23
Local
85.85 + 0.93
38.40 + 0.69
16.20 + 0.30
FedAMP
88.88 + 0.83
38.36 + 0.79
16.13 + 0.55
FedPer
87.51 + 0.95
41.54 0.74
20.25 + 0.65
FedRep
87.10 0.91
40.63 + 0.74
19.24 + 0.33
FedBN
87.02 1.41
47.75 1.03
24.91 0.48
FedRoD
88.06 + 1.70
52.55 0.92
32.25 1 0.80
pFedSD
89.97 1.45
52.30 1.18 1.18
30.27 + 0.78
pFedGate
43.73 0.76
22.42 1 0.83
FedCAC
89.77 1.14
49.07 0.87
30.83 = 0.42
PFedPT
86.29 + 1.11
39.92
21.38 = 0.98
FedPFT
89.67 + 1.96
57.62 = 1.18
36.13 = 1.32
FedPFT+Con
90.55 + 1.35
58.14 = 0.71
37.59 = 0.39
Results in Pathological non-III scenario. This is an extreme setting where each client has data from only subset of classes. This scenario is particularly pronounced in the CIFAR-I0 dataset, where each client essentially performs simple binary classification task. Here, clients can achieve decent performance by solely focusing on their local tasks ("Local'), even without collaboration with other clients. As such, methods that prioritize local data distribution, such as FedAMP,, pFedSD, and pFedGate, perform well. In contrast, on CIFAR- 100 and Tiny ImageNet datasets, as clients have more local classes with fewer samples per class, local tasks become more challenging, Effective collaboration with other clients becomes crucial. Consequently, methods such as FedRoD, which emphasize client collaboration, exhibit increasingly significant performance. FedAMP and pFedGate show considerable performance degradation. FedPer, FedRep, FedBN, and FedCAC, by personalizing certain parameters of FedAvg, enhance local performance by indirectly aligning local features with classifiers to some extent. However, as they do not address the mismatch issue, they compromise the performance of feature extractors to some extent, thereby limiting their performance to moderate level across the three datasets. FedPFT aligns local features with the global feature space using classification prompts, enhancing both local feature-classiffe alignment and inter-client collaboration effectiveness. achieves competitive performance on CIFAR-10 10 and surpasses existing SOTA methods CIFAR- 100 and Tiny ImageNet. FedPFT+Con further incorporates contrastive learning tasks to enhance feature extractor performance, outperforming SOTA methods significantly across all datasets.
H Addressing Mismatch by Inserting Linear Layer
As discussed in Section of the main text, a straightforward approach to address the mismatch problem is to insert a personalized linear transformation layer between the global feature extractor and the global classifier FredAvg+Lineall In this section, we validate this method through experiments, with the results shown in Table [1]
By combining the results in Table 1 of the main text, we observe that FedAvg+Linear outper- forms most SOTA methods on the CIFAR-10 dataset, demonstrating the effectiveness of addressing the mismatch problem during training. However, on the more challenging CIFAR- 100 dataset, FedAvg+Ling underperforms several SOTA methods. This illustrates that simple linear transfor- mation is insufficient for complex datasets. Notably, on CIFAR-100 with a = 1.0, FedAvg+Linear underperforms FedAvg, highlighting that FedAvg+Linear tends overfit the limited local training data due to the large number of personalized parameters introduced.
In comparison, FedPFT demonstrates superior performance across all scenarios. Leveraging the flexibility of the FTM, Fon further enhances model performance, significantly outperforming FedAvg+Lines
26

---

Table 11: Test accuracy (%) of FedAvg+Linear under Dirichlet non-IlI on CIFAR-10 and CIFAR- 100.
CIFAR-10
CIFAR-100
Methods
Q 0.1
= 0.5
0= 1.0
0.1
= 01.5
0: = 1.0
FedAvg
60.39:1144
10.41:21.36
00.71
34.9110.88
32.78:10.33
33.94:10.39
Local
81.91
60.15:10.86
52.24:1.41
47.61
22.6510.51
18.76:10.63
FedAvg+Linear
85.96112.23
71.17:11.28
67.63:10.83
58.0710.41
31.23:10.24 37.09
FedPFT
87.29:12.69
1..10
69.23
60.9810.39
41.83110.37 44.87
FedPFT+Con
88.60.12.19
1.85
7..81
62.03.11441
47.98.10.78
44.29.10.74
I Learned Features of Different Methods
In this section, we visually compare the feature quality extracted by different methods. We conduct experiments on the CIFAR-10 dataset with 10 clients and the data distribution is illustrated in Fig, at).. For each method, visualize the feature vectors of testing data from different clients using t-SNE 47 The results are shown in Fig. b)-((f).
Ula
Clac
Ulm
Clos atds
Clees
Uket
(a) Data Distribution
(b) FedAvg
(c) FedCAC
(d) FedPer
(e) FedPFT
(I) FedPFT w/o Alter:
Figure 9: t-SNE visualization of features extracted by different methods on the CIFAR-10 10 dataset. Different colors indicate various data categories, while distinct markers represent different clients.
FedAvg and FedCAC exhibit noticeable cluster structures of features but lack clear discriminative boundaries. FedPer shows overlapping features across various classes.
FedPFT exhibits clearer discriminative boundaries, which can be attributed to the alignment of local features with the global classifier achieved during local training. We also observe that data from the same class across different clients are mapped to the same positions in the feature space, indicating that our approach not only resolves the mismatch issue but also addresses client feuture discrepancies, as emphasized many prior works 27 54). 'FedPFT w/o Alter.' denotes the variant without alternating training. Although it shows improved clustering compared FedAvg,
27

---

its boundary discriminative quality is weaker than that of FedPFT, as it does not fully resolve the mismatch problem.
Feature Separability of Different Methods
In this section, we delve deeper into the linear separability of features extracted by various PFL methods. We conduct linear probing experiments on the CIFAR-10 10 and CIFAR-106 assess this property, with results detailed in Table
Table 12: Linear probe accuracy (%) of different methods.
CIFAR-10
CIFAR-100
Methods
0.1
0.5
1.0
0.1
0.5
1.0
FedAvg 85.01% %
72.52%
68.38%
59.50%
37.40%
32.33%
FedPer
84.44%
71.07%
66.51%
52.09%
26.61%
20.51% %
FedBN
84.52%
70.15%
66.51%
57.86%
35.24%
30.28%
FedCAC
85.22%
71.56%
66.98%
56.86%
34.64%
29.35%
FedRoD
82.79%
67.07%
63.12%
56.88%
33.99%
29.22%
pFedSD 85.86%
72.42%
68.12%
60.07%
37.33%
31.99%
FedPFT
85.52%
72.59%
69.57 %
51.60%
43.14%
38.47 %
It can be observed that the feature linear separability of most PFL methods is inferior to FedAvg. This suggests that although they partially alleviate the mismatch issue and achieve better model performance, the quality of the feature extractor is inevitably compromised due to their design, thereby limiting the full potential of PFL.
In stark contrast, FedPFT significantly improves the linear separability of features compared to FedAvg. It accomplishes this by fundamentally addressing the mismatch issue during the training process, thereby reducing model divergence among clients and the mutual interference between the feature extractor and classifier caused by misalignment.
K Comparison of FedPFT+Con with Two-stage Approach
In FedPFT+Con, we propose an FTM to coordinate the joint training contrastive learning and classification tasks. To illustrate the superiority of this design, introduce baseline called 'Two- stage,' similar to [48] where contrastive learning training is conducted first, followed by classification task training after convergence. For fairness, in the two-stage method, first perform 1000 rounds of contrastive learning training, followed by 1000 rounds of classification task training. The experimental results are depicted in Figure I0
Twe-stge d 0.1 Okm, a 0.1
Oun,
Twoosigo a 0.5
28
Two-sigo
Okns 0.5
Methods a 0.1 0
0.5
Ckn, ir
2o6
2d77
Two-stage Ours 62.03
53.43
43.87
14
47.98
30
180
-80
120
120
140
(a) T.on
(b) Accuracy (%)
(c) LCE
Figure 10: Comparison with two-stage approach on training Lcons LCE, and testing accuracy.
Firstly, from the perspective of the contrastive learning loss (Tcon), FedPFT+Con registers lower loss values compared to the Two-stage approach, suggesting that simultaneous training with the classifi cation task enhances the efficacy of contrastive learning. Secondly, considering both Figure 10
28

---

and Figure c0 our method exhibits significantly higher accuracy compared to the Two-stage approach. However, LCF converges higher training loss value, suggesting that in our design. contrastive learning tasks can alleviate overfitting issues in the classification task during training. These experiments demonstrate that our proposed approach can effectively coordinate both tasks, allowing them to assist each other. Importantl, these experiments also indicate that the significant performance improvement brought by contrastive learning in our method largely attributed to the design of our FTM and training approach.
L Attention Weight Visualization
In the FTM of FedPFT and FedPFT+Con, selffattention mechanisms are employed facilitate the integration of prompts with sample features. This section visualizes the attention weights to reveal how prompts influence the transformation process. We analyze 20 test samples from a single client on the CIFAR-10 10 dataset, with results depicted in Figure 1T Each row in the figure corresponds the attention weights for the output feature f' of single sample. Columns represent the input dimensions of the FTM: the first column corresponds to the original input feature f. while subsequent columns relate to different prompts from the sets Px,i or Pp.i-
1408
Libate 882h mach sindee
Input Postion
INDu Postion
11put Posk.o
(a) a 0.1
(b) Pp.
0.5
(c) Pps a = 1.0
1000
140M
Input Postion
Input Position
Input Posttion
(d) Pk, a = 0.1
(c) Pk. a 0.5
( Pa. a = 1.0
Figure 11: Visualize attention weights for different prompts client in the CIFAR-10 dataset under the Dirichlet non-IID scenario.
can be observed that when a = 0.1, indicating severe local class imbalances, each client has data from only few classes. In this case, the feature transformation task is relatively simple, and the influence of different prompts on a sample similar. As a increases, indicating more complex local tasks, the influence prompts becomes more intricate. Particularly at a 1.0, it can be seen that each sample is affected differently by different prompts. This also indicates that our approach performs sample-level feature transformation.
M Expanding to more clients and More Complex Models
To demonstrate the scalability of FedPFT, we further conduct experiments on CIFAR-10 in Dirichlet non-IIID scenarios with more clients and more complex models.
Scaling more clients. We perform experiments with 100 clients, and the test accuracies of various methods are presented in Table B3
29

---

Table 13: Accuracy (%) of different methods with 100 clients.
Scenarios
FedPer
FedBN
FedRol
FedCAC
FedPFT
a 0.1
84.68
85.51
87.58
87.40
88.43
0.5
71.40
70.85
75.23
72.82
76.49
a 1.0
66.92
67.18
70.99
69.06
72.53
Scaling to more complex models. We conduct experiments using ResNet-18, with the test accuracies of different methods shown in T4
Table 14: Accuracy (%) different methods with ResNet-18.
Scenarios
FedPer
FedBN
FedRoD
FedCAC
FedPFT
0.1
82.80
78.10
83.50
82.60
83.52
a = 0.5
65.35
55.38
67.55
63.27
68.93
1.0
61.83
56.88
62.05
60.53
64.95
In both experiments, FedPFT significantly outperforms state-of-the-art methods, highlighting the scalability of our approach.
N Partial Client Participation
As discussed in Section 4.7 we evaluated the robustness of FedPFT under partial client participation. In this section, we further consider more challenging scenario where only 20% of clients participate in each round, and compare our method against state-of-the totte-ff-the-rt baselines. The results are reported Table T5 As shown, FedPFT consistently outperforms all competing methods across different settings, demonstrating its strong robustness.
Table 15: Test accuracy (%) on CIFAR-100 with Dirichlet a 0.5 under different client participation ratios.
Methods
90%
70%
50%
20%
FedAvg-FT
41.50
40.88
41.95
40.23
FedBABU-FT
40.33
41.28
41.65
38.80
FedRoD
39.75
40.55
38.37
38.23
FedPFT
45.45
46.08
46.22
46.10
O
Effeet of Hyperparameters
In the previous experiments, we utilize the default hyperparameter combination. In this section, we verify how variations in these hyperparameters influence the performance of FedPFT.
O.1 Effect of nh.
72 represent the number of prompts in Px.: for each client. We examine the impact of this hyperpa- rameter on the performance of FedPFT on CIFAR-10 and CIFAR-100 datasets. The experimental results are depicted in Figure I2
FedPFT shows considerable robustness to variations in On the CIFAR-10 dataset, changes пı have minimal impact on performance, suggesting that the model can effectively handle simpler data distributions even with fewer prompts. In contrast, on the more complex CIFAR-100 dataset, performance initially limited by small number of prompts, which may not sufficiently cover the diverse feature space required for effective feature transformation. As the number of
30

---

(a) CIFAR-10, a 0.5 (b) CIFAR-10, a == 1.0 (c) CIFAR-100, 0.5 (d) CIFAR- 100, 1.0
Figure 12: The effect of hyperparameter 7u on CIFAR-10 and CIFAR-100 in the Dirichlet non-III scenario.
prompts increases, the model's ability transform and adapt features improves, leading enhanced performance.
O.2 Effect of Ry and Ra
R and R. used to control the number of training epochs for the two training stages. Since we set Ry + R. = R, in this experiment, we only adjust R to examine the impact these two hyperparameters model performance. The experimental results are illustrated in Figure [3
60
40
80
55
Acsurcy
-35
30
45
30
60
40
35
(a) CIFAR-IO,
(b) CIFAR-100, a 0.1
(c) Tiny ImageNet. 0.1
Figure 13: The effect of hyperparameter Ry CIFAR-10, CIFAR-100, and Tiny ImageNet in the Dirichlet non-IID scenario with a 0.1.
When Ry = 0, it indicates that local features are not aligned with the global classifier before training the model parameters. Under this condition, the model performance observed to be very poor. As Ry gradually increases, the model performance initially improves but then declines in some scenarios, suggesting that Ry balances the trade-off between the training stages. When Ry is too small, local features are not sufficiently transformed to match the classifier, resulting in the model being affected by the mismatch during the model training phase, which reduces the synergy between the feature extractor and classifier, On the other hand, when R is too large, the model parameters are insufficiently trained, limiting the learning of local knowledge from clients.
In general, R2 and R. are two hyperparameters that need careful adjustment, as they have a significant impact the performance of FedPFT. Typically, in scenarios where clients' local tasks simple, it may be appropriate to decrease the value of Ry. In other cases, we recommend using a larger R₃ value fully align the local features with the global classifier.
0.3 Impact of Local Epoch E
In this section, we extend our evaluation to explore the impact of increasing local update epochs (E {10, 15, 20}). The results shown in the Table 16 As expected, larger E values lead to performance degradation across all methods due to exacerbated client drift under non-IID settings. However, FedPFT exhibits smaller performance drops, highlighting its robustness.
This robustness stems from our two-phase training design: by resolving the feature-classifie mis- match before model updates, FedPFT ensures that training begins from aligned feature space and
31

---

Table 16: Test accuracy (%) CIFAR-I0 with Dirichlet 1.0 under different E.
Methods
20
FedAvg-FT
67.73
67.55
67.23
FedBABU-FT
68.30
68.13
67.85
FedRoD
67.23
67.05
66.80
FedPFT
69.43
69.32
69.10
mitigates interference between the feature extractor and classifier, thus reducing the adverse effects of drift.
P Computation Efficiency
We empirically evaluate the computational efficiency of our method on CIFAR-100 using ResNet- 10, with results displayed in Table [7 We run each method for 100 rounds and calculate their average runtime per round. Each method exclusively utilizes a single machine during runtime. All experiments are conducted on four NVIDIA RTX 2080 GPUs.
Table 17: The average computation time per round for different methods on CIFAR-100.
Methods
FedAvg
FedPer
FedBN
FedRoD
FedCAC
FedPFT
Time per round (s)
50.10
51.81
52.65
54.76
55.35
51.61
FedPFT introduces the FTM, which adds some computational overhead. However, since it does not require updating the feature extractor during the feature learning phase, this helps reduce training costs to some extent. Overall, from the results in Table 17] FedPFT takes slightly longer to compute than FedAvg but more time-efficient compared SOTA methods.
Q Limitations and Future Work
In this paper, we primarily investigate PFL methods that derive personalized models based on a global model. We analyze the essential reasons these methods enhance performance from the perspective of mismatches between local features and classifiers. Although such methods occupy the mainstream in the current PFL. field, it is necessary to admit that there are some PFI. methods that are not based on global models, such as seooaalieed-eiihh--gggree in this study. Additionally, while this paper observes that personalizing a subset of parameters degrades the quality of the feature extractor, the underlying reasons for this phenomenon require further investigation.
methods, which are not explored
R Theoretical Analysis
Since the main problem in Eq. 3.2) is non-convex, we focus on the factors affecting convergence in the non-convex setting.
R.1
Problem Setup
Non-convex case analyses are provided, because model is multi-layer transformer. Analyses as follows.
We transform the problem into an unconditional bi-level optimization problem:
min EF(w) E, (.(F(() min Ed,, (ww(w, Pr,3i d:))
Prc.
32

---

Table 18: The glossary of notations used in the theoretical analysis.
Implication
Notation
Global / Local loss
L L FIF
Global / Local problem
Local Dataset on 2th client
d E d; Φ T
Feature extractor
Feature transformation module
Classification Contrastive learning prompts
Ps I Pp w
Feature extractor & Feature transformation module & Classifier
Classification Contrastive learning task head
bg / hp V (.()))
Global Local problem's 's gradient
4 (F(w)
Local gradient approximation
Client number
N
Local update epoch
R S S
The number of clients sampled each global epoch
The set of clients sampled at global epoch &
The actual learning rate of global problem
TJ T
The learning rate of local problem
Approximated local gradient error's 's upper- upper-bound Local-global gradient error's upper-bound Index of client, local epoch and global epoch
bound
OF i [N], E [R] E [T]
where represents the expectation of all random variables, E; means the expectation of client
sampling, Es, is the local data sampling expectation, and we use w {Ø, r, h.,) for simplification, based on the equivalence of block coordinate descent and gradient descent.
With contrastive learning the problem could be transformed into a similar problem with constrain. By Lagrange duality, the main problem is transformed as follows:
min min E,E,, Leu(d, T, has Pr,3; di) hr (pe,he
s.t. E,E,, Loon (0, (;, r;d,) s Hcon
R.2 Propositions
Proposition R.1 (L-smooth). If f is L-smooth, Vx, we have:
(v)(x) Vf()), v) < L||x y/1
[V/((x) - V/(y)]1 L|| yll
|V / (2(((=))=--)))))((_1(()) - f(y)
- -
x1|
Proposition R.2 (Jensen's inequality). If is convex, we have the following inequality:
(x) 2 ()
A variant of the general one shown above, given a group {z, be(N
Σ 2.1 N Σ ||..
EE[M]
EE[N]
Proposition R.3 (Triangle inequality). The triangle inequality, where ||- || the norm, and B the elements in the corresponding norm space:
|A + Bl|
||||| + ||B||
33

---

Proposition R.4 (Matrix norm compatibility). The matrix nomm compatibility, A E Raxb B E R U R
|| mm S |m//||||BB] || S ||A||m|lo||
Proposition R.5 (Peter Paul inequality). b Vx, y and Ve we have the following inequality:
2(r,0) - | + cllvl1
R.3 Assumptions
Assumption R.1 (L-smooth local objectives), Vi, F; is LE-Smooth, the main proposition is shown in Prop. R.T. Notice that the F; is assumed be L-smooth and non-convex, which matches the problem and neural network architecture setting in the main paper.
Assumption R.2 (Bounded local variance). The local problem's gradient is assumed not to be too fur from the global problem's gradient.
Yw, E, [ -VF(())-F(uull)) VF(w)ll OF
Assumption R.3 (Bounded approximated gradient). The first-order approximation of the local problem's 's gradient ghr should not be too far from the ground truth V F, (w( In this assumption, the approximated error of the block coordinate descent in Algorithm iis bounded.
((,,r,)), 1l9,, V.((w, S
R.4 Lemmas
Lemma R.1 (Bounded local approximation error). =f nR S of client drift error:
I
we have the following bound
2Lp
1 N.R
2672
Σ
Bllg
10
V
(wo(0)
20 2 (3r
Σ E (w))
NR
Proof. The client drift error given 2th client and its upper bound are as follows:
(
(1)
2 Mi.
V F; w (1) 11 VF(
(1))
266 + 2L El w0
11
where the first inequality is by Proposition R.3 and the second one is by Assumption R.1
For the last term in the upper bound, we have the iterative formulation as follows:
E | - w 11
E | [ - (t) 112
<2E|| [B Wi
(t)
- (1) 1VVF(( (1) 27 19 (1)
<2(1 + 2R
E
(1)
- 11 2(1 + 2R) 2 [
1--
+ 144 182
+ 1
w
(t)
w [1
-2(1 +
1
2n
13
- (t) 2n
(8)
2R
+ 2(1 + 2R) 'E V
(w
where the two inequalities are by Proposition R.3 Proposition R.5 and Eq. 11
Take 7j n ZLE recursively unroll the incquality follows:
34

---

E w
(1) 11
E 1 [ < <
<2 +
(1)
417 + 2 4
VE
27j
12842
R
where the inequality unrolled and we use
R S 1. Thus, we have:
52 s 28 2 11² I 3F +
R
R.5
Theorem and Discussion
Theorem R.2 (Non-convex and smooth convergence of FedPFT), Let Assumption R.
Assump-
tion R.2 and Assumption R.3 hold, if == nR z min ni} taken, where ij == 24 1122 07 we have the following bound:
N75
1)28
AF
28/3 1.F
R. 82/1/3
+ JFV L 1)Ar
OE (F O
+
+ +
1/T
T2/3 R2/3
VTN
Proof.
EF(w - -FF(w(O)
<E(VF(w(() (1+1) - (w)
LF
E ([ - 20
2
TE(FF(W (1)) LF E Ellg 11)
2
LF
TE V )1 VV n00 VP(w(()))
E 10
N.
LF
EIVF
||
El
9
(2)
- ())
+
E g 10
NR
2
N.
<
"
(1)
7 1|1
E
Σ
Dir
VF((w(9)11
2
NR
4r
+ 2
3
LF
- (10
Elllg(
+
1
Σ - +
iESIO
7(1 - 3hLp)
(
71 3dLF)
1
Σ
Ellg (t) vF(
2
2
NR
37 LF
Σ
())(w(()) (1(
+
2
5
(ES())
< 7i(1 3T)FF)
(F + 37 LF N/S - 103 + [VF((W)))
2
82
+ 71 3
+ 2 R+3 LF 33 3E V (w((t) 2|12
where the four inequalities are respectively by Lp-smooth of F = F, Proposition R.5 Lemma R. and the similar classic Lemma 4 in [1
Let C₁ = 382 o2 3 N/S N-I
C3 = R
3o + 1
35

---

(F(wllt))) EF(w!!) S 7 3 N/S
- oF +2 2d i V/
2
+ C3T + c27 + cc
S 7 El| Fl (1) )| + c 021 + ci
where let ij S (22 f. where ;= 3x2 N-1 and accumulating, we have:
N/8-1 d 13. Re-arranging the inequality above
E ( (F( EF(w + caff² c2n c
t=T
V (F( E)( oz/ + c24 + c
27
Let Ap = F(w")- - F(w"), where w" the minimum of the main problem arg min EF(w). To measure the exact term of the bounds, we consider the following cases:
Tg S 8 or
Ar
Ar T
S Sf. let vj = min fiin(( (7
1//3,)(3333)/) ( )
1/2), we have:
1/3 A/3
E (VF
C33
C2 F
1/2 + c1
72/3
T
Ag
and
Ag c
2 let कj = f7. we have:
of
E V 11
AF
33 F
1/3
2/3
2/3
(c2Ar
1/2 + 11
T2/3
T
Uniformly sample afe E [T] 1, we have the upper bound as follows:
VF (0 113) OEE F( (11))
AF
2R/3 (Ro + +53) 2/3
FFVLE(N/S - 1)AF + +8)
nT
T2/3
VTN
Remark R.2.1. According to Theorem R.2 our proposed FedPFT converges at a sub-linear level. The linear term O Ar) is affected by ij and the initialization gap P. The sub-linear term (1/T22/3) is affected by R, especially when R is large due to the exponential factor 28, As the local approximation ervor of the gradient 8 grows, both the convergence radius (6)) and the sub-linear term (1/T22/3) are affected by the local optimizer selection significantly. Another sub-linear term O(VT) is eliminated if N/S- affected by dr-
O when all the clients are sampled. Otherwise, the sub-linear rate is mainly
FedPFT addresses the mismatch and aligns the training objectives across clients by introducing Pe, Our design can effectively reduce differences in local gradients among clients during training, thereby reducing OF and subsequently lowering the upper bound. During training, Px,s incorporate information from the local datasets. By using them as part of the input, FedPFT effectively reduces the randomness in gradient computation, thereby lowering d and consequently reducing the upper bound.
Empirical Support for Remark R.2.1
To validate Remark R.2.1
we measure the standard deviation of local gradients within clients (8) and the gradient difterences across clients (ar) in each training round. As shown in the Fig. [4 our proposed prompts significantly reduce both metries, directly supporting the theoretical insights.
36

---

0.009
1.5
0.005
FedPFT w/o prompt
FedPFT w/o prompt
FedPFT
1.1
FedPFT
0.001
0.7
0
100
200
300
400
500
100
200
300
400
500
(a) Local Gradient Deviation
(b) Client Gradient Difference
Figure 14: Comparison of local gradient deviation and inter-client gradient difference with and without prompts FedPFT on CIFAR-100.
37


Reviews:
[{"review_id": "ou6stjPp2I", "reviewer": "Reviewer_zsuL", "summary": "This paper addresses data heterogeneity in federated learning by proposing a novel personalized federated learning framework. Unlike existing methods that localize or fine-tune certain model parameters, the proposed FedPFT framework identifies the mismatch between local features and the global classifier as the main performance bottleneck. It uses personalized prompts combined with local features, processed through a shared self-attention transformation module, to align them with the global classifier during training. Experimental results demonstrate that FedPFT outperforms state-of-the-art methods by up to 7.08%.", "questions": "1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.", "limitations": "**Limitations:**\n\n1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. This paper is well-written and easy to follow. Exploiting a prompt-driven transformation module to achieve alignment between local features and the global classifier for improved personalization in federated learning is both interesting and promising.\n\n2. The transformation experiments effectively explain the motivations and insights behind the approach, making them clear and accessible to readers.\n\n3. Extensive experiments validate the effectiveness of the proposed method across various datasets, demonstrating its superiority over existing competitors.\n\n**Weaknesses:**\n\n1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.\n\n2. The evaluation section can be improved to strengthen the applicability of the proposed method to real-world scenarios:\n\n   **(1) Scalability:** The current evaluation only considers 40 clients. The performance with a larger number of participating clients should be assessed to demonstrate scalability.\n\n   **(2) Participation mode:** Only the full participation setting is considered in the experiments. In practice, partial participation is more common in federated learning systems and should be evaluated.\n\n   **(3) Local update epochs:** The impact of the number of local update epochs should be evaluated to demonstrate the robustness of the proposed algorithm.\n\n   I would be glad to increase my rating if you can validate the scalability and robustness of the proposed method, particularly with a large number of participating clients and under partial participation settings.", "was_revised": true, "final_justification": "The authors' clarifications and additional experimental results has effectively addressed most of my concerns. After carefully considering the authors' responses and other reviews' comments, I have decided to raise my original rating to 4.", "rebuttal": "We sincerely thank the reviewer for valuable feedback. We try our best to carefully address each of the raised weaknesses and questions. Additionally, we would like to emphasize that the code for our method was submitted along with the original submission, and we commit to open-sourcing the additional experiments conducted during the rebuttal phase if this paper is accepted.\n\n## Weakness 1 & Question 1 & Limitation 1: Regarding the distinction from FedPAC\n\nWe appreciate the reviewer’s thoughtful comparison. We would like to clarify that our method differs significantly from FedPAC in terms of **objective, methodological design, and practical advantage**.\n\n### 1. Difference in Objective\n\n- **FedPAC aims to align features across clients**, using global prototypes to reduce inter-client feature divergence and performing classifier collaboration via weighted aggregation of personalized heads.\n- In contrast, **our method aims to align local features and the shared global classifier within each client**. This is a new problem that has not been studied in federated learning but is critical for stable training and accurate inference under heterogeneous data.\n\n### 2. Differences in Methodological Design and Our Advantage\n\nWhile our method, like FedPAC, enables feature alignment and classifier collaboration, **it further addresses the alignment between local features and the global classifier**, **which FedPAC is not designed to handle**:\n\n- Before local training, we freeze the feature extractor and classifier, and optimize lightweight, client-specific prompts to **first resolve the mismatch** between local features and the global classifier (see Section 3.5). This ensures that training starts from an aligned feature space and avoids mutual interference between modules during local updates.\n- Although FedPAC promotes inter-client alignment, it **does not address the feature–classifier mismatch within each client**, which we identify as a core bottleneck. As shown in the table below, FedPAC still suffers from a noticeable gap between Match Accuracy and Origin Accuracy, indicating residual mismatch.\n- In contrast, our method **resolves this mismatch** and **implicitly enables inter-client feature alignment** by aligning local features to a shared semantic space defined by the global classifier. This leads to class-wise feature consistency across clients, as visualized in Appendix Fig. 9(e), where features from different clients cluster naturally by class.\n\nThese results suggest that **FedPFT not only achieves inter-client consistency comparable to methods explicitly designed for it, but also addresses intra-client mismatch effectively**.\n\n|        | Dir 0.5 (Match Acc.) | Dir 0.5 (Origin Acc.) | Dir 1.0 (Match Acc.) | Dir 1.0 (Origin Acc.) |\n| ------ | -------------------- | --------------------- | -------------------- | --------------------- |\n| FedPAC | 70.79                | 69.59                 | 67.12                | 66.14                 |\n\n\n\n## Weakness 2: Evaluation Scope – Scalability, Partial Participation, and Local Epochs\n\nWe thank the reviewer for this thoughtful suggestion. In response, we have conducted additional experiments to address these concerns and summarize our findings below.\n\n### (1) Scalability to more clients\n\nAs part of our original submission, we have already included experiments involving **100 clients**, comparing FedPFT with multiple state-of-the-art baselines under varying non-IID settings. The detailed results are provided in **Appendix K, Table 12**, and for the reviewer’s convenience, we summarize them below. As shown, FedPFT consistently outperforms competing methods, even as the number of clients increases, demonstrating strong scalability.\n\n| Scenarios    | FedPer | FedBN | FedRoD | FedCAC | FedPFT    |\n| ------------ | ------ | ----- | ------ | ------ | --------- |\n| $\\alpha=0.1$ | 84.68  | 85.51 | 87.58  | 87.40  | **88.43** |\n| $\\alpha=0.5$ | 71.40  | 70.85 | 75.23  | 72.82  | **76.49** |\n| $\\alpha=1.0$ | 66.92  | 67.18 | 70.99  | 69.06  | **72.53** |\n\nIn the rebuttal phase, we additionally conducted experiments with 200 clients. As shown below, our method consistently outperforms SOTA methods.\n\n| 200 clients  | FedAvg-FT | FedBABU-FT | FedRoD | FedPFT |\n| ------------ | --------- | ---------- | ------ | ------ |\n| $\\alpha=0.1$ | 86.20     | 87.04      | 85.68  | 87.90  |\n| $\\alpha=0.5$ | 72.82     | 73.80      | 72.60  | 74.81  |\n| $\\alpha=1.0$ | 70.04     | 71.31      | 69.80  | 72.32  |\n\n### (2) Participation Mode\n\nWe have also evaluated **partial client participation** in our original submission (see **Section 4.7, Table 7**), under **90%, 70%, and 50%** participation rates on multiple datasets. For ease of reference, we summarize those results here. The results indicate that FedPFT maintains strong robustness and performance under varying participation ratios.\n\n| Datasets  | 90%            | 70%            | 50%            |\n| --------- | -------------- | -------------- | -------------- |\n| CIFAR-10  | 73.88$\\pm$1.84 | 74.21$\\pm$1.45 | 74.33$\\pm$1.38 |\n| CIFAR-100 | 45.74$\\pm$0.32 | 45.46$\\pm$1.14 | 45.87$\\pm$0.81 |\n| Tiny      | 28.53$\\pm$0.62 | 29.24$\\pm$0.16 | 29.90$\\pm$0.10 |\n\nTo further test under a more challenging setting, we conducted new experiments during the rebuttal phase with only **20%** client participation, and compared against SOTA methods. As shown in the table below, FedPFT continues to achieve superior performance across all settings.\n\n| CIFAR100, Dir 0.5 | 90%       | 70%       | 50%       | 20%       |\n| ----------------- | --------- | --------- | --------- | --------- |\n| FedAvg-FT         | 41.50     | 40.88     | 41.95     | 40.23     |\n| FedBABU-FT        | 40.33     | 41.28     | 41.65     | 38.80     |\n| FedRoD            | 39.75     | 40.55     | 38.37     | 38.23     |\n| **FedPFT**        | **45.45** | **46.08** | **46.22** | **46.10** |\n\n### (3) Impact of Local Update Epochs\n\nIn the rebuttal phase, we extended our evaluation to explore the impact of increasing local update epochs ($E \\in {10, 15, 20}$). As expected, larger $E$ values lead to performance degradation across all methods due to exacerbated client drift under non-IID settings. However, **FedPFT exhibits smaller performance drops**, highlighting its robustness. \n\nThis robustness stems from our two-phase training design: by resolving the feature–classifier mismatch before model updates, FedPFT ensures that training begins from an aligned feature space and **mitigates interference between the feature extractor and classifier**, thus reducing the adverse effects of drift.\n\n| CIFAR10, Dir 1.0 | 10        | 15        | 20        |\n| ---------------- | --------- | --------- | --------- |\n| FedAvg-FT        | 67.73     | 67.55     | 67.23     |\n| FedBABU-FT       | 68.30     | 68.13     | 67.85     |\n| FedRoD           | 67.23     | 67.05     | 66.80     |\n| **FedPFT**       | **69.43** | **69.32** | **69.10** |\n\nWe will include all of the above results and analyses in the revised version of the paper. Thank you again for the helpful and constructive suggestions."}, {"review_id": "CivNN6O5Hc", "reviewer": "Reviewer_3APB", "summary": "The paper addresses a challenge in Federated Learning (FL) caused by data heterogeneity, which leads to suboptimal performance of global models across diverse client distributions. \nThey identify that the primary cause of suboptimal performance is the mismatch between local features and the global classifier, rather than the inadequacy of the feature extractor or classifier themselves. \n\nTo address this issue, they propose FedPFT, a new Personalized Federated Learning (PFL) framework that uses personalized prompts and a shared self-attention-based feature transformation module (FTM) to align local features with the global classifier during training. The paper demonstrates through experiments that FedPFT outperforms state-of-the-art methods by up to 5.07%, and when combined with contrastive learning, the improvement can reach up to 7.08%.", "questions": "- Could you formally define “feature–classifier mismatch” and clarify how it differs from discrepancies investigated in prior works [1,2,3,4,5,6,7]? \n\n- Please report a quantitative metric—such as cosine similarity between local features, classifier weights, or average gradient-angle divergence—and show how FedPFT and existing baselines change that metric during training.\n\n- The robustness exp in Section 4.7 should be  compared  with other baselilnes, and more sparse client (eg, 10%, 20%) should be considered for the realistic problem setting.\n\n- Have you measured communication cost or accuracy when applying FedPFT to out-of-distribution (OOD) clients?\n\n- It is highly recommended to follow the experimental settings of large benchmarks, such as those in [8, 9], to ensure a fair comparison. Could you describe the problem setting more detaily in comparison to these benchcmarks? \n\n[8] https://github.com/KarhouTam/FL-bench\n\n[9] https://github.com/alibaba/FederatedScope", "limitations": "The authors have not discussed limitations of the work in the conclusion. I think the method is purely theoretical and there is no potential negative societal impact.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n1. The paper offers a fresh perspective on the feature–classifier mismatch problem in FL.\n\n2. FedPFT introduces a concrete framework—personalized prompts plus a self-attention FTM—that explicitly targets this mismatch.\n\n3. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet under both Dirichlet and Pathological non-IID settings demonstrate clear performance gains over existing methods.\n\n### Weaknesses \n1. Although the authors promote feature–classifier mismatch as a new explanation for performance degradation, it is unclear precisely what this mismatch is or why it constitutes a novel, critical problem. Prior studies have already examined gaps between global and personalized models in feature, classifier, or weight space side [1, 2, 3, 4, 5, 6, 7].\n\n2. To claim that FedPFT solves the mismatch problem, the paper should provide more evidence—e.g., extend Figure 1 (Gradient Difference) to include other SOTA models, or introduce quantitative metrics beyond pre-/post-accuracy. Performance gains alone are insufficient supports.\n\n3. Showing that FedPFT pairs well with contrastive learning is useful but not, by itself, a novel contribution; it is incremental.\n\n4. The personalized prompts and FTM increase both model complexity and communication cost. Figure 1 and Table 1 suggest that simple post-training or fine-tuning (FT) yields nearly the same accuracy boost with far less overhead.\n\n\n[1] FedBABU: Toward Enhanced Representation For Federated Image Classification, ICLR 2022\n\n[2] No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data, NeurIPS 2021\n\n[3]  Guiding The Last Layer in Federated Learning with Pre-Trained Models, NeurIPS 2023\n\n[4] Personalized Fedrated Learning with Ehroretical Guarantees: A Model-Agnostic Meta-L earning Approach, NeurIPS 2020\n\n[5] Personalized Federated Learning using Hypernetworks, ICML 2021\n\n[6] Federated Learning with Partial Model Personalization, PMLR 2022\n\n[7] Personalized Federated Learning via Variational Bayesian Inferennce, PMLR 2022", "was_revised": true, "final_justification": "Based on the discussion, I have decided to increase my score slightly.", "rebuttal": "We sincerely thank the reviewer for valuable feedback. We try our best to carefully address each of the raised weaknesses and questions. Additionally, we would like to emphasize that the code for our method was submitted along with the original submission, and we commit to open-sourcing the additional experiments conducted during the rebuttal phase if this paper is accepted.\n\n## Weakness 1 & Question 1: Clarifying the Definition and Novelty of Feature–Classifier Mismatch\n\nWe thank the reviewer for this important question. Below we provide a formal definition of feature–classifier mismatch, and explain in detail how it differs from the discrepancies studied in prior works [1–7].\n\n### (1) Formal Definition of Feature–Classifier Mismatch\n\nWe define feature–classifier mismatch as:\n$$\n\\text{Mismatch}(f, W) = \\mathbb{E}_{(x,y)\\sim D_i}[\\mathcal{L}(f', W) - \\mathcal{L}(f, W)]\n$$\nwhere:\n\n- $D_i$ is the data distribution of client $i$.\n- $f$ denotes the original local features extracted by the feature extractor,\n- $W$ is the classifier,\n- $f'$ is the transformed feature produced by feature transformation module (e.g., linear transformation layer),\n- $\\mathcal{L}(\\cdot, \\cdot) $ is the loss function (e.g., cross-entropy).\n\nThis quantity measures how much the performance (i.e., classification loss) improves when we align the local features with the classifier. Since raw loss values are often not intuitive, we use the accuracy gap between **Origin Acc.**  and **Match Acc.** (as we defined in the caption of Figure 1) to reflect the mismatch degree throughout the paper.\n\n### (2) Key Difference from Prior Work\n\nWe highlight the **fundamental difference** between our work and existing literature.\n\n**1. Objective: Inter-Client vs. Intra-Client Alignment**\n\n- Prior works on feature/classifier [1, 2, 3] and weight space alignment [4, 5, 6, 7] mainly focus on **inter-client consistency**, aiming to reduce **client drift** by aligning models, features, or classifiers **across clients**.\n- In contrast, **our method explicitly targets intra-client mismatch**—that is, the misalignment between a client’s own local features and the shared global classifier. This is a new problem that has not been studied in federated learning but is critical for stable training and accurate inference under heterogeneous data.\n\n**2. Our Advantage**\n\n**Existing methods fail to address intra-client mismatch, whereas our approach simultaneously resolves intra-client mismatch and enhances inter-client consistency to mitigate client drift.** This joint resolution is crucial, because even with successful inter-client alignment, feature–classifier mismatches can still persist within individual clients, leading to degraded performance (as illustrated on the left side of Figure 1).\n\n- Our method tackles this issue at its root by first resolving the mismatch before local updates. This enables training to begin from an aligned feature space and avoids interference between the feature extractor and classifier—reducing client drift (Figure 1, bottom right).\n- Furthermore, by transforming local features to match the shared semantic space defined by the global classifier, our method **implicitly promotes inter-client feature consistency**, as evidenced by class-wise clustering across clients (Appendix Figure 9(e)).\n\nWe thank the reviewer once again for the insightful question. We will incorporate the above clarification, formal definition, and comparative discussion into the revised version of the paper to better highlight the novelty and significance of our work.\n\n\n\n## Weakness 2 & Question 2: Quantitative Evidence for Mismatch\n\nWe thank the reviewer for this insightful suggestion. In response, we conducted two additional experiments during the rebuttal phase to provide more **quantitative evidence** beyond accuracy:\n\n### (1) Gradient Divergence\n\nWe measured the gradient divergence of the feature extractor across clients throughout training, and reported the **minimum, maximum, and average** divergence observed over the entire training process. (Figures are omitted per rebuttal policy.)\n\n|            | Min  | Max  | Average |\n| ---------- | ---- | ---- | ------- |\n| FedAvg     | 0.93 | 3.38 | 1.62    |\n| FedBABU    | 0.85 | 2.81 | 1.45    |\n| FedRoD     | 0.83 | 3.20 | 1.43    |\n| **FedPFT** | 0.67 | 1.82 | 1.36    |\n\nThe results show that **FedPFT achieves consistently lower gradient divergence** compared to other methods. This suggests that resolving the feature–classifier mismatch prior to model updates leads to more consistent optimization directions, thereby mitigating client drift caused by non-IID data.\n\n\n\n### (2) Feature–Classifier Cosine Similarity\n\nWe also computed the **cosine similarity between local features and their corresponding classifier weights**. As shown in the table below, FedPFT achieves comparable values to strong baselines like FedBABU-FT.\n\n| CIFAR 10  | FedAvg | FedAvg-FT | FedBABU | FedBABU-FT | FedPFT |\n| --------- | ------ | --------- | ------- | ---------- | ------ |\n| Dir (0.5) | 0.20   | 0.26      | 0.28    | 0.39       | 0.42   |\n| Dir (1.0) | 0.22   | 0.24      | 0.31    | 0.37       | 0.36   |\n\nHowever, we note that **this similarity metric is not strongly correlated with accuracy or loss**. It can be influenced by confounding factors such as feature norm and classifier bias terms. This is why we **do not adopt this metric to define mismatch** in our paper.\n\nInstead, we define mismatch based on the **performance gap before and after feature transformation**, which directly reflects model effectiveness. Our method is also designed accordingly—to **learn a transformation module that minimizes classification loss**, rather than enforcing explicit angular constraints.\n\n## Weakness 3: On the Role of Contrastive Learning\n\nWe appreciate the reviewer’s observation. We clarify that the integration of contrastive learning is **not presented as a standalone contribution**, but rather intended to illustrate the **flexibility and extensibility** of our framework. In many real-world applications, achieving high accuracy is critical. Our method’s modular prompt-based design naturally supports the integration of complementary techniques—such as contrastive learning—to further boost performance when necessary. This makes FedPFT practically valuable and well-suited for a wide range of deployment scenarios.\n\n\n\n## Weakness 4: On Overhead\n\nWe thank the reviewer for raising this important point. We agree that FedPFT introduces additional model complexity and communication overhead due to the inclusion of the feature transformation module. This results in approximately 20% more resource consumption compared to baseline methods.\n\nWhile such overhead is relatively minor in cross-silo settings where clients typically have sufficient computational and communication resources, we acknowledge that it may pose challenges in resource-constrained cross-device scenarios. We sincerely appreciate the reviewer’s attention to this aspect.\n\nAs part of future work, we plan to investigate more lightweight feature transformation modules to further reduce the overhead while preserving the effectiveness of our approach.\n\n## Question 3: On Robustness Comparison\n\nWe have conducted additional experiments under a **20% client participation** setting and compared FedPFT with SOTA methods. The results show that FedPFT maintains strong robustness and consistently outperforms baselines, even under sparse participation.\n\n| CIFAR100, Dir 0.5 | 90%       | 70%       | 50%       | 20%       |\n| ----------------- | --------- | --------- | --------- | --------- |\n| FedAvg-FT         | 41.50     | 40.88     | 41.95     | 40.23     |\n| FedBABU-FT        | 40.33     | 41.28     | 41.65     | 38.80     |\n| FedRoD            | 39.75     | 40.55     | 38.37     | 38.23     |\n| **FedPFT**        | **45.45** | **46.08** | **46.22** | **46.10** |\n\n## Question 4: Communication and Accuracy on OOD Clients\n\nWe thank the reviewer for raising this important point. In our current work, our primary focus is on achieving strong performance in in-distribution (ID) settings, and we have not conducted evaluations on out-of-distribution (OOD) clients. We agree that evaluating performance on out-of-distribution (OOD) clients is an interesting and meaningful extension. We consider this an important direction and plan to leave a comprehensive OOD evaluation to future work.\n\n## Question 5: Comparison with FL Benchmarks [8, 9]\n\nWe thank the reviewer for this valuable suggestion. We have carefully reviewed both **FederatedScope [9]** and **FL-Bench [8]**, and provide the following clarification:\n\n- **FederatedScope** is a highly modular FL platform that emphasizes system-level efficiency and rapid deployment. However, it currently provides limited support for state-of-the-art personalization algorithms, which are the focus of our work. As our contribution is algorithmic in nature, we chose not to base our implementation on FederatedScope.\n- **FL-Bench**, on the other hand, is a well-designed research benchmark that includes many strong personalized FL methods, standardized datasets, model architectures, and non-IID settings. Our experimental setup is **highly aligned with FL-Bench**, including:\n  - Use of public datasets (CIFAR-10, CIFAR-100, TinyImageNet);\n  - Adoption of Dirichlet and feature-shift non-IID partitions;\n  - Use of standard model architectures (e.g., ResNet);\n  - Client counts of 40 and 100 for small- and large-scale evaluations;\n  - Including many baselines that are also part of FL-Bench, such as FedRep, FedPer, FedBN, FedRoD, and FedBABU.\n\nAlthough our experiments are not strictly based on the FL-Bench codebase, our design choices are largely consistent with its experimental protocols, enabling fair and meaningful comparisons."}, {"review_id": "bxNGznFevb", "reviewer": "Reviewer_JoZC", "summary": "This paper proposes FedPFT, a framework that uses personalized prompts to address the mismatch between local features and the global classifier during training in heterogeneous FL. FedPFT improves feature alignment and supports additional training objectives like contrastive learning, achieving up to 5.07% better performance and 7.08% further gains with collaborative contrastive learning.", "questions": "1) Why does a smaller parameter in the Dirichlet distribution lead to a more heterogeneous setting? In most previous papers, a smaller Dirichlet parameter typically implies a more heterogeneous data distribution.", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1) The paper is well-motivated with clear and effective representations. \n2) The experiments include extensive baselines for thorough evaluation. \n\nWeaknesses:\n1) The proposed method appears somewhat straightforward and lacks theoretical motivation and analysis. \n2) The experiments are conducted on a relatively small scale, limiting the evaluation's generalizability. \n\nConsidering the efforts in Table 1, I tend to give a marginaly positive score.\n\nMinor: There are also several works addressing feature and classifier misalignment, which are recommended to be included in the related works section for discussion. Some of these methods are designed for G-FL, while others are applicable to P-FL. To name a few:\n* FedAWS:Federated learning with only positive labels \n* FedMR: Federated learning under partially class-disjoint data via manifold reshaping \n* FedRS: Fedrs: Federated learning with restricted softmax for label distribution non-iid data \n* FedGELA: Federated learning with bilateral curation for partially class-disjoint data", "was_revised": true, "final_justification": "The author solved most of my questions. And I found a typo and misundertanding in my proposed question. After discussed with the authors, the question is solved. Therefore I kept the positive rating.", "rebuttal": "We sincerely thank the reviewer for valuable feedback. We try our best to carefully address each of the raised weaknesses and questions. Additionally, we would like to emphasize that the code for our method was submitted along with the original submission, and we commit to open-sourcing the additional experiments conducted during the rebuttal phase if this paper is accepted.\n\n## Weakness 1: Lack of Theoretical Motivation and Analysis\n\nWe thank the reviewer for the thoughtful comment. We acknowledge the importance of conducting thorough theoretical analysis to strengthen the foundation of our work. While our current motivation is primarily empirical, grounded in practical observations, we have included a preliminary theoretical convergence analysis in Appendix O of the original manuscript. We recognize that a more comprehensive theoretical study would further enhance the rigor of our approach, and we leave this as an important direction for future work.\n\n## Weakness 2: Limited Scale of Experiments\n\nWe appreciate the reviewer’s concern. We have provided **larger-scale experiments** in **Appendix K (Table 12 and Table 13)** of the original manuscript, including:\n\n- **100 participating clients**,\n- **Deeper model architectures**, such as **ResNet-18**\n\nFor the reviewer’s convenience, we highlight these results here. As shown, FedPFT consistently outperforms SOTA baselines even under these more demanding configurations, demonstrating its robustness and scalability.\n\n| 100 clients  | FedPer | FedBN | FedRoD | FedCAC | FedPFT    |\n| ------------ | ------ | ----- | ------ | ------ | --------- |\n| $\\alpha=0.1$ | 84.68  | 85.51 | 87.58  | 87.40  | **88.43** |\n| $\\alpha=0.5$ | 71.40  | 70.85 | 75.23  | 72.82  | **76.49** |\n| $\\alpha=1.0$ | 66.92  | 67.18 | 70.99  | 69.06  | **72.53** |\n\n| ResNet-18    | FedPer | FedBN | FedRoD | FedCAC | FedPFT    |\n| ------------ | ------ | ----- | ------ | ------ | --------- |\n| $\\alpha=0.1$ | 82.80  | 78.10 | 83.50  | 82.60  | **83.52** |\n| $\\alpha=0.5$ | 65.35  | 55.38 | 67.55  | 63.27  | **68.93** |\n| $\\alpha=1.0$ | 61.83  | 56.88 | 62.05  | 60.53  | **64.95** |\n\nIn the rebuttal phase, we further conducted larger-scale experiments with **200 clients**. As shown below, FedPFT continues to outperform SOTA methods under these settings as well.\n\n| 200 clients  | FedAvg-FT | FedBABU-FT | FedRoD | FedPFT |\n| ------------ | --------- | ---------- | ------ | ------ |\n| $\\alpha=0.1$ | 86.20     | 87.04      | 85.68  | 87.90  |\n| $\\alpha=0.5$ | 72.82     | 73.80      | 72.60  | 74.81  |\n| $\\alpha=1.0$ | 70.04     | 71.31      | 69.80  | 72.32  |\n\n## Minor: Related Works on Feature–Classifier Misalignment\n\nWe thank the reviewer for pointing out these additional related works.\n\nWe have carefully reviewed the references provided by the reviewer and found that these works primarily address **feature or classifier misalignment across clients**. In contrast, our method focuses on the **intra-client mismatch between local features and the global classifier**, which we believe is an orthogonal but equally important problem.\n\nWe will include a dedicated discussion in the revised version comparing our approach with feature/classifier alignment methods, incorporating the references suggested by you and other reviewers.\n\n## Question 1: Dirichlet Parameter and Data Heterogeneity\n\nWe thank the reviewer for the question. In federated learning, a smaller Dirichlet parameter ($\\alpha$) leads to sparser and more concentrated class distributions on each client—that is, clients are more likely to receive data from only a few classes, resulting in stronger inter-client heterogeneity. Conversely, larger $\\alpha$ values yield more balanced and overlapping distributions across clients.\n\nWe also note that more details about the Dirichlet-based non-IID setting are provided in Appendix D.1, and Figure 8 provides a visual illustration to help readers better understand how varying $\\alpha$ influences the data distribution."}]

Author rebuttal:

Meta-review: This paper presents FedPFT, a personalized federated learning framework that addresses data heterogeneity by mitigating the mismatch between local features and the global classifier. The work makes valuable contributions: 1) it offers a targeted solution to feature-classifier mismatch in federated learning; 2) it introduces a practical prompt-driven framework with strong empirical performance; and 3) it demonstrates compatibility with auxiliary objectives like contrastive learning, enhancing its versatility. Key concerns raised by reviewers were effectively addressed during the rebuttal process. Thus, it is strongly recommended that the authors integrate the constructive comments from reviewers to further enhance the work’s quality.
Decision: Accept (poster)