🩻 M3D-CLIP-CrossModal-MRI-CT

Contrastively Aligned 3D Vision Transformer for Cross-Modal MRI-to-CT Representation Learning

---

📌 Overview

M3D-CLIP-CrossModal-MRI-CT is a specialized 3D Vision Transformer encoder trained to bridge the cross-modal domain gap between Magnetic Resonance Imaging (MRI) and Computed Tomography (CT) in 3D medical volumes.

While the foundation model GoodBaiBai88/M3D-CLIP was primarily pre-trained on CT volumes and paired radiological text reports, its pre-trained visual encoder struggles when applied directly to MRI. This failure stems from fundamental physical differences:

1. Voxel Intensity Discrepancies: CT represents calibrated tissue electron density in fixed Hounsfield Units (HU), whereas MRI reflects qualitative proton relaxation times ($T_1, T_2, T_2^*$) lacking absolute physical scales.
2. Inverted Tissue Contrast: Cortical bone appears as near-maximum intensity in CT ($\ge +1000\text{ HU}$) but as signal void (black) in standard MRI. Air cavities appear identical to cortical bone on MRI, causing severe feature ambiguity.
3. Scanner Inhomogeneities: MRI is subject to RF coil bias fields, $B_0$ field distortions, and sequence variations.

To solve this, M3D-CLIP-CrossModal-MRI-CT was trained using Teacher-Student Cross-Modal Contrastive Distillation (InfoNCE):

• CT Teacher (Frozen): The pre-trained M3D-CLIP vision encoder generates reference semantic embeddings from patient CT scans.
• MRI Student (Trained): The MRI encoder adapts its 3D ViT attention blocks to project 3D MRI volumes into the exact same latent metric space as the CT teacher.

---

⚡ Key Highlights & Dual Representation

Unlike standard vision encoders that only output a pooled 1D vector, M3D-CLIP-CrossModal-MRI-CT provides two complementary levels of semantic representation:

                       Input 3D MRI (1, 32, 256, 256)
                                     │
                     ┌───────────────▼───────────────┐
                     │     3D Patch Embedding        │ (Patch: 4 × 16 × 16)
                     │   12-Layer 3D-ViT Encoder     │
                     └───────────────┬───────────────┘
                                     │
                     Output Tensor: (B, 2049, 768)
                                     │
             ┌───────────────────────┴───────────────────────┐
             ▼                                               ▼
   Index 0: [CLS] Token                            Indices 1..2048: Patch Tokens
      Shape: (B, 768)                                   Shape: (B, 2048, 768)
             │                                               │
             ▼                                               ▼
  【Global Cross-Modal Semantics】                【Dense 3D Spatial Guidance】
  • Patient-level cross-modal retrieval           • Reshaped to (B, 768, 8, 16, 16)
  • Anatomical cohort classification              • Preserves 3D spatial voxel coordinates
  • Diffusion FiLM / AdaLN modulation             • Disentangles bone cortical shells vs air
  • Cross-patient metric search                   • Dense semantic guidance for generative models

---

📊 Benchmark & Validation Results

Evaluated on held-out paired patient volumes across multiple anatomical regions (Head & Neck, Thorax, Abdomen, Pelvis):

1. Cross-Modal MRI-to-CT Patient Retrieval

Given a held-out MRI scan, does the encoder identify the patient's own CT scan among all cohort CT scans?

Model	Top-1 Accuracy ↑	Top-5 Accuracy ↑	Median Rank ↓	Same-Patient Cosine Sim ↑
Base M3D-CLIP (Zero-Shot on MRI)	0.033 (1/30)	0.200 (6/30)	14.5	0.128 ± 0.04
M3D-CLIP-CrossModal-MRI-CT (Ours)	0.353 (Val) / 0.800 (Test)	0.725 (Val) / 0.967 (Test)	1.0	0.654 ± 0.07
Relative Improvement	+17.65× chance	+3.6×	14.5 → 1.0	+410%

2. Anatomical Cohort Separation (t-SNE)

In t-SNE latent manifold projections:

• Base M3D-CLIP: MRI embeddings from different anatomies (Head & Neck, Thorax, Abdomen) collapse together with no coherent cluster boundaries.
• M3D-CLIP-CrossModal-MRI-CT: Features form cleanly separated, compact clusters matching the CT latent geometry, correctly separating pelvic/abdominal soft tissues from thoracic lungs and cranial structures.

---

🛠️ Usage Quickstart

Installation

pip install torch transformers monai safetensors SimpleITK

1. Load Model with AutoModel

import torch
import torch.nn.functional as F
from transformers import AutoModel

# Load the aligned MRI encoder
model = AutoModel.from_pretrained(
    "anuragpradhan/M3D-CLIP-CrossModal-MRI-CT",
    trust_remote_code=True
).eval()

# Move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

2. Extract Global Semantic Vector ([CLS])

