Research map / Federated learning and privacy / Personalized and heterogeneous FL
Watch Your Head: Assembling Projection Heads to Save the Reliability of Federated Models
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.
Shows federated models are poorly calibrated and overconfident on OOD data due to biased projection heads, and proposes APH, an ensemble of fine-tuned projection heads sampled around the federated head to restore reliability.
Problem
Federated learning models trained on heterogeneous (non-IID) data are unreliable: they are miscalibrated on in-distribution test data and assign low uncertainty to out-of-distribution samples, and standard uncertainty methods (MC-Dropout, deep ensembles) are impractical in federated settings.
Approach
Treat the federated projection head as a prior, sample several perturbed head initializations around it, fine-tune each head on local data with varying learning rates, and average their predictions over a frozen feature extractor, adding parameter diversity at small inference cost.
Key results
- On CIFAR-100 with ResNet-50, FedAvg + APH improves F-ECE from 0.458 to 0.091 and accuracy from 0.284 to 0.484 with 10 heads.
- APH applied to FedProx on CIFAR-10 reduces F-ECE from 0.226 to 0.066 and NLL from 1.662 to 0.344.
- With 100 projection heads, APH adds about 24.49% computation cost on ResNet-50 on CIFAR-100, below 30%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add APH, assembling multiple fine-tuned projection heads | FedAvg · Cifar10, Cifar100, Tiny-ImageNet | F-ECE: 0.168 -> 0.036; 0.458 -> 0.091; 0.395 -> 0.089 | Helped |
| Add APH, assembling multiple fine-tuned projection heads | FedAvg · Cifar10, Cifar100, Tiny-ImageNet | Accuracy: 0.615 -> 0.902; 0.284 -> 0.484; 0.098 -> 0.294 | Helped |
| Add APH to FedNTD | FedNTD · Cifar10, Cifar100, Tiny-ImageNet | F-ECE: 0.263 -> 0.033; 0.449 -> 0.128; 0.351 -> 0.149 | Helped |
| Add APH to FedNTD | FedNTD · Cifar100 | Accuracy: 0.298 -> 0.485 | Helped |
| Add APH to FedFOMO | FedFOMO · Cifar100 | Accuracy: 0.315 -> 0.303 | Hurt |
| Add APH to FedDyn | FedDyn · Tiny-ImageNet | NLL: 10.496 -> 13.769 | Hurt |
| Fine-tune the whole projection head once while freezing the feature extractor | Generic federated model · Cifar10 in-domain test data | F-ECE | Helped |
| Fine-tune only the last fully connected layer | Generic federated model · Cifar10 in-domain test data | F-ECE | Helped |
| Fine-tune projection heads without adding parameter diversity | Generic federated model · Cifar10 OOD test data | Predictive entropy | No clear effect |
| Increase the number of assembled heads from 10 to 50 (ablation) | APH on FedAvg · Cifar100 | F-ECE: 0.091 -> 0.077 | Helped |
| Increase the number of assembled heads from 10 to 100 (ablation) | APH on FedAvg · Cifar100 | F-ECE: 0.091 -> 0.071 | Helped |
Surprising findings
- Personalized federated models are more calibrated than centralized training but still show low uncertainty on OOD data.
- Fine-tuning only the projection head improves calibration but leaves OOD predictive entropy almost unchanged.
Limitations
- Ensemble-based baselines such as MC-Dropout and deep ensembles are only partly compared because other calibration methods were deemed inapplicable to federated training.
- Evaluation limited to image classification on CIFAR-10/100 and Tiny-ImageNet with a CNN and ResNet-50.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 Tiny-ImageNet
Benchmarks: CIFAR-10 CIFAR-100 Tiny-ImageNet SVHN ImageNet-O
Metrics: accuracy F-ECE NLL predictive entropy additional computation cost
Baselines: FedAvg FedDropout Client Ensembles FedAvg + FineTune MC-Dropout FedProx FedDyn FedNTD FedALA FedFOMO FedPAC
Builds on: FedAvg deep ensembles MC-Dropout CCVR FedRoD
Closest papers by meaning
- Federated Learning with Sample-level Client Drift Mitigation (AAAI 2025 · similarity 0.87)
- Unlocking the Potential of Model Calibration in Federated Learning (ICLR 2025 · similarity 0.87)
- FedGPS: Statistical Rectification Against Data Heterogeneity in Federated Learning (NeurIPS 2025 · similarity 0.87)
- FLHetBench: Benchmarking Device and State Heterogeneity in Federated Learning (CVPR 2024 · similarity 0.87)
- FedAdaVR: Adaptive Variance Reduction for Robust Federated Learning under Limited Client Participation (NeurIPS 2026 · similarity 0.87)
- FedQS: Optimizing Gradient and Model Aggregation for Semi-Asynchronous Federated Learning (NeurIPS 2025 · similarity 0.87)
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.