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

Pretraining scheme of baselines (dataset and computational budget).

Thank you for your inquiry regarding the pretraining process. We confirm that for self-supervised pretraining baselines like BrainLM, we used the same pretraining dataset and computational budget, specifically the UKB dataset on 4 A100 GPUs (40GB each).

Scaling of dataset size.

Thank you for your suggestion on exploring the impact of dataset size scaling.

We have compared the performance of Brain-JEPA trained with varying portions of the UKB pretraining dataset: 25%, 50%, 75%, and 100%. As shown in the table below, the performance improves as the dataset size increases, highlighting the scalability of Brain-JEPA in relation to the pretraining dataset size.

Combine contributions with BrainLM.

Thank you for your insightful suggestion.

We further compared the performance of BrainLM combined with our contributions to vanilla BrainLM. As shown in the global author rebuttal Table 4, BrainLM combined with our contributions outperforms vanilla BrainLM consistently, demonstrating that our contributions (gradient positioning and spatiotempotal masking) could benefit the training of BrainLM as well.

Full open-source code, preprocessing.

We appreciate your important suggestions. We have supplemented the following materials/information:

Code and Data Source: We have now supplemented the codebase to include all downstream tasks on public datasets mentioned in the paper. The complete codebase, along with the list of subject IDs for all datasets used, is available through an anonymous link. In accordance with the author guidelines, we have provided this link to the Area Chair in a separate comment.

Data Preprocessing: We utilized open-source preprocessing pipelines as outlined in [1] and [2]. These pipelines are fully open-sourced and have been referenced in our paper. For data normalization, we implemented robust scaling as practiced in BrainLM.

By sharing the complete codebase, the list of subject IDs, and using open-source preprocessing pipelines, we believe that our results are highly reproducible.

Data splitting of downstream datasets.

Thank you for your inquiry regarding the splitting of our downstream datasets. We confirm that the data is divided at the participant level, ensuring that each participant appears only once in either the training, validation, or test set. It ensures that all participants in the test set are unseen during training, thereby demonstrating the excellent generalizability of Brain-JEPA. Additionally, the same division is applied consistently across all models tested.

Different TR for different datasets.

Thank you for the insightful question regarding the temporal resolution of datasets used. For fMRI in UKB, we used multi-band data with a high temporal resolution (TR is ~0.7s). In our downstream datasets, HCP-Aging also uses multi-band data with a TR of ~0.7s, while both ADNI and the Asian cohort use single-band data with a TR of ~2s.

To address the differences in TR, we downsampled the multi-band data with a temporal stride of 3, aligning the TR to ~2s. This ensures consistency across different datasets.

For future work, we plan to use learnable TR embeddings to enable the model to adapt to different TRs dynamically.

Dimension of brain gradient embedding.

Thank you for your insightful feedback on the dimensionality of the brain gradient. We compared the model performance between 3-dimensional (3-dim) and 30-dim brain gradient positioning, shown in the global author rebuttal Table 5. The 30-dim model consistently outperformed the 3-dim model by a large margin. This indicates that higher-dimensional brain gradients may encapsulate finer-grained information on brain network organization, which benefits the learning of brain dynamics.

We leave the investigation of the mapped gradient embedding in the future work. Specifically, we will conduct comprehensive experiments to explore how these gradients are organized in the latent space and their relationships with brain network organization.

Explain to neuroscience community.

We appreciate your valuable suggestion. We will include short explanatory paragraphs on concepts like JEPA and linear probing in the supplementary materials to make the paper more accessible to the neuroscientific community.

Fig7.

Thank you for your insightful suggestion regarding Figure 7. However, the absolute attention values are not meaningful in isolation, so it may not be appropriate to directly compare attention values across different cohorts. Our focus is on whether the rankings among different networks are aligned between the two cohorts (Caucasian versus Asian). From the current Figure, we can clearly see the top 4 networks are the same across the two cohorts. Please let us know if you have any further concerns regarding this.

Training time.

Thank you for your inquiry regarding the training time.

For pre-training (UKB) with the ViT-base model on 4 A100 GPUs (40GB each), using a batch size of 16x4x8, the training time per batch is ~0.3s. The total training time is ~16.5 hours.

For fine-tuning, taking HCP-Aging as an example, with the ViT-base model on 1 A100 GPU (40GB), using a batch size of 16. The total training time is ~14 mins.

Fig1.

Thank you for the valuable feedback on Figure 1. The positional embedding part is intended to be a 2D schematic view of brain gradient positioning along with temporal positioning. We will revise the figure to explicitly highlight the inputs to the predictor $g_{\phi}$.

Related work.

Thank you for highlighting other related works. While it is true that the mentioned studies have trained models on large fMRI datasets, their downstream applications are limited. [3] and [4] focus on mental state classification, and [5] on brain decoding. A brain foundation model should exhibit broad applicability across diverse brain-related tasks, such as demographic prediction, trait prediction, and disease diagnosis/prognosis. We have added additional comparisons with CSM [4], together with most recent work suggested by other reviewers. Shown in the global author rebuttal Table 1 and 2, Brain-JEPA outperforms them on most tasks. We highlight that Brain-JEPA demonstrates the most diverse range of downstream applications, showcasing its powerful generalizability across different cohorts and tasks. We will include the mentioned works in our revised version to better illustrate these distinctions.

