Large-scale modality-invariant foundation models for brain MRI analysis: Application to lesion segmentation

Technical Summary: Large-Scale Modality-Invariant Foundation Models for Brain MRI Analysis

Problem Statement
The field of medical image analysis is shifting toward large-scale foundation models pre-trained via self-supervised learning (SSL) to overcome the scarcity of clinical-grade annotated cohorts. While SSL has shown promise in natural images, its adaptation to brain MRI remains underexplored, particularly regarding multi-modal data. Brain MRI comprises multiple modalities (e.g., T1, T2, FLAIR, DWI) that offer complementary views of biological structures. The central hypothesis of this work is that subject-specific, modality-invariant representations can be learned by mapping these diverse input modalities into a shared latent space. Such models could theoretically capture discriminative structural features, mitigate the impact of missing modalities in retrospective cohorts, and reduce future acquisition demands. However, it remains unclear whether enforcing modality invariance improves downstream tasks like lesion segmentation, which often rely on fine-grained, modality-specific contrasts.

Methodology
The authors propose a large-scale pre-training framework utilizing the FOMO60k dataset, comprising 60,529 brain MRI scans from 11,187 subjects across 16 sources. The pre-processing pipeline involved skull-stripping (HD-BET), resampling to 1 mm³ isotropic resolution, and intensity normalization.

The core methodological contribution is a Modality-Invariant Contrastive Learning (MCL) setup integrated with Masked Image Modeling (MIM):

• Modality-Invariant Contrastive Learning (MCL): Based on a MoCo v2 architecture, this branch defines positive pairs as co-localized 3D patches from different modalities within the same acquisition session of the same subject. Negative pairs consist of patches from different subjects. This contrasts with default SSL approaches where positive pairs are augmented views of the same modality.
• Masked Image Modeling (MIM): A masked reconstruction objective was integrated, where patches are masked at the bottleneck and interpolated back to the patch embedding space. The target for reconstruction preceded intensity augmentations to encourage noise-invariant representations.
• Architecture: A hybrid Swin encoder with residual convolutions was used. A lightweight Feature Pyramid Network (FPN) decoder mapped features from the MIM branch to voxel space for segmentation.
• Training: Pre-training was conducted for 200k steps, followed by 5k steps of fine-tuning for downstream tasks.

Experimental Setup
The framework was evaluated on downstream lesion segmentation tasks for stroke and epilepsy across three cohorts:

1. ATLAS v2: Multi-center post-stroke cases (T1-weighted).
2. ISLES 2022: Ischemic stroke cases (DWI, ADC, FLAIR).
3. FCD BONN: Focal cortical dysplasia patients (FLAIR, T1).

Performance was measured using the Dice Similarity Coefficient (DSC) and compared against training from scratch and standard SSL baselines (default Contrastive Learning and MIM).

Key Results

• Cross-Modality Alignment: The MCL approach successfully aligned cross-modality embeddings. The mean cosine distance between embeddings of different modalities for the same anatomical region dropped significantly (MCL: 0.0066) compared to default contrastive learning (0.5995), indicating effective alignment in the latent space.
• Segmentation Performance: Despite successful alignment, the modality-invariant approach did not improve per-modality lesion segmentation performance compared to the default contrastive learning baseline.
  • Performance was primarily driven by the presence of the MIM branch, which encourages highly localized embeddings beneficial for segmentation.
  • In many cases, integrating MCL with MIM slightly deteriorated performance compared to default CL + MIM. The authors attribute this to the MCL objective potentially limiting MIM effectiveness by dominating the optimization process (higher gradient norms) and eliminating modality-specific texture features necessary for accurate reconstruction.
• Modality Specificity: Downstream performance remained strongly coupled to the specific imaging modality (e.g., FLAIR for epilepsy, DWI for stroke), regardless of the pre-training strategy. This pattern persisted even in models trained from scratch, suggesting that modality-specific contrasts are essential for segmentation and cannot be fully inferred from a modality-invariant representation.

Significance and Conclusions
The paper concludes that while large-scale pre-training can successfully learn modality-invariant representations that align cross-modality embeddings, this invariance does not translate to improved lesion segmentation performance. The authors argue that accurate lesion segmentation relies on fine-grained spatial details tied to modality-specific contrasts and pathology, which are lost when forcing a unified representation.

Consequently, the proposed modality-invariant framework may be more effective for global classification or regression tasks (e.g., brain age prediction) where patient-level anatomical features are paramount. For segmentation, the study suggests that preserving discriminative, modality-specific features is more beneficial than unifying representations, though multi-modal fusion remains a viable alternative. The authors have made all pre-trained foundation models and code publicly available to facilitate further research in this direction.