Technical Summary: MultiMAE for Brain MRIs

Problem Statement

In medical imaging, particularly with brain MRI, deep learning models often face the challenge of missing input sequences. Clinical environments frequently lack complete multi-modal data (e.g., missing T1, T1c, T2, or FLAIR sequences) due to acquisition issues or patient constraints. Standard approaches typically employ an early-fusion strategy, stacking all available modalities along the channel dimension (analogous to RGB channels). While simple, this approach requires all modalities to be present during both training and inference. Consequently, models trained with early fusion suffer sharp performance drops when inputs are missing and fail to learn distinct, modality-specific representations that could be leveraged for cross-modal reasoning.

Methodology

The authors propose a Multi-Modal Masked Autoencoder (MultiMAE) framework adapted for 3D brain MRI. Inspired by the MultiMAE architecture for natural images, this method shifts from early fusion to a modality-specific encoding strategy combined with late-fusion via a transformer encoder.

Architecture and Pretraining

• Backbone: The model utilizes a modified Vision Transformer (ViT-B/16) as the encoder.
• Modality-Specific Adapters: Instead of stacking channels, each MRI sequence (t1, t1c, t2, fla) is processed through its own linear adapter to project non-overlapping $16 \times 16 \times 16$ patches into a shared token space.
• Masking Strategy: A global masking ratio of 75% is applied. Crucially, masking is sampled per-modality from a Dirichlet distribution ($\alpha=1$), allowing up to 100% of a specific modality to be masked while maintaining the global rate. This forces the model to learn cross-modal dependencies.
• Encoder-Decoder Design:
  • Encoder: A standard ViT-B/16 (12 layers) processes the unmasked tokens from all modalities along with a [cls] token.
  • Decoders: Four lightweight, modality-specific decoders (3 layers each) are employed. They utilize a cross-attention mechanism where queries (learnable [mask] tokens for missing patches) attend to the shared context provided by the encoder.
• Objective: The model is pretrained to reconstruct masked patches using Mean Squared Error (MSE) loss, encouraging the learning of rich representations that can infer missing sequences from available ones.

Downstream Adaptation

The pretrained encoder is adapted to two downstream tasks:

1. Segmentation: Using the UNETR architecture. To handle the variable sequence length caused by separate tokenization per modality, tokens at each spatial location are averaged to maintain a consistent sequence length.
2. Classification: A single linear layer replaces the decoders. Patch tokens and the [cls] token are averaged to form a feature vector for glioma subtype classification (GBM, Astro, Oligo).

Key Contributions

1. MultiMAE Pretraining for 3D MRI: The work investigates the application of MultiMAE to 3D brain MRIs, enabling modality-specific encoding and masked modeling across multiple sequences.
2. Robustness to Missing Modalities: The authors demonstrate that this pretraining strategy significantly improves performance and robustness in both segmentation (glioma and metastasis) and classification tasks when input modalities are missing, outperforming standard MAE-ViT baselines.
3. Synthesis of Missing Modalities: The framework demonstrates the ability to synthesize entirely missing MRI sequences at inference time by leveraging cross-modal relationships learned during pretraining.

Experimental Results

The method was evaluated on a large combined dataset (BraTS 2021/2023, UPENN-GBM, UCSF-PGDM, etc.) and tested on internal and external datasets.

• Missing Modality Generation: The model successfully reconstructed missing modalities with high fidelity (PSNR ~20–21, SSIM ~0.71–0.76), significantly outperforming a standard MAE-ViT baseline (PSNR ~7–10, SSIM ~0.52–0.61).
• Segmentation Performance:
  • The pretrained MultiMAE consistently outperformed the MAE-ViT baseline on both internal and external test sets for glioma and metastasis segmentation.
  • In missing-modality scenarios, the MultiMAE showed superior resilience. For instance, when the T1c modality was missing, the MultiMAE maintained a Dice score of ~51.55 (frozen encoder) for the tumor core, whereas the ViT baseline dropped to ~47.91.
  • The model achieved an absolute improvement of 10.1 in overall Dice score and 0.46 in MCC over baselines when inputs were missing.
• Classification Performance:
  • The model achieved strong multi-class accuracy and MCC (up to 0.70) for glioma subtyping.
  • Unlike the segmentation task, freezing the encoder was less effective for classification; full fine-tuning yielded the best results.
  • Even with missing modalities (e.g., no T1c), the MultiMAE maintained high accuracy (0.89) compared to the ViT baseline (0.41), demonstrating effective cross-modal reasoning.

Significance and Claims

The paper claims that the MultiMAE paradigm offers a flexible and generalizable encoder for brain MRIs that addresses the critical clinical challenge of incomplete data. By treating each sequence as a separate modality and enforcing cross-modal reconstruction, the model learns to infer missing sequences from available inputs.

The authors emphasize that while the approach introduces minor changes to the ViT backbone, it yields substantial gains in robustness. The pretraining strategy allows the model to handle "unseen input configurations" without retraining, making it suitable for resource-constrained clinical environments where data completeness cannot be guaranteed. However, the authors modestly note that:

• Reconstructions exhibit some blurring due to the limitations of pixel-wise MSE loss.
• A domain gap persists between internal and external datasets, reflecting the broader challenge of cross-site adaptation.
• The added complexity of cross-modal integration requires significant pretraining compute and data to converge effectively, as randomly initialized versions of the model performed worse in some cases.

Ultimately, the work presents a reformulation of the MAE pretraining paradigm that enables more elegant handling of multi-modal brain MRI data, prioritizing robustness in the face of missing inputs.