Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (2024)

Xinghao Wu^1,∗, Jianwei Niu^1,2, Xuefeng Liu ${}^{1,2,\text{\textdagger}}$ , Mingjia Shi³,
Guogang Zhu¹, and Shaojie Tang⁴
¹School of Computer Science and Engineering, Beihang University, Beijing, China
²Zhongguancun Laboratory, Beijing, China
³ Sichuan University, Sichuan, China
⁴Department of Management Science and Systems, University at Buffalo

Abstract

In traditional Federated Learning approaches like FedAvg, the global model underperforms when faced with data heterogeneity. Personalized Federated Learning (PFL) enables clients to train personalized models to fit their local data distribution better. However, we surprisingly find that the feature extractor in FedAvg is superior to those in most PFL methods. More interestingly, by applying a linear transformation on local features extracted by the feature extractor to align with the classifier, FedAvg can surpass the majority of PFL methods. This suggests that the primary cause of FedAvg’s inadequate performance stems from the mismatch between the locally extracted features and the classifier. While current PFL methods mitigate this issue to some extent, their designs compromise the quality of the feature extractor, thus limiting the full potential of PFL.In this paper, we propose a new PFL framework called FedPFT to address the mismatch problem while enhancing the quality of the feature extractor. FedPFT integrates a feature transformation module, driven by personalized prompts, between the global feature extractor and classifier. In each round, clients first train prompts to transform local features to match the global classifier, followed by training model parameters. This approach can also align the training objectives of clients, reducing the impact of data heterogeneity on model collaboration. Moreover, FedPFT’s feature transformation module is highly scalable, allowing for the use of different prompts to tailor local features to various tasks. Leveraging this, we introduce a collaborative contrastive learning task to further refine feature extractor quality. Our experiments demonstrate that FedPFT outperforms state-of-the-art methods by up to 7.08%.

^†^†footnotetext: ^∗wuxinghao@buaa.edu.cn^†^†footnotetext:

{}^{\text{\textdagger}}

Corresponding author

1 Introduction

Federated Learning (FL) allows all clients to train a global model collaboratively without sharing their raw data. A key challenge in FL is data heterogeneity, meaning the data across different clients is not independently and identically distributed (non-IID). This issue results in degraded performance of the global model trained in conventional FL methods such as FedAvg [27].

To address this issue, Personalized Federated Learning (PFL) has been proposed, which allows clients to train personalized models to fit their local data distribution better. Many current PFL methods achieve personalization by personalizing some parameters of the global model. For example, FedPer [2] personalizes classifiers, FedBN [20] personalizes BN layers, AlignFed [39] personalizes feature extractors, and FedCAC [35] selects parameters susceptible to non-IID effect for personalization.

	CIFAR-10, $\alpha=0.5$			CIFAR-10, $\alpha=1.0$
Methods	Probe Acc.	Origin Acc.	Match Acc.	Probe Acc.	Origin Acc.	Match Acc.
FedAvg	72.52%	59.66%	72.60%	68.38%	60.33%	68.37%
FedPer	71.07%	68.86%	71.03%	66.51%	64.83%	66.75%
FedBN	70.15%	66.20%	70.60%	66.51%	62.97%	66.80%
FedCAC	71.56%	68.71%	71.63%	66.98%	64.90%	67.11%
Ours	77.25%	77.06%	77.68%	74.02%	73.88%	74.75%

Although the above methods demonstrate significant performance improvements over the global model, an interesting observation emerged from our experiments: the feature extractor derived from FedAvg outperforms those in most PFL methods. Specifically, we conduct linear probe experiments in which each client employs a randomly initialized linear classifier (probe) behind the feature extractor, and this classifier is subsequently retrained. As evident from Table1, the Probe Acc. of FedAvg exceeds that of the PFL methods, indicating that the features extracted by FedAvg exhibit superior linear separability. This suggests that FedAvg has greater potential to outperform PFL methods.

These findings prompt us to further explore why FedAvg underperforms on client-local data compared to PFL methods. To unveil this puzzle, we introduce a linear layer between the global feature extractor and the classifier on each client, training this layer with local data to align the features with the classifier. According to the Match Acc. in Table1, applying a linear transformation to local features significantly improves accuracy over the original model (Origin Acc.), even exceeding the Origin Acc. of current PFL methods. This indicates that the fundamental reason for the global model’s inadequate performance lies in the mismatch between local features and the global classifier.

Further experiments with PFL methods demonstrate that while they somewhat mitigate the mismatch issue, their design inadvertently degrades the quality of the feature extractor, leading to a lower Match Acc. compared to FedAvg. More importantly, current PFL methods still face issues of mismatch. This problem not only diminishes model accuracy during inference but also affects the synergy between the feature extractor and the classifier during training, ultimately impacting the feature extractor’s quality. These observations suggest that significant untapped potential remains within PFL.

In PFL, targeted designs are imperative to tackle the mismatch problem during training and improve the quality of the feature extractor. Hence, we introduce a novel PFL method called FedPFT. Drawing inspiration from prompt technology [11], which utilizes prompts as inputs to guide a model’s behavior, FedPFT integrates a vision-prompt-driven feature transformation module between the global feature extractor and classifier. In each iteration, FedPFT initially trains prompts to guide local feature transformation to align with the global classifier. This process aligns the local features with the global feature space partitioned by the classifier, thereby achieving alignment of training objectives among clients. Subsequently, training the model parameters based on this alignment can alleviate the impact of non-IID data on client collaboration and enhance the quality of the feature extractor.

Furthermore, our proposed framework exhibits notable scalability. Clients’ local features can be transformed by different task-specific prompts to accommodate various tasks. Leveraging this capability, we introduce a collaborative contrastive learning task among clients to further enhance the quality of the feature extractor. As evidenced in Table 1, our method not only resolves the mismatch issue but also significantly improves the quality of the feature extractor.

Our main contributions can be summarized as follows:

•
We identify the root cause of the inadequate performance of the global model stemming from the mismatch between local features and the classifier. The reason personalizing some parameters can improve performance is that it alleviates the impact of this issue. This provides a new perspective for future PFL approaches to better address the non-IID problem.
•
We propose a new PFL framework, which incorporates a feature transformation module to align local features with the global classifier. This approach not only resolves the mismatch problem but also significantly enhances the performance of the feature extractor.
•
Our experiments on multiple datasets and non-IID scenarios demonstrate the superiority of FedPFT, outperforming state-of-the-art methods by up to 7.08%.

2 Related Work

PFL is a kind of effective approach to address the challenges of non-IID data in FL. There is a surge of methodologies within PFL, with parameter decoupling methods gaining significant attention due to their simplicity and effectiveness, thus becoming one of the mainstream research directions in PFL. For a more detailed summary of other categories of PFL methods, please refer to AppendixA.

Parameter decoupling methods aim to decouple a subset of parameters from the global model for personalization. Approaches such as FedPer [2], FedRep [5], and GPFL [38] focus on personalizing the classifier. In contrast, methods like LG-FedAvg [22], and AlignFed [39] advocate for the personalization of the feature extractor. Additionally, FedBN [20] and MTFL [28] propose personalizing batch normalization (BN) layers within the feature extractor. Techniques employing deep reinforcement learning [31] or hypernetworks [26] have been used to determine which specific layers to personalize. The recent FedCAC [35] method advances this by introducing a metric for parameter-wise selection.

These decoupling methods help alleviate the mismatch issue within the global model by allowing local parameter adjustments. For instance, personalized classifiers involve local adjustments to the classifier to match it with the local features extracted by the global feature extractor. However, these methods do not completely resolve the mismatch issue during training. Personalizing parameters often reduce the extent of client information exchange, which can diminish the overall quality of the feature extractor, thus limiting the potential benefits of PFL.

3 Methodology

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (1)

3.1 Overview of FedPFT

Before delving into the details of FedPFT, we first provide an overview, as illustrated in Figure1(a). Each training round in FedPFT includes five key steps:

(1) Clients download the global models, which include the feature extractor $\phi$ , feature transformation module $\tau$ , classifier $h_{\kappa}$ , and feature projection head $h_{\rho}$ . These models serve to initialize the corresponding local models $\{\phi_{i},\tau_{i},h_{\kappa,i},h_{\rho,i}\}$ .

(2) Each client $i$ updates the $\phi_{i}$ , $\tau_{i}$ and $h_{\rho,i}$ to minimize contrastive learning loss $L_{\text{Con}}$ , aiming to enhance the generalization of the feature extractor. It also updates $\tau_{i}$ and classification prompts $p_{\kappa,i}$ with the cross-entropy loss $L_{\text{CE}}$ to align local features with the global classifier.

(3) Each client $i$ freezes the prompts $p_{\kappa,i}$ and trains $\phi_{i}$ , $\tau_{i}$ , and $h_{\kappa,i}$ using $L_{\text{CE}}$ to adapt the model to the classification task. It also makes the contrastive learning prompts $p_{\rho,i}$ trainable to align features with the contrastive learning task.

(4) Clients upload $\{\phi_{i},\tau_{i},h_{\kappa,i},h_{\rho,i}\}$ to the server while retaining $\{p_{\kappa,i},p_{\rho,i}\}$ locally.

(5) The server aggregates the models uploaded by the clients.

3.2 Problem Formulation

In PFL, $N$ clients train their personalized models $w_{i},i\in[N]$ under the coordination of a server, aiming for each $w_{i}$ to perform well on client data distribution $\mathbb{D}_{i}$ . This objective can be formalized as $\min_{\{w_{i}\}_{i\in[N]}}\frac{1}{N}\sum_{i=1}^{N}L_{i}(w_{i};\mathbb{D}_{i})$ , where $L_{i}$ represents the loss function of the $i$ -th client.

In this paper, our goal is to enhance personalized models by addressing the mismatch problem between local features and the classifier in the global model and improving the quality of the feature extractor. Thus, the training objective of FedPFT can be formulated as:

\min_{\phi,\tau,h_{\kappa}}\min_{\{p_{\kappa,i}\}_{i\in[N]}}\mathbf{E}_{i}\{L_%{i}(\phi,\tau,h_{\kappa},p_{\kappa,i};d_{i}):=\mathbf{E}_{d_{i}}[L_{\text{CE}}%(\phi,\tau,h_{\kappa},p_{\kappa,i};d_{i})+L_{\text{Con}}(\phi,\tau;d_{i})]\},