# MRI input shape: (B, 1, D, H, W) = (1, 1, 32, 256, 256), min-max normalized to [0, 1]
dummy_mri = torch.randn(1, 1, 32, 256, 256).to(device)

with torch.no_grad():
    all_features = model.encode_image(dummy_mri)  # (1, 2049, 768)
    
    # Global [CLS] semantic embedding for retrieval, FiLM, or classification
    cls_embedding = all_features[:, 0]  # Shape: (1, 768)
    print("Global [CLS] embedding shape:", cls_embedding.shape)

3. Extract Dense 3D Spatial Feature Volume

with torch.no_grad():
    all_features = model.encode_image(dummy_mri)
    
    # 2048 spatial patch tokens (patch size 4x16x16 on volume 32x256x256 -> 8x16x16 grid)
    spatial_tokens = all_features[:, 1:]  # Shape: (B, 2048, 768)
    
    # Reshape back to 3D spatial grid: (B, C, D_grid, H_grid, W_grid)
    B = spatial_tokens.shape[0]
    dense_3d_volume = spatial_tokens.transpose(1, 2).view(B, 768, 8, 16, 16)
    print("Dense 3D volume shape:", dense_3d_volume.shape)
    
    # Interpolate to high-resolution voxel grid (e.g. for dense conditioning in diffusion models)
    high_res_guidance = F.interpolate(dense_3d_volume, size=(32, 64, 64), mode="trilinear", align_corners=False)
    print("High-res 3D guidance shape:", high_res_guidance.shape)

4. Cross-Modal MRI-CT Cosine Similarity

# Compute cosine similarity between an MRI embedding and a CT embedding
cosine_sim = F.cosine_similarity(mri_cls_embedding, ct_cls_embedding, dim=-1)
print(f"Cross-modal alignment score: {cosine_sim.item():.4f}")

---

🔬 Model Architecture & Training Details

Parameter	Value
Architecture	3D Vision Transformer (ViT-Base adapted for 3D volumes)
Input Shape	$(B, 1, 32, 256, 256)$
Patch Resolution	$(4, 16, 16)$ voxels per patch
Latent Patch Grid	$8 \times 16 \times 16 = 2,048$ spatial tokens
Hidden Dimension ($d$)	768
Transformer Layers	12 blocks (self-attention + MLP)
Attention Heads	12 heads
MLP Intermediate Dim	3,072
Total Parameters	198,052,609 (~198M)
Trainable Parameters	63,771,648 (~63.8M, last 9 transformer blocks + LayerNorm + Projection)
Loss Function	InfoNCE Contrastive Loss with MoCo-style Memory Bank ($K=1024$)
Temperature ($\tau$)	0.07
Optimizer	AdamW ($\text{lr} = 10^{-5}$, weight decay $10^{-4}$, Cosine Annealing)
Training Cohorts	SynthRAD2025 (Task 1: HN, TH, AB) + Gold Atlas (Male Pelvis)
Total Patient Pairs	295 training pairs, 51 validation pairs

---

🎯 Intended Applications

1. Medical Generative Diffusion Models: Conditioning 3D MRI-to-sCT synthesis pipelines (via FiLM AdaLN modulation using [CLS] or dense spatial concatenation).
2. Cross-Modal Patient Alignment & Retrieval: Locating corresponding CT scans for unlabelled or anonymized MRI cases in multi-centre databases.
3. Pre-trained Backbone for 3D Segmentation / Registration: Utilizing dense 3D ViT patch representations as rich geometric priors for multiorgan segmentation.
4. Scanner Invariance & Domain Generalization: Normalizing heterogeneous MRI protocols (T1, T2, FLAIR, bSSFP) into a canonical CT-aligned feature space.

---

⚠️ Limitations & Ethical Considerations

• Research Use Only: This checkpoint is intended solely for scientific research and evaluation in medical image computing. It is not cleared for direct clinical diagnosis or patient treatment planning.
• Language Encoder: The text encoder (bert-base-uncased) retains its original CT-report weights and was not fine-tuned on MRI radiology reports. For text queries, refer to the base M3D-CLIP model.
• Input Dimensions: Voxel inputs should be pre-interpolated to $(32, 256, 256)$ with intensity min-max percentile normalization $[0, 1]$.

---

📜 Citation

If you use this model or its pre-trained representations in your research, please cite:

@misc{pradhan2026m3dclipcrossmodal,
  title={M3D-CLIP-CrossModal-MRI-CT: Contrastively Aligned 3D Vision Transformer for Cross-Modal MRI-to-CT Representation Learning},
  author={Pradhan, Anurag},
  year={2026},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/anuragpradhan/M3D-CLIP-CrossModal-MRI-CT}}
}

Base model acknowledgment:

@article{wang2024m3d,
  title={M3D: Advancing 3D Medical Image Analysis with Multi-Modal Large Language Models},
  author={Wang, Fan and others},
  journal={arXiv preprint arXiv:2404.00578},
  year={2024}
}