Skip to content

CheXWorld: Exploring Image World Modeling for Radiograph Representation Learning

🕒 Published (v1): 2025-04-18 17:50 UTC · Source: Arxiv · Venue: CVPR 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

CheXWorld is the first self-supervised world model for radiographic images, built on a joint-embedding predictive architecture (JEPA) that simultaneously encodes local anatomical microstructures, global organ layouts, and cross-domain appearance variations. Pre-trained on ~0.5M chest X-rays using a unified latent-space prediction objective, it outperforms existing SSL methods and large-scale medical foundation models across eight classification and segmentation benchmarks while using 20× less compute than the closest competitor.

Problem

Existing SSL methods for radiology either focus on pixel-level reconstruction (masked autoencoders) or invariant contrastive representations, neither of which explicitly models the hierarchical anatomical knowledge and cross-domain adaptability that radiologists rely on. Large foundation models like Rad-DINO achieve competitive performance but require massive compute budgets and private data. No prior work has applied the world-modeling paradigm to medical imaging.

Method

CheXWorld extends the Joint-Embedding Predictive Architecture (I-JEPA) to radiology via three complementary tasks unified in a single forward pass:

  1. Local anatomical structure modeling: Masks four random rectangular regions; the context encoder (ViT-B) + predictor reconstructs target patch embeddings from surrounding visible patches (L2 loss in latent space). Encourages capture of fine-grained tissue textures and microstructures.

  2. Global anatomical layout modeling: Crops two spatially disjoint patches from the same image; predictor uses sinusoidal positional encoding of the relative displacement Δx→y as a latent variable to predict one crop's embeddings from the other. Captures long-range organ topology.

  3. Domain variation modeling: Constructs context-target pairs by applying augmentation pipelines (Gaussian blur, brightness, contrast, gamma) parameterized by vector a; predictor conditioned on a via a lightweight policy network π must predict the un-augmented target embedding. Learns equivariant (not invariant) representations encoding predictable domain transitions.

The unified framework samples two crops (y₁, y₂), applies independent augmentation+mask configurations (a₁,M₁) and (a₂,M₂), and computes four cross-prediction losses (L₁→₁, L₂→₂, L₁→₂, L₂→₁) in one forward pass. Target encoder is an EMA of the context encoder. Backbone: ViT-Base; pre-training: 300 epochs, 8× RTX 4090 (~16 GPU-hours).

Key Contributions

  • First world-modeling framework for self-supervised radiograph representation learning.
  • Three-task unified JEPA pipeline that models local anatomy, global layout, and domain shift simultaneously without increased forward passes.
  • Equivariant (vs. invariant) latent space via domain variation modeling, preserving cross-domain information rather than discarding it.
  • State-of-the-art results on 8 benchmarks using only 0.5M public X-rays and ~128 GPU-hours, vs. Rad-DINO's 2560 GPU-hours + private data.
  • Qualitative verification via RCDM diffusion decoding, anatomical correspondence heatmaps, and a domain sensitivity recall test.

Results

Classification (AUROC / Accuracy): - VinDr-CXR: 95.24 vs. Adam-v2 91.46, SimMIM 92.81, Rad-DINO 95.16 (20× more compute) - ShenZhen-CXR: 98.88 vs. Adam-v2 97.80, SimMIM 98.09 - ChestX-ray14: 83.58 vs. Adam-v2 83.4, SimMIM 83.04 - RSNA Pneumonia (accuracy): 75.03 vs. SimMIM 74.09, Rad-DINO 74.51 - CheXpert: 89.63 vs. SimMIM 89.14, Adam-v2 88.90

Segmentation (Dice): - SIIM-ACR Pneumothorax: 84.58 vs. Adam-v2 83.60, MAE 83.01 - RSNA Pneumonia (seg): 79.02 vs. Adam-v2 78.53, MAE 77.39

Few-shot (MedFMC-ChestDR AUROC): - 1-shot: 64.60 vs. MAE 61.31; 5-shot: 75.19 vs. MAE 74.03; 10-shot: 76.40 vs. MAE 75.26

Sample efficiency: CheXWorld with 10% training data outperforms all baselines trained on 100% data on VinDr-CXR.

Domain sensitivity: Top-5 recall = 77.67 vs. 44.37 without domain condition and 7.81 random baseline.

Ablation: All three tasks are individually additive; unified formulation outperforms naive loss sum (89.91 vs. 89.12 on VinDr 1%); latent variables Δx→y and a are each critical—removing either causes significant performance drops.

Limitations

  • Pre-training restricted to frontal chest X-rays (~0.5M images); generalizability to other modalities (CT, MRI, pathology) is untested.
  • Domain variation is simulated via photometric augmentations rather than real multi-scanner domain shifts, which may not fully capture real distribution shifts.
  • Evaluation confined to 2D radiographs; extension to 3D volumetric data requires non-trivial architectural adaptation.
  • The unified framework requires careful hyperparameter balancing across four loss terms; sensitivity to this is not fully ablated.
  • Rad-DINO comparison is acknowledged as not fully controlled (private data, 20× compute), leaving headroom in the strongest baseline unclear.

Relevance to Foundation Models in Medicine

CheXWorld demonstrates that world-modeling—predicting latent outcomes of "actions" (masking, spatial displacement, augmentation)—is a principled alternative to masked reconstruction or contrastive learning for building radiology foundation models, and achieves SOTA with dramatically fewer resources than models like Rad-DINO. Its equivariant representation design directly addresses domain shift, a pervasive bottleneck in deploying medical foundation models across hospitals and devices. The unified JEPA framing also offers a compositional template for extending world-model pre-training to other imaging modalities by swapping in domain-relevant "actions." For the foundation-models-in-medicine research agenda, this paper positions world modeling as a data-efficient and compute-efficient path toward general-purpose medical vision encoders.