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

Summary

This paper introduces Brain-JEPA, a self-supervised learning approach that leverages joint-predictive architecture to learn representations of brain fMRI images. The authors introduce two novel components on top of the JEPA architecture to adapt it to brain images: 1) Brain Gradient Positioning, which encodes the functionality of each ROI into the patch positional encoding, and 2) an fMRI-specific masking strategy. The authors pretrain a family of ViT encoders using Brain-JEPA on the UK Biobank dataset and then explore downstream task performance using both fine-tuning and linear probing protocols to evaluate the quality of their learned representations. They look into trait prediction of a held-out set from the same UK Biobank dataset and three other datasets: HCP-Aging, ADNI, and another resting state fMRI data source. The methods achieve strong empirical results, outperforming previous approaches such as BrainLM with a significant margin across evaluation protocols (fine-tuning, linear eval).

Strengths

- The proposed approach demonstrates strong empirical results compared to previous work in the field. - The authors provide a clear ablation study that highlights the significance of the proposed contributions, specifically the brain-gradient position embedding and masking strategy, for applying JEPA to fMRI data. - The methods exhibit good scaling properties, indicating their potential for broader applicability.

Weaknesses

- It is unclear from the empirical evaluation how the pretraining data affects downstream performance. Are the baselines, such as BrainLM, using the same pretraining dataset and computational budget? - Similarly, it would be informative to explore the scaling properties of BrainJEPA with respect to dataset size. Does performance improve as the size of the pretraining dataset increases? - What is the impact of some of the contributions (brain gradient positioning, masking strategy) compared to other design choices, such as predictive in latent space? Would the former contributions also benefit a BrainLM baseline?

Questions

See weaknesses.

Rating

8

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Authors adequately addressed the limitations

Summary

In this paper, the authors train a foundation model on fMRI data. To this end, they combine multiple deep learning techniques in a novel way: - they rely on a Joint-Embedding Predictive Architecture, and devise a specific masking strategy for brain data (referred to as spatio-temporal masking) - they make use of pre-trained embeddings containing functional information (as opposed to anatomical only) of fMRI data as positional embeddings for transformers comprised in their JEPA model Overall, this study shows that the foundation model obtained yields state-of-the-art performance in a variety of downstream tasks (some of which test the obtained embedding without further tuning through linear probing).

Strengths

In my opinion, this paper tackles an important problem: large collections of neuro data have been collected in the past decade, but inter-individual variability makes it hard to derive meaningful models of the brain. Many recent endeavours have sought to show how deep-learning can help alleviate this issue and provide meaningful embeddings of brain data that can be used in downstream tasks. I believe the current work would be of interest to a growing number of computational neuroscience researchers. Moreover, I find the writing to be clear and rather easy to follow.

Weaknesses

I find the paper interesting and well written. I think it brings valuable information to the community. Methodological contributions could be deemed poor compared to other submissions, but I think the kind of benchmarks featured in this paper are challenging to implement and represent a great amount of work ; in this regard, I would encourage the authors to make their code public so that other research teams can potentially reproduce this benchmark. However, some important pieces of information are missing at this stage, notably concerning how data was processed for downstream tasks.

Questions

