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

Paper | Code

Architecture

Rad-JEPA 3D learns volumetric CT representations through masked latent prediction. Its H-Mamba encoder combines a Mamba state-space branch for inter-slice continuity and a grouped-query attention (GQA) branch for spatial context, fused by a learned router.

Important: The MLLM checkpoint already contains its fine-tuned 384-d vision encoder. Do not substitute the separate 768-d encoder/encoder.pt checkpoint for it.

Model usage

Use the code from the RadJepa repository. The Python examples assume cd RadJepa/src; replace <repo_dir> with the directory containing the downloaded weights. The encoder requires CUDA for the Mamba selective-scan kernel.

1. Encoder feature extraction

import torch
from load_checkpoint import load_encoder

enc, cfg, meta = load_encoder("<repo_dir>/encoder/encoder.pt", router_mode="layer")
enc = enc.cuda().eval()
volume = torch.rand(1, 1, 32, 256, 256, device="cuda")  # (B, C, D, H, W), [0, 1]

with torch.no_grad():
    tokens = enc(volume, indices=None)  # (1, 1024, 768)
    embedding = tokens.mean(dim=1)      # (1, 768)

2. Medical VQA with the MLLM

The Qwen2.5-7B base model is downloaded separately when loading the checkpoint.

import torch
from load_checkpoint import load_mllm

ckpt_path = "<repo_dir>/mllms_qwen25_7b/stage2_best.pt"
overrides = torch.load(ckpt_path, map_location="cpu", weights_only=False)["config"]
model = load_mllm(ckpt_path, config_overrides=overrides, dtype=torch.bfloat16)

volume = torch.rand(1, 1, 32, 256, 256, device="cuda")
answer = model.generate(volume, "Which imaging plane is shown in this CT scan?")
print(answer)

For benchmark evaluation:

python mllms/benchmark.py --checkpoint <repo_dir>/mllms_qwen25_7b/stage2_best.pt --mode closed \
    --vision-checkpoint "" --llm-name Qwen/Qwen2.5-7B --llm-dim 3584 --use-hybrid \
    --vision-dim 384 --vision-n-layers 12 --vision-upsample-dim 768 --projector-hidden 2048 \
    --vision-pool-tokens 256 --vision-router-mode layer

3. TotalSegmentator evaluation and inference

python test/eval_totalseg.py \
    --checkpoint <repo_dir>/segmentation_totalseg/decoder_best.pt \
    --vision-checkpoint <repo_dir>/encoder/encoder.pt \
    --data-dir /path/to/Totalsegmentator_dataset_v201 --split val

To export the masks for one scan:

python test/eval_totalseg.py \
    --checkpoint <repo_dir>/segmentation_totalseg/decoder_best.pt \
    --vision-checkpoint <repo_dir>/encoder/encoder.pt \
    --data-dir /path/to/Totalsegmentator_dataset_v201 \
    --case-id s0000 --output-dir outputs/totalseg/s0000

Input: Encoder and MLLM accept normalized CT volumes of shape (B, 1, 32, 256, 256) in [0, 1]. Segmentation expects CT scans in HU and applies the [-1024, 3071] window and sliding-window inference in its evaluation script.

Results

M3D-VQA

SpatialMed VQA (10,047 questions)

TotalSegmentator (validation, 123 cases)

Per-class metrics are available in segmentation_totalseg/val_metrics.json.

Citation

@article{trinh2026rad,
  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={2026}
}