References:

[1] Spatial topography of individual-specific cortical networks predicts human cognition, personality, and emotion. Cerebral Cortex 2019.

[2] Global signal regression strengthens association between resting-state functional connectivity and behavior. Neuroimage 2019.

[3] Extracting Representations of Cognition across Neuroimaging Studies Improves Brain Decoding. PLoS Computational Biology

[4] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.

[5] Seeing Beyond the Brain: Conditional Diffusion Model with Sparse Masked Modeling for Vision Decoding. CVPR 2023.

Ablation study on framework.

Thank you for your feedback on our ablation study.

To thoroughly compare the performance between JEPA with anatomical locations (AL) and BrainLM (MAE), we have extended our comparison to include all the tasks except for the three in the current version, as well as two newly added datasets, OASIS-3 and CamCAN. The results shown in the table below, demonstrates that JEPA w AL outperforms BrainLM in seven out of eleven tasks, demonstrating the superiority of prediction in latent space. For the tasks where BrainLM performs better, it is likely that JEPA requires gradient positioning for precise ROI placement to achieve optimal performance. In future work, we will further investigate the possible interactions between the self-supervised learning framework and brain gradient positioning.

We further compared Brain-JEPA without JEPA architecture (i.e., BrainLM w contributions) to Brain-JEPA. As shown in the global author rebuttal Table 4, Brain-JEPA (JEPA framework) outperforms BrainLM (MAE framework) w contributions consistently, demonstrating the superiority of JEPA framework.

Reconstruction of time series.

Thank you for pointing out the underlying difference between the MAE and JEPA frameworks. We would like to emphasize that the primary goal of reconstructing masked fMRI time series is to facilitate self-supervised training rather than for a specific clinical application. Our primary focus is on learning representations that directly enhance performance across diverse downstream tasks. By achieving better latent performance for the samples, we can utilize powerful decoders such as diffusion models to reconstruct input signals for various applications. We leave this as the future work.

Italic lines 71-73.

Thank you for your comment. Regarding lines 71 to 73, we intentionally italicized this part to emphasize the importance of these questions. The italicization was meant to highlight the significance of developing a functional coordinate system and a masking strategy for large-scale pretraining with fMRI data, which are crucial interdisciplinary challenges at the intersection of AI and neuroscience.

Line 154, the value of $c$.

Thank you for your attention to detail. The functional connectivity is represented by an adjacency matrix. Given $N$ ROIs, the adjacency matrix is an $N×N$ symmetric matrix. Each row/column $i$ of this matrix represents the connectivity of ROI $i$ with other regions. In our paper, the brain is parcellated into $N=450$ regions, therefore, the dimension of connectivity for each ROI is 450.

Downsampling.

We thank the reviewer for the insightful question on temporal downsampling. We performed temporal downsampling for some of our datasets because of the variable temporal resolutions among different datasets. For the fMRI data in the UKB (our pretraining dataset), we used multi-band data with a high temporal resolution, where the TR is ~0.7s. In our downstream datasets, HCP-Aging also uses multi-band data with a TR of ~0.7s, while both ADNI and the Asian cohort use single-band data with a TR of ~2s.

To address the differences in temporal resolution, we downsampled the multi-band data with a temporal stride of 3, effectively aligning the TR to approximately 2s. This ensures consistency across datasets.

For future work, we plan to incorporate learnable TR embeddings, enabling the model to handle different temporal resolutions in a learning-driven manner. This approach will enhance the model's flexibility across varied datasets.

Other self-supervised study on fMRI.

Thank you for your suggestion on discussing other self-supervised studies on fMRI.

We have incorporated CSM [1], SwiFT [2], and BrainMass[3] (one concurrent work suggested by another reviewer) into our comparisons on the HCP-Aging and ADNI datasets, as well as two newly added datasets, OASIS-3 and CamCAN. As shown in the global author rebuttal Table 1-3, our proposed Brain-JEPA consistently outperforms these models on most tasks.

CSM is domain-specific, trained exclusively for mental state decoding. SwiFT specializes in demographic prediction, and BrainMass is limited to disease diagnosis and prognosis. Brain-JEPA, however, has demonstrated versatility across a broader range of downstream applications, including demographic prediction, trait prediction, cognitive score prediction, and disease prognosis/diagnosis, showcasing its extensive potential.

In our revised version, we will include comparisons with additional models to further demonstrate the outstanding performance of Brain-JEPA. We will also incorporate a discussion of the above-mentioned works into the related work section to better position our contributions.

References:

[1] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.

[2] SwiFT: Swin 4D fMRI Transformer. NeurIPS 2023.

[3] BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning. TMI 2024.

Code, data source, and preprocessing.

Thank you for your valuable feedback.

