huyquoctrinh/Rad-Jepa-3D

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.

Model	Architecture	Checkpoint
Rad-JEPA 3D encoder	8 H-Mamba blocks, width 768; 8 query heads, 2 KV heads, 3D RoPE; layerwise routing	encoder/encoder.pt
Rad-JEPA 3D MLLM	12 H-Mamba blocks, width 384; upsampling to 768, pooling to 256 tokens, projection to Qwen2.5-7B; LoRA adapters	mllms_qwen25_7b/stage2_best.pt
Rad-JEPA 3D segmentation	Frozen 768-d encoder and residual decoder with five anatomy-group heads for 117 structures	segmentation_totalseg/decoder_best.pt

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

Benchmark	Metric	Score
Closed-ended (4,995 questions)	Overall accuracy	80.1%
Open-ended	TopExact	78.5%
Open-ended	TopMatch	80.2%

Plane	Phase	Organ	Abnormality	Location
98.3	75.3	79.0	81.2	66.9

SpatialMed VQA (10,047 questions)

Overall accuracy	Direction	Distance	Extent	Volume	Comparison
62.7%	73.8	20.7	73.9	45.9	67.2

TotalSegmentator (validation, 123 cases)

Anatomy group	Macro Dice	Macro IoU
All 117 structures	0.697	0.569
Organs	0.652	0.567
Vertebrae	0.648	0.489
Cardiac	0.709	0.563
Muscles	0.890	0.805
Ribs	0.607	0.448

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}
}