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

AI summary

Overview

Research area: Medical computer vision, specifically self-supervised representation learning for chest X-ray radiography (CXR), combining world modeling, knowledge distillation from 3D CT, and domain adaptation.

Technical level: Advanced. The paper assumes familiarity with self-supervised learning objectives (InfoNCE, masked image modeling), joint-embedding predictive architectures with exponential moving average teachers, vision transformers, and tomographic reconstruction.

Scope: The paper introduces X-WIN, a CXR world model that distills 3D spatial knowledge from chest CT volumes by learning to predict CT projections in latent space, and demonstrates its performance on six CXR benchmarks plus 3D CT reconstruction.

What This Paper Is About

Chest X-rays are 2D projection images, so overlapping anatomy hides the 3D structure that radiologists mentally reconstruct when they read a frontal and lateral view. X-WIN aims to give a CXR model that same kind of internal 3D understanding by training it to predict what a chest X-ray would look like after the X-ray source moves to a new angle, using real chest CT volumes to supply the target views. The goal is better learned representations for downstream disease diagnosis on ordinary 2D radiographs.

Key Contributions

1. X-WIN framework for radiograph world modeling. The authors describe it as the first approach to integrate 3D spatial knowledge into a CXR world model, by distilling volumetric knowledge from chest CT through latent projection prediction.

2. Affinity-guided contrastive alignment objective. A softened contrastive loss that uses mutual similarities between projections from the same CT volume instead of a one-hot identity label matrix, so that correlated cross-view information is not treated as purely negative.

3. Structure-preserving domain adaptation. Real CXRs are incorporated into pretraining through masked image modeling in both domains, plus a domain classifier loss that encourages representations of the view predictor to be statistically similar to real-domain representations while retaining patch-level structural consistency.

4. Demonstrated downstream performance and 3D reconstruction. State-of-the-art linear-probing results on six CXR benchmarks and an ability to render 2D projections whose FDK reconstruction recovers a 3D CT volume.

Main Findings

• Linear probing outperforms prior CXR foundation models. With a ViT-Base backbone, X-WIN reached an average AUROC of 0.883 across VinDr, CheXpert, NIH-CXR, RSNA, and JSRT, and with ViT-Large 0.892. The next-best comparison model, Ark+, averaged 0.863, followed by CheXFound at 0.858 and CheXWorld at 0.844. General-domain models trailed furthest: I-JEPA at 0.736 and DINOv2 at 0.728.

• Per-dataset linear probing numbers (ViT-Large). VinDr 0.925, CheXpert 0.908, NIH-CXR 0.843, RSNA 0.929, JSRT 0.857. The corresponding ViT-Base values were 0.905, 0.902, 0.841, 0.917, and 0.848.

• Few-shot COVID-19 pneumonia detection. X-WIN achieved AUROC of 0.868 at 4 shots, 0.924 at 8 shots, 0.939 at 16 shots, and 0.993 with all data, exceeding CheXFound (0.823 / 0.883 / 0.897 / 0.977), CheXWorld (0.843 / 0.893 / 0.902 / 0.981), and RAD-DINO (0.793 / 0.852 / 0.884 / 0.971). Ark+ was higher at 4 shots (0.872) but lower at 8, 16, and all-data settings (0.908 / 0.921 / 0.985), and the paper notes Ark+ additionally used disease labels during pretraining.

• Direct yaw rotation is the best action design. Direct rotation (yaw) gave 0.925 AUROC on VinDr and 0.908 on CheXpert, versus 0.857 and 0.843 for step-wise rotation (yaw) and 0.871 and 0.865 for direct rotation with yaw, pitch, and roll.

• A 3° rotation step size works best. Step sizes tested were 1°, 3°, 5°, 10°, and 15°, corresponding to 180, 60, 36, 18, and 12 potential projections for sampling. The 3° setting performed best; larger steps degraded performance through increased sparsity, and 1° underperformed 3° likely because of reduced target diversity.

• Best loss weights from the ablation. The authors report that λ_affinity of 0.4 and λ_domain of 0.6 gave the best performance; λ_cls was fixed at 1.0 because that loss does not update the encoder. Decreasing any loss weight toward 0.0 consistently degraded performance, and applying the InfoNCE or MIM losses alone served as strong baselines while their combination improved results.

• Domain adaptation measurably closes the sim-to-real gap. Between simulated-domain and real-domain cluster centers, cosine similarity rose from 0.845 to 0.967 and L2 distance fell from 3.973 to 1.825 when the domain adaptation loss was applied.

• 3D CT reconstruction is feasible from learned representations. A VQ-GAN decoder trained on the frozen encoder and view predictor achieved 81.6% codebook usage, rendered 2D projections with PSNR 30.23 dB and SSIM 0.888, and reconstructed 3D volumes with PSNR 27.87 dB and SSIM 0.789. Reconstructions preserved overall structure and reasonable local detail, with some blurriness attributed to the abstraction of representation learning.

