⚡ electrical engineering

Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

This paper introduces Probabilistic Feature Imputation Network (P-FIN), a federated learning framework that generates calibrated uncertainty estimates for missing modalities to enable reliability-aware local feature gating and global aggregation, significantly improving multimodal chest X-ray classification performance in privacy-preserving healthcare settings.

Original authors: Nafis Fuad Shahid, Maroof Ahmed, Md Akib Haider, Saidur Rahman Sagor, Aashnan Rahman, Md Azam Hossain

Published 2026-04-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Nafis Fuad Shahid, Maroof Ahmed, Md Akib Haider, Saidur Rahman Sagor, Aashnan Rahman, Md Azam Hossain

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a team of doctors from different hospitals trying to build a super-smart AI assistant to help diagnose chest X-rays. They want to work together, but they have a big problem: privacy. They can't share patient files because of strict laws. So, they use a technique called Federated Learning, where the AI travels to each hospital, learns from the local data, and brings back only its "lessons" (mathematical updates) without ever seeing the actual patient records.

However, there's a second, trickier problem: Missing Information.

The Problem: The "Half-Empty" Puzzle

Some big hospitals have everything: the X-ray image and the radiologist's written report. But smaller clinics might only have the X-ray image and no written notes.

In the past, when the AI tried to learn from these smaller clinics, it had to guess what the missing report would say.

The Old Way (Deterministic): Imagine a student taking a test who doesn't know the answer. The old AI would just pick an answer and write it down with 100% confidence, even if it was a complete guess. If the guess was wrong, the whole team would learn the wrong lesson, and the AI would get worse. It was like a confident liar.
The Risk: In medicine, a confident wrong guess can lead to a misdiagnosis.

The Solution: P-FIN (The "Honest" AI)

The authors of this paper created a new system called P-FIN (Probabilistic Feature Imputation Network). Instead of just guessing a single answer, P-FIN acts like an honest student.

When it sees an X-ray and has to guess the missing report, it doesn't just say, "The patient has pneumonia." Instead, it says:

"Based on this image, I think it might be pneumonia, but I'm only 60% sure. My confidence is low because the image is blurry."

It outputs two things:

The Guess: What it thinks the missing text says.
The Confidence Score: A number telling everyone how shaky that guess is.

How It Uses This "Confidence Score"

The paper uses this confidence score in two clever ways, like a smart traffic control system:

1. Local Gating (The "Filter" at the Clinic)

Before a single hospital uses its own guess to make a diagnosis, the system checks the confidence score.

High Confidence: "Okay, this guess is solid. Let's use it."
Low Confidence: "Whoa, this guess is shaky. Let's mute this part of the information so it doesn't confuse the diagnosis."
Analogy: Imagine a chef tasting a soup. If the ingredient is fresh (high confidence), they add a lot. If the ingredient is old and questionable (low confidence), they add a tiny pinch or none at all, so it doesn't ruin the whole pot.

2. Global Aggregation (The "Team Meeting")

After all hospitals finish their local learning, they send their updates to a central server to combine them into one big model.

The Old Way (FedAvg): The server treated every hospital equally. If a small clinic had 1,000 patients but made terrible guesses, their bad advice would still count just as much as a big hospital's good advice.
The New Way (Fed-UQ-Avg): The server looks at the confidence scores.
- If a hospital's guesses were confident and reliable, the server says, "Great job! We'll listen to you more."
- If a hospital's guesses were full of uncertainty, the server says, "We hear you, but we'll weigh your advice lightly so you don't drag the whole team down."
Analogy: Imagine a group project. If one student submits a paper they wrote while half-asleep (low confidence), the teacher (the server) grades it but doesn't let that one bad grade ruin the whole group's average. They listen to the student who stayed up all night studying (high confidence) more.

The Results

The researchers tested this on real chest X-ray data from three different datasets.

The Outcome: In the hardest scenarios (where most hospitals were missing text data), their new method improved the AI's accuracy by 5.36% compared to the old "confident guesser" methods.
The Takeaway: By teaching the AI to say "I don't know" (or "I'm not sure") when it's guessing, the whole system becomes safer, more accurate, and more trustworthy.

Why This Matters

In the real world, we don't want AI that is confidently wrong. We want AI that knows its limits. This paper gives us a blueprint for building medical AI that can collaborate across different hospitals, handle missing data gracefully, and always tell the doctor, "Here is my best guess, and here is how much you should trust it."

1. Problem Statement

The paper addresses a critical challenge in Multimodal Federated Learning (FL) for healthcare: Modality Heterogeneity.

Context: In real-world medical settings, institutions vary in resources. Large hospitals often possess both radiological images (X-rays) and textual reports, while smaller clinics may only have images.
The Gap: Existing FL approaches attempt to handle missing modalities (e.g., missing text reports) using deterministic feature imputation networks. These networks synthesize missing data (e.g., predicting a text report from an image) as a single "point estimate."
The Risk: Deterministic models produce confident predictions even when the input is ambiguous or the imputation is poor. In safety-critical medical applications, this leads to "silent failures" where the model hallucinates features with high confidence, potentially causing erroneous diagnoses. Current methods lack Uncertainty Quantification (UQ), preventing downstream classifiers from distinguishing between reliable and unreliable imputed data.

