BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Technical Summary of "BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability"

Problem Statement

Clinical diagnostic workflows for neurodegenerative diseases, particularly Alzheimer's Disease (AD), follow a "modality escalation pathway." Clinicians typically begin with routine structural imaging (e.g., T1-weighted MRI), selectively add sequences (e.g., FLAIR, T2) to refine differential diagnoses, and reserve molecular imaging (e.g., amyloid-PET) for uncertain cases. Consequently, patient data in real-world cohorts is characterized by heterogeneous and often incomplete modality subsets.

Current AI models generally assume fixed input modalities, requiring architectural switching or retraining when data availability changes. Furthermore, existing large-scale brain imaging pretraining models often suffer from three limitations in this context:

1. They are typically MRI-only, failing to learn explicit bridges to amyloid-PET, a critical biomarker.
2. They assume fixed modality combinations, mismatching the stage-dependent missingness of real-world data.
3. Standard Masked Autoencoder (MAE) pretraining relies on spatially uniform masking, ignoring the fact that AD exhibits selective regional vulnerability (e.g., the hippocampus) rather than uniform pathology.

Methodology

The authors propose BrainAnytime, a unified pretraining framework designed to support brain image analysis under arbitrary modality availability. The framework is pretrained on 34,899 3D brain scans from five datasets (A4, DOD-ADNI, AIBL, BraTS23, NACC) spanning multi-sequence MRI and amyloid-PET. It integrates three core components within a shared 3D masked autoencoder architecture:

1. Multi-MAE3D (3D Multi-Modal Masked Autoencoder)

• Architecture: A shared Transformer encoder processes any subset of available modalities (T1, T2, FLAIR, PET).
• Tokenization: Each modality uses a modality-specific 3D convolutional patch embedding.
• Handling Missing Data: Missing modalities are fully masked. To maintain a constant token count regardless of input availability, the visible token budget is redistributed among observed modalities using a Dirichlet distribution (Dir($\alpha$)).
• Decoding: Each modality has a lightweight decoder that reconstructs all patches. The training loss ($L_{MAE}$) is computed only on masked patches using mean squared error with per-patch normalization.

2. Reciprocal Cross-Modal Distillation (RCMD)

• Objective: To learn correspondences between structural MRI and molecular PET.
• Mechanism: An Exponential Moving Average (EMA) teacher network shares the encoder architecture. For paired MRI-PET samples, group-level representations ($z_{MRI}$ and $z_{PET}$) are computed via average pooling.
• Loss: Two symmetric MLP predictors ($g_{M\to P}$ and $g_{P\to M}$) are trained to predict cross-modal teacher features. The loss ($L_{RCMD}$) maximizes the cosine similarity between the student's predicted features and the teacher's features in the opposite modality space.

3. Pathology-Aware Curriculum Masking (PACM)

• Objective: To prioritize disease-vulnerable anatomy during reconstruction rather than masking uniformly.
• Importance Scoring: A static score is derived from neuropathological priors (Braak staging) mapped via the AAL116 atlas. A dynamic score is obtained by aggregating attention maps from an EMA anatomy teacher. These are combined to generate patch importance scores.
• Curriculum Schedule: A three-phase strategy is employed:
  1. Early stage: Uniform random masking.
  2. Middle stage: The temperature parameter ($\tau$) anneals via cosine decay, progressively sharpening the mask distribution to favor high-importance regions for masking (i.e., forcing the model to reconstruct critical regions).
  3. Late stage: Fixed temperature.
• Selection: Patches are selected using the Gumbel-top-k trick, ensuring high-importance patches are more likely to be masked and thus reconstructed.

Training and Fine-tuning:
The total loss is $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[paired]} \cdot L_{RCMD}$. For downstream tasks, only the student encoder is used with a lightweight task head. Missing modalities remain zeroed and attention-blocked during inference, ensuring the model handles arbitrary subsets without architectural changes.

Key Contributions

1. Unified Framework: A single model capable of accepting any subset of modalities (from a lone T1 scan to a full multimodal workup) without retraining or architectural switching.
2. Cross-Modal Learning: Explicitly learns structural-molecular correspondences between MRI and PET via reciprocal distillation, bridging the gap between structural and pathological evidence.
3. Anatomy-Aware Pretraining: Introduces a curriculum masking strategy that directs the model's reconstructive effort toward AD-salient neuroanatomy, addressing the limitations of uniform masking in standard MAEs.
4. Scalable Pretraining: Successfully pretrained on a large-scale, heterogeneous dataset (34,899 scans) mimicking clinical missingness patterns.

Experimental Results

The model was evaluated on four downstream tasks (CN vs. AD, CN vs. MCI, MMSE regression, Age prediction) across five clinically motivated modality settings (ranging from T1-only to T1+T2+FLAIR+PET).

• Performance: BrainAnytime outperformed modality-specific models (e.g., 3D ResNet, M3T), dual-modality fusion methods, large-scale brain foundation models (BrainIAC, SAM-Brain3D), and incomplete multimodal learning approaches (Flex-MoE, FuseMoE).
• Key Metrics:
  • CN vs. AD: Achieved an average accuracy of 89.4% (vs. 84.2% for the next best baseline, Flex-MoE), representing a relative improvement of 6.2% over the strongest missing-modality baseline.
  • CN vs. MCI: Achieved an average accuracy of 67.2% (vs. 60.6% for MoE-retriever), a relative improvement of 7.0%.
  • Regression: Demonstrated superior performance in MMSE and age prediction across all modality settings.
• Robustness: A within-subject robustness study confirmed that performance consistently improved as more modalities were added (e.g., CN vs. AD F1 increased from 65.0% with T1-only to 80.4% with full input), demonstrating the model's ability to leverage available data effectively.
• Ablation: Removing RCMD or PACM individually resulted in performance drops, confirming the necessity of both cross-modal distillation and anatomy-aware masking.

Significance

The paper claims that BrainAnytime addresses a critical gap between clinical reality (staged, incomplete data acquisition) and current AI paradigms (fixed-input models). By enabling a single model to flexibly accept whatever imaging is available, it aligns with the clinical modality escalation pathway. The framework demonstrates that pretraining with cross-modal distillation and pathology-aware masking yields a robust foundation model that significantly outperforms existing baselines, particularly in scenarios where data is incomplete or modalities are missing. This approach offers a practical solution for deploying AI in real-world clinical settings where full multimodal data is rarely available for every patient.