Rad-Jepa-3D

Rad-JEPA 3D: Radiology Joint-Embedding Predictive Model for 3D Computed Tomography

<p align="center"> <img src="imgs/architecture.png" alt="Rad-JEPA 3D Architecture" width="90%"> </p>

Overview

Rad-JEPA 3D is a joint-embedding predictive framework that learns volumetric CT representations by predicting the latent features of a complete scan from a masked view. At its core is a hybrid H-Mamba encoder that fuses a Mamba state-space branch (inter-slice continuity) with a grouped-query attention branch (cross-plane spatial context), combined through a lightweight per-token router. Hidden States Orthogonal Regularization (HSOR) aligns student-teacher hidden states and reduces feature redundancy across encoder layers.

Pretrained on ~120,000 CT scans, Rad-JEPA 3D achieves state-of-the-art results with only 4.0B total parameters: competitive closed-ended VQA and the best average spatial-reasoning score on the Spatial-Med benchmark.

Model Checkpoints

Encoder Details

• Architecture: 12 H-Mamba blocks with per-layer routing between Mamba SSM and grouped-query attention (GQA) with 3D RoPE
• Input: 3D CT volume (1, 32, 256, 256) — patchified via Conv3d(8,16,16) into 1024 tokens at 384-d
• Output: 1024 patch tokens at 384-d; global-pool to a single 384-d embedding
• Pretraining: V-JEPA objective (L1 loss in representation space) + HSOR on ~120k CT volumes
• Scan order: Raster (not Morton)
• Epoch: 252

MLLM Details

• Vision encoder: Same H-Mamba encoder as above (LoRA-adapted during stage 2)
• Projector: 2-layer MLP (384 → 1024 → 2560)
• LLM: Qwen3-4B with LoRA adapters (r=16, alpha=32)
• Training: Stage 1 (projector + CLIP contrastive) → Stage 2 (vision + projector + LoRA)

Quick Start

Install the code and pull the weights from the Hub:

git clone https://github.com/huyquoctrinh/RadJepa.git
cd RadJepa/src
pip install huggingface_hub

Download the Weights

from huggingface_hub import snapshot_download

repo_dir = snapshot_download("huyquoctrinh/Rad-JEPA-3D")
# repo_dir now contains config.json, encoder/encoder.pt, mllms_qwen3/stage2_best.pt

Or fetch a single file:

from huggingface_hub import hf_hub_download

config_path  = hf_hub_download("huyquoctrinh/Rad-JEPA-3D", "config.json")
encoder_path = hf_hub_download("huyquoctrinh/Rad-JEPA-3D", "encoder/encoder.pt")

Load the Encoder

import os, torch
from load_checkpoint import load_encoder

enc, cfg, meta = load_encoder(
    os.path.join(repo_dir, "encoder/encoder.pt"),
    router_mode="layer",
)

# Encode a 3D CT volume
volume = torch.randn(1, 1, 32, 256, 256).cuda()  # or load a real .npy
with torch.no_grad():
    tokens = enc(volume, indices=None)       # (1, 1024, 384)
    embedding = tokens.mean(dim=1)           # (1, 384)

Load the MLLM

from load_checkpoint import load_mllm

model = load_mllm(
    os.path.join(repo_dir, "mllms_qwen3/stage2_best.pt"),
    config_overrides={
        "use_hybrid": True,
        "vision_use_morton": False,
        "vision_checkpoint": os.path.join(repo_dir, "encoder/encoder.pt"),
    },
)

volume = torch.randn(1, 1, 32, 256, 256).cuda()
answer = model.generate(volume, "What organ is shown in this CT scan?")
print(answer)

Extract Frozen Embeddings for kNN

from load_checkpoint import load_encoder, encoder_extract_fn

enc, cfg, meta = load_encoder(
    os.path.join(repo_dir, "encoder/encoder.pt"), router_mode="layer"
)
extract = encoder_extract_fn(enc, batch_size=32)

npy_paths = ["volume_001.npy", "volume_002.npy", ...]
embeddings = extract(npy_paths)  # (N, 384) float32

Input Format

Volumes should be (1, 32, 256, 256) float32 tensors normalized to [0, 1]. For M3D-Cap data this is per-volume min-max normalization. See src/infer/extract_m3d.py and src/infer/extract_inspect.py in the codebase for preprocessing scripts.

Citation

@article{trinh2025radjepa3d,
  title={Rad-JEPA 3D: Radiology Joint-Embedding Predictive Model for 3D Computed Tomography},
  author={Trinh, Quoc-Huy and Nguyen, Minh-Van and Bagci, Ulas},
  journal={arXiv preprint arXiv:2607.26196},
  year={2025}
}