Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

2026-04-14

AI Summary

This work addresses the challenge of unreliable features in multimodal federated learning caused by missing modalities across institutions, a critical issue in high-stakes domains such as healthcare where existing methods lack proper uncertainty quantification for imputed features. To mitigate this risk, the authors propose Probabilistic Feature Imputation Network (P-FIN), the first approach in this setting to incorporate calibrated uncertainty estimation. P-FIN employs a local sigmoid-gated mechanism and a global Fed-UQ-Avg aggregation algorithm to enable uncertainty-aware bi-level optimization. Evaluated on federated chest X-ray classification tasks using CheXpert, NIH Open-I, and PadChest datasets, the method achieves a 5.36% AUC improvement over deterministic baselines under the most challenging modality-missing configuration, demonstrating its dual advantage in both reliability and predictive performance.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Computer Vision: Multi-modal Vision

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• multimodal federated learning
• modality heterogeneity
• feature imputation
• uncertainty quantification
• medical AI

Innovation

• Probabilistic Feature Imputation
• Uncertainty Quantification
• Multimodal Federated Learning
• Fed-UQ-Avg
• Missing Modality