huyquoctrinh/Rad-Jepa-3D

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

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

File	Description	Size	Params
encoder/encoder.pt	H-Mamba encoder (hybrid Mamba + GQA, 12 layers, 384-d)	149M	18.3M
mllms_qwen3/stage2_best.pt	Full MLLM (vision encoder + projector + Qwen3-4B LoRA adapters)	615M	4.0B

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

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

Load the Encoder

from load_checkpoint import load_encoder

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

# Encode a 3D CT volume
import torch
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(
    "path/to/mllms_qwen3/stage2_best.pt",
    config_overrides={
        "use_hybrid": True,
        "vision_use_morton": False,
        "vision_checkpoint": "path/to/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("path/to/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}
}