Title: Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Personalized Federated Learning

Feature-Classifier Mismatch

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


Peer reviews:
comments: SUMMARY:
This paper proposes FedPFT, a framework that uses personalized prompts to address the mismatch between local features and the global classifier during training in heterogeneous FL. FedPFT improves feature alignment and supports additional training objectives like contrastive learning, achieving up to 5.07% better performance and 7.08% further gains with collaborative contrastive learning.

STRENGTHS AND WEAKNESSES:
Strengths:
1) The paper is well-motivated with clear and effective representations. 
2) The experiments include extensive baselines for thorough evaluation. 

Weaknesses:
1) The proposed method appears somewhat straightforward and lacks theoretical motivation and analysis. 
2) The experiments are conducted on a relatively small scale, limiting the evaluation's generalizability. 

Considering the efforts in Table 1, I tend to give a marginaly positive score.

Minor: There are also several works addressing feature and classifier misalignment, which are recommended to be included in the related works section for discussion. Some of these methods are designed for G-FL, while others are applicable to P-FL. To name a few:
* FedAWS:Federated learning with only positive labels 
* FedMR: Federated learning under partially class-disjoint data via manifold reshaping 
* FedRS: Fedrs: Federated learning with restricted softmax for label distribution non-iid data 
* FedGELA: Federated learning with bilateral curation for partially class-disjoint data

QUESTIONS:
1) Why does a smaller parameter in the Dirichlet distribution lead to a more heterogeneous setting? In most previous papers, a smaller Dirichlet parameter typically implies a more heterogeneous data distribution.

LIMITATIONS:
NA
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses a challenge in Federated Learning (FL) caused by data heterogeneity, which leads to suboptimal performance of global models across diverse client distributions. \nThey identify that the primary cause of suboptimal performance is the mismatch between local features and the global classifier, rather than the inadequacy of the feature extractor or classifier themselves. \n\nTo address this issue, they propose FedPFT, a new Personalized Federated Learning (PFL) framework that uses personalized prompts and a shared self-attention-based feature transformation module (FTM) to align local features with the global classifier during training. The paper demonstrates through experiments that FedPFT outperforms state-of-the-art methods by up to 5.07%, and when combined with contrastive learning, the improvement can reach up to 7.08%.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n1. The paper offers a fresh perspective on the feature–classifier mismatch problem in FL.\n\n2. FedPFT introduces a concrete framework—personalized prompts plus a self-attention FTM—that explicitly targets this mismatch.\n\n3. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet under both Dirichlet and Pathological non-IID settings demonstrate clear performance gains over existing methods.\n\n### Weaknesses \n1. Although the authors promote feature–classifier mismatch as a new explanation for performance degradation, it is unclear precisely what this mismatch is or why it constitutes a novel, critical problem. Prior studies have already examined gaps between global and personalized models in feature, classifier, or weight space side [1, 2, 3, 4, 5, 6, 7].\n\n2. To claim that FedPFT solves the mismatch problem, the paper should provide more evidence—e.g., extend Figure 1 (Gradient Difference) to include other SOTA models, or introduce quantitative metrics beyond pre-/post-accuracy. Performance gains alone are insufficient supports.\n\n3. Showing that FedPFT pairs well with contrastive learning is useful but not, by itself, a novel contribution; it is incremental.\n\n4. The personalized prompts and FTM increase both model complexity and communication cost. Figure 1 and Table 1 suggest that simple post-training or fine-tuning (FT) yields nearly the same accuracy boost with far less overhead.\n\n\n[1] FedBABU: Toward Enhanced Representation For Federated Image Classification, ICLR 2022\n\n[2] No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data, NeurIPS 2021\n\n[3]  Guiding The Last Layer in Federated Learning with Pre-Trained Models, NeurIPS 2023\n\n[4] Personalized Fedrated Learning with Ehroretical Guarantees: A Model-Agnostic Meta-L earning Approach, NeurIPS 2020\n\n[5] Personalized Federated Learning using Hypernetworks, ICML 2021\n\n[6] Federated Learning with Partial Model Personalization, PMLR 2022\n\n[7] Personalized Federated Learning via Variational Bayesian Inferennce, PMLR 2022\n\nQUESTIONS:\n- Could you formally define “feature–classifier mismatch” and clarify how it differs from discrepancies investigated in prior works [1,2,3,4,5,6,7]? \n\n- Please report a quantitative metric—such as cosine similarity between local features, classifier weights, or average gradient-angle divergence—and show how FedPFT and existing baselines change that metric during training.\n\n- The robustness exp in Section 4.7 should be  compared  with other baselilnes, and more sparse client (eg, 10%, 20%) should be considered for the realistic problem setting.\n\n- Have you measured communication cost or accuracy when applying FedPFT to out-of-distribution (OOD) clients?\n\n- It is highly recommended to follow the experimental settings of large benchmarks, such as those in [8, 9], to ensure a fair comparison. Could you describe the problem setting more detaily in comparison to these benchcmarks? \n\n[8] https://github.com/KarhouTam/FL-bench\n\n[9] https://github.com/alibaba/FederatedScope\n\nLIMITATIONS:\nThe authors have not discussed limitations of the work in the conclusion. I think the method is purely theoretical and there is no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses data heterogeneity in federated learning by proposing a novel personalized federated learning framework. Unlike existing methods that localize or fine-tune certain model parameters, the proposed FedPFT framework identifies the mismatch between local features and the global classifier as the main performance bottleneck. It uses personalized prompts combined with local features, processed through a shared self-attention transformation module, to align them with the global classifier during training. Experimental results demonstrate that FedPFT outperforms state-of-the-art methods by up to 7.08%.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1. This paper is well-written and easy to follow. Exploiting a prompt-driven transformation module to achieve alignment between local features and the global classifier for improved personalization in federated learning is both interesting and promising.\n\n2. The transformation experiments effectively explain the motivations and insights behind the approach, making them clear and accessible to readers.\n\n3. Extensive experiments validate the effectiveness of the proposed method across various datasets, demonstrating its superiority over existing competitors.\n\n**Weaknesses:**\n\n1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.\n\n2. The evaluation section can be improved to strengthen the applicability of the proposed method to real-world scenarios:\n\n   **(1) Scalability:** The current evaluation only considers 40 clients. The performance with a larger number of participating clients should be assessed to demonstrate scalability.\n\n   **(2) Participation mode:** Only the full participation setting is considered in the experiments. In practice, partial participation is more common in federated learning systems and should be evaluated.\n\n   **(3) Local update epochs:** The impact of the number of local update epochs should be evaluated to demonstrate the robustness of the proposed algorithm.\n\n   I would be glad to increase my rating if you can validate the scalability and robustness of the proposed method, particularly with a large number of participating clients and under partial participation settings.\n\nQUESTIONS:\n1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.\n\nLIMITATIONS:\n**Limitations:**\n\n1. The advantage of the proposed prompt-driven feature transformation over existing personalized federated learning methods based on feature alignment and classifier collaboration (e.g., FedPAC [1]) remains unclear to me. As shown in Figure 3, if $(p_{\\kappa,i}, \\tau_i)$ and $h_{\\kappa}$ are regarded as a unified \"classifier\" module, the proposed framework essentially also performs global feature alignment and local classifier collaboration, similar to FedPAC. Therefore, the core advantage and key contribution of the proposed framework are not clearly differentiated.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}