Masked and Predictive Self-Supervised Foundation Models for 3D Brain MRI

🕒 Published (v1): 2026-06-11 13:09 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper systematically compares two self-supervised pretraining paradigms—Masked Autoencoders (MAE) and Joint Embedding Predictive Architectures (JEPA)—for 3D structural brain MRI disease detection. The authors introduce a spectral-domain auxiliary loss for MAE and variance–covariance regularization (VCR) for JEPA, then evaluate both paradigms across five neurological disease classification tasks. MAE with spectral supervision consistently outperforms JEPA for MRI-based disease detection, and the benefit of each auxiliary objective is conditioned on the frequency structure of the downstream task's pathological signal.

Problem

Existing MRI self-supervised learning (SSL) literature overwhelmingly targets segmentation and dense prediction; systematic investigation of transferable representations for MRI-based disease detection is limited. No controlled side-by-side comparison of MAE vs. JEPA existed for 3D structural MRI. Additionally, it was unclear how auxiliary SSL objectives (spectral, regularization) interact with the pathological signal structure of different downstream tasks, and most prior work relied on spatial/intensity normalization or multi-contrast concatenation that limits clinical applicability.

Method

Pretraining data: 58,781 3D MRI volumes (~12,000 subjects) from seven datasets (ADNI, NACC/SCAN, PPMI, OASIS-3, IXI, MOOD, BraTS-2024) spanning T1, T2, FLAIR, T2* contrasts; single-contrast inputs without modality concatenation, spatial normalization, or intensity harmonization. Native in-plane orientation is preserved for anisotropic scans; isotropic volumes receive random axial/coronal/sagittal orientation sampling. Volumes are cropped to 160×160 in-plane and depth-interpolated to 64 slices.

MAE variant (SL-MAE): Standard ViT-based MAE (75% masking, spatiotemporal block masking from V-JEPA) augmented with a spectral auxiliary loss. Both reconstructed and ground-truth volumes are 3D-FFT-transformed; high-pass-filtered log-magnitude spectra are compared via MSE. The spectral loss is computed every n iterations (amortized) with a cached gradient buffer. Total loss: L_Total = L_MAE + λ_FFT · L_spec.

JEPA variant (VCR-JEPA): Student-teacher architecture where the student encoder predicts the EMA-teacher's latent representations of masked target tubelets. Loss is L1 prediction with stop-gradient on the teacher. Variance–covariance regularization (VCR) is added to the concatenated context+target encoder representations to penalize feature collapse (variance term) and feature redundancy (off-diagonal covariance term). Covariance computed batch-wide across all token positions for efficiency.

Downstream evaluation: Pretrained encoder + lightweight Attentive Classifier; evaluated under frozen and fine-tuned settings on five tasks: AD vs. NC (ADNI, NACC/SCAN), MCI vs. NC (ADNI, NACC/SCAN), tumor grade prediction (UCSF-PDGM), and autism classification (ABIDE). Metrics: AUROC, accuracy, F1. Strict subject-level splits released publicly.

Key Contributions

• First controlled side-by-side comparison of MAE and JEPA as SSL pretraining paradigms for 3D structural MRI disease detection.
• Novel spectral-domain auxiliary loss for MAE using 3D FFT + high-pass filtering of log-magnitude spectra, with amortized computation for efficiency.
• Integration of variance–covariance regularization (VCR) into a 3D MRI JEPA framework to encourage decorrelated latent representations.
• Empirical finding that auxiliary objective benefit is task-conditioned: spectral loss helps most for tasks with high-frequency anatomical discriminative signal; VCR helps most when discriminative information spans multiple decorrelated feature dimensions.
• Contrast-agnostic pretraining framework with native orientation sampling, no spatial/intensity normalization, compatible with incomplete contrast availability.
• Public release of all subject-level pretraining and downstream train/val/test splits for reproducibility (github.com/ituvisionlab/mjepa).

Results

• SL-MAE (MAE + spectral loss) consistently achieves superior downstream AUROC over baseline MAE, VCR-JEPA, and baseline JEPA across the five disease detection tasks.
• JEPA underperforms MAE-based methods on MRI disease detection despite strong performance in natural image modeling.
• Spectral loss yields the largest gains on tasks whose discriminative signal is characterized by high-frequency anatomical structure (e.g., fine cortical/lesion features).
• VCR is most beneficial for JEPA on tasks where discriminative signal spans multiple decorrelated dimensions; it provides marginal or no benefit when the signal is dominated by a single high-frequency component.
• Frozen encoder evaluation confirms that SL-MAE representations transfer better without fine-tuning, indicating stronger semantic richness of the pretrained features.
• (Exact AUROC numbers for individual tasks were not present in the provided excerpt; only qualitative performance rankings are reported above as they appear in the text.)

Limitations

• Pretraining dataset (~58k volumes, ~12k subjects) is smaller than CT foundation model counterparts (e.g., 361k head CT scans in comparable work).
• Evaluation is restricted to neurological/brain classification tasks; generalizability to other organs or modalities is untested.
• The spectral loss requires amortized computation to remain tractable, adding implementation complexity and potential gradient lag.
• Contrast-agnostic single-channel design may not exploit complementary information available across contrasts for subjects with complete multi-contrast acquisitions.
• No direct comparison to contrastive SSL methods (SimCLR, DINO) under the same 3D MRI disease-detection setup.
• Exact pretraining compute costs and model parameter counts are not detailed in the provided text.

Relevance to Foundation Models in Medicine

This paper directly advances the understanding of how SSL pretraining objective design translates to clinical utility in medical foundation models, establishing that the choice between pixel-reconstruction (MAE) and latent-prediction (JEPA) paradigms is not task-agnostic but governed by the frequency structure of pathological signals. For researchers tracking medical foundation models, the key insight is that domain-motivated auxiliary losses (spectral, decorrelation) can substitute for scale when curated labeled data is scarce, and that the frozen-vs-fine-tuned transfer gap is objective-dependent rather than a fixed property of SSL. The contrast-agnostic, normalization-free pretraining framework also provides a practically deployable template for clinical heterogeneity challenges endemic to MRI. This work fills the gap noted by Ulrich et al. (2026) that reconstruction-based classification benchmarking in 3D MRI had been systematically underpowered.