We have supplemented code of all downstream tasks on public datasets to the codebase. The code and subject IDs are available via an anonymous link provided to AC in a separate comment, per author guidelines.

We utilized open-source preprocessing pipelines as outlined in [1] and [2]. We implemented robust scaling for normalization as practiced in BrainLM.

By sharing the complete codebase, subject IDs, and using open-source preprocessing pipelines, we believe that our results are highly reproducible.

Other related work.

Thank you for your important suggestions.

We have incorporated SVM/SVR, BrainMass [3], CSM [4], and SwiFT [5] (as suggested by another reviewer) into comparisons. As shown in the global author rebuttal Table 1 and 2, Brain-JEPA outperforms these models on most tasks.

Although SVM/SVR may perform well in demographic prediction (e.g., sex), previous research has indicated that SVMs perform worse than fMRI deep learning models [6] [7]. Additionally, note that BrainMass is a concurrent work to ours (published after our initial submission). We have demonstrated a broader range of downstream applications, including demographic prediction, trait prediction, cognition score prediction, and disease prognosis/diagnosis. In contrast, experiments in BrainMass are limited to disease diagnosis/prognosis only. Moreover, the work mentioned [8] lacks available code and is applied to a multi-modal setting, making it difficult to reproduce and not suitable for our context.

In revised version, we will include comparisons with additional models to further demonstrate the outstanding performance of Brain-JEPA. We will also incorporate a discussion of the above works to better position our contributions.

More downstream datasets.

Thank you for the valuable suggestion.

Note that datasets such as ABIDE I, ABIDE II, and ADHD200 are all children cohorts. Given that our model was pretrained on UKB, which primarily includes middle-aged to elderly participants, it may not generalize well to younger populations yet (same as BrainLM).

To further demonstrate the diversity of our downstream applications, we have conducted additional experiments using two aging-related public datasets: OASIS-3 and CamCAN, for AD conversion prediction and depression classification, respectively (global author rebuttal Table 3). By applying Brain-JEPA to five downstream datasets across eight distinct tasks totally, we have demonstrated its versatility in a wider range of applications compared to the existing models. Specifically, Brain-JEPA excels in demographic prediction, trait prediction, and disease diagnosis and prognosis. This stands in contrast to experiments done in BrainLM, which is limited to demographic and clinical score prediction, and BrainMass, which focuses solely on disease diagnosis and prognosis.

Temporal resolution.

We thank the reviewer for the insightful question. For fMRI in UKB, we used multi-band data with a high temporal resolution (TR is ~0.7s). In our downstream datasets, HCP-Aging also uses multi-band data with a TR of ~0.7s, while both ADNI and the Asian cohort use single-band data with a TR of ~2s.

To address the differences in TR, we downsampled the multi-band data with a temporal stride of 3, aligning the TR to ~2s. This ensures consistency across different datasets.

For future work, we plan to use learnable TR embeddings to enable the model to adapt to different TRs dynamically.

Clarification on Fig 7.

Brain networks are systems of interconnected regions that work together to perform specific functions, such as the Default Mode Network (DMN) for self-referential and memory tasks and the Control Network (CN) for executive functions [9] [10]. Figure 7 shows attention values among networks. Higher attention values in DMN, CN, and SAN indicate their significant involvement in MCI, consistent across different ethnic groups, showcasing Brain-JEPA's robustness and generalizability.

Claim in abstract.

Thank you for your valuable feedback. While we recognize the contributions of recent works, our aim was to highlight specific advancements introduced by Brain-JEPA. These include the integration of Brain Gradient Positioning and Spatiotemporal Masking. Furthermore, the downstream applications of Brain-JEPA are exceptionally diverse. We will remove the “unprecedented chapter” in the abstract for our revision to better place our work in the context of recent developments in the field.

UKB has three scanners.

In revised version, we will revise "one scanner" to "three scanners" in the UKB introduction of our revised version.

IRB approval.

We have obtained all necessary IRB approvals for the datasets used for analysis.

References:

[1] Spatial topography of individual-specific cortical networks predicts human cognition, personality, and emotion. Cerebral Cortex 2019.

[2] Global signal regression strengthens association between resting-state functional connectivity and behavior. Neuroimage 2019.

[3] BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning. TMI 2024.

[4] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.

[5] SwiFT: Swin 4D fMRI Transformer. NeurIPS 2023.

[6] Interpretable Graph Neural Networks for Connectome-Based Brain Disorder Analysis. MICCAI 2022.

[7] Beyond the Snapshot: Brain Tokenized Graph Transformer for Longitudinal Brain Functional Connectome Embedding. MICCAI 2023.

[8] TOWARDS NEURAL FOUNDATION MODELS FOR VISION : ALIGNING EEG, MEG AND FMRI REPRESENTATIONS TO PERFORM DECODING , ENCODING AND MODALITY CONVERSION. ICLR 2024 Workshop Re-Align.

[9] The organization of the human cerebral cortex estimated by intrinsic functional connectivity. J Neurophysiol 2011.

[10] Correspondence of the brain's functional architecture during activation and rest. PNAS 2009.