(1)

where $\phi$ and $h_{\kappa}$ represent the feature extractor and classifier of the global model, respectively. $\tau$ is the newly introduced global feature transformation module. This module, along with the classification prompt $p_{\kappa,i}$ , transforms local features to align with the global classifier. $L_{\text{CE}}$ denotes the cross-entropy loss for classification tasks, while $L_{\text{Con}}$ represents the contrastive learning loss designed to enhance the feature extractor’s quality. $d_{i}$ represents the local data of the client.

3.3 Feature Transformation Module

In FedPFT, we introduce a global feature transformation module $\tau$ , along with a set of prompts $p_{\kappa,i}$ for each client $i$ , to align the features extracted by the feature extractor $\phi$ with the global classifier.

Formally, given a sample $x_{j}\in d_{i}$ , extracted by the feature extractor $\phi$ , the obtained feature is $f_{j}\in\mathbb{R}^{m}$ , where $m$ is the feature dimension. A collection of $n$ prompts is denoted as $p=\{\boldsymbol{p}^{k}\in\mathbb{R}^{m}|k\in\mathbb{N},1\leq k\leq n\}$ . The operation of the feature transformation module is formulated as

[f_{j}^{\prime},p^{\prime}]=\tau([f_{j},p]),

(2)

where $[\cdot,\cdot]$ signifies stacking and concatenation along the sequence length dimension, yielding $[f_{j}^{\prime},p^{\prime}]\in\mathbb{R}^{(1+n)\times m}$ . The $f_{j}^{\prime}$ represents the transformed feature. An example of the feature transform module is illustrated in Figure1(b).

The feature transformation module essentially adapts local features for downstream tasks, providing good scalability. We can introduce tasks beneficial for client collaboration using different task-specific prompts $p$ . Leveraging this, FedPFT additionally introduces a contrastive learning task and utilizes contrastive learning prompts $p_{\rho,i}$ for feature transformation. We denote $n_{\kappa}$ and $n_{\rho}$ as the number of prompts contained in $p_{\kappa,i}$ and $p_{\rho,i}$ , respectively.

3.4 Classification Task with Personalized Prompts

The classifier is highly susceptible to the influence of non-IID data, leading to a mismatch between the global classifier and local features. Different from the previous methods, which personalize the classifier to match local features, we find that using a global classifier provides clients with a unified feature partition space. Clients aligning features with this space not only solves the mismatch problem but also aligns training objectives among clients, reducing the impact of non-IID on collaboration.

To implement this, we retain the global feature extractor and classifier while employing a set of personalized classification prompts $p_{\kappa,i}$ to transform each client $i$ ’s local features to better align with the global classifier. Specifically, the classification loss in each client $i$ is defined as:

L_{\text{CE}}(\phi,\tau,p_{\kappa,i},h_{\kappa};x,y)=-\log\sum_{c=1}^{C}y_{c}%\log(o_{i,c}),\text{where }x,y\sim d_{i}.

(3)

$C$ is the number of classes, and $o_{i}=\text{Softmax}(h_{\kappa}\circ\tau([\phi(x),p_{\kappa,i}]))$ represents the predicted probabilities, with $o_{i,c}$ being the ones of class $c$ . Details on coordinating the training of the model and prompts to achieve feature and classifier alignment are discussed in Section3.6.

3.5 Contrastive Learning Task

Contrastive learning tasks have shown robustness to the challenges posed by non-IID data distributions [34]. To further enhance the quality of the model’s feature extractor, we introduce a contrastive learning task using the Momentum Contrast (MoCo) [8] framework. The associated contrastive loss function is defined as:

L_{\text{Con}}(\phi,\tau,p_{\rho,i},h_{\rho};x)=-\log\frac{\exp\left(q\cdot k_%{+}/\beta\right)}{\sum_{j=0}^{K}\exp\left(q\cdot k_{j}/\beta\right)},\text{%where }x\sim d_{i}.

(4)

$h_{\rho}$ is the projection head used for contrastive learning. In this formula, $q=h_{\rho}\circ\tau([\phi(x^{\prime}),p_{\rho,i}])$ represents the query vector, and $k_{+}=\tilde{h}_{\rho}\circ\tau([\tilde{\phi}(x^{\prime\prime}),p_{\rho,i}])$ denotes the positive key vector. Here, $x^{\prime}$ and $x^{\prime\prime}$ are augmented versions of the sample $x$ , $\tilde{\phi}$ and $\tilde{h}_{\rho}$ refer to the momentum-updated encoder and projection head, respectively. $\beta$ is a temperature hyperparameter, and $K$ is the number of negative samples drawn from MoCo’s queue, comprising the set $\{k_{j}\}_{j=0}^{K}$ .

3.6 Alternating Training Strategy

To effectively coordinate the training of different modules in FedPFT, we propose an alternating training strategy, which partitions each local training round into two distinct phases: the feature learning phase and the task adaptation phase.

Feature Learning Phase.

In this phase, the training objective can be formulated as

\mathop{\min}\limits_{p_{\kappa,i},\phi_{i},\tau_{i},h_{\rho,i}}\left\{L_{%\text{CE}}(\phi_{i},\tau_{i},p_{\kappa,i},h_{\kappa};d_{i})+L_{\text{Con}}(%\phi_{i},\tau_{i},p_{\rho,i},h_{\rho,i};d_{i})\right\}.

(5)

$L_{\text{CE}}$ trains the classification prompts $p_{\kappa,i}$ , aligning local features with the global classifier $h_{\kappa}$ , while $L_{\text{Con}}$ is aimed at training the feature extractor $\phi_{i}$ to derive general feature representations and mitigate the impact of non-IID data. Notably, during this phase, $\phi_{i}$ is exclusively updated by $L_{\text{Con}}$ .

Task Adaptation Phase.

Following the above phase, this phase refines the previously learned features for the classification task and further enhances the classifier. The training objective is

\mathop{\min}\limits_{\phi_{i},\tau_{i},h_{\kappa,i},p_{\rho,i}}\left\{L_{%\text{CE}}(\phi_{i},\tau_{i},p_{\kappa,i},h_{\kappa,i};d_{i})+L_{\text{Con}}(%\phi_{i},\tau_{i},p_{\rho,i},h_{\rho,i};d_{i})\right\}.

(6)

The second item in Eq.(6) focuses on training the contrastive learning prompts $p_{\rho,i}$ , aiming to align the feature with the contrastive learning task and mitigate the impact of the classification task on the contrastive learning task. In this phase, $\phi_{i}$ is updated solely by $L_{\text{CE}}$ .

Let $R$ represent the total number of local epochs in one training round. We divide into $R_{f}$ epochs for the feature learning phase and $R_{a}$ epochs for the task adaptation phase, where $R_{f}+R_{a}=R$ . It is crucial that $R_{f}$ is always larger than $R_{a}$ to ensure: 1) $\phi_{i}$ is predominantly trained by the contrastive learning task, reducing the impact of the non-IID problem on collaboration in feature extraction; 2) Improved alignment of features utilized for classification at the client side with the global classifier, thereby achieving better alignment of training objectives across clients.

Upon completing local training, the parameters $\phi_{i}$ , $\tau_{i}$ , $h_{\kappa,i}$ , and $h_{\rho,i}$ are aggregated at the server to foster client collaboration, while $p_{\kappa,i}$ and $p_{\rho,i}$ remain locally. We simply adopt the aggregation method used in FedAvg. The pseudo-code of FedPFT is summarized in Algorithm1.

Input: Each client’s initial personalized prompts $p_{\kappa,i}^{(0)}$ and $p_{\rho,i}^{(0)}$ ; The initial global models $\{\phi^{(0)},\tau^{(0)},h_{\kappa}^{(0)},h_{\rho}^{(0)}\}$ ;Client Number $N$ ;Total round $T$ ; Epochs of two learning phases $R_{f}$ and $R_{a}$ .

Output: Personalized model $\{\phi^{(T)},\tau^{(T)},h_{\kappa}^{(T)},p_{\kappa,i}^{(T)}\}$ for each client.

for $t=0$ to $T-1$ do

Client-side:

for $i=1$ to $N$ in paralleldo

Initializing $\{\phi_{i}^{(t)},\tau_{i}^{(t)},h_{\kappa,i}^{(t)},h_{\rho,i}^{(t)}\}$ with $\{\phi^{(t)},\tau^{(t)},h_{\kappa}^{(t)},h_{\rho}^{(t)}\}$ .

4 Experiments

4.1 Experimental Setup

Datasets.

We employ three datasets for experimental validation: CIFAR-10 [14], CIFAR-100 [13], and Tiny ImageNet [15]. We utilize two scenarios: Dirichlet non-IID and Pathological non-IID.

In our experiments, each client is assigned 500 training samples. For CIFAR-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 nine state-of-the-art (SOTA) methods: FedAMP [10], Fedper [2], FedRep [5], FedBN [20], FedRoD [4], pFedSD [12], pFedGate [3], FedCAC [35], and pFedPT [17]. These methods cover the advancements in mainstream PFL research directions.

Hyperparameter Settings. For the general hyperparameters of FL, we set the number of clients $N=40$ , Batch Size $B=100$ , and local update rounds $R=5$ . Across all datasets, we set the total rounds $T=1000$ in each experiment to ensure convergence and select the highest average accuracy achieved by all clients across all rounds as the result. Each experiment is 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 and ResNet-10 for CIFAR-100 and Tiny ImageNet.

For more details on the experimental setup, please refer to AppendixB.

4.2 Comparison with State-of-the-art Methods

	CIFAR-100			Tiny ImageNet
