Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation
2026-04-14 • Computer Vision and Pattern Recognition
Computer Vision and Pattern Recognition
AI summaryⓘ
The authors address a problem in healthcare where different hospitals have only some types of medical data, making it hard to train shared AI models. They propose a new method, P-FIN, that not only fills in missing data but also tells how confident it is in those guesses. This confidence helps improve the AI's decisions locally and also helps decide which hospitals' data updates to trust more when combining models. Their tests on chest X-ray data show better performance compared to older methods, especially when data is more incomplete.
federated learningmultimodal datafeature imputationuncertainty estimationchest X-ray classificationprivacy-preserving AImodel aggregationhealthcare data heterogeneityAUC (Area Under Curve)CheXpert
Authors
Nafis Fuad Shahid, Maroof Ahmed, Md Akib Haider, Saidur Rahman Sagor, Aashnan Rahman, Md Azam Hossain
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.