2. Methodology

The authors propose P-FIN (Probabilistic Feature Imputation Network) combined with Fed-UQ-Avg, a novel aggregation strategy. The framework operates in two main stages:

A. Probabilistic Feature Imputation Network (P-FIN)

Instead of outputting a fixed vector for the missing modality, P-FIN models the conditional distribution $p(z_T | z_I)$ as a heteroscedastic Gaussian.

Architecture:
- Encoders: Uses ResNet-50 for images and BERT-base for text, projecting features into a shared 256-dimensional space.
- Imputation Core: A Transformer-based encoder takes image features and a learnable query token.
- Dual Output Heads: Two Multi-Layer Perceptrons (MLPs) output the mean ( $\mu$ ) (the imputed feature) and the variance ( $\sigma^2$ ) (the uncertainty) for each dimension of the feature vector.
Training (Calibration): To prevent the model from minimizing loss by simply predicting infinite uncertainty (a common issue in Gaussian NLL), the authors use $\beta$ -NLL loss. This loss function applies a stop-gradient operation to the variance term, forcing the model to balance reconstruction quality with calibrated uncertainty estimation.

B. Local Uncertainty-Aware Fusion

On clients with missing modalities, the imputed features are not used directly. Instead, an Uncertainty Gating mechanism is applied:

Gating: A sigmoid gate is computed from the predicted variance: $g = \text{sigmoid}(-\log \sigma^2)$ $g = sigmoid (- lo g σ^{2})$ .
- High uncertainty $\rightarrow$ Gate approaches 0 (suppresses the feature).
- Low uncertainty $\rightarrow$ Gate approaches 1 (retains the feature).
Fusion: The gated imputed features are fused with the original image features using Cross-Modal Attention, allowing the model to dynamically weigh contributions based on reliability.

C. Global Aggregation: Fed-UQ-Avg

Standard Federated Averaging (FedAvg) weights clients solely by dataset size, which can be detrimental if a large client has poor imputation quality. The authors introduce Fed-UQ-Avg:

Weighting Scheme: The aggregation weight $W_k$ $W_{k}$ for client $k$ $k$ is a blend of two components:
1. Data Weight ( $w_{data}$ ): Proportional to the number of samples (standard FedAvg).
2. Confidence Weight ( $w_{conf}$ ): Based on the client's mean imputation uncertainty ( $\bar{\sigma}^2$ ), calculated via a temperature-scaled exponential: $\exp(-\bar{\sigma}^2 / T)$ .
Mechanism: Clients with high uncertainty (poor imputation quality) receive lower weights, reducing their negative impact on the global model.

3. Key Contributions

Probabilistic Imputation: Introduction of P-FIN, the first framework to output calibrated uncertainty estimates alongside imputed features in a federated multimodal setting.
Dual-Level Uncertainty Utilization:
- Local: Uncertainty gating suppresses unreliable synthetic features before classification.
- Global: Fed-UQ-Avg prioritizes updates from clients with high-confidence imputations during model aggregation.
Calibrated Training: Adoption of $\beta$ -NLL loss to ensure the predicted uncertainty is meaningful and correlated with actual reconstruction error, preventing "variance explosion."
Robustness in Heterogeneity: Demonstrated that explicitly modeling uncertainty allows the system to "learn to distrust" poor data, significantly outperforming deterministic baselines in scenarios with high modality imbalance.

4. Experimental Results

Datasets: Experiments were conducted on three chest X-ray datasets: CheXpert, NIH Open-I, and PadChest.
Setup: Simulated federated environments with 10 clients and varying ratios of unimodal (image-only) to multimodal (image+text) clients (8:2, 6:4, 4:6).
Baselines: Compared against Zero-filling, Uniform-filling, deterministic FIN + FedAvg, and an ablation of their method (P-FIN + FedAvg).
Performance:
- P-FIN + Fed-UQ-Avg achieved the highest AUC across all configurations.
- In the most challenging scenario (80% unimodal clients), the proposed method achieved a 5.36% absolute AUC gain over the deterministic baseline (FIN + FedAvg).
- Ablation Study: Confirmed that both the probabilistic imputation (P-FIN) and the uncertainty-aware aggregation (Fed-UQ-Avg) contribute independently to performance, with the aggregation strategy providing an additional ~2.4–2.8% gain.
Calibration: The model demonstrated excellent calibration (Expected Calibration Error = 0.0422), with predicted uncertainty strongly correlating with actual imputation error.

5. Significance

Safety-Critical AI: The work directly addresses regulatory requirements (e.g., FDA) for uncertainty quantification in medical AI. It ensures that clinicians are warned when the system is "guessing" missing data, rather than presenting hallucinated features as fact.
Practical Deployment: By enabling effective collaboration between resource-rich and resource-poor hospitals without sharing raw data, the framework facilitates more inclusive and robust medical AI development.
Paradigm Shift: Moves the field from "blind imputation" to "uncertainty-aware synthesis," establishing a new standard for handling missing modalities in federated learning.