Methods	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$
FedAvg	34.91 $\pm$ 0.86	32.78 $\pm$ 0.23	33.94 $\pm$ 0.39	21.26 $\pm$ 1.28	20.32 $\pm$ 0.91	17.20 $\pm$ 0.54
Local	47.61 $\pm$ 0.96	22.65 $\pm$ 0.51	18.76 $\pm$ 0.63	24.07 $\pm$ 0.62	8.75 $\pm$ 0.30	6.87 $\pm$ 0.28
FedAMP	46.68 $\pm$ 1.06	24.74 $\pm$ 0.58	18.22 $\pm$ 0.41	27.85 $\pm$ 0.71	10.70 $\pm$ 0.32	7.13 $\pm$ 0.21
FedPer	51.38 $\pm$ 0.94	28.25 $\pm$ 1.03	21.53 $\pm$ 0.50	32.33 $\pm$ 0.31	12.69 $\pm$ 0.42	8.67 $\pm$ 0.40
FedRep	51.25 $\pm$ 1.37	26.97 $\pm$ 0.33	20.63 $\pm$ 0.42	30.83 $\pm$ 1.05	12.14 $\pm$ 0.28	8.37 $\pm$ 0.25
FedBN	54.35 $\pm$ 0.63	36.94 $\pm$ 0.94	33.67 $\pm$ 0.12	33.34 $\pm$ 0.71	19.61 $\pm$ 0.35	16.57 $\pm$ 0.44
FedRoD	60.17 $\pm$ 0.48	39.88 $\pm$ 1.18	36.80 $\pm$ 0.56	41.06 $\pm$ 0.77	25.63 $\pm$ 1.11	22.32 $\pm$ 1.13
pFedSD	54.14 $\pm$ 0.77	41.06 $\pm$ 0.83	38.27 $\pm$ 0.20	39.31 $\pm$ 0.19	19.25 $\pm$ 1.80	15.91 $\pm$ 0.33
pFedGate	48.54 $\pm$ 0.39	27.47 $\pm$ 0.79	22.98 $\pm$ 0.03	37.59 $\pm$ 0.39	24.09 $\pm$ 0.67	19.69 $\pm$ 0.14
FedCAC	57.22 $\pm$ 1.52	38.64 $\pm$ 0.63	32.59 $\pm$ 0.32	40.19 $\pm$ 1.20	23.70 $\pm$ 0.28	18.58 $\pm$ 0.62
pFedPT	43.21 $\pm$ 1.66	35.23 $\pm$ 0.87	36.25 $\pm$ 0.37	23.55 $\pm$ 0.68	22.35 $\pm$ 0.49	21.69 $\pm$ 0.24
FedPFT w/o $L_{\text{Con}}$	60.98 $\pm$ 0.39	44.87 $\pm$ 0.76	41.83 $\pm$ 0.37	41.49 $\pm$ 0.10	28.61 $\pm$ 0.40	25.10 $\pm$ 0.59
FedPFT	62.03 $\pm$ 1.41	47.98 $\pm$ 0.78	44.29 $\pm$ 0.74	43.42 $\pm$ 1.62	32.44 $\pm$ 0.58	27.84 $\pm$ 0.41

In this section, we compare our proposed FedPFT with two baseline methods and nine SOTA methods across three datasets and two non-IID scenarios. We also introduce ‘FedPFT w/o $L_{\text{Con}}$ ,’ which solely addresses the mismatch problem without contrastive learning. The experimental results on CIFAR-100 and Tiny ImageNet in Dirichlet non-IID scenario are presented in Table2.Please refer to the AppendixC for experimental results in Pathological non-IID scenarios and the CIFAR-10 dataset.

Results in Dirichlet non-IID scenario.

In this setting, by varying $\alpha$ , we can evaluate the performance of methods under different non-IID degrees. The results, as detailed in Table2, demonstrate that performance varies significantly depending on the underlying design principles of each method. Among all methods, FedRoD demonstrates robust performance across all datasets and non-IID degrees. This is attributed to its design of two classifiers: a personalized classifier for local feature alignment and a global classifier for assistance from other clients to improve generalization. ‘FedPFT w/o $L_{\text{Con}}$ ’ addresses the mismatch issue specifically and achieves competitive or superior results across all scenarios. FedPFT further improves feature extractor quality and outperforms SOTA methods significantly across all scenarios, achieving up to a 7.08% improvement.

4.3 Ablation Study

In this section, we validate the effectiveness of each component of FedPFT on the CIFAR-100 dataset under two non-IID degrees. The experimental results are illustrated in Table3.

	$\alpha=0.1$					$\alpha=0.5$
Settings	$p_{\kappa}$	Alter.	$L_{\text{Con}}$	$p_{\rho}$	Accuracy (%)	$p_{\kappa}$	Alter.	$L_{\text{Con}}$	$p_{\rho}$	Accuracy (%)
I					33.87 $\pm$ 1.35					30.09 $\pm$ 0.31
II	✓				40.97 $\pm$ 1.28	✓				31.45 $\pm$ 1.35
III	✓	✓			60.98 $\pm$ 0.39	✓	✓			44.87 $\pm$ 0.76
IV	✓	✓	✓		61.13 $\pm$ 0.50	✓	✓	✓		47.67 $\pm$ 1.42
V	✓	✓	✓	✓	62.03 $\pm$ 1.41	✓	✓	✓	✓	47.98 $\pm$ 0.78
VI			✓		36.24 $\pm$ 1.10			✓		34.70 $\pm$ 1.33
VII	✓		✓		53.17 $\pm$ 0.58	✓		✓		38.90 $\pm$ 0.91
VIII	✓		✓	✓	53.76 $\pm$ 0.35	✓		✓	✓	39.29 $\pm$ 1.00

Setting I represents FedAvg. Setting II incorporates classification prompts $p_{\kappa}$ to allow each client to adjust the global model individually to obtain a personalized model, resulting in a performance improvement. Setting III incorporates alternating training, where prompts are firstly updated to align local features with the global classifier, followed by training model parameters. This approach essentially aligns training objectives among clients. This effectively mitigates the impact of non-IID data on model collaboration, thus further enhancing the quality of the global model.

Setting IV adds contrastive learning loss to Setting III, , focusing primarily on enhancing the feature extractor’s performance through contrastive learning techniques. Setting V incorporates specific prompts $p_{\rho}$ for the contrastive learning task. This reduces mutual interference between the two tasks during training, especially effective when non-IID is strong (e.g., $\alpha=0.1$ ).

Setting VI illustrates that adding contrastive learning alone brings very limited improvements. Settings VII and VIII partially achieve feature-classifier alignment by introducing $p_{\kappa}$ , greatly enhancing model performance. However, without using alternating training, local features cannot adapt well to the classifier. This leads to a significant performance gap between Settings VII, VIII, and Setting V.

This ablation study underlines the importance of each module in FedPFT. It confirms that aligning local features with the global classifier and enhancing the feature extractor’s quality are both crucial for optimizing model performance, aligning with the core motivations behind our methodology.

4.4 Separability of Features

In this section, we assess the effectiveness of FedPFT in enhancing the quality of the feature extractor by conducting linear probing experiments on CIFAR-10 and CIFAR-100. The results are shown in Table4. Higher accuracy means that the extracted features have better linear separability.

Compared to FedAvg, features extracted by ‘FedPFT w/o $L_{\text{Con}}$ ’ demonstrate superior linear separability. This improvement is attributable to the alignment of local features with the global classifier in the feature learning phase, which synchronizes client training objectives and mitigates the adverse effects of non-IID data. FedPFT further improves the quality of the feature extractor by integrating a collaborative contrastive learning task. For more experimental results, please refer to AppendixD.

	CIFAR-10			CIFAR-100
Methods	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$
FedAvg	85.01%	72.52%	68.38%	59.50%	37.40%	32.33%
FedPFT w/o $L_{\text{Con}}$	85.52%	72.59%	69.57%	61.60%	43.14%	38.47%
FedPFT	87.83%	77.25%	74.02%	64.12%	46.43%	40.95%

4.5 Learned Features of Different Methods

In this subsection, we visually compare the quality of features extracted by different methods and highlight the impact of different modules in FedPFT on feature extraction. We conduct experiments on the CIFAR-10 dataset with 10 clients, each allocated 1000 training images and 500 testing images. The data distribution is shown in Figure2(a). For each method, we visualize the feature vectors of testing data from different clients using t-SNE [33]. The visualization results are depicted in Figure2(b)-(h), where colors represent different data classes, and markers represent different clients, as detailed in Figure2(a).

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (2)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (3)

FedAvg and FedCAC exhibit noticeable cluster structures of features but lack strong discriminative boundaries. FedPer displays overlapping features across various classes, attributable to the use of personalized classifiers that create different local feature spaces for each client. Consequently, data from different classes across different clients are mapped to similar positions. This interference between clients reduces the quality of the global feature extractor.

‘FedPFT w/o $L_{\text{Con}}$ ’ shows clearer discriminative boundaries, which is 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 the global classifier provides a unified feature space for all clients. Adapting local features to this space essentially aligns the training objectives among clients in non-IID scenarios, promoting collaboration among clients. FedPFT further enhances feature separability by incorporating $L_{\text{Con}}$ .

‘FedPFT w/o Alter’ represents not using alternating training. While it shows better clustering than FedAvg, the discriminative quality of the boundaries is weaker compared to ‘FedPFT w/o $L_{\text{Con}}$ .’ This configuration shows increased interference among client models, lacking alignment to the common global feature space. ‘FedPFT_p_classifier’ indicates using personalized classifiers. In this case, the feature space becomes highly scattered, similar to FedPer’s issue. Since we train prompt $p_{\kappa}$ to adapt to personalized classifiers first, this exacerbates the variability in feature spaces across clients

4.6 Effect of Different Prompts

In this section, we delve into the role of prompts in FedPFT. We visualize the features transformed by different prompts using t-SNE. The experimental setup is consistent with Section4.5. The results are depicted in Figure3. Larger markers in the figures represent feature centroids of corresponding classes for each client.

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (10)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (11)

It is evident that features obtained from classification prompts $p_{\kappa}$ are not significantly correlated with image similarity but rather with the distribution of client data. For example, two classes within a client are close together. Conversely, features transformed by contrastive learning prompts are more related to image similarity. For instance, in Figure3(b), the feature centroids of ‘cat’ and ‘dog’ are closer, as are ‘truck’ and ‘automobile,’ which aligns with the principles of contrastive learning.

	CIFAR-10			CIFAR-100
Prompt Type	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$	$\alpha=0.1$	$\alpha=0.5$	$\alpha=0.5$
None	87.69%	77.12%	73.93%	64.08%	46.50%	40.79%
$p_{\kappa}$	87.83%	77.25%	74.02%	64.12%	46.43%	40.95%
$p_{\rho}$	87.82%	77.25%	74.02%	64.18%	46.40%	40.95%

