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

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.

Weakness 1: Lack of Theoretical Motivation and Analysis

We 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.

Weakness 2: Limited Scale of Experiments

We appreciate the reviewer’s concern. We have provided larger-scale experiments in Appendix K (Table 12 and Table 13) of the original manuscript, including:

• 100 participating clients,
• Deeper model architectures, such as ResNet-18

For 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.

In 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.

Minor: Related Works on Feature–Classifier Misalignment

We thank the reviewer for pointing out these additional related works.

We 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.

We 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.

Question 1: Dirichlet Parameter and Data Heterogeneity

We 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.

We 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.

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.

Weakness 1 & Question 1: Clarifying the Definition and Novelty of Feature–Classifier Mismatch

We 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].

(1) Formal Definition of Feature–Classifier Mismatch

We define feature–classifier mismatch as:

$$\text{Mismatch}(f, W) = \mathbb{E}_{(x,y)\sim D_i}[\mathcal{L}(f', W) - \mathcal{L}(f, W)]$$

where:

• $D_i$ is the data distribution of client $i$.
• $f$ denotes the original local features extracted by the feature extractor,
• $W$ is the classifier,
• $f'$ is the transformed feature produced by feature transformation module (e.g., linear transformation layer),
• $\mathcal{L}(\cdot, \cdot)$ is the loss function (e.g., cross-entropy).

This 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.

(2) Key Difference from Prior Work

We highlight the fundamental difference between our work and existing literature.

1. Objective: Inter-Client vs. Intra-Client Alignment

• 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.
• 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.

2. Our Advantage

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).

• 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).
• 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)).

We 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.

Weakness 2 & Question 2: Quantitative Evidence for Mismatch

We 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:

(1) Gradient Divergence

We 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.)

The 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.

(2) Feature–Classifier Cosine Similarity

We 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.

However, 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.

Instead, 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.

Weakness 3: On the Role of Contrastive Learning

We 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.

Weakness 4: On Overhead

We 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.

While 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.

As 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.

Question 3: On Robustness Comparison

We 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.

Question 4: Communication and Accuracy on OOD Clients

We 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.

Question 5: Comparison with FL Benchmarks [8, 9]

We thank the reviewer for this valuable suggestion. We have carefully reviewed both FederatedScope [9] and FL-Bench [8], and provide the following clarification:

• 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.
• 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:
  • Use of public datasets (CIFAR-10, CIFAR-100, TinyImageNet);
  • Adoption of Dirichlet and feature-shift non-IID partitions;
  • Use of standard model architectures (e.g., ResNet);
  • Client counts of 40 and 100 for small- and large-scale evaluations;
  • Including many baselines that are also part of FL-Bench, such as FedRep, FedPer, FedBN, FedRoD, and FedBABU.

Although 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.

I sincerely appreciate the authors’ detailed response to my questions. However, I remain unconvinced that the proposed approach introduces a fundamentally new perspective. While the authors frame the concept of feature–classifier mismatch as a novel problem, many prior works in personalized federated learning (PFL) have naturally addressed similar issues—though perhaps under different terminology. In particular, methods like FedBABU, MAML-based approaches, and FedRep all incorporate mechanisms that implicitly or explicitly align features and classifiers at the client level.

For instance, FedBABU separates the model into a shared feature extractor and a local classifier, updating only the extractor during training and personalizing the head on each client. This design directly targets the alignment between client-specific features and the classifier, effectively resolving what the authors now describe as a mismatch. Similarly, MAML-based methods such as Per-FedAvg allow rapid adaptation of the global model to local data, which inherently reduces any initial misalignment between the shared model and client-specific distributions. FedRep also optimizes the classifier locally before updating the global feature extractor, again implicitly correcting mismatches between features and classifiers.

Moreover, even simple baselines like FedAvg with local classifier fine-tuning demonstrate substantial improvements in personalization performance, highlighting that the challenge often lies in aligning the classifier with well-learned features. In light of this, I believe the feature–classifier mismatch is not a newly identified problem, but rather a reframing of a widely recognized challenge. The true contribution of FedPFT lies in its specific solution—prompt-based feature transformation—rather than in the novelty of the problem itself.

That said, I find the proposed methodology itself to be a promising direction. However, the benefits gained from the Prompt-based Feature Transformation (PFT) appear to be offset by the additional communication and model complexity it introduces. As a result, the practical advantage of FedPFT over simpler adaptation-based methods remains unclear. To justify a higher rating, I believe the authors need to provide stronger theoretical justification for why PFT offers a fundamentally superior solution compared to existing personalization approaches.

Based on the authors' final clarification—specifically, the explicit formulation of the feature–classifier mismatch and the training-time mechanism designed to address it—I am inclined to raise my score slightly in recognition of the paper’s contribution. However, I still have some reservations.

In particular, it remains unclear whether the observed performance improvements stem primarily from the proposed mismatch correction mechanism or simply from increased model complexity, which could introduce additional computational and communication overhead. Moreover, if the approach relies on learning client-specific transformation modules, it is important to consider how the model would generalize to previously unseen users without local adaptation.

That said, the authors have addressed many of my concerns to a reasonable extent, and I appreciate their thoughtful response. As such, I have decided to increase my score slightly (based on the assumption that authors will update the contribution more clearly).

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.

Weakness 1 & Question 1 & Limitation 1: Regarding the distinction from FedPAC

We 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.

1. Difference in Objective

• 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.
• 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.

2. Differences in Methodological Design and Our Advantage

While 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:

• 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.
• 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.
• 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.

These results suggest that FedPFT not only achieves inter-client consistency comparable to methods explicitly designed for it, but also addresses intra-client mismatch effectively.

Weakness 2: Evaluation Scope – Scalability, Partial Participation, and Local Epochs

We thank the reviewer for this thoughtful suggestion. In response, we have conducted additional experiments to address these concerns and summarize our findings below.

(1) Scalability to more clients

As 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.

In the rebuttal phase, we additionally conducted experiments with 200 clients. As shown below, our method consistently outperforms SOTA methods.

(2) Participation Mode

We 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.

To 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.

(3) Impact of Local Update Epochs

In 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.

This 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.

We 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.

Dear Reviewer,

As the Area Chair handling this submission, I would like to express my appreciation for the thorough and insightful comments you have provided. The authors have now submitted a comprehensive rebuttal addressing each of the points raised in your reviews. To facilitate a fair and rigorous decision-making process, I kindly urge each reviewer to carefully assess whether the authors' rebuttal adequately addresses your specific concerns and technical questions. Your active participation and constructive input during the upcoming discussion phase will be crucial in arriving at a well-informed decision.

Please feel free to share any additional thoughts or clarifications you may have as we proceed with our deliberations.

Best regards,

Your AC