Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking
Authors: Zijian Dong, Ruilin Li, Yilei Wu, Thuan Tinh Nguyen, Joanna Chong, Fang Ji, Nathanael Tong, Christopher Chen, Juan Helen Zhou
NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | 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 f MRI 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. Code is available at: https://github.com/Eric-LRL/Brain-JEPA. |
| Researcher Affiliation | Academia | National University of Singapore |
| Pseudocode | No | The paper describes the methodology in text and with diagrams, but does not include any explicit pseudocode blocks or algorithm listings. |
| Open Source Code | Yes | Code is available at: https://github.com/Eric-LRL/Brain-JEPA. |
| Open Datasets | Yes | We leveraged the large-scale public dataset UK Biobank (UKB) [44, 45] for the self-supervised pretraining of Brain-JEPA... We used three datasets for external evaluation: HCP-Aging, as a segment of the public Human Connectome Project (HCP) [46]... The Alzheimer s Disease Neuroimaging Initiative (ADNI) [47]... Memory, Ageing and Cognition Centre (MACC)... All f MRI data was parcellated into n = 450 ROIs, using Schaefer-400 [48] for cortical regions and Tian-Scale III [49] for subcortical regions. |
| Dataset Splits | Yes | During the fine-tuning and linear probing stage, all the downstream datasets were divided into a 6:2:2 ratio for training, validation, and testing. |
| Hardware Specification | Yes | The pre-training process utilized four A100 GPUs, each with 40GB of memory. |
| Software Dependencies | No | The paper mentions software components like 'Adam W' and 'Flash Attention' and cites their respective papers, but it does not specify exact version numbers for these libraries or other critical software dependencies like Python or PyTorch versions. |
| Experiment Setup | Yes | The default settings are detailed in Table 4. We initialized all transformer blocks using the Xavier uniform method, as described in [19]... The default settings for end-to-end fine-tuning and linear probing are detailed in Table 5... The range ratios for obtaining the observation block and three target blocks introduced in Section 3.2 are presented in Tables 6. |