Methodology in Plain English

X-WIN borrows the idea of a "world model": a model that learns how its sensory input changes when the world changes, so that its internal state captures the underlying structure of that world. Here the "world" is the 3D chest, the "senses" are X-ray projections, and the "action" is rotating the X-ray source around the patient.

The setup uses two networks. A student encoder reads one routine projection, either frontal or lateral, and a small action-conditioned view predictor is asked to produce the latent representation of the same chest volume seen from a rotated angle. A teacher encoder, updated as an exponential moving average of the student, sees the actual projections at those angles and supplies the targets. Because a model that can predict novel views has to have internalized the 3D anatomy that produced them, the latent space is pushed toward volumetric understanding rather than surface appearance.

Three additional losses shape the training. An affinity-guided contrastive loss aligns predicted and target representations, but instead of treating every non-matching pair as equally wrong, it uses a similarity matrix between the target projections to soften the repulsion, which is appropriate since all projections come from one volume and share anatomy. A masked image modeling loss asks the model to reconstruct randomly masked patches, applied to both real CXRs and simulated projections, which builds local and contextual features. A domain classifier is trained to distinguish real from simulated images, and the view predictor is then pushed to produce representations that the classifier labels as real, while a patch-level term keeps those representations faithful to the targets. The overall loss is a weighted sum of these terms.

For evaluations, the authors froze the encoder and trained a linear layer for classification, and separately tested few-shot fine-tuning. For the reconstruction demonstration, they trained a VQ-GAN decoder on top of the frozen encoder and view predictor and used the FDK algorithm in the ASTRA toolbox to reconstruct a volume from rendered projections.

Why This Matters

Impact on research. The paper argues that a CXR world model limited to 2D leaves a capability gap versus radiologists, and it offers a concrete route to embedding 3D knowledge in a 2D imaging model without requiring CT at inference. It also extends the world-modeling paradigm from general vision and video into medical imaging, and shows that simulated projections from CT can act as a supervision signal for real radiographs.

Real-world applications.

• Better disease detection on routine chest X-rays, where only frontal and lateral views are available, by transferring spatial priors learned from CT.
• Few-shot adaptation to new diseases or new scanners, since the paper shows strong performance with 4, 8, or 16 labeled samples per class.
• Reducing reliance on CT for volumetric information, relevant where CT is roughly five times the cost of CXR and delivers notably more radiation.
• Improving access in settings where CT installation and maintenance costs restrict availability, by strengthening what can be extracted from widely available radiography.

Industry relevance. Hospital systems, imaging device manufacturers, and clinical AI vendors can use the pretrained representations for downstream diagnostic models without needing paired CT at deployment. The pretrained model checkpoint is released at the project's GitHub repository. The approach is also relevant to radiology workflow tools that must work from the most common and cheapest thoracic imaging exam.

Future Directions

• Further addressing the sim-to-real domain adaptation challenge, for example through invariant feature learning or domain mapping strategies, with the aim of improving data efficiency and generalizability.
• Comparing X-WIN directly against models trained on CT, to test how far CXR-based models can go for disease detection and diagnosis.
• Extending the world model beyond yaw rotation, since direct 3D rotation with yaw, pitch, and roll underperformed yaw alone in the reported ablation, suggesting more work is needed on action design.
• Closing the remaining gap in 3D reconstruction fidelity, since reconstructed volumes showed a certain level of blurriness attributed to the abstraction of visual representation learning.

Target Audience

Researchers and engineers working on medical image representation learning, self-supervised pretraining, or foundation models for radiology, who already know terms like InfoNCE, masked image modeling, and EMA teachers. It is also relevant to clinical AI teams evaluating how to use CT-derived knowledge to improve CXR models, and to world-model researchers interested in an application outside video and robotics. Readers without a machine learning background will find the method sections difficult, though the motivation section and the results tables are accessible.

Authors’ abstract

Chest X-ray radiography (CXR) is an essential medical imaging technique for disease diagnosis. However, as 2D projectional images, CXRs are limited by structural superposition and hence fail to capture 3D anatomies. This limitation makes representation learning and disease diagnosis challenging. To address this challenge, we propose a novel CXR world model named X-WIN, which distills volumetric knowledge from chest computed tomography (CT) by learning to predict its 2D projections in latent space. The core idea is that a world model with internalized knowledge of 3D anatomical structure can predict CXRs under various transformations in 3D space. During projection prediction, we introduce an affinity-guided contrastive alignment loss that leverages mutual similarities to capture rich, correlated information across projections from the same volume. To improve model adaptability, we incorporate real CXRs into training through masked image modeling and employ a domain classifier to encourage statistically similar representations for real and simulated CXRs. Comprehensive experiments show that X-WIN outperforms existing foundation models on diverse downstream tasks using linear probing and few-shot fine-tuning. X-WIN also demonstrates the ability to render 2D projections for reconstructing a 3D CT volume.