Research map / Federated learning and privacy / Personalized and heterogeneous FL
Tackling Feature-Classifier Mismatch in Federated Learning via Prompt-Driven Feature Transformation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FedPFT, which resolves feature-classifier mismatch during FL training by alternating prompt-driven self-attention feature transformation with model updates, plus an optional contrastive-learning extension (FedPFT+Con).
Problem
Personalized federated learning under non-IID data: the global model's feature extractor and classifier are individually strong, but mismatch between local features and the global classifier persists during training, causing gradient divergence and suboptimal aggregation.
Approach
Insert a shared self-attention feature transformation module (FTM) between the feature extractor and heads, and train per-client personalized prompts (with the extractor frozen) to align local features to the global classifier before updating model parameters each round, so mismatch is resolved during training rather than post-hoc; task-specific prompts also let auxiliary objectives such as MoCo c
Key results
- FedPFT outperforms state-of-the-art methods by up to 5.07% under Dirichlet non-IID, and FedPFT+Con by up to 7.08%, per the abstract.
- On CIFAR-100 with alpha=0.1, FedPFT reaches 60.98% vs 34.47% for FedAvg (Dirichlet non-IID table).
- Linear probe on CIFAR-100 (alpha=0.1) gives FedPFT 61.60% vs FedAvg 59.50%, indicating improved feature separability.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add classification prompts (ablation) | FedAvg on CIFAR-100, α=0.1 · CIFAR-100, α=0.1 | accuracy: 33.87$7 1.35 -> 40.977 1.28 | Helped |
| add alternating training to classification prompts (ablation) | classification-prompt variant · CIFAR-100, α=0.1 | accuracy: 40.97± 1.28 -> 60.98± 0.39 | Helped |
| add contrastive loss (ablation) | FedPFT · CIFAR-100, α=0.1 | accuracy: 60.98± 0.39 -> 61.13± 0.50 | Helped |
| add contrastive-task prompts (ablation) | FedPFT with contrastive loss · CIFAR-100, α=0.1 | accuracy: 61.13± 0.50 -> 62.03± 1.41 | Helped |
| use contrastive learning without mismatch alignment (ablation) | personalized FL · CIFAR-100, α=0.1 | accuracy: 60.98± 0.39 -> 53.76± 0.35 | Hurt |
| add FedPFT prompt-driven feature transformation | FedAvg · CIFAR-100, Pathological non-IID | accuracy: 34.27 ± 0.44 -> 57.62 ± 1.18 | Helped |
| add collaborative contrastive learning | FedPFT · CIFAR-100, Pathological non-IID | accuracy: 57.62 ± 1.18 -> 58.14 ± 0.71 | Helped |
| use classification prompts for linear-probe features (ablation) | FedPFT feature representation · CIFAR-100, Dirichlet α=0.1 | linear probe accuracy: 64.08% -> 64.12% | No clear effect |
Surprising findings
- Localizing or personalizing parts of the model degrades the global feature extractor's linear-probe quality, even though it improves personalized accuracy.
- Simple prompt-only fine-tuning after FedAvg, or a plain personalized linear layer, does not match FedPFT's gains, suggesting the training-time alignment matters more than added capacity.
Limitations
- Adds communication overhead of roughly 20% per round due to uploading the FTM.
- Experiments use small ResNet backbones (ResNet-8/10) and small per-client data (500 samples), limiting evidence on larger models.
- Hyperparameters (prompt count, R_f/R_a split) are chosen by default and may need per-scenario tuning.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 Tiny ImageNet PACS DomainNet
Benchmarks: CIFAR-10 CIFAR-100 Tiny ImageNet PACS DomainNet
Metrics: test accuracy
Baselines: FedAvg FedBABU FedETF FedFA FedAMP FedPer FedRep FedBN FedRoD pFedSD
Builds on: FedAvg FedBABU MoCo FedRoD
Closest papers by meaning
- Personalized Federated Learning via Feature Distribution Adaptation (NeurIPS 2024 · similarity 0.91)
- Unlocking the Potential of Prompt-Tuning in Bridging Generalized and Personalized Federated Learning (CVPR 2024 · similarity 0.90)
- Harmonizing Generalization and Personalization in Federated Prompt Learning (ICML 2024 · similarity 0.90)
- FedPHA: Federated Prompt Learning for Heterogeneous Client Adaptation (ICML 2025 · similarity 0.90)
- PFedCS: A Personalized Federated Learning Method for Enhancing Collaboration among Similar Classifiers (AAAI 2025 · similarity 0.89)
- FedLoGe: Joint Local and Generic Federated Learning under Long-tailed Data (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.