Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation
Nafis Fuad Shahid, Maroof Ahmed, Md Akib Haider, Saidur Rahman Sagor, Aashnan Rahman + 1 more
TLDR
P-FIN introduces probabilistic feature imputation with uncertainty estimates for robust multimodal federated learning, improving medical image classification.
Key contributions
- Proposes P-FIN for probabilistic feature imputation, providing calibrated uncertainty estimates with imputed features.
- Leverages uncertainty locally via sigmoid gating to attenuate unreliable feature dimensions before classification.
- Employs Fed-UQ-Avg, an aggregation strategy prioritizing updates from clients with reliable imputation.
- Achieves +5.36% AUC gain in federated chest X-ray classification over deterministic baselines.
Why it matters
Existing multimodal federated learning methods lack reliability measures for imputed features, posing significant risks in safety-critical medical applications. P-FIN addresses this by providing uncertainty estimates, making collaborative AI safer and more robust for healthcare.
Original Abstract
Multimodal federated learning enables privacy-preserving collaborative model training across healthcare institutions. However, a fundamental challenge arises from modality heterogeneity: many clinical sites possess only a subset of modalities due to resource constraints or workflow variations. Existing approaches address this through feature imputation networks that synthesize missing modality representations, yet these methods produce point estimates without reliability measures, forcing downstream classifiers to treat all imputed features as equally trustworthy. In safety-critical medical applications, this limitation poses significant risks. We propose the Probabilistic Feature Imputation Network (P-FIN), which outputs calibrated uncertainty estimates alongside imputed features. This uncertainty is leveraged at two levels: (1) locally, through sigmoid gating that attenuates unreliable feature dimensions before classification, and (2) globally, through Fed-UQ-Avg, an aggregation strategy that prioritizes updates from clients with reliable imputation. Experiments on federated chest X-ray classification using CheXpert, NIH Open-I, and PadChest demonstrate consistent improvements over deterministic baselines, with +5.36% AUC gain in the most challenging configuration.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.