Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking

Zijian Dong,Ruilin Li,Yilei Wu,Thuan Tinh Nguyen,Joanna Su Xian Chong,Fang Ji,Nathanael Ren Jie Tong,Christopher Li Hsian Chen,Juan Helen Zhou

2024-09-29

Abstract:We introduce Brain-JEPA, a brain dynamics foundation model with the Joint-Embedding Predictive Architecture (JEPA). This pioneering model achieves state-of-the-art performance in demographic prediction, disease diagnosis/prognosis, and trait prediction through fine-tuning. Furthermore, it excels in off-the-shelf evaluations (e.g., linear probing) and demonstrates superior generalizability across different ethnic groups, surpassing the previous large model for brain activity significantly. Brain-JEPA incorporates two innovative techniques: Brain Gradient Positioning and Spatiotemporal Masking. Brain Gradient Positioning introduces a functional coordinate system for brain functional parcellation, enhancing the positional encoding of different Regions of Interest (ROIs). Spatiotemporal Masking, tailored to the unique characteristics of fMRI data, addresses the challenge of heterogeneous time-series patches. These methodologies enhance model performance and advance our understanding of the neural circuits underlying cognition. Overall, Brain-JEPA is paving the way to address pivotal questions of building brain functional coordinate system and masking brain activity at the AI-neuroscience interface, and setting a potentially new paradigm in brain activity analysis through downstream adaptation.

Neurons and Cognition,Artificial Intelligence,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### Problems the Paper Attempts to Solve

The paper primarily aims to address the following issues:

1. **Improving Model Generalization**: Existing fMRI analysis models for specific tasks (such as BrainNetCNN, BrainGNN, BNT, etc.) have limited generalization capabilities across different tasks and cannot fully utilize large amounts of unlabeled fMRI data.

2. **Enhancing Pre-training Strategies**: Current foundational fMRI models (such as BrainLM) use self-supervised pre-training methods, but directly reconstructing masked time series segments may not be suitable for noisy fMRI data, making it difficult for the model to distinguish between signal and noise.

3. **Enhancing Downstream Task Performance**: Existing models perform poorly in downstream tasks (such as linear probing) and are limited to specific populations (such as Caucasians), restricting their clinical application scope.

To address these issues, the paper proposes the Brain-JEPA model, which enhances model performance through the following two innovative techniques:
- **Brain Gradient Positioning**: Introduces a functional coordinate system based on functional connectivity gradients to enhance the positional encoding of different regions of interest (ROIs).
- **Spatiotemporal Masking**: Designs a specialized masking strategy tailored to the unique characteristics of fMRI data to handle heterogeneous time series segments.

Through these techniques, Brain-JEPA demonstrates excellent performance in various downstream tasks, including demographic prediction, disease diagnosis/prognosis, and trait prediction, and shows outstanding generalization capabilities across different ethnic groups.