# Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

arXiv: https://arxiv.org/abs/2604.12970 · HTML: https://mlsanity.com/paper/2604.12970/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 35 references that are in the map.
- First posted on arXiv 14 Apr 2026.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [CheXpert: A Large Chest Radiograph Dataset with Uncertainty Labels and Expert Comparison](https://mlsanity.com/paper/1901.07031/) (2019)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Contrastive Learning of Medical Visual Representations from Paired Images and Text](https://mlsanity.com/paper/2010.00747/) (2020)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)
- [CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning](https://mlsanity.com/paper/1711.05225/) (2017)
- [Stacked Cross Attention for Image-Text Matching](https://mlsanity.com/paper/1803.08024/) (2018)
- [Making the Most of Text Semantics to Improve Biomedical Vision--Language Processing](https://mlsanity.com/paper/2204.09817/) (2022)
- [PadChest: A large chest x-ray image dataset with multi-label annotated reports](https://mlsanity.com/paper/1901.07441/) (2019)
- [Multimodal Deep Learning](https://mlsanity.com/paper/2301.04856/) (2023)
- [Gated Multimodal Units for Information Fusion](https://mlsanity.com/paper/1702.01992/) (2017)
- [SMIL: Multimodal Learning with Severely Missing Modality](https://mlsanity.com/paper/2103.05677/) (2021)
- [Exploring Uncertainty Measures in Deep Networks for Multiple Sclerosis Lesion Detection and Segmentation](https://mlsanity.com/paper/1808.01200/) (2018)
- [On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks](https://mlsanity.com/paper/2203.09168/) (2022)
- [Assessing Reliability and Challenges of Uncertainty Estimations for Medical Image Segmentation](https://mlsanity.com/paper/1907.03338/) (2019)
- [Recalibration of Aleatoric and Epistemic Regression Uncertainty in Medical Imaging](https://mlsanity.com/paper/2104.12376/) (2021)
- [Multimodal Federated Learning With Missing Modalities through Feature Imputation Network](https://mlsanity.com/paper/2505.20232/) (2025)

## Built on by (most cited first, 1 of 1)

- [OmniMed-FL: A Robust Multimodal Federated Learning Framework for Clinical Diagnosis](https://mlsanity.com/paper/2609.10364/) (2026)

## Abstract (arXiv)

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.