BrainMAE: A Region-aware Self-supervised Learning Framework for Brain Signals

TL;DR

Brain Masked Auto-Encoder (BrainMAE) is proposed for learning representations directly from fMRI time-series data that combines a region-aware graph attention mechanism designed to capture the relationships between different brain ROIs, and a novel self-supervised masked autoencoding framework for effective model pre-training.

Abstract

The human brain is a complex, dynamic network, which is commonly studied using functional magnetic resonance imaging (fMRI) and modeled as network of Regions of interest (ROIs) for understanding various brain functions. Recent studies utilize deep learning approaches to learn the brain network representation based on functional connectivity (FC) profile, broadly falling into two main categories. The Fixed-FC approaches, utilizing the FC profile which represents the linear temporal relation within the brain network, are limited by failing to capture informative brain temporal dynamics. On the other hand, the Dynamic-FC approaches, modeling the evolving FC profile over time, often exhibit less satisfactory performance due to challenges in handling the inherent noisy nature of fMRI data. To address these challenges, we propose Brain Masked Auto-Encoder (BrainMAE) for learning representations directly from fMRI time-series data. Our approach incorporates two essential components: a region-aware graph attention mechanism designed to capture the relationships between different brain ROIs, and a novel self-supervised masked autoencoding framework for effective model pre-training. These components enable the model to capture rich temporal dynamics of brain activity while maintaining resilience to inherent noise in fMRI data. Our experiments demonstrate that BrainMAE consistently outperforms established baseline methods by significant margins in four distinct downstream tasks. Finally, leveraging the model's inherent interpretability, our analysis of model-generated representations reveals findings that resonate with ongoing research in the field of neuroscience.

Figures (16)

• Figure 1: Overview of the Proposed BrainMAE Method. (A). Overall pre-training procedures for BrainMAE. (B). Region-aware Graph Attention. (C). Architecture of proposed TSE modules.
• Figure 2: Validation on synthetic dataset. (A) Ground truth connectivity matrix. (B) Example synthetic signals and their reconstruction using pre-trained BrainMAE. (C) Similarity matrix between learned ROI embeddings. (D) t-SNE visualization of the ROI embeddings revealing the real network.
• Figure 3: Evaluation of pre-trained BrainMAE. (A) Reconstruction examples on unseen fMRI datasets. (B) Reconstruction performance across various mask ratios. (C-D) t-SNE plot of the ROI embeddings pre-trained with HCP-3T Rest Dataset, revealing Yeo-17 network (C) and brain cortical hierarchy (D). (E) First principal component of ROI embeddings aligns with the cortical hierarchy. (F-G) ROI embedding similarity matrix shows high degrees of similarity to the FC matrix (F) and is consistent across models pretrained with different datasets (G). (H-I) t-SNE plot of fMRI representation ($h_{[CLS]}$) on unseen NSD dataset to reveal subject identity (H) and differentiate between task and rest states (I).
• Figure 4: Model "pays more attention" on task blocks and the attention score correlates with brain arousal state change (measured by 1/RT).
• Figure 5: A. BrainMAE fine-tuning framework. B. vanilla-TSE modules utilizing only the self-attention in the transformer blocks.