We also investigate whether different types of prompts influence feature separability. We conduct linear probe experiments using the CIFAR-10 and CIFAR-100 datasets. The results are detailed in Table5. In these experiments, we compare three conditions: ‘None’ (no prompts used), ‘ $p_{\kappa}$ ’ (using classification prompts), and ‘ $p_{\rho}$ ’ (using contrastive learning prompts). Interestingly, the accuracies across different prompt conditions are generally similar, suggesting that the use of either type of prompt does not significantly impact the overall quality of the features extracted.

The above experiments demonstrate that prompts work by transforming features into the required format for downstream tasks using task-specific prompts. This also indicates the scalability and adaptability of our designed feature transformation module. It can incorporate various client-collaborative tasks beneficial for enhancing the performance of personalized models through task-specific prompts.

5 Conclusion and Discussion

We observe that the feature extractor from FedAvg surpasses those in most PFL methods, yet it suffers from inadequate performance due to a mismatch between the local features and the classifier. This mismatch issue not only impacts the performance during model inference but also affects the synergy between the feature extractor and the classifier during training. We propose a new PFL method called FedPFT with a prompt-driven feature transform module to address these issues during training. Our experiments demonstrate that FedPFT not only resolves the mismatch issue but also significantly improves the quality of the feature extractor, achieving substantial performance gains compared to state-of-the-art methods. We discuss the limitations and our future work in AppendixJ.

References

[1]D.A.E. Acar, Y.Zhao, R.Zhu, R.Matas, M.Mattina, P.Whatmough, and V.Saligrama.Debiasing model updates for improving personalized federated training.In International Conference on Machine Learning, pages 21–31. PMLR, 2021.
[2]M.G. Arivazhagan, V.Aggarwal, A.K. Singh, and S.Choudhary.Federated learning with personalization layers.arXiv preprint arXiv:1912.00818, 2019.
[3]D.Chen, L.Yao, D.Gao, B.Ding, and Y.Li.Efficient personalized federated learning via sparse model-adaptation.arXiv preprint arXiv:2305.02776, 2023.
[4]H.-Y. Chen and W.-L. Chao.On bridging generic and personalized federated learning for image classification.In International Conference on Learning Representations, 2022.
[5]L.Collins, H.Hassani, A.Mokhtari, and S.Shakkottai.Exploiting shared representations for personalized federated learning.pages 2089–2099, 2021.
[6]W.Deng, C.Thrampoulidis, and X.Li.Unlocking the potential of prompt-tuning in bridging generalized and personalized federated learning.arXiv e-prints, pages arXiv–2310, 2023.
[7]A.Fallah, A.Mokhtari, and A.Ozdaglar.Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.Advances in Neural Information Processing Systems, 33:3557–3568, 2020.
[8]K.He, H.Fan, Y.Wu, S.Xie, and R.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]K.He, X.Zhang, S.Ren, and J.Sun.Deep residual learning for image recognition.In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
[10]Y.Huang, L.Chu, Z.Zhou, L.Wang, J.Liu, J.Pei, and Y.Zhang.Personalized cross-silo federated learning on non-iid data.In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
[11]M.Jia, L.Tang, B.-C. Chen, C.Cardie, S.Belongie, B.Hariharan, and S.-N. Lim.Visual prompt tuning.In European Conference on Computer Vision, 2022.
[12]H.Jin, D.Bai, D.Yao, Y.Dai, L.Gu, C.Yu, and L.Sun.Personalized edge intelligence via federated self-knowledge distillation.IEEE Transactions on Parallel and Distributed Systems, 34(2):567–580, 2022.
[13]A.Krizhevsky, G.Hinton, etal.Learning multiple layers of features from tiny images.2009.
[14]A.Krizhevsky, V.Nair, and G.Hinton.Cifar-10 (canadian institute for advanced research).URL http://www. cs. toronto. edu/kriz/cifar. html, 5, 2010.
[15]Y.Le and X.Yang.Tiny imagenet visual recognition challenge.CS 231N, 7(7):3, 2015.
[16]B.Lester, R.Al-Rfou, and N.Constant.The power of scale for parameter-efficient prompt tuning.arXiv preprint arXiv:2104.08691, 2021.
[17]G.Li, W.Wu, Y.Sun, L.Shen, B.Wu, and D.Tao.Visual prompt based personalized federated learning.Transactions on Machine Learning Research, 2023.
[18]H.Li, W.Huang, J.Wang, and Y.Shi.Global and local prompts cooperation via optimal transport for federated learning.arXiv preprint arXiv:2403.00041, 2024.
[19]T.Li, S.Hu, A.Beirami, and V.Smith.Ditto: Fair and robust federated learning through personalization.In International Conference on Machine Learning, pages 6357–6368. PMLR, 2021.
[20]X.Li, M.JIANG, X.Zhang, M.Kamp, and Q.Dou.FedBN: Federated learning on non-IID features via local batch normalization.In International Conference on Learning Representations, 2021.
[21]Z.Li, X.Shang, R.He, T.Lin, and C.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.
[22]P.P. Liang, T.Liu, L.Ziyin, N.B. Allen, R.P. Auerbach, D.Brent, R.Salakhutdinov, and L.-P. Morency.Think locally, act globally: Federated learning with local and global representations.arXiv preprint arXiv:2001.01523, 2020.
[23]H.Liu, C.Li, Q.Wu, and Y.J. Lee.Visual instruction tuning.Advances in neural information processing systems, 36, 2024.
[24]X.Liu, K.Ji, Y.Fu, W.L. Tam, Z.Du, Z.Yang, and J.Tang.P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks.arXiv preprint arXiv:2110.07602, 2021.
[25]J.Luo and S.Wu.Adapt to adaptation: Learning personalization for cross-silo federated learning.In L.D. Raedt, editor, Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22, pages 2166–2173. International Joint Conferences on Artificial Intelligence Organization, 7 2022.Main Track.
[26]X.Ma, J.Zhang, S.Guo, and W.Xu.Layer-wised model aggregation for personalized federated learning.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10092–10101, 2022.
[27]B.McMahan, E.Moore, D.Ramage, S.Hampson, and B.A. yArcas.Communication-efficient learning of deep networks from decentralized data.In Artificial Intelligence and Statistics, pages 1273–1282. PMLR, 2017.
[28]J.Mills, J.Hu, and G.Min.Multi-task federated learning for personalised deep neural networks in edge computing.IEEE Transactions on Parallel and Distributed Systems, 33(3):630–641, 2021.
[29]M.Shi, Y.Zhou, K.Wang, H.Zhang, S.Huang, Q.Ye, and J.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, volume36, pages 28378–28392. Curran Associates, Inc., 2023.
[30]S.Su, M.Yang, B.Li, and X.Xue.Federated adaptive prompt tuning for multi-domain collaborative learning.In Proceedings of the AAAI Conference on Artificial Intelligence, volume38, pages 15117–15125, 2024.
[31]B.Sun, H.Huo, Y.Yang, and B.Bai.Partialfed: Cross-domain personalized federated learning via partial initialization.Advances in Neural Information Processing Systems, 34:23309–23320, 2021.
[32]C.TDinh, N.Tran, and J.Nguyen.Personalized federated learning with moreau envelopes.Advances in Neural Information Processing Systems, 33:21394–21405, 2020.
[33]L.Vander Maaten and G.Hinton.Visualizing data using t-sne.Journal of machine learning research, 9(11), 2008.
[34]L.Wang, K.Zhang, Y.Li, Y.Tian, and R.Tedrake.Does learning from decentralized non-IID unlabeled data benefit from self supervision?In The Eleventh International Conference on Learning Representations, 2023.
[35]X.Wu, X.Liu, J.Niu, G.Zhu, and S.Tang.Bold but cautious: Unlocking the potential of personalized federated learning through cautiously aggressive collaboration.In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 19375–19384, October 2023.
[36]X.Wu, J.Niu, X.Liu, T.Ren, Z.Huang, and Z.Li.pfedgf: Enabling personalized federated learning via gradient fusion.In 2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pages 639–649. IEEE, 2022.
[37]F.-E. Yang, C.-Y. Wang, and Y.-C.F. Wang.Efficient model personalization in federated learning via client-specific prompt generation.In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 19159–19168, 2023.
[38]J.Zhang, Y.Hua, H.Wang, T.Song, Z.Xue, R.Ma, J.Cao, and H.Guan.Gpfl: Simultaneously learning global and personalized feature information for personalized federated learning.In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 5041–5051, October 2023.
[39]G.Zhu, X.Liu, S.Tang, and J.Niu.Aligning before aggregating: Enabling communication efficient cross-domain federated learning via consistent feature extraction.IEEE Transactions on Mobile Computing, 23(5):5880–5896, 2024.

Appendix A Related Work

Current PFL methods can primarily be categorized into several major types: meta-learning-based methods [7; 1], model-regularization-based methods [32; 19], fine-tuning-based methods [12; 3; 21], personalized-weight-aggregation-based methods [10; 25], and parameter-decoupling-based methods. This paper delves into the issues inherent in the global model of FedAvg and primarily discusses parameter-decoupling methods that rely on the global model.

In addition to the aforementioned methods, a new category based on prompts has recently emerged.

Prompt-based methods.

Recently, prompt technology has garnered widespread attention in the fields of computer vision [11; 23] and natural language processing [16; 24]. This technology involves using prompts as inputs to guide the behavior or output of models, typically for fine-tuning purposes. The domain of PFL has also seen the emergence of prompt-based approaches. Most of these are based on pre-trained models, aiming to train prompts to fine-tune the pre-trained models to fit client-local data, as seen in pFedPG [37], SGPT [6], FedOTP [18], and FedAPT [30]. pFedPT [17] trains both the model and prompts, using prompts at the input level to learn personalized knowledge for fine-tuning the global model to adapt to the client’s local distributions. Our FedPFT fundamentally differs from these methods in its training objective. Rather than fine-tuning, we introduce prompts to guide feature transformations to align with the global classifier, thereby addressing the mismatch issue inherent in the global model during the training process.

Appendix B Experiment Setup

B.1 Introduction to non-IID Scenarios

Pathological non-IID.

In this setting, each client is randomly assigned data from a 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 to each client, respectively.

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (12)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (13)

Dirichlet non-IID.

