Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

Nafis Fuad Shahid, Maroof Ahmed, Md Akib Haider, Saidur Rahman Sagor, Aashnan Rahman, Md Azam Hossain
Proceedings of The 9th International Conference on Medical Imaging with Deep Learning, PMLR 315:3593-3607, 2026.

Abstract

Multimodal federated learning enables privacy-preserving collaborative model training healthcare applications. 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v315-shahid26a, title = {Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation}, author = {Shahid, Nafis Fuad and Ahmed, Maroof and Haider, Md Akib and Sagor, Saidur Rahman and Rahman, Aashnan and Hossain, Md Azam}, booktitle = {Proceedings of The 9th International Conference on Medical Imaging with Deep Learning}, pages = {3593--3607}, year = {2026}, editor = {Huo, Yuankai and Gao, Mingchen and Kuo, Chang-Fu and Jin, Yueming and Deng, Ruining}, volume = {315}, series = {Proceedings of Machine Learning Research}, month = {08--10 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v315/main/assets/shahid26a/shahid26a.pdf}, url = {https://proceedings.mlr.press/v315/shahid26a.html}, abstract = {Multimodal federated learning enables privacy-preserving collaborative model training healthcare applications. 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.} }
Endnote
%0 Conference Paper %T Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation %A Nafis Fuad Shahid %A Maroof Ahmed %A Md Akib Haider %A Saidur Rahman Sagor %A Aashnan Rahman %A Md Azam Hossain %B Proceedings of The 9th International Conference on Medical Imaging with Deep Learning %C Proceedings of Machine Learning Research %D 2026 %E Yuankai Huo %E Mingchen Gao %E Chang-Fu Kuo %E Yueming Jin %E Ruining Deng %F pmlr-v315-shahid26a %I PMLR %P 3593--3607 %U https://proceedings.mlr.press/v315/shahid26a.html %V 315 %X Multimodal federated learning enables privacy-preserving collaborative model training healthcare applications. 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.
APA
Shahid, N.F., Ahmed, M., Haider, M.A., Sagor, S.R., Rahman, A. & Hossain, M.A.. (2026). Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation. Proceedings of The 9th International Conference on Medical Imaging with Deep Learning, in Proceedings of Machine Learning Research 315:3593-3607 Available from https://proceedings.mlr.press/v315/shahid26a.html.

Related Material