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

1. Problem Statement

Chest X-ray (CXR) is the most widely used imaging modality for thoracic disease detection. However, as 2D projectional images, CXRs suffer from structural superposition, where overlapping organs obscure anatomical details, making it difficult to capture true 3D structures.

• The Gap: While Chest CT provides detailed 3D volumetric information, it is significantly more expensive, involves higher radiation doses, and has lower accessibility than CXRs.
• The Challenge: Existing deep learning models and "World Models" (like CheXWorld) for CXRs are limited to 2D representations. They lack the internalized 3D anatomical knowledge that radiologists use to mentally reconstruct 3D volumes from 2D frontal and lateral views. This limits their ability to reason about spatial relationships and occlusions.

2. Methodology: The X-WIN Framework

The authors propose X-WIN (X-ray World Intelligence Network), a novel world model that distills 3D volumetric knowledge from CT scans into a model capable of processing 2D CXRs. The core philosophy is Predictive Sensing: the model learns to predict how an X-ray image changes under 3D transformations, thereby internalizing 3D anatomical structure.

A. Core Architecture

The framework consists of two synergistic networks:

1. Context Encoder ($f_\theta$): Takes a routine frontal or lateral CXR as input.
2. Target Encoder ($f_{\theta'}$): An Exponential Moving Average (EMA) encoder that processes multiple projections from a CT volume.
3. View Predictor ($g_v$): A lightweight network conditioned on an "action" (rotation) to predict the latent representation of a new projection view.
4. Mask Predictor ($g_m$): Used for Masked Image Modeling (MIM) to reconstruct local patches.

Action Design: The "action" is defined as rotating the X-ray source (Yaw, Pitch, Roll) around the patient. The model is trained to predict the latent representation of a new projection ($z_{patch}$) given a context projection and a specific rotation action, supervised by the actual projection from the CT volume.

B. Key Loss Functions

The training objective combines four distinct losses to ensure robust representation learning:

1. Affinity-Guided Contrastive Alignment Loss ($\mathcal{L}_{align}$):

   • Standard contrastive learning (InfoNCE) treats all non-matching pairs as negatives. However, projections from the same CT volume share rich anatomical correlations.
   • X-WIN introduces an affinity matrix based on mutual similarities between target projections. This creates a "soft" alignment, allowing the model to learn that different views of the same volume are related, not just repelling them.
   • Formula: $\mathcal{L}_{align} = \mathcal{L}_{InfoNCE} + \lambda_{affinity} \mathcal{L}_{affinity}$.

2. Masked Image Modeling Loss ($\mathcal{L}_{MIM}$):

   • Applied to both real CXRs and simulated CT projections.
   • Randomly masks patch tokens and forces the model to reconstruct them. This ensures the model learns fine-grained local features and contextual information, enhancing adaptability to downstream tasks.

3. Structure-Preserving Domain Adaptation Loss ($\mathcal{L}_{domain}$):

   • Addresses the Sim-to-Real gap between simulated CT projections and real clinical CXRs.
   • A domain classifier is trained to distinguish between real and simulated domains.
   • The loss encourages the simulated domain representations to be statistically indistinguishable from real domain representations (minimizing the gap) while preserving the structural information required for projection prediction.

4. Classification Loss ($\mathcal{L}_{cls}$):

   • Used to train the domain classifier itself to ensure it can effectively differentiate domains during the adaptation process.

3. Key Contributions

1. First 3D-Integrated CXR World Model: X-WIN is the first approach to integrate 3D spatial knowledge from CT volumes into a CXR world model via predictive sensing, bridging the gap between 2D imaging and 3D anatomy.
2. Affinity-Guided Contrastive Learning: Introduces a novel loss that leverages mutual similarities across projections from the same volume, capturing rich anatomical correspondences that standard contrastive losses miss.
3. Robust Domain Adaptation: Successfully aligns simulated (CT-derived) and real (clinical) CXR domains, enabling the model to generalize effectively to real-world data despite being trained heavily on synthetic projections.
4. 3D Reconstruction Capability: Demonstrates that the learned latent space contains sufficient 3D information to reconstruct 3D CT volumes from rendered 2D projections.

4. Experimental Results

The model was evaluated on six standard CXR benchmarks (VinDr, CheXpert, NIH-CXR, RSNA, JSRT, COVIDx) using Linear Probing and Few-Shot Fine-Tuning.

• Performance: X-WIN achieved State-of-the-Art (SOTA) performance across all benchmarks.
  • Linear Probing: Achieved an average AUROC of 0.892 (ViT-Large), outperforming previous CXR foundation models (e.g., CheXFound, Ark+) and vision-language models.
  • Few-Shot Learning: In the 4-shot setting for COVID-19 detection, X-WIN achieved an AUROC of 0.868, significantly outperforming competitors. With 16 shots, it reached 0.993.
• Ablation Studies:
  • Removing the affinity-guided loss or domain adaptation resulted in significant performance drops.
  • Direct rotation (Yaw) was found to be more effective than step-wise rotation or full 3D Euler angle rotation.
  • A rotation step size of 3° yielded optimal results.
• 3D Reconstruction: The model successfully reconstructed 3D CT volumes from rendered projections, achieving a PSNR of 27.87 dB and SSIM of 0.789, proving the latent space encodes valid 3D structural information.

5. Significance and Impact

• Bridging the 2D-3D Gap: X-WIN demonstrates that it is possible to imbue 2D radiography models with 3D anatomical reasoning without requiring 3D inputs during inference. This mimics the cognitive ability of radiologists to reconstruct 3D models from 2D views.
• Cost-Effective Medical AI: By distilling knowledge from expensive CT scans into a model that runs on cheap, accessible CXRs, X-WIN offers a pathway to improve diagnostic accuracy in resource-limited settings where CT is unavailable.
• Explainability: The ability to reconstruct 3D volumes from the learned representations provides a new form of interpretability, allowing clinicians to visualize the underlying 3D structure the model "sees."
• Future Directions: The work opens avenues for further sim-to-real adaptation research and comparing CXR-based 3D reasoning against models trained directly on CT data.

In summary, X-WIN represents a paradigm shift in medical imaging representation learning, moving from static 2D feature extraction to dynamic, 3D-aware predictive modeling, resulting in superior diagnostic performance and deeper anatomical understanding.