This is a commonly used setting in current FL research [36; 35; 29] In this scenario, the data for each client is generated from a Dirichlet distribution denoted as $Dir(\alpha)$ . As the value of $\alpha$ increases, the class imbalance within each client’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 a clearer, more intuitive understanding, we involve 20 clients with 10-class and 50-class datasets to visualize the data distribution among clients with varying $\alpha$ values. As depicted in Figure 4, 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 a client, with larger dots signifying a higher volume of data in that class.

B.2 Introduction to Comparative Methods

FedAMP [10] is a weighted-aggregation-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 from the design point of view. FedPer [2], FedRep [5], FedBN [20], FedRoD [4], and FedCAC [35] are parameter-decoupling-based methods, which personalize the global model by retaining certain parameters locally based on FedAvg. FedRoD additionally introduces a balanced global classifier to obtain assistance from other clients, alleviating the overfitting issue caused by personalized classifiers alone. pFedSD [12] and pFedGate [3] are fine-tuning-based 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 it to the global model. pFedPT [17], a prompt-based method, can also be viewed as a fine-tuning approach, enhancing the global model’s adaptation to local data distributions by adding prompts to images.

B.3 Hyperparameter Settings in Different Methods

For the unique hyperparameters of each baseline method, we utilize the optimal parameter combinations reported in their respective papers. For learning rates, we adjust within {1e-1, 1e-2, 1e-3}.

B.4 Compute Resources

All the experiments are implemented using PyTorch and conducted on NVIDIA V100 GPUs. For the methods we compared, as well as ‘FedPFT w/o $L_{\text{Con}}$ ,’ a single training session requires 24-48 hours. For FedPFT, the training process takes longer due to the use of the MoCo algorithm, which requires data augmentation that can only be executed on the CPU. Consequently, a single training session for FedPFT requires 48-72 hours.

Appendix C Comparison with State-of-the-art Methods

We present the comparative results of FedPFT against established methods on CIFAR-10, CIFAR-100, and Tiny ImageNet datasets under Pathological non-IID scenarios, as well as CIFAR-10 under Dirichlet non-IID scenarios in Tables6 and 7.

Methods	CIFAR-10	CIFAR-100	Tiny ImageNet
FedAvg	54.33 $\pm$ 3.03	34.27 $\pm$ 0.44	18.05 $\pm$ 0.23
Local	85.85 $\pm$ 0.93	38.40 $\pm$ 0.69	16.20 $\pm$ 0.30
FedAMP	88.88 $\pm$ 0.83	38.36 $\pm$ 0.79	16.13 $\pm$ 0.55
FedPer	87.51 $\pm$ 0.95	41.54 $\pm$ 0.74	20.25 $\pm$ 0.65
FedRep	87.10 $\pm$ 0.91	40.63 $\pm$ 0.74	19.24 $\pm$ 0.33
FedBN	87.02 $\pm$ 1.41	47.75 $\pm$ 1.03	24.91 $\pm$ 0.48
FedRoD	88.06 $\pm$ 1.70	52.55 $\pm$ 0.92	32.25 $\pm$ 0.80
pFedSD	89.97 $\pm$ 1.45	52.30 $\pm$ 1.18	30.27 $\pm$ 0.78
pFedGate	89.15 $\pm$ 0.76	43.73 $\pm$ 0.14	22.42 $\pm$ 0.83
FedCAC	89.77 $\pm$ 1.14	49.07 $\pm$ 0.87	30.83 $\pm$ 0.42
pFedPT	86.29 $\pm$ 1.11	39.92 $\pm$ 0.33	21.38 $\pm$ 0.98
FedPFT w/o $L^{Con}$	89.67 $\pm$ 1.96	57.62 $\pm$ 1.18	36.13 $\pm$ 1,32
FedPFT	90.55 $\pm$ 1.35	58.14 $\pm$ 0.71	37.59 $\pm$ 0.39

Methods	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$
FedAvg	60.39 $\pm$ 1.46	60.41 $\pm$ 1.36	60.91 $\pm$ 0.72
Local	81.91 $\pm$ 3.09	60.15 $\pm$ 0.86	52.24 $\pm$ 0.41
FedAMP	84.99 $\pm$ 1.82	68.26 $\pm$ 0.79	64.87 $\pm$ 0.95
FedPer	84.43 $\pm$ 0.47	68.80 $\pm$ 0.49	64.92 $\pm$ 0.66
FedRep	84.59 $\pm$ 1.58	67.69 $\pm$ 0.86	60.52 $\pm$ 0.72
FedBN	83.55 $\pm$ 2.32	66.79 $\pm$ 1.08	62.20 $\pm$ 0.67
FedRoD	86.23 $\pm$ 2.12	72.34 $\pm$ 1.77	68.45 $\pm$ 1.94
pFedSD	86.34 $\pm$ 2.61	71.97 $\pm$ 2.07	67.21 $\pm$ 1.89
pFedGate	87.25 $\pm$ 1.91	71.98 $\pm$ 1.61	67.85 $\pm$ 0.87
FedCAC	86.82 $\pm$ 1.18	69.83 $\pm$ 0.46	65.39 $\pm$ 0.51
pFedPT	82.38 $\pm$ 2.91	67.33 $\pm$ 1.33	64.37 $\pm$ 1.22
FedPFT w/o $L^{Con}$	87.23 $\pm$ 2.69	74.10 $\pm$ 1.95	69.23 $\pm$ 0.76
FedPFT	88.60 $\pm$ 2.19	77.54 $\pm$ 1.88	74.81 $\pm$ 0.77

Results in Pathological non-IID scenario.

This is an extreme setting where each client has data from only a subset of classes. This scenario is particularly pronounced in the CIFAR-10 dataset, where each client essentially performs a 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 a moderate level across the three datasets. ‘FedPFT w/o $L_{\text{Con}}$ ’ aligns local features with the global feature space using classification prompts, enhancing both local feature-classifier alignment and inter-client collaboration effectiveness. It achieves competitive performance on CIFAR-10 and surpasses existing SOTA methods on CIFAR-100 and Tiny ImageNet. FedPFT further incorporates contrastive learning tasks to enhance feature extractor performance, outperforming SOTA methods significantly across all datasets.

Appendix D Feature Separability of Different Methods

	CIFAR-10			CIFAR-100
Methods	$\alpha=0.1$	$\alpha=0.5$	$\alpha=1.0$	$\alpha=0.1$	$\alpha=0.5$	$\alpha=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 w/o $L_{\text{Con}}$	85.52%	72.59%	69.57%	61.60%	43.14%	38.47%
FedPFT	87.83%	77.25%	74.02%	64.12%	46.43%	40.95%

In this section, we delve deeper into the linear separability of features extracted by various PFL methods. Linear separability is a critical measure of feature quality, indicating the ability of a model to distinguish between classes using simple linear classifiers. We conduct linear probing experiments on the CIFAR-10 and CIFAR-100 datasets to assess this metric, with results detailed in Table8.

It can be observed that the feature linear separability of most PFL methods is inferior to FedAvg. This indicates 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, constraining the full potential of PFL.

In stark contrast, FedPFT significantly improves the linear separability of features compared to FedAvg. Our method accomplishes this by fundamentally addressing the mismatch issue during the training process rather than merely adapting the model post hoc. This proactive approach ensures that the feature extractor not only aligns more closely with the global classifier but also preserves its ability to generalize across diverse data distributions. Consequently, FedPFT enhances both the performance and the utility of the feature extractor.

Appendix E Comparison with Two-stage Approach

In FedPFT, we propose using a feature transformation module to coordinate the joint training of contrastive learning and classification tasks. To illustrate the superiority of this design, we introduce a baseline called ‘Two-stage,’ similar to [34], where contrastive learning training is conducted first, followed by classification task training after convergence. For fairness, in the two-stage method, we first perform 1000 rounds of contrastive learning training, followed by 1000 rounds of classification task training. The experimental results are depicted in Figure5.

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (17)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (18)

Firstly, from the perspective of the contrastive learning loss ( $L_{\text{Con}}$ ), FedPFT registers lower loss values compared to the Two-stage approach, suggesting that simultaneous training with the classification task enhances the efficacy of contrastive learning. Secondly, considering both Figure5(b) and Figure5(c), our method exhibits significantly higher accuracy compared to the Two-stage approach. However, $L_{\text{CE}}$ converges to a 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. Importantly, these experiments also indicate that the significant performance improvement brought by contrastive learning in our method is largely attributed to the design of our feature transformation module and training approach.

Appendix F Attention Weight Visualization

In the feature transformation module of FedPFT, self-attention mechanisms are employed to 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 dataset, with results depicted in Figure6. Each row in the figure corresponds to the attention weights for the output feature $f^{\prime}$ of a single sample. Columns represent the input dimensions of the transformation module: the first column corresponds to the original input feature $f$ , while subsequent columns relate to different prompts from the sets $p_{\kappa,i}$ or $p_{\rho,i}$ .

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (19)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (20)

It can be observed that when $\alpha=0.1$ , indicating severe local class imbalances, each client has data from only a few classes. In this case, the feature transformation task is relatively simple, and the influence of different prompts on a sample is similar. As $\alpha$ increases, indicating more complex local tasks, the influence of prompts becomes more intricate. Particularly at $\alpha=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.

Appendix G Partial Client Participation

In FL, challenges such as offline clients and unstable communication may result in only a subset of clients participating in training each round, posing a challenge to the robustness of FL algorithms. In this section, we investigate whether FedPFT is robust to this issue. We conduct experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet, considering scenarios where only a random 50%, 70%, and 90% of clients participate in training each round. The experimental results are presented in Table9.

Datasets	100%	90%	70%	50%
CIFAR-10	88.60 $\pm$ 2.19	88.50 $\pm$ 2.01 (-0.10)	88.57 $\pm$ 2.53 (-0.03)	88.69 $\pm$ 1.83 (+0.09)
CIFAR-100	62.03 $\pm$ 1.41	61.65 $\pm$ 0.41 (-0.38)	63.54 $\pm$ 0.55 (+1.51)	63.97 $\pm$ 0.10 (+1.94)
Tiny	43.42 $\pm$ 1.62	43.03 $\pm$ 2.09 (-0.39)	44.59 $\pm$ 1.19 (+1.17)	45.81 $\pm$ 1.02 (+2.39)

It can be observed that compared to scenarios where all clients participate in training, FedPFT’s accuracy is not significantly reduced when only a subset of clients participate. Furthermore, in CIFAR-100 and Tiny ImageNet, the performance of FedPFT may even be improved. This is because reducing the number of participating clients each round may mitigate the impact of non-IID data distribution on the global model. These experiments demonstrate the robustness of FedPFT to scenarios where only a subset of clients participate.

