Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

1. Problématique et Contexte

L'apprentissage fédéré (Federated Learning - FL) permet de former des modèles collaboratifs en préservant la confidentialité des données médicales, ce qui est crucial dans le secteur de la santé. Cependant, une contrainte majeure persiste : l'hétérogénéité des modalités.

• Le défi : Dans la réalité, les sites cliniques disposent de ressources variables. Les grands centres médicaux possèdent souvent des données complètes (images radiologiques + rapports textuels), tandis que les petites cliniques ou hôpitaux ruraux n'ont souvent accès qu'à un sous-ensemble de modalités (par exemple, uniquement les images X-ray, sans rapports textuels structurés).
• La limitation des approches existantes : Pour combler ce manque, des réseaux d'imputation de caractéristiques sont utilisés pour synthétiser les modalités manquantes (ex: générer un texte à partir d'une image). Cependant, ces méthodes actuelles sont déterministes : elles produisent une estimation ponctuelle unique sans mesure de fiabilité. En cas d'ambiguïté, le modèle peut « halluciner » des caractéristiques avec une confiance excessive, ce qui pose un risque critique pour le diagnostic médical.
• Besoin : Il est nécessaire de quantifier l'incertitude pour que le système puisse alerter les cliniciens lorsque les données synthétisées sont peu fiables, conformément aux réglementations (ex: FDA).

2. Méthodologie Proposée

Les auteurs proposent un cadre intégré appelé P-FIN (Probabilistic Feature Imputation Network) couplé à une nouvelle stratégie d'agrégation Fed-UQ-Avg.

A. Réseau d'Imputation Probabiliste (P-FIN)

Au lieu de prédire une valeur fixe pour la modalité manquante, P-FIN modélise la distribution conditionnelle $p(z_T | z_I)$ comme une Gaussienne hétéroscédastique.

• Architecture :
  • Un encodeur d'images (ResNet-50) et un encodeur de texte (BERT) projettent les données dans un espace commun.
  • Un module d'imputation basé sur un Transformateur (2 couches, 4 têtes d'attention) prend les caractéristiques d'image et un token d'interrogation apprenable pour prédire deux sorties : la moyenne ($\mu$) et la variance ($\sigma^2$) de la modalité textuelle manquante.
• Fonction de perte ($\beta$-NLL) : Pour éviter que le modèle ne minimise la perte en prédissant une incertitude infinie (un raccourci classique), les auteurs utilisent une perte de vraisemblance négative logarithmique (NLL) modifiée par un opérateur stop-gradient ($\beta$-NLL). Cela force le modèle à apprendre des représentations significatives tout en calibrant correctement l'incertitude.

B. Fusion Locale Consciente de l'Incertitude

Sur les clients unimodaux (sans texte), les caractéristiques imputées ne sont pas utilisées aveuglément.

• Gating par incertitude : Une porte (gate) est calculée via une fonction sigmoïde appliquée à la variance prédite : $g = \text{sigmoid}(-\log \sigma^2)$.
  • Si l'incertitude est élevée ($\sigma^2$ grand), la porte tend vers 0, atténuant fortement la contribution des caractéristiques imputées.
  • Si l'incertitude est faible, la porte reste proche de 1.
• Fusion par attention : Les caractéristiques d'image et les caractéristiques textuelles (filtrées par la porte) sont fusionnées via un mécanisme d'attention croisée bidirectionnelle, permettant au modèle de pondérer dynamiquement la fiabilité de chaque modalité.

C. Agrégation Fédérée Globale (Fed-UQ-Avg)

La stratégie d'agrégation standard (FedAvg) pondère les mises à jour des clients uniquement par la taille de leurs données, ce qui est dangereux si un client a beaucoup de données mais une mauvaise qualité d'imputation.

• Nouvelle pondération : Fed-UQ-Avg introduit un poids combiné basé sur deux facteurs :
  1. Poids des données ($w_{data}$) : Proportionnel à la taille du jeu de données local.
  2. Poids de confiance ($w_{conf}$) : Dépend de l'incertitude moyenne d'imputation du client ($\bar{\sigma}^2$). Les clients avec une faible incertitude (imputation de haute qualité) reçoivent un poids plus élevé.
• Formule : $W_k = (1 - \alpha)w_{data} + \alpha w_{conf}$, où $\alpha$ est un hyperparamètre de balance (fixé à 0.6). Cela permet de privilégier les mises à jour provenant de clients capables d'imputer fiablement leurs données manquantes.

3. Contributions Clés

1. P-FIN : Un réseau d'imputation probabiliste qui génère des distributions calibrées (moyenne + variance) plutôt que des estimations ponctuelles, utilisant la perte $\beta$-NLL.
2. Mécanisme de Gating Local : Utilisation de la variance prédite pour supprimer dynamiquement les caractéristiques imputées non fiables avant la classification locale.
3. Fed-UQ-Avg : Une stratégie d'agrégation fédérée novatrice qui intègre la qualité de l'imputation (via l'incertitude) dans le processus de mise à jour du modèle global, réduisant l'impact des clients peu fiables.
4. Validation Rigoureuse : Démonstration que la quantification de l'incertitude améliore la robustesse dans des scénarios de FL avec hétérogénéité sévère des modalités.

4. Résultats Expérimentaux

Les expériences ont été menées sur trois ensembles de données de radiographie thoracique (CheXpert, NIH Open-I, PadChest) dans un environnement fédéré simulé avec 10 clients et des ratios variables de clients unimodaux/multimodaux (jusqu'à 80% de clients sans texte).

• Performance : La méthode proposée (P-FIN + Fed-UQ-Avg) surpasse systématiquement les baselines déterministes (Zero-filling, FIN + FedAvg).
  • Dans le scénario le plus difficile (80% de clients unimodaux), l'approche proposée obtient une amélioration absolue de +5,36% de l'AUC (Area Under the Curve) par rapport à la méthode déterministe standard.
  • L'ajout de Fed-UQ-Avg à P-FIN apporte un gain supplémentaire de ~2,4% à 2,8% par rapport à l'utilisation de P-FIN avec FedAvg standard, confirmant l'importance de l'agrégation consciente de l'incertitude.
• Calibration : L'analyse montre que les incertitudes prédites sont bien calibrées (Erreur de Calibration Attendue - ECE = 0,0422) et corrélées positivement avec l'erreur de reconstruction réelle. Les cas à haute incertitude correspondent effectivement à des imputations de mauvaise qualité.

5. Signification et Impact

Ce travail adresse un problème critique de sécurité dans l'IA médicale : la confiance aveugle dans les données synthétisées.

• Sécurité Clinique : En permettant au système de « dire ce qu'il ne sait pas » (via l'incertitude), P-FIN permet aux radiologues de filtrer les prédictions peu fiables, alignant l'IA sur les exigences réglementaires (FDA) concernant la quantification de l'incertitude.
• Robustesse du FL : La méthode démontre que l'on peut construire des modèles collaboratifs performants même lorsque la majorité des participants ont des données incomplètes, à condition de gérer intelligemment la fiabilité de ces données manquantes.
• Paradigme : L'article promeut un paradigme d'apprentissage où le modèle apprend à « se méfier » (learning to distrust) de ses propres imputations incertaines, offrant une base pour des systèmes de santé plus résilients et dignes de confiance.