At this stage, it is not clear to me from line 235 how the downstream dataset were divided into training, validation and test sets. For instance, can data from a given participant appear in different categories? In my understanding, the division is the same across all models tested (Brain-JEPA, BrainLM, etc), is that indeed the case? I believe it is crucial to evaluate how models like Brain-JEPA generalise to unseen participants. The authors write that the temporal resolution of the dataset used for the pre-training phase is 0.735s (line 216). Classical repetition times in fMRI are usually higher. Do the datasets used in downstream tasks have the same TR as the one used for pre-training? If not, how did the authors adapt to this change? In particular, since $p$, the number of brain volumes concatenated to form patches is set to 16 (Table 4), patches are approximately 10 seconds long during the pre-training phase be would be about 30 seconds long with a classical TR of 2 seconds in downstream tasks, which seems pretty long. This brings the question: how was the value of $p$ chosen? Does it seem possible to the authors that different downstream tasks would have different optimal values for $p$? I am curious about the size of the brain gradients used to derive the positional embeddings. The authors indicate in Table 4 that $m$ the dimension of the brain gradients is 30, which seems pretty high to me. I would expect that only the first few dimensions are actually useful to the model (maybe the first 3, as illustrated in Figure 2, are enough to distinguish the most important networks cortical networks). Moreover, these vectors are actually mapped to a higher-dimensional space of size $d/2$ (where $d$ depends on the size of the ViT used). I am curious as to what the final functiono-anatomical embedding obtained here looks like. Can the authors try to give more intuitions about it? **Additional suggestions:** In my humble opinion, this paper mostly targets the neuroscientific community and should therefore help members of this community who are not deep-learning experts to dive into the current work. In particular, I think the authors could add short paragraphs (maybe in the supplementary materials if they cannot fit in the main text) to explain concepts behind JEPA, linear probing, with simple words. The left and middle panes of Figure 7 could be merged together so that bars for the caucasian and asian participants for each network would be side by side. It would be nice to indicate how much time the pre-training and fune-tuning phases required (lines 556-560). In my opinion, the positional embedding part of Figure 1 (top right) is rather unclear. I think the figure would benefit from highlighting clearly what the inputs of $g_{\phi}$ are (embeddings, positional embeddings) The authors write: "In the field of brain activity analysis, brain language model (brainLM) is the first and only foundation model to date" (line 32). I would personally be less assertive about this. In my understanding, many existing works have trained models on large fMRI datasets to extract meaningful representations (before these would be called "foundation models"). In the list below, the last two examples are rather close to what is being done in the current paper (self-supervised settings), and the last example even uses masking strategies close to that of BrainLM: - Mensch, Arthur, Julien Mairal, Bertrand Thirion, and Gaël Varoquaux. ‘Extracting Representations of Cognition across Neuroimaging Studies Improves Brain Decoding’. _PLoS Computational Biology_ 17, no. 5 (2021): e1008795. - Thomas, Armin, Christopher Ré, and Russell Poldrack. ‘Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data’. _Advances in Neural Information Processing Systems_ 35 (6 December 2022): 21255–69. - Chen, Zijiao, Jiaxin Qing, Tiange Xiang, Wan Lin Yue, and Juan Helen Zhou. ‘Seeing Beyond the Brain: Conditional Diffusion Model with Sparse Masked Modeling for Vision Decoding’. arXiv, 14 November 2022. [https://doi.org/10.48550/arXiv.2211.06956](https://doi.org/10.48550/arXiv.2211.06956).

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

The study introduces Brain-JEPA, an fMRI foundation model that leverages joint-embedding predictive architecture and a novel position-embedding approach based on brain gradient. This model achieves state-of-the-art performance in demographic prediction, disease diagnosis/prognosis, and trait prediction, and excels in off-the-shelf evaluations like linear probing.

Strengths

1. The idea of using fMRI gradient information to guide the position encoding is very interesting and is proven to be effective. 2. Overall the paper is nicely organized and written. 3. The model shows great generalizability.

Weaknesses

1. I think the effectiveness of JEPA, i.e., predicting the representations of target blocks rather than constructing the masked input like MAE, is not well supported by the ablation study: * BrainLM uses the AAL-424 atlas instead of the Schaefer functional atlas. Since these two studies use different atlases, it is not a very fair comparison. * I think using different atlases is totally fine if there is a very significant performance difference between the current model with anatomical position embedding (incorporating the same position settings as BrainLM) and BrainLM. However, from the ablation study, the performances with anatomical position embedding on three downstream tasks are 0.716, 78.79%, and 74.74%, respectively, while for BrainLM, they are 0.832, 74.39%, and 75.79%, with two of them higher than Brain-JEPA (with anatomical position embedding). With these results, I feel the main contribution would be the novel position embedding based on the fMRI gradient rather than the JEPA framework. Additionally, the author didn’t report other ablation results on downstream tasks like Neuroticism, Flanker, Amyloid, and NC>MCI (Asian), so there is no further evidence that JEPA consistently performs well on other downstream tasks. Nor did the author explore Brain-JEPA without the JEPA framework, i.e., predicting the original signals rather than representations. 2. Since the model performs prediction at the latent embedding level rather than at the original data level, the current model is unable to reconstruct the original missing time series. However, reconstructing the unseen missing brain time series is also an important and valuable application of the brain foundation model.

Questions

1. In lines 71 to 73, I wonder if the author made it italics on purpose or if it’s a formatting mistake. 2. In line 154, what’s the shape/size of functional connectivity $c_i$ in ROI i? 4. Why does the author choose to sample 160 frames from the original data with a temporal stride of 3, instead of using the original data frame? 3. I think BrainLM is not the only self-supervised study on fMRI; there are other works like “Thomas et al., Neurips 2022” and “SwiFT” by Kim et al., Neurips 2023, etc, that could be discussed in the introduction or even compared with the current model.

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The author has adequately discussed the limitations.

Summary

This paper introduces a foundation model for fMRI time series, using a classic vision transformer backbone. It incorporates two main original developments: (1) a positional encoding for brain regions, using a “functional gradient” analysis (aka PCA on a Jacobian matrix derived from a temporal correlation matrix), (2) a masking strategy that explicitly masks different type of spatial and/or temporal interactions. The model is trained on a large public popular resource (UK biobank), and evaluated on downstream supervised learning in the same resource (age and sex prediction). Then the model is used for transfer learning in three independent datasets (including both north american and asian participants) for various supervised downstream tasks, predicting either demographic data (age, sex), clinical diagnosis (mild cognitive impairment) and AD biomarker status (amyloid beta deposition). The performance of the proposed model Brain-JEPA is contrasted with other models from the literature with different kinds of architectures, and in particular another “fMRI foundation” model called BrainLM. Brain-JEPA outperforms other models on most tasks, and also retains good performance when a simple linear layer is used for transfer, unlike BrainLM. These results are very promising, as they show self-supervised pretraining can successfully be applied to fMRI time series, and may lead to more robust brain biomarkers for brain disorders.

Strengths

The work is very clear and well constructed. It adapts a now classic vision transformer architecture by proposing compelling solutions for the two main domain-specific ingredients: positional encoding and masking. It includes a fair number of baseline models, and tasks for evaluation. The code and the checkpoints of the models are shared. I have not done an in depth review of these, but the code is clear and is meant to replicate one of the tasks. The comparison of transfer with fine-tuning vs linear probing as well as the improvement in downstream accuracy across training epochs all improve confidence in the results and the generalizability of the learned representations by brain-JEPA. The ablation study is also great.

Weaknesses

Although the code is shared, it does not appear to cover all of the downstream tasks in the paper though. Importantly, I could not find detailed information on data sources. ADNI for example is a massive dataset, and depending on the details of data release and other criteria, the sample size and exact list of subjects for “NC vs MCI” may vary dramatically. I am fairly confident I would not be able to reproduce all the results in this paper based on the provided information, and it’s not clear to me which parts of the paper would be easy to reproduce. There are also no details in fMRI data preprocessing, and the normalization applied across subjects is not standard. Were all datasets preprocessed with similar tools? EDIT: the authors shared anonymized code and a list of subject IDs for their analysis. The level of reproducibility is thus adequate. Surprisingly, the paper lacks a good SVM/SVR or linear regression baseline. Given the regime of limited data size in fMRI, these vanilla models are really tough to beat. Based on reported accuracy, SVM would likely beat brain-JEPA prediction on sex prediction on uk biobank. https://www.nature.com/articles/s41467-020-18037-z EDIT: the authors added an SVM baseline (along with other methods). More of a stylistic weakness: some of the text uses excessively strong languages, in particular the abstract. “This pioneering model sets an unprecedented chapter in brain activity analysis with fMRI, achieving state-of-the-art performance in demographic prediction, disease diagnosis/prognosis, and trait prediction through fine-tuning.” This is not accurate: several similar papers have come out in the past two years (see below), at least the brainLM paper. EDIT: the authors have toned down claims of novelty in the abstract. The paper fails to acknowledge several closely related works. BrainLM is not the only fMRI foundational model. I have listed a few below. The brainMASS paper in particular is an impressive work very relevant to this submission, using 30 different datasets and seven downstream tasks. It would be important to position the paper compared to some of these models, and tone down the claims of novelty. EDIT: the authors have added more recent models to their evaluation and added a discussion of other models that could not be directly compared. Yang et al., 2024 BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning Ferrante et al., 2024 TOWARDS NEURAL FOUNDATION MODELS FOR VISION : ALIGNING EEG, MEG AND FMRI REPRESENTATIONS TO PERFORM DECODING , ENCODING AND MODALITY CONVERSION https://openreview.net/forum?id=nxoKCdmteM Thomas et al., 2023 Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data https://arxiv.org/abs/2206.11417 Finally, and related to that last point, there are many public dataset to benchmark downstream tasks. ABIDE I, ABIDE II and ADHD200 in particular are readily available. It is surprising to see “only” three datasets used for downstream. EDIT: the authors have added several datasets for downstream tasks.

Questions

How do the authors handle the massive difference in fMRI temporal resolution between UK biobank and ADNI? EDIT: all required details were provided in the rebuttal below. Positional encoding would be critical with variable brain region location and temporal sampling, but fMRI only has variable temporal sampling, and it is not clear from the text how this is handled by the model. The brain representation in Figure 7 is hard to decipher. What are the readers supposed to observe? EDIT: not really resolved, see my comment to the rebuttal. I believe UK biobank uses three Siemens 3T scanners, not one. EDIT: manuscript was amended to reflect this. Could the authors please double check that IRB approval is not required for secondary analysis of human neuroimaging data in their institution.EDIT: authors have required IRB approvals.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The number of downstream tasks is limited, and the authors did not use simple baseline models (such as SVM on connectomes) despite the known excellent performance of these models at the scale of fMRI datasets. EDIT: the number of baseline models and downstream tasks have been expanded. It is also unclear based on the current set of results how the brain-JEPA model handles a diversity of scanner and image acquisition characteristics, considering it was trained on a single protocol. EDIT: dowstream tasks include datasets with various protocols, demonstrating robustness.

Thank you for the rebuttal.

The rebuttal effectively addressed my primary concerns about the size of the pretraining dataset and the evaluation of each contribution in the BrainLM framework. As a result, I have updated my score to 8, as I believe this paper will be a valuable addition for both the SSL and neuroscience communities.

Thank you for rebuttal

Thank you for the detailed explanations. The additional experiments have addressed most of my concerns, and I have increased my score accordingly. However, regarding Concern 1, I actually still think that BrainLM and Brain-JEPA with AL perform comparably, given the large standard deviations observed in some tasks and the absence of statistical testing here. Therefore, the effectiveness of JEPA remains uncertain from my perspective.