Appendix H Effect 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.

H.1 Effect of $n_{\kappa}$ and $n_{\rho}$

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (25)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (26)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (29)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (30)

$n_{\kappa}$ and $n_{\rho}$ respectively represent the number of prompts in $p_{\kappa,i}$ and $p_{\rho,i}$ for each client. We examine the impact of these two hyperparameters on the performance of FedPFT on CIFAR-10 and CIFAR-100 datasets. When assessing the effect of $n_{\kappa}$ , we hold $n_{\rho}=20$ constant. Similarly, when evaluating the impact of $n_{\rho}$ , $n_{\kappa}$ is fixed at 10. The experimental results are depicted in Figure7 and Figure8.

FedPFT shows considerable robustness to variations in these hyperparameters. On the CIFAR-10 dataset, changes in $n_{\kappa}$ and $n_{\rho}$ 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 is initially limited by a smaller number of prompts, which may not sufficiently cover the diverse feature space required for effective feature transformation. As the number of prompts increases, the model’s ability to transform and adapt features improves, leading to enhanced performance.

H.2 Effect of $R_{f}$ and $R_{a}$

$R_{f}$ and $R_{a}$ are used to control the number of training epochs for the two training stages. Since we set $R_{f}+R_{a}=R$ , in this experiment, we only adjust $R_{f}$ to examine the impact of these two hyperparameters on model performance. The experimental results are illustrated in Figure9.

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (33)

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (34)

When $R_{f}=0$ , it indicates that contrastive learning is not used to train the feature extractor, and local features are not aligned with the global classifier before training model parameters. It can be observed that the model performance is very poor under this condition. As $R_{f}$ gradually increases, the model performance shows a trend of initially increasing and then decreasing. This suggests that $R_{f}$ essentially balances the trade-off between the two training stages. When $R_{f}$ is small, the feature extractor is predominantly trained by $L_{\text{CE}}$ , and the classifier undergoes more training epochs. At this point, the model pays more attention to the local data distribution of clients, but collaboration among clients is also more susceptible to non-IID effects. Conversely, when $R_{f}$ is large, the feature extractor is primarily trained by $L_{\text{Con}}$ , focusing more on general features, and collaboration among clients is less affected by non-IID issues. However, because the model is rarely trained with $L_{\text{CE}}$ , it pays less attention to the local data distribution of clients, resulting in poorer performance on local data.

In general, $R_{f}$ and $R_{a}$ are two hyperparameters that need to be carefully adjusted, as they have a significant impact on the performance of FedPFT. Typically, in scenarios where the local tasks of clients are simple, it may be appropriate to decrease the value of $R_{f}$ . In other cases, we recommend using a larger value of $R_{f}$ to enhance the degree of collaboration among clients.

Appendix I Communication Cost

In this section, we calculate the communication overhead of one client in FedAvg and FedPFT in each communication round.

Model	$\phi_{i}$	$\tau_{i}$	$h_{\kappa,i}$	$h_{\rho,i}$	FedAvg	FedPFT	Incre. Ratio
ResNet-8	1.24M	0.26M	25.70K	32.90K	1.27M	1.56M	23.14%
ResNet-10	4.91M	1.05M	51.30K	65.66K	4.96M	6.08M	22.49%

In FedAvg, each communication round involves uploading the feature extractor $\phi_{i}$ and the classifier $h_{\kappa,i}$ . FedPFT adds the feature transformation module $\tau_{i}$ and the feature projection layer $h_{\rho,i}$ , thereby increasing the volume of parameters transmitted per round. According to the results presented in Table10, the communication overhead for FedPFT using ResNet-8 and ResNet-10 architectures is increased by 23.14% and 22.49%, respectively, relative to FedAvg.

While FedPFT brings additional communication cost, it is important to weigh it against the performance enhancements and scalability offered by $\tau_{i}$ , as discussed in earlier sections of this paper. The improved model accuracy and robustness to non-IID data might justify the additional costs in scenarios where model performance is critical.

Moving forward, considering the increase in communication cost is primarily due to the additional components $\tau_{i}$ , we aim to develop a more efficient and lightweight feature transformation module to reduce communication demands without compromising model effectiveness in our future work.

Appendix J 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 PFL methods that are not based on global models, such as personalized-weight-aggregation-based methods, which are not explored 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.

Appendix K Theoretical Analysis

Since the main problem in Eq.(1) is non-convex, we focus on the factors affecting convergence in the non-convex setting.

Implication	Notation
Global / Local loss	$L$ / $L_{i}$
Global / Local problem	$F$ / $F_{i}$
Local Dataset on $i^{\text{th}}$ client	$\tilde{d}_{i}\in d_{i}$
Feature extractor	$\phi$
Feature transformation module	$\tau$
Classification / Contrastive learning prompts	$p_{\kappa}$ / $p_{\rho}$
Feature extractor & Feature transformation module & Classifier	$w$
Classification / Contrastive learningtask head	$h_{\kappa}$ / $h_{\rho}$
Global / Local problem’s gradient	$\nabla F(w)$ / $\nabla F_{i}(w)$
Local gradient approximation	$g_{i,r}^{t}$
Client number	$N$
Local update epoch	$R$
The number of clients sampled at each global epoch	$S$
The set of clients sampled at global epoch $t$	$\mathcal{S}_{t}$
The actual learning rate of global problem	$\tilde{\eta}$
The learning rate of local problem	$\eta$
Approximated local gradient error’s upper-bound	$\delta$
Local-global gradient error’s upper-bound	$\sigma_{F}$
Index of client, local epoch and global epoch	$i\in[N],r\in[R],t\in[T]$

K.1 Problem Setup

Non-convex case analyses are as follows. By Lagrange duality, the main problem is transformed as follows:

	$\displaystyle\min_{\phi,\tau,h_{\kappa}}$	$\displaystyle\min_{\{p_{\kappa,i}\}_{i\in[N]}}\mathbf{E}_{i}\mathbf{E}_{d_{i}}%L_{\text{CE}}(\phi,\tau,h_{\kappa},p_{\kappa,i};d_{i})$
		$\displaystyle\text{s.t. }\mathbf{E}_{i}\mathbf{E}_{d_{i}}L_{\text{Con}}(\phi,%\tau;d_{i})\leq H_{\text{Con}}$

We transform the problem into an unconditional bi-level optimization problem:

\displaystyle\min_{w}\mathbf{E}F(w)=\mathbf{E}_{i}\{F_{i}(w):=\min_{p_{\kappa,%i}}\mathbf{E}_{d_{i}}L_{\text{CE}}(w,p_{\kappa,i};d_{i})\}

where $\mathbf{E}$ represents the expectation of all random variables, $\mathbf{E}_{i}$ means the expectation of client sampling, $\mathbf{E}_{d_{i}}$ is the local data sampling expectation, and we use $w=\{\phi,\tau,h_{\kappa}\}$ for simplification, based on the equivalence of block coordinate descent and gradient descent.

K.2 Propositions

Proposition K.1 ( $L$ -smooth).

If $f$ is $L$ -smooth, $\forall x,y$ we have:

	$\displaystyle\langle\nabla f(x)-\nabla f(y),x-y\rangle$	$\displaystyle\leq L\|\|x-y\|\|^{2}$
	$\displaystyle\|\|\nabla f(x)-\nabla f(y)\|\|$	$\displaystyle\leq L\|\|x-y\|\|$
	$\displaystyle\|\|\nabla f(x)-\nabla f(y)\|\|^{2}$	$\displaystyle\leq 2L[f(x)-f(y)]$
	$\displaystyle f(y)-f(x)-\langle\nabla f(x),y-x\rangle$	$\displaystyle\leq\frac{L}{2}\|\|y-x\|\|^{2}$

Proposition K.2 (Jensen’s inequality).

If $f$ is convex, we have the following inequality:

\displaystyle\mathbf{E}_{X}f(X)\geq f(\mathbf{E}_{X}X).

A variant of the general one shown above, given a group $\{x_{i}\}_{i\in[N]}$ :

\displaystyle||\sum_{i\in[N]}x_{i}||^{2}\leq N\sum_{i\in[N]}||x_{i}||^{2}.

Proposition K.3 (Triangle inequality).

The triangle inequality, where $||\cdot||$ is the norm, and $A$ , $B$ is the elements in the corresponding norm space:

||A+B||\leq||A||+||B||

Proposition K.4 (Matrix norm compatibility).

The matrix norm compatibility, $A\in\mathbf{R}^{a\times b},B\in\mathbf{R}^{b\times c},v\in\mathbf{R}^{b}$ :

	$\displaystyle\|\|AB\|\|_{m}\leq\|\|A\|\|_{m}\|\|B\|\|_{m}$
	$\displaystyle\|\|Av\|\|_{m}\leq\|\|A\|\|_{m}\|\|v\|\|$

Proposition K.5 (Peter Paul inequality).

$\forall x,y$ and $\forall\epsilon>0$ , we have the following inequality:

2\langle x,y\rangle\leq\frac{1}{\epsilon}||x||^{2}+\epsilon||y||^{2}

K.3 Assumptions

Assumption K.1 (L-smooth local objectives).

$\forall i$ , $F_{i}$ is $L_{F}$ -Smooth, the main proposition is shown in Prop.K.1. Notice that the $F_{i}$ is assumed to be L-smooth and non-convex, which matches the problem and neural network architecture setting in the main paper.

Assumption K.2 (Bounded local variance).

The local problem’s gradient is assumed not to be too far from the global problem’s gradient.

\forall w,\mathbf{E}_{i}||\nabla F_{i}(w)-\nabla F(w)||\leq\sigma_{F}

Assumption K.3 (Bounded approximated gradient).

The first-order approximation of the local problem’s gradient $g_{i,r}^{t}$ should not be too far from the ground truth $\nabla F_{i}(w_{i,r}^{t})$ . In this assumption, the approximated error of the block coordinate descent in Algorithm1 is bounded.

\forall\{(i,r,t)\},||g_{i,r}^{t}-\nabla F_{i}(w_{i,r}^{t})||\leq\delta

K.4 Lemmas

Lemma K.1 (Bounded local approximation error).

