X-WIN: Building Chest Radiograph World Model via Predictive Sensing

Probleemstelling

Borstfoto's (Chest X-ray of CXR) zijn de meest gebruikte beeldvormingstechniek voor de diagnose van thoracale ziekten. Een fundamentele beperking van CXR is echter dat het 2D-projectieafbeeldingen zijn. Hierdoor treedt er structurele superpositie op (organen overlappen elkaar), wat het vastleggen van de 3D-anatomie bemoeilijkt. Dit maakt representatieleren en ziektediagnose uitdagend.

Aan de andere kant biedt Computertomografie (CT) gedetailleerde 3D-informatie, maar is deze veel duurder, minder toegankelijk (vooral in ontwikkelingslanden) en blootstelt de patiënt aan een aanzienlijk hogere stralingsdosis. Bestaande "wereldmodellen" voor medische beeldvorming (zoals CheXWorld) zijn beperkt tot 2D en missen de interne kennis van 3D-structuren, wat een groot gat laat tussen de prestaties van AI en die van radiologen (die wel in staat zijn om cognitief een 3D-model te reconstrueren op basis van 2D-foto's).

Methodologie: X-WIN

De auteurs stellen X-WIN (X-ray World Intelligence Network) voor, een nieuw wereldmodel dat volumetrische kennis distilleert uit CT-scans door te leren om 2D-projecties te voorspellen in een latente ruimte. Het kernidee is dat een model met geïnternaliseerde kennis van 3D-anatomie nieuwe projecties kan voorspellen onder verschillende transformaties in de 3D-ruimte.

Het framework bestaat uit de volgende componenten:

1. Architectuur:

   • Het systeem gebruikt twee netwerken: een encoder ($f_\theta$) die routine X-ray projecties (frontaal of lateraal) verwerkt, en een EMA-encoder (Exponential Moving Average, $f_{\theta'}$) die meerdere projecties van een CT-volume verwerkt als "target".
   • Een View Predictor ($g_v$), geconditioneerd op een "actie" (rotatie van de X-ray bron), voorspelt de latente representatie van een nieuwe projectie.
   • Een Mask Predictor ($g_m$) wordt gebruikt voor Masked Image Modeling (MIM).

2. Actie-ontwerp (Predictive Sensing):

   • In plaats van statische beelden te analyseren, simuleert het model het draaien van de X-ray bron rondom een patiënt (gebaseerd op CT-data).
   • De "actie" ($a_i$) definieert een rotatiehoek (yaw). Het model leert hoe de contextuele X-ray verandert bij deze rotatie.

3. Verliesfuncties (Loss Functions):

   • Affinity-guided Contrastive Alignment Loss: Een zachte contrastieve loss die niet alleen gepaarde projecties (voorspeld vs. target) dichter bij elkaar brengt, maar ook rekening houdt met de onderlinge gelijkenis tussen ongepaarde projecties uit hetzelfde volume. Dit vangt rijke anatomische correlaties op.
   • Masked Image Modeling (MIM) Loss: Wordt toegepast op zowel echte CXR's als gesimuleerde projecties. Dit dwingt het model om lokale en contextuele kenmerken te leren en verbetert de aanpasbaarheid aan downstream taken.
   • Structure-preserving Domain Adaptation Loss: Om de kloof tussen de gesimuleerde domein (CT-projecties) en het echte domein (echte CXR's) te overbruggen, wordt een domeinklassificator gebruikt. Deze dwingt het model om representaties van gesimuleerde data statistisch gelijk te maken aan die van echte data, terwijl de structurele informatie behouden blijft.

Belangrijkste Bijdragen

1. Eerste integratie van 3D-kennis in een CXR-wereldmodel: X-WIN is het eerste model dat 3D-spatiale kennis uit CT-scans distilleert om een wereldmodel voor 2D-radiografie te bouwen.
2. Affinity-geleide contrastieve uitlijning: Een nieuwe objectief die de encoding van discriminatieve kenmerken verbetert door gebruik te maken van rijke correspondenties tussen verschillende projecties van hetzelfde volume.
3. Superieure prestaties: Het model bereikt state-of-the-art resultaten op standaard CXR-taken (zoals ziektedetectie) met lineaire probing en few-shot fine-tuning.
4. 3D-reconstructiecapaciteit: Het bewijs dat het model daadwerkelijk 3D-structuren begrijpt, door het renderen van 2D-projecties die gebruikt kunnen worden voor de reconstructie van een 3D CT-volume.

Resultaten

De auteurs hebben X-WIN getest op zes populaire CXR-benchmarks (o.a. CheXpert, NIH-CXR, VinDr) en vergeleken met bestaande foundation modellen (zoals I-JEPA, DINOv2, RAD-DINO, CheXWorld).

• Lineaire Probing: X-WIN presteerde significant beter dan alle vergelijkende modellen. Bijvoorbeeld, op de VinDr-dataset behaalde X-WIN (ViT-Large) een AUROC van 0,925, vergeleken met 0,906 voor de tweede beste (Ark+) en 0,844 voor CheXWorld.
• Few-Shot Fine-tuning: Op de COVIDx-dataset (met slechts 4 tot 16 voorbeelden per klasse) behaalde X-WIN een AUROC van 0,993, wat alle andere modellen overtrof. Dit toont een sterke aanpasbaarheid aan nieuwe taken.
• 3D Reconstructie: Het model kon succesvol 2D-projecties renderen die gebruikt werden om een 3D CT-volume te reconstrueren (PSNR 27,87 dB, SSIM 0,789), wat aantoont dat de latente ruimte zinvolle 3D-informatie bevat.

Betekenis en Conclusie

X-WIN markeert een belangrijke doorbraak in medische beeldvorming door de beperkingen van 2D-radiografie te overwinnen zonder de patiënt bloot te stellen aan de kosten en risico's van CT-scans. Door 3D-kennis te "distilleren" uit CT-data en deze toe te passen op een 2D-wereldmodel, creëren de auteurs een systeem dat beter begrijpt hoe anatomie in de ruimte is opgebouwd.

Dit leidt tot:

• Hogere diagnose-accuraatheid bij ziekteherkenning.
• Betere interpretatie van beelden door het model (via 3D-reconstructie).
• Een nieuwe richting voor "sim-to-real" transfer learning in de medische AI, waarbij schaarse en dure 3D-data wordt gebruikt om goedkopere 2D-modellen te versterken.

De auteurs wijzen erop dat toekomstig werk gericht zal zijn op het verder verbeteren van de domeinadaptatie en het vergelijken van modellen die direct op CT-data zijn getraind versus deze distillatie-aanpak.