Spatio-Temporal Encoding of Brain Dynamics with Surface Masked Autoencoders

TL;DR

The paper tackles robust encoding of the brain's spatio-temporal dynamics amid substantial individual variability by introducing surface Masked AutoEncoders (sMAE) and video MAE (vsMAE) that operate on icosahedral cortical grids. By pretraining Surface Vision Transformers (SiTs) with MAE-style self-supervision on large-scale datasets (UKB) and fine-tuning on neonatal (dHCP) and HCP data, the authors achieve at least a $\ge 26\%$ improvement in downstream phenotype predictions and faster convergence. vsMAE further enables reconstruction of cortical dynamics with up to $75\%$ missing data and demonstrates transfer-learning benefits to data-scarce regimes, including improved fluid intelligence predictions (correlation $\approx 0.39$) compared to scratch models. Overall, the work provides a scalable, geometry-aware pretraining paradigm that enhances both static phenotyping and dynamic brain-function modelling, with potential clinical translational impact in low-data neuroimaging settings.

Abstract

The development of robust and generalisable models for encoding the spatio-temporal dynamics of human brain activity is crucial for advancing neuroscientific discoveries. However, significant individual variation in the organisation of the human cerebral cortex makes it difficult to identify population-level trends in these signals. Recently, Surface Vision Transformers (SiTs) have emerged as a promising approach for modelling cortical signals, yet they face some limitations in low-data scenarios due to the lack of inductive biases in their architecture. To address these challenges, this paper proposes the surface Masked AutoEncoder (sMAE) and video surface Masked AutoEncoder (vsMAE) - for multivariate and spatio-temporal pre-training of cortical signals over regular icosahedral grids. These models are trained to reconstruct cortical feature maps from masked versions of the input by learning strong latent representations of cortical structure and function. Such representations translate into better modelling of individual phenotypes and enhanced performance in downstream tasks. The proposed approach was evaluated on cortical phenotype regression using data from the young adult Human Connectome Project (HCP) and developing HCP (dHCP). Results show that (v)sMAE pre-trained models improve phenotyping prediction performance on multiple tasks by $\ge 26\%$, and offer faster convergence relative to models trained from scratch. Finally, we show that pre-training vision transformers on large datasets, such as the UK Biobank (UKB), supports transfer learning to low-data regimes. Our code and pre-trained models are publicly available at https://github.com/metrics-lab/surface-masked-autoencoders .

Figures (7)

• Figure 1: [A] (v)sMAE partitioning and learning pipelines; [B] Sequence Masking and [C] Unmasking strategies.
• Figure 2: (a) sMAE sulcal depth reconstruction results on a UKB test subject ($\rho=75\%$). (b) vsMAE reconstruction ($\rho=50\%$) results with 3 7T HCP consecutive frames.
• Figure 3: (a) Sex classification 7T HCP - comparing SiT models - trained from scratch (with 3 frames) (grey), fine-tuned (with 1 frame) from sMAE (1 frame)(orange) and fine-tuned (with 3 frames) from vsMAE (3 frames). (b) dHCP transfer learning experiment from sMAE (UKB) pre-training against SiT-tiny trained from scratch
• Figure 4: Sulcal depth reconstruction from sMAE pre-training at different masking ratio ($25\%$, $50\%$, $75\%$, $90\%$). Results are shown for the same validation subject.
• Figure 5: Validation accuracy and loss curves for the sec classification task with (v)sMAE models. (a) Three different training regimes are compared: training from scratch (with 3 frames) (gray), fine-tuned (with 1 frame) from sMAE (1 frame) (orange) and fine-tuned (with 3 frames) from vsMAE (3 frames). (b) validation loss curves for the same three training schemes. Incorporating spatio-temporal information via vsMAE pre-training and fine-tuning leads to the best results.