If $\tilde{\eta}:=\eta R\leq\frac{1}{2L_{F}}$ , we have the following bound of client drift error:

\frac{1}{NR}\sum_{i,r}^{N,R}\mathbf{E}||g_{i,r}^{(t)}-\nabla F_{i}(w^{(t)})||^%{2}\leq 2\delta^{2}+2^{R+3}L_{F}[3\tilde{\eta}^{2}\sum_{i}^{N}\mathbf{E}||%\nabla F_{i}(w^{(t)})||^{2}+\frac{2\tilde{\eta}^{2}\delta^{2}}{R}]

Proof.

The client drift error on given $i^{\text{th}}$ client and its upper bound are as follows:

	$\displaystyle\mathbf{E}\|\|g_{i,r}^{(t)}-\nabla F_{i}(w^{(t)})\|\|^{2}$	(7)
$\displaystyle\leq$	$\displaystyle 2\mathbf{E}\|\|g_{i,r}^{(t)}-\nabla F_{i}(w_{i,r}^{(t)})\|\|^{2}+2%\mathbf{E}\|\|\nabla F_{i}(w^{(t)})-\nabla F_{i}(w_{i,r}^{(t)})\|\|^{2}$
$\displaystyle\leq$	$\displaystyle 2\delta^{2}+2L_{F}\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|\|^{2}$

where the first inequality is by PropositionK.3 and the second one is by AssumptionK.1.

For the last term in the upper bound, we have the iterative formulation as follows:

		$\displaystyle\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}-g_{i,r-1}^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}-\eta\nabla F_{i}(w^{(t)})\|\|%^{2}+2\eta^{2}\mathbf{E}\|\|g_{i,r-1}^{(t)}-\nabla F_{i}(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2(1+\frac{1}{2R})\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}\|\|^{2}+2(1+2%R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+4\eta^{2}[\delta^{2}+L_{F}^{2}\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|%\|^{2}]$
	$\displaystyle=$	$\displaystyle 2(1+\frac{1}{2R}+2\eta^{2}L_{F}^{2})\mathbf{E}\|\|w_{i,r-1}^{(t)}-%w^{(t)}\|\|^{2}+4\eta^{2}\delta^{2}$
		$\displaystyle+2(1+2R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$

where the two inequalities are by PropositionK.3, PropositionK.5 and Eq.(7).

Take $\tilde{\eta}:=\eta R\leq\frac{1}{2L_{F}}$ , we recursively unroll the inequality as follows:

		$\displaystyle\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2(1+\frac{1}{R})\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}\|\|^{2}+4\eta^%{2}\delta^{2}+2(1+2R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle[3\tilde{\eta}^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}+\frac{2%\tilde{\eta}^{2}\delta^{2}}{R}]2^{R+2}$

where the inequality is unrolled and we use $\frac{1}{R}\leq 1$ . Thus, we have:

\displaystyle\mathbf{E}||g_{i,r}^{(t)}-\nabla F_{i}(w^{(t)})||^{2}\leq 2\delta%^{2}+2^{R+4}\tilde{\eta}^{2}L_{F}[3\sigma_{F}^{2}+3\mathbf{E}||\nabla F(w^{(t)%})||^{2}+\frac{\delta^{2}}{R}]

∎

K.5 Theorem and Discussion

Theorem K.2 (Non-convex and smooth convergence of FedPFT).

Let AssumptionK.1, AssumptionK.2 and AssumptionK.3 hold, if $\tilde{\eta}:=\eta R\leq\min\{\frac{1}{2L_{F}},\hat{\eta}\}$ is taken, where $\hat{\eta}:=\frac{N/S-1}{24(N-1)2^{R}}\sigma_{F}^{2}-1$ , we have the following bound:

\displaystyle\mathcal{O}(\mathbf{E}||\nabla F(w^{(\bar{t})})||^{2}):=\mathcal{%O}(\frac{\Delta_{F}}{\hat{\eta}T}+\frac{2^{R/3}L_{F}^{1/3}(R\sigma_{F}^{2}+%\delta^{2})^{1/3}\Delta_{F}^{2/3}}{T^{2/3}R^{1/3}}+(\frac{\sigma_{F}\sqrt{L_{F%}(N/S-1)\Delta_{F}}}{\sqrt{TN}})+\delta^{2})

Proof.

		$\displaystyle\mathbf{E}F(w^{(t+1)})-\mathbf{E}F(w^{(t)})$
	$\displaystyle\leq$	$\displaystyle\mathbf{E}\langle\nabla F(w^{(t)}),w^{(t+1)}-w^{(t)}\rangle+\frac%{L_{F}}{2}\mathbf{E}\|\|w^{(t+1)}-w^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle-\tilde{\eta}\mathbf{E}\langle\nabla F(w^{(t)}),g^{(t)}\rangle+%\frac{\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle-\tilde{\eta}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}-\tilde{\eta}%\mathbf{E}\langle\nabla F(w^{(t)}),g^{(t)}-\nabla F(w^{(t)})\rangle+\frac{%\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{%\tilde{\eta}}{2}\mathbf{E}\|\|\frac{1}{NR}\sum_{i,r}^{N,R}g_{i,r}^{(t)}-\nabla F%_{i}(w^{(t)})\|\|^{2}+\frac{\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{%\tilde{\eta}}{2}\mathbf{E}\|\|\frac{1}{NR}\sum_{i,r}^{N,R}g_{i,r}^{(t)}-\nabla F%_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+\frac{3\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}[\|\|g^{(t)}-\nabla F_{i%}(w^{(t)})\|\|^{2}+\|\|\frac{1}{S}\sum_{i\in\mathcal{S}^{(t)}}\nabla F_{i}(w^{(t)}%)-\nabla F(w^{(t)})\|\|^{2}+\|\|\nabla F(w^{(t)})\|\|^{2}]$
	$\displaystyle=$	$\displaystyle-\frac{\tilde{\eta}(1-3\tilde{\eta}L_{F})}{2}\mathbf{E}\|\|\nabla F%(w^{(t)})\|\|^{2}+\frac{\tilde{\eta}(1+3\tilde{\eta}L_{F})}{2}\frac{1}{NR}\sum_{%i,r}^{N,R}\mathbf{E}\|\|g_{i,r}^{(t)}-\nabla F_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+\frac{3\tilde{\eta}^{2}L_{F}}{2}\|\|\frac{1}{S}\sum_{i\in\mathcal{%S}^{(t)}}\nabla F_{i}(w^{(t)})-\nabla F(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}(1-3\tilde{\eta}L_{F})}{2}\mathbf{E}\|\|\nabla F%(w^{(t)})\|\|^{2}+3\tilde{\eta}^{2}L_{F}\frac{N/S-1}{N-1}[\sigma_{F}^{2}+\|\|%\nabla F(w^{(t)})\|\|^{2}]$
		$\displaystyle+\tilde{\eta}(1+3\tilde{\eta}L_{F})[\delta^{2}+2^{R+3}\tilde{\eta%}^{2}L_{F}[3\sigma_{F}^{2}+3\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{\delta^{%2}}{R}]]$

where the four inequalities are respectively by $L_{F}$ -smooth of $F:=\mathbf{E}_{i}F_{i}$ , PropositionK.5, LemmaK.1 and the similar classic Lemma 4 in [29].

Let $c_{1}:=3\delta^{2}$ , $c_{2}:=3L_{F}\sigma_{F}^{2}\frac{N/S-1}{N-1}$ , $c_{3}:=2^{R+3}L_{F}[3\sigma_{F}^{2}+\frac{\delta^{2}}{R}]$ ,

	$\displaystyle\mathbf{E}F(w^{(t+1)})-\mathbf{E}F(w^{(t)})\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\{1-[\frac{3}{2}-3\frac{N/S-1}{N-1}\sigma_%{F}^{2}+72\times 2^{R}\tilde{\eta}]\}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}$
		$\displaystyle+c_{3}\tilde{\eta}^{3}+c_{2}\tilde{\eta}^{2}+c_{1}\tilde{\eta}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+c_{3}%\tilde{\eta}^{3}+c_{2}\tilde{\eta}^{2}+c_{1}\tilde{\eta}$

where let $\tilde{\eta}\leq\min\{\frac{1}{2L_{F}},\hat{\eta}$ , where $\hat{\eta}:=\frac{2}{3\times 2^{R+4}}\frac{N/S-1}{N-1}\sigma_{F}^{2}-1\}$ . Re-arranging the inequality above and accumulating, we have:

	$\displaystyle\frac{1}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}$	$\displaystyle\leq\mathbf{E}F(w^{(t+1)})-\mathbf{E}F(w^{(t)})+c_{3}\tilde{\eta}%^{2}+c_{2}\tilde{\eta}+c_{1}$
	$\displaystyle\frac{1}{2T}\sum_{t=0}^{t=T-1}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}$	$\displaystyle\leq\mathbf{E}F(w^{(T)})-\mathbf{E}F(w^{(0)})+c_{3}\tilde{\eta}^{%2}+c_{2}\tilde{\eta}+c_{1}$

Let $\Delta_{F}=F(w^{0})-F(w^{*})$ , where $w^{*}$ is the minimum of the main problem $\arg\min_{w}\mathbf{E}F(w)$ . To measure the exact term of the bounds, we consider the following cases:

•

$\frac{\Delta_{F}}{c_{3}T}\leq\tilde{\eta}^{3}$ or $\frac{\Delta_{F}}{c_{2}T}\leq\tilde{\eta}^{2}$ , let $\tilde{\eta}=\min\{(\frac{\Delta_{F}}{c_{3}T})^{1/3},(\frac{\Delta_{F}}{c_{2}T%})^{1/2}\}$ ,we have:

\frac{1}{2}\mathbf{E}||\nabla F(w^{(t)})||^{2}\leq\frac{c_{3}^{1/3}\Delta_{F}^%{2/3}}{T^{2/3}}+(\frac{c_{2}\Delta_{F}}{T})^{1/2}+c_{1}

•

$\frac{\Delta_{F}}{c_{3}T}\geq\tilde{\eta}^{3}$ and $\frac{\Delta_{F}}{c_{2}T}\geq\tilde{\eta}^{2}$ , let $\tilde{\eta}=\hat{\eta}$ ,we have:

