Probabilistic Feature Imputation and Uncertainty-Aware Multimodal Federated Aggregation

Probleemstelling

In de moderne gezondheidszorg is multimodaal leren (het combineren van beelden, zoals röntgenfoto's, met tekstuele rapporten) essentieel voor nauwkeurige diagnoses. Echter, bij **Federated Learning **(FL)—een privacybehoudende methode waarbij modellen lokaal worden getraind zonder dat patiëntdata wordt gedeeld—ontstaat een fundamenteel probleem: modale heterogeniteit.

Veel klinische locaties hebben niet dezelfde infrastructuur. Grote academische centra hebben vaak zowel röntgenfoto's als gedetailleerde tekstuele rapporten, terwijl kleinere ziekenhuizen of plattelandsklinieken vaak alleen beeldmateriaal hebben. Bestaande oplossingen voor dit ontbrekende data-probleem gebruiken deterministische netwerken om de ontbrekende modus (bijv. de tekst) te imputeren (samenstellen) op basis van de aanwezige modus (het beeld).

De kritieke tekortkoming van deze bestaande methoden is dat ze punt-schattingen produceren zonder maatstaven voor betrouwbaarheid. Zelfs als het model "raadt" of hallucineert, geeft het een zelfverzekerde output. In veiligheidskritieke medische toepassingen kan dit leiden tot foutieve diagnoses, omdat downstream-classificatoren alle geïmputeerde features als even betrouwbaar behandelen.

Methodologie

De auteurs stellen **P-FIN **(Probabilistic Feature Imputation Network) voor, een framework dat onzekerheid expliciet modelleert tijdens het imputeren van ontbrekende modaliteiten. De aanpak bestaat uit drie kerncomponenten:

**1. Probabilistische Feature Imputation **(P-FIN)

In plaats van een vast vector te voorspellen, modelleert P-FIN de conditionele verdeling $p(z_T | z_I)$ als een heteroscedastische Gaussische verdeling.

• Architectuur: Het netwerk gebruikt een Transformer-encoder die beeldfeatures (uit ResNet-50) omzet naar een verdeling van tekstfeatures (uit BERT).
• Output: Het netwerk voorspelt niet alleen een gemiddelde ($\mu$) voor de ontbrekende features, maar ook een variantie ($\sigma^2$) die de onzekerheid per dimensie aangeeft.
• Training: Om te voorkomen dat het model de loss minimaliseert door simpelweg oneindige onzekerheid te voorspellen, wordt gebruikgemaakt van de $\beta$-NLL loss (Negative Log-Likelihood). Deze lossfunctie gebruikt een stop-gradient operatie om de variantie te regulariseren en zorgt voor een goed gekalibreerde onzekerheidsschatting.

2. Lokale Onzekerheidsbewuste Fusie

Op clients met slechts één modus (bijv. alleen beelden) wordt de geïmputeerde tekst niet blindelings gebruikt.

• Gating Mechanisme: De voorspelde variantie ($\sigma^2$) wordt omgezet in een "gate" via een sigmoid-functie: $g = \text{sigmoid}(-\log \sigma^2)$.
  • Bij hoge onzekerheid nadert de gate naar 0, waardoor de onbetrouwbare features worden onderdrukt.
  • Bij lage onzekerheid blijft de gate dicht bij 1.
• Fusie: De gefilterde tekstfeatures worden via cross-modal attention (Multi-Head Attention) gecombineerd met de beeldfeatures. Dit zorgt ervoor dat het model dynamisch weegt welke modus betrouwbaar is voordat de classificatie plaatsvindt.

3. Globale Aggregatie: Fed-UQ-Avg

De standaard FedAvg-algoritme weegt client-bijdragen puur op basis van datasetgrootte. Dit is gevaarlijk als clients met grote datasets maar slechte imputatiekwaliteit het globale model verstoren.

• Fed-UQ-Avg: Dit nieuwe aggregatiestrategie introduceert een vertrouwensgewicht gebaseerd op de gemiddelde onzekerheid van de client.
• Formule: Het gewicht is een combinatie van het data-gewicht (grootte) en een confidence-gewicht: $w_{conf} \propto \exp(-\bar{\sigma}^2 / T)$.
• Effect: Clients met lage onzekerheid (hoge kwaliteit imputatie) krijgen een hoger gewicht in de globale update, terwijl clients met hoge onzekerheid minder invloed hebben, zelfs als ze veel data hebben.

Belangrijkste Bijdragen

1. P-FIN Architectuur: Het eerste framework dat probabilistische imputatie toepast in een federale multimodale setting, waarbij onzekerheid wordt gemodelleerd als een verdeling in plaats van een punt.
2. Dubbele Onzekerheidsbenutting: Een uniek tweeledig mechanisme dat onzekerheid lokaal gebruikt om features te filteren (gating) en globaal om client-bijdragen te wegen (aggregatie).
3. Fed-UQ-Avg: Een nieuwe aggregatiestrategie die de kwaliteit van imputatie expliciet meenemt in het federale leerproces, wat cruciaal is voor heterogene omgevingen.
4. Regelgeving en Veiligheid: De methode adresseert direct de eisen van toezichthouders (zoals de FDA) voor onzekerheidskwantificatie in medische AI-systemen.

Resultaten

De methode werd getest op drie grote chest X-ray datasets: CheXpert, NIH Open-I en PadChest. De experimenten simuleerden federale omgevingen met verschillende verhoudingen van unimodale (alleen beeld) tot multimodale clients (8:2, 6:4, 4:6).

• Prestatieverbetering: P-FIN in combinatie met Fed-UQ-Avg overtrof consistent deterministische baselines (zoals standaard imputatie met FedAvg).
• Kerncijfer: In de meest uitdagende configuratie (80% unimodale clients) werd een absolute verbetering van 5,36% in AUC (Area Under the Curve) behaald ten opzichte van de beste deterministische methode.
• Ablatie-studie: Het tonen van de meerwaarde van beide componenten:
  • P-FIN + FedAvg (alleen lokale onzekerheid) verbeterde al ten opzichte van deterministische methoden.
  • De toevoeging van Fed-UQ-Avg (globale onzekerheid) leverde een extra verbetering van ongeveer 2,4% tot 2,8% op.
• Kalibratie: De onzekerheidsschattingen waren goed gekalibreerd (Expected Calibration Error = 0,0422) en correleerden sterk met de daadwerkelijke reconstructiefouten: hoge onzekerheid betekende consistent slechte imputatie.

Significantie

Dit werk biedt een robuust antwoord op het probleem van data-fragmentatie in de federale gezondheidszorg. Door het paradigma te verschuiven van "blind vertrouwen in geïmputeerde data" naar "leren wantrouwen" bij hoge onzekerheid, verhoogt P-FIN de veiligheid en betrouwbaarheid van AI-systemen.

De methode zorgt ervoor dat:

1. Lokale artsen niet worden misleid door hallucinaties van het model.
2. Het globale model niet wordt "vergiftigd" door clients die slechte imputaties genereren.
3. Medische AI-systemen transparant zijn over hun beperkingen, wat essentieel is voor klinische adoptie en naleving van regelgeving.

De code is beschikbaar gesteld, wat bijdraagt aan de reproduceerbaarheid en verdere ontwikkeling van veilige multimodale federale leer-systemen.