\frac{1}{2}\mathbf{E}||\nabla F(w^{(t)})||^{2}\leq\frac{\Delta_{F}}{\hat{\eta}%T}+\frac{c_{3}^{1/3}\Delta_{F}^{2/3}}{T^{2/3}}+(\frac{c_{2}\Delta_{F}}{T})^{1/%2}+c_{1}

Uniformly sample a $\bar{t}\in[T]-1$ , we have the upper bound as follows:

	$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}$	$\displaystyle\mathbf{E}{\|\|\nabla F(w^{(t)}\|\|^{2})}=\mathcal{O}(\mathbf{E}\|\|F(w%^{(\bar{t})})\|\|^{2})$
	$\displaystyle:=$	$\displaystyle\mathcal{O}(\frac{\Delta_{F}}{\hat{\eta}T}+\frac{2^{R/3}L_{F}^{1/%3}(R\sigma_{F}^{2}+\delta^{2})^{1/3}\Delta_{F}^{2/3}}{T^{2/3}R^{1/3}}+(\frac{%\sigma_{F}\sqrt{L_{F}(N/S-1)\Delta_{F}}}{\sqrt{TN}})+\delta^{2})$

∎

Remark K.2.1.

According to TheoremK.2, our proposed FedPFT converges at a sub-linear level. The linear term $\mathcal{O}(\frac{\Delta_{F}}{\hat{\eta}T})$ is affected by $\hat{\eta}$ and the initialization gap $\Delta_{F}$ . The sub-linear term $\mathcal{O}(1/T^{2/3})$ is affected by $R$ , especially when $R$ is large due to the exponential factor $2^{R}$ . As the local approximation error of the gradient $\delta$ grows, both the convergence radius $\mathcal{O}(\delta)$ and the sub-linear term $\mathcal{O}(1/T^{2/3})$ are affected by the local optimizer selection significantly. Another sub-linear term $\mathcal{O}(\sqrt{T})$ is eliminated if $N/S-1=0$ when all the clients are sampled. Otherwise, the sub-linear rate is mainly affected by $\sigma_{F}$ .

FedPFT aligns the training objectives across clients by introducing $p_{\kappa,i}$ and reduces the impact of non-IID data on the feature extractor $\phi$ through contrastive learning. Both of these designs can effectively reduce differences in local gradients among clients during training, thereby reducing $\sigma_{F}$ and subsequently lowering the upper bound. During training, $p_{\kappa,i}$

	$\displaystyle\langle\nabla f(x)-\nabla f(y),x-y\rangle$	$\displaystyle\leq L\|\|x-y\|\|^{2}$
	$\displaystyle\|\|\nabla f(x)-\nabla f(y)\|\|$	$\displaystyle\leq L\|\|x-y\|\|$
	$\displaystyle\|\|\nabla f(x)-\nabla f(y)\|\|^{2}$	$\displaystyle\leq 2L[f(x)-f(y)]$
	$\displaystyle f(y)-f(x)-\langle\nabla f(x),y-x\rangle$	$\displaystyle\leq\frac{L}{2}\|\|y-x\|\|^{2}$

		$\displaystyle\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}-g_{i,r-1}^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}-\eta\nabla F_{i}(w^{(t)})\|\|%^{2}+2\eta^{2}\mathbf{E}\|\|g_{i,r-1}^{(t)}-\nabla F_{i}(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2(1+\frac{1}{2R})\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}\|\|^{2}+2(1+2%R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+4\eta^{2}[\delta^{2}+L_{F}^{2}\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|%\|^{2}]$
	$\displaystyle=$	$\displaystyle 2(1+\frac{1}{2R}+2\eta^{2}L_{F}^{2})\mathbf{E}\|\|w_{i,r-1}^{(t)}-%w^{(t)}\|\|^{2}+4\eta^{2}\delta^{2}$
		$\displaystyle+2(1+2R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$

		$\displaystyle\mathbf{E}\|\|w_{i,r}^{(t)}-w^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle 2(1+\frac{1}{R})\mathbf{E}\|\|w_{i,r-1}^{(t)}-w^{(t)}\|\|^{2}+4\eta^%{2}\delta^{2}+2(1+2R)\eta^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle[3\tilde{\eta}^{2}\mathbf{E}\|\|\nabla F_{i}(w^{(t)})\|\|^{2}+\frac{2%\tilde{\eta}^{2}\delta^{2}}{R}]2^{R+2}$

		$\displaystyle\mathbf{E}F(w^{(t+1)})-\mathbf{E}F(w^{(t)})$
	$\displaystyle\leq$	$\displaystyle\mathbf{E}\langle\nabla F(w^{(t)}),w^{(t+1)}-w^{(t)}\rangle+\frac%{L_{F}}{2}\mathbf{E}\|\|w^{(t+1)}-w^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle-\tilde{\eta}\mathbf{E}\langle\nabla F(w^{(t)}),g^{(t)}\rangle+%\frac{\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle=$	$\displaystyle-\tilde{\eta}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}-\tilde{\eta}%\mathbf{E}\langle\nabla F(w^{(t)}),g^{(t)}-\nabla F(w^{(t)})\rangle+\frac{%\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{%\tilde{\eta}}{2}\mathbf{E}\|\|\frac{1}{NR}\sum_{i,r}^{N,R}g_{i,r}^{(t)}-\nabla F%_{i}(w^{(t)})\|\|^{2}+\frac{\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}\|\|g^{(t)}\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}}{2}\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{%\tilde{\eta}}{2}\mathbf{E}\|\|\frac{1}{NR}\sum_{i,r}^{N,R}g_{i,r}^{(t)}-\nabla F%_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+\frac{3\tilde{\eta}^{2}L_{F}}{2}\mathbf{E}[\|\|g^{(t)}-\nabla F_{i%}(w^{(t)})\|\|^{2}+\|\|\frac{1}{S}\sum_{i\in\mathcal{S}^{(t)}}\nabla F_{i}(w^{(t)}%)-\nabla F(w^{(t)})\|\|^{2}+\|\|\nabla F(w^{(t)})\|\|^{2}]$
	$\displaystyle=$	$\displaystyle-\frac{\tilde{\eta}(1-3\tilde{\eta}L_{F})}{2}\mathbf{E}\|\|\nabla F%(w^{(t)})\|\|^{2}+\frac{\tilde{\eta}(1+3\tilde{\eta}L_{F})}{2}\frac{1}{NR}\sum_{%i,r}^{N,R}\mathbf{E}\|\|g_{i,r}^{(t)}-\nabla F_{i}(w^{(t)})\|\|^{2}$
		$\displaystyle+\frac{3\tilde{\eta}^{2}L_{F}}{2}\|\|\frac{1}{S}\sum_{i\in\mathcal{%S}^{(t)}}\nabla F_{i}(w^{(t)})-\nabla F(w^{(t)})\|\|^{2}$
	$\displaystyle\leq$	$\displaystyle-\frac{\tilde{\eta}(1-3\tilde{\eta}L_{F})}{2}\mathbf{E}\|\|\nabla F%(w^{(t)})\|\|^{2}+3\tilde{\eta}^{2}L_{F}\frac{N/S-1}{N-1}[\sigma_{F}^{2}+\|\|%\nabla F(w^{(t)})\|\|^{2}]$
		$\displaystyle+\tilde{\eta}(1+3\tilde{\eta}L_{F})[\delta^{2}+2^{R+3}\tilde{\eta%}^{2}L_{F}[3\sigma_{F}^{2}+3\mathbf{E}\|\|\nabla F(w^{(t)})\|\|^{2}+\frac{\delta^{%2}}{R}]]$

Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation (2024)

Abstract

1 Introduction

2 Related Work

3 Methodology

3.1 Overview of FedPFT

3.2 Problem Formulation

3.3 Feature Transformation Module

3.4 Classification Task with Personalized Prompts

3.5 Contrastive Learning Task

3.6 Alternating Training Strategy

Feature Learning Phase.

Task Adaptation Phase.

4 Experiments

4.1 Experimental Setup

Datasets.

4.2 Comparison with State-of-the-art Methods

Results in Dirichlet non-IID scenario.

4.3 Ablation Study

4.4 Separability of Features

4.5 Learned Features of Different Methods

4.6 Effect of Different Prompts

5 Conclusion and Discussion

References

Appendix A Related Work

Prompt-based methods.

Appendix B Experiment Setup

B.1 Introduction to non-IID Scenarios

Pathological non-IID.

Dirichlet non-IID.

B.2 Introduction to Comparative Methods

B.3 Hyperparameter Settings in Different Methods

B.4 Compute Resources

Appendix C Comparison with State-of-the-art Methods

Results in Pathological non-IID scenario.

Appendix D Feature Separability of Different Methods

Appendix E Comparison with Two-stage Approach

Appendix F Attention Weight Visualization

Appendix G Partial Client Participation

Appendix H Effect of Hyperparameters

H.1 Effect of nκsubscript𝑛𝜅n_{\kappa}italic_n start_POSTSUBSCRIPT italic_κ end_POSTSUBSCRIPT and nρsubscript𝑛𝜌n_{\rho}italic_n start_POSTSUBSCRIPT italic_ρ end_POSTSUBSCRIPT

H.2 Effect of Rfsubscript𝑅𝑓R_{f}italic_R start_POSTSUBSCRIPT italic_f end_POSTSUBSCRIPT and Rasubscript𝑅𝑎R_{a}italic_R start_POSTSUBSCRIPT italic_a end_POSTSUBSCRIPT

Appendix I Communication Cost

Appendix J Limitations and Future Work

Appendix K Theoretical Analysis

K.1 Problem Setup

K.2 Propositions

Proposition K.1 (L𝐿Litalic_L-smooth).

Proposition K.2 (Jensen’s inequality).

Proposition K.3 (Triangle inequality).

Proposition K.4 (Matrix norm compatibility).

Proposition K.5 (Peter Paul inequality).

K.3 Assumptions

Assumption K.1 (L-smooth local objectives).

Assumption K.2 (Bounded local variance).

Assumption K.3 (Bounded approximated gradient).

K.4 Lemmas

Lemma K.1 (Bounded local approximation error).

Proof.

K.5 Theorem and Discussion

Theorem K.2 (Non-convex and smooth convergence of FedPFT).

Proof.

Remark K.2.1.

H.1 Effect of $n_{\kappa}$ and $n_{\rho}$

H.2 Effect of $R_{f}$ and $R_{a}$

Proposition K.1 ( $L$ -smooth).