Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking
Authors: Zijian Dong
Li Ruilin
Yilei Wu
Thuan Tinh Nguyen
Joanna Su Xian Chong
Fang Ji
Nathanael Ren Jie Tong
Christopher Li Hsian Chen
Juan Helen Zhou
Venue: NeurIPS 2024 spotlight (2024)
Keywords: foundation model
fMRI
Primary area: neuroscience_and_cognitive_science
TL;DR: Brain-JEPA is a state-of-the-art brain dynamics foundation model enhancing brain activity analysis, it achieves superior performance on different downstream tasks with broad applicability.

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.

Brain-JEPA: Brain Dynamies 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
zijian daaauaudd.u.....aaaa nus edu, fli.rl, helen. helen.zhoul wnu..
National University of Singapore
Abstract
We introduce Brain-JEPA, 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 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 het- erogeneous time-series patches. These methodologies enhance model performance and advance our understanding the neural circuits underlying cognition. Overall, Brain-JEPA is paving the way to address pivotal questions of building brain func- tional coordinate system and masking brain activity at the Al-ncuroscience interface, and setting potentially new paradigm in brain activity analysis through downstream adaptation. Code available at: https://gittuu.. wow.e//re/r/r--L//riiiin... LRL/Brain-JEPA
Introduction
Understanding large-scale brain activity data is crucial for deciphering the complex mechanisms underlying cognitive processes and human behavior. Functional magnetic resonance imaging (fMRI) captures frood-oxygen-leva dependent (BOLD) signals that reflect regional brain activity. It emerges as an indispensable tool in neuroscience for identifying the neural bases of cognitive processes 22 Deep learning approaches have been developed for fMRI analysis, improving brain disease diagnosis and deepening insights into cognition and behavior S 6 g 9 TO 13. Despite notable
advances, these task-specific models suffer from limited generalizability and adaptability to other tasks. In addition, they fail to leverage the vast amounts of unlabeled fMRI data available [4
Artificial intelligence (AI) is experiencing a paradigm shift from task-specific training to building foundation models that are trained on extensive data using self-supervision scale (16). Unlike the models with singular functions, foundation models can be adapted to diverse array of downstream tasks. I Large language models such as GPT (17) and LLaMM (18) have shown significant potential in natural language processing, with expansive applications in healthcare, biomedicine, and beyond [16].
* Equal contribution Corresponding author
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

In the field of fMRI time series analysis, brain language model (brainLM) is one of the most representa- tive foundation models [4. is a masked autoencoder (MAE) [9 trained to reconstruct masked fMRI time series. However, as an indirect measure of neuronal activity, the BOLD signal has a relatively low signal-to-noise ratio (SNR), which is influenced by a mixture of factors and distorted by non-neuronal fluctuations [20]. Filling every bit of the fMRI time series as in MrainI. can hinder the model's 's ability to distinguish between noise and actual signals. This can result in either amplifying noise or missing critical subtle variations in brain activity. Unlike natural images, which have high information density with structures such as edges and colors, fMRI data has spatiotemporally sparser signals distributed across brain volumes without clear boundaries, making it difficult to accurately reconstruct signal of masked regions of interest (ROIs). Furthermore, it has been widely shown that masked pretraining in generative architectures such MAE leads to suboptimal performance in off-the-sheli evaluations (e.g., linear probing) 21] Because of that, BrainI .M requires computationally intensive end-to-end finetuning, with three-layer MIP attached to the pretrained encoder, to achieve optimal perfor- mance. Furthermore, the absence of comparisons with state-of-the-ee methods for downstream task performance and the focus only on Caucasian cohorts limit BrainLM'ss applicability in clinical settings.
Therefore, rather than focusing on the original brain activity time series, the inherent noise and sparse information density of iMRI lead us to explore the latent space of IMRI time series extracted from a strong encoder (e.g.. Vision Transformer (ViT) (22)). It potentially offers a higher SNR after "compression" achieving a greater level of abstraction that captures subtle yet crucial patterns 23 Recently, Imaged- based Joint- Embedding Predictive Architecture (I-JEPA) has been proposed as non- -on-generative architecture for self-supervised learning from images (2I. It predicts the representations of various target blocks rather than reconstructing the masked input like MAF during pretraining. By predicting representations in the latent space, I-JEPA enhances the semantic quality of learned representations and boosts scalability and efficiency.
Training a brain dynamics foundation model using JEPA-like architecture might offer advantages over the MAE approach. However, the distinct spatiotemporal characteristics of fMRI data make direct application of the JEPA architecture suboptimal:
Positional embeddings in transformer play a crucial role by incorporating information about the order or position of tokens in the input data (e.g.. the order of different words in a sentence the locations of pixels in an image) (24] However, there is no such natural "order" for different ROIs across the 3D brain volume in fMRL. BrainLM utilizes anatomical positions to label each ROI [14), yet it does not account for brain functional parcellation, where nearby anatomical ROIS might exhibit rather different brain activation patterns represented by lack of local coherence in fMRI data 25 I-JEPA employs random multi- selection of context and target. However, unlike images, IMRI presents complex patterns across both spatial and temporal domains. Given the smaller sample size and sparser information density in fMRI datasets compared to datasets like ImageNet (26). learning in f'MII requires a stronger inductive bias. This would enhance the efficiency of training models by better capturing the underlying patterns specific to brain activity. Given the unique challenges presented by fMRI data, there is a pressing need to develop a functinal coordinate system and a tailored masking strategy for large-scale pretraining on fMRI data.
These neglected yet crucial questions of developing a functional coordinate system and a masking strategy for large-scale pretraining with fMRI data, lie at the intersection of Al and neuroscience, highlighting important interdisciplimary challenges.
To address these gaps, here we introduce Brain-JEPA, a brain dynamics foundation model with the Doint-Embedding Predictive Architecture (JEPA). Instead of reconstructing masked inputs during pretraining, Brain-JEPA predicts abstract representations of sampled targets from the observation. We propose two innovative techniques to enhance model performance and address key questions in AI for neuroscience; First, Brain Gradient Positioning provides brain functional coordinate system for positional embedding of brain functional parcellation (Section 3.1 Second, Spatiotemporal Masking offers tailored masking strategy for the heterogeneous time-series patches inherent in fMRI (Section 3.2). Moreover, in downstream experiments, our proposed Brain-JEPA achieves tfat-oof-the-ar art results in demographic prediction, disease diagnosis/pronnosis, and trait prediction through fine-tuning. It also excels in off-the-shelf evaluations (e.g., linear probing), and shows superior generalizability across different ethnic groups. Brain-JEPA enhances brain activity analysis and deepens our understanding of critical Al-neuroscience questions related to constructing functional coordinate systems and developing spatiotemporal masking strategies.
2

---

Impat Data
Brain Gradient Posttooning
Pareta
8
Rotel
BO112
Eanoder
3
:
fe
ROLNN
Timsstry
Paichity Shutfle ROO,
Spatiotemporal Mosting
g+
Temporal
loriing
f
Lg (MSE)
Chos-36 3.
Figure I: Brain-JEPA. With a Vision Transformer (ViT) as the observation encoder fo. Brain-JEPA
employs single observation block predict the representations of target blocks. (1) The input
fMRI data is initially segmented into patches for subsequent processing. (2) Through Spatiotemporal Masking, the input excluding the observation block-is divided into three distinet regions: Cross-ROI (a)), Cross-Time Time (3), and Double-Cross (7). The target blocks are sampled from different regions separately. (3) A narrower VIT, serving as the predictor Jé, takes the output 8, from fo. predicts the representations of a target block 8y conditioned on positional embedding (brain gradient positioning for ROI locations and sine and cosine functions for temporal positioning). predicted representations align with those s" from the target encoder fz. whose parameters are incrementally updated through an Exponential Moving Average (EMA) of the observation encoder's parameters.
(4) These
2 Related Work
Task-specific Models for fMRI (state-of-the-arrt) SVR and have been used in fMRI analysis,
utilizing Pearson correlation matrices derived from fMRI time series as input 8 9). Deep learning models have substantially advanced fMRI analysis recent years. BrainNetCNN (6) introduces convolutional neural network (CNN) with specialized convolutional filters tailored for brain network. BrainGNN [S) utilizes ROI-aware graph neural networks (GNNs) to effectively harness functional brain network information, incorporating a pooling operator to highlight key ROIs. More recently, Brain network transformer (BNT) 4 employs transformer encoders to generate embeddings for ROIs based on Pearson correlation matrices, alongside readout layer designed to identify clusters within the brain. Swift [27] applies Swin Transformer architecture 28 to process brain functional data. As noted in Section I these task-specific models have limited generalizability and adaptability across different tasks, and fail utilize extensive unlabeled fMRI data.
14 stands out as the first fMRI foundation model, employing MAE for self-supervised pretraining of fMRI data. In this approach. 1MRI time series are treated as images and patchified. The training goal is to reconstruct the masked patches of the time series. As outlined in Section BrainLM exhibits several limitations: 1) Direct reconstruction of masked input may not be suitable for inherently noisy data with low information density, such as fMRI. It complicates the differentiation between noise and signal, making difficult capture underlying patterns. 2) Generative architectures like MAE result in suboptimal performance in linear probing, critical method for evaluating learned representations. 3) The absence of comparisons with state-of-the--rt models and evaluations limited to Caucasian cohorts restricts its broader applicability. Brain 229]. concurrent work in large-scale self-supervised learning for neuroimaging. focuses on brain network analysis rather than brain dynamics, distinguishing it from research.
The fMRT Foundation Model. BrainIM
3

---

3 Method
In this section, we outline the methodology of Brain-JEPA. Instead of reconstructing masked patches of fMRI time series, Brain-JEPA operates in the latent space, depicted in Figure With the observation block excluded, the input data is divided into three non-overlapping regions: Cross-Time (a), Cross- ROI (8), and Double-Cross (7). This division forces the model engage in forecasting time series, generalizing across unseen ROIs, and predicting time series for unseen ROIs. Section [3. details the Brain Gradient Positioning we proposed, which encodes the functional relationships among different ROIs, serving as brain functional coordinate system in the brain's 's functional organization. In Section 3.2
we introduce Spatiotemporal Masking, which injects a strong inductive bias during the masking process, leading to faster convergence during pretraining and superior performance in downstream tasks.
3.1 Brain Gradient Positioning
We propose Brain Gradient Positioning, which provides brain functional coordi- nates system based on the functional con- nectivity gradient. Positional embeddings are crucial in transformer architectures, they encode information about the posi- tions of tokens in sequence. These em- beddings can be implemented using fixed sine and cosine functions across various fre- quencies [24] or through learnable embed- dings that adapt during training [30). How- ever, the integration of positional informa- tion into IMRI time series has long been ne- glected. FMRI data, incorporating complex spatiotemporal information, requires sepa- rate consideration of its temporal and spatial domains. The temporal domain, represent- ing timesteps during scanning, is well-suited for conventional sine and cosine positional embeddings, as the time series in each ROI is sequentially ordered by time. However, this method not appropriate for the spatial domain, where ROIs across brain volumes lack a simple, inherent order, making sine and cosine embeddings unsuitable for capturing spatial relationships. Anatomical locations of ROIs offer an alternative sine and cosine functions [14] but fall short in capturing functional parcellation. Spatially adjacent ROIs can exhibit significantly different brain activation patterns, reflecting the inherent lack
Gradient Positioning
Grodient
Grad
Gradient
Colour Proyection
(A) Brain Surface Space
(B) Brain Gradient Space
Brain gradient positioning. Brain cortical regions are situated in the top gradient axes and colored based on their positions. These colors are then projected back into the brain surface.
Figure 2:
of local coherence in fMRI data 25
The functional connectivity gradient is continuous measure that captures the functional relations among different ROIs Each attribute in the gradient represents an axis in the latent space of brain regions and networks. The relative distance between different ROIs indicates the similarity in their connectivity (i.e., shorter distance means higher similarity in connectivity). The concept of spatial gradient as conceptualized by Mesulam in 1998 entailed a synaptic hierarchy that supports cognitive processes [31]. Recent studies have built upon this concept, revealing that brain networks in adult humans and macaques exhibit linear distributions across different gradient axes [32] Using this oethodology, it has been shown that these gradients reflect the functional changes related to age
33 34 35 36 37], cognition 33 38 and brain diseases 39 40. These gradients together provide
a framework to assess the relationship between brain regions based on their relative positioning across different gradient axes.
Before deriving the gradients, we first calculate a non-negative affinity matrix A(ij) (a graph Lapla- cian) as follows:
C
A(i.j) - [
1(
(I)
|
π
where c; and ey represents the features (functional connectivity) across the ROI and j. respectively. Gradients are then derived using diffusion map [2). a nonlinear dimension reduction method used to identify the underlying manifold structure of the data. We can obtain the diffusion matrix Lş and the diffusion operator M, from A as follows:

---

M L8 D AD
where D is the degree matrix of A. Here is set to 0.5 to maintain the global relations between ROIs in the embedding space.
Finally, we can compute the eigenvectors and eigenvalues of Má, and stack the column vectors to formulate the diffusion map D E =R= (n ROIs in total, with m gradients for each) at time t and the gradient matrix G with the same dimension:
4 N2
Wm). G= G-1111.12 (1.../2.....//
where A are the eigenvalues and the are the corresponding eigenvectors (gradients) of the graph Laplacian. The parameter t represents the diffusion time, which controls the scale of the diffusion process. Here we estimated the eigenvalues A at time by dividing it by 1 X to enhance robustness against noisy eigenvalues.
In Brain-JEPA, we leverage G the spatial positioning ROIs. Specifically, the gradient G E R transformed into G E Rndd/2 through a trainable linear layer, where represents the embedding dimension of the VIT backbone. The predefined temporal positioning T =R is obtained using sine and cosine functions [43] The final positional emhedding can then be formulated as P= [T, G] Raa Figure Figure[]provviess visualization of the top 3 gradients in Euclidean space with each ROI color coded by their locations. As shown, the brain gradient positioning reflects functional network architecture, such as the somatomotor, default mode and visual networks, consistent with previous literature 32 37
3.2 Spatiotemporal Masking
Observation. Brain-JEPP aims to predict representations of multiple target blocks based on the representation of a single observation block. For an input fMRI time series, the temporal signal for each parcel is divided into patches after shuffling ROIs, each containing p time points (dash boxes in Figure(T). The observation block x is obtained by randomly sampling a block within the range (i/e nf.)). VR specifies the range ratio along the ROI dimension, and 7/ pertains to the timestep patches (10 in total). Subsequently, x is fed through the observation encoder fe. generating corresponding patch-level representation 8g:
8, - {82, JJeB.
where B, represents the mask associated with the observation block 2; Sxy, is the representation of the j"h patch.
Targets. Given a single observation, the model is trained to predict other parts of the IMRI within the latent space. Random sampling of targets like MAE 19 might allow the model to learn shortcuts (e.g., interpolation of time series) or rely heavily on simpler, more frequent patterns in the data, which could limit its generalizability. It crucial to recognize that patches in fMRI vary spatially depending on the positions in their brain functional organization, and temporally regarding brain states and task conditions. The nonlinear relationship among brain networks further complicates the interactions between different brain patches.
As shown in Figure distinct and non-overlapping regions: Cross- ROI (a), Cross-Tine (3), and Double-Croo Cross (7). For targets in the and 3 regions, the model should generalize the observation across different ROIS spatially or timesteps temporally. For targets in the y regions, which are the most challenging, the model should generalize unseen ROIs unencountered timesteps. We randomly sample K blocks from each of the three types of regions as targets, forcing the model to handle a variety of prediction tasks with a stronger inductive bias. We denote the mask corresponding of the region (r E (a,B,n)) as By.
categorize the remaining parts (with the observation excluded) into three
Overlapped sampling. It has been shown in 21 that sufficiently large dynamic range of masking ratio could benefit pretraining. To effectively adjust the observation-to-inge ratio during pretraining, we implement an overlapped sampling strategy that allows for flexible, rather than fixed, ratio. When sampling the target block sy from region r, for r == a or B, we sample the target from the union of the observation mask and region/ mask; while for r we directly sample the target from the y region mask. Formally, the overlapped sampling strategy is defined as:
5 BBBBU s - BB 8 B3
(5)
5

---

Afterwards, part of the observation region might overlap with some a and 3 targets. We remove any ROIs in the observation that overlap with the a targets. Additionally, we eliminate all timesteps for ROIs that show overlap with the / targets. Refer to Table the block sizes.
Training. Given the output 8, from the observation encoder fo, the predictor do is trained predict the three kinds of targets sy conditioned on the positional embedding P (Figure I. The training loss C the average L2 distance between By and its corresponding prediction:
-=
1 3K
E
6 By ==99 (8(
(6)
4 Experiments
4.1 Datasets
We leveraged the large-scale public dataset UK Biobank (UKB)
44
45) for the self-supervised
pretraining of Brain-JEPA, It includes resting-state fMRI recordings with medical records from 10, participants aged 44 to 83. Multi-site recordings were acquired with the temporal resolution of 0.735s. We allocated 80% of this dataset for pretraining (of which calculated the group- gradients as well), with the 20% held-out for downstream evaluation (internal tasks of age and sex prediction).
We used three datasets for external evaluation: HHCP-Aging as a segment of the public Human Connectome Project (HCP) [46,, includes resting- fMRI from 656 healthy elderly participants. It was used to predict traits (Neuroticism and Flanker score) and demographics (age and sex). The Alzheimer's s Disease Neuroimaging Initiative (ADNI) 47 was used for the early diagnosis and prognosis of neurodegenerative diseases, with fMRI from 189 participants for normal control (NC) v.s. mild cogmitive impairment (MCT) classification, and 100 cognitively normal participants for amyloid positive v.s. negative classification. Moreover, assess generalizability across different ethnic groups and real-world clinical applications, we included resting-state IMRI of Asian participants recruited by Memory, Ageing and Cognition CenCre((MACC) with 539 participants for NC v.s. MCI classification. More details of the downstream tasks performed can be found in the Appendix A
All fMRI data was parcellated into = 450 ROIs, using Schaefer-400
48 for cortical regions and
Tian-Scale III [49 for subcortical regions. Robust scaling was implemented by subtracting the median and dividing by the interquartile range, calculated across participants for each ROI [14). Our default input size 160 timesteps for each the 450 ROIS (i.e., 450x160) UKB and ICi--gging used multi- band acquisition with high temporal resolution (TR x 0.7 seconds), while ADNI and MACC used single-band acquisition with a lower resolution (TR se 2 seconds). To ensure consistency across datasets, we standardized the temporal resolution by downsampling the multi-band data using a temporal stride of 3, aligning the TR of all datasets to approximately 2 seconds. During the fine-tuning and linear probing stage, all the downstream datasets were divided into 6:2:2 ratio for training, validation, and testing.
4.2 Implementation details
For Brain-JEPA pretraining, we utilized ViT architectures for the observation encoder, target encoder, and predictor. We employed FltshAttention [50] 51) in our self-attention implementation to improve computational efficiency and reduce memory usage. Balancing the trade-off between data quantity and the model complexity, we experimented with ViT-Small (ViT-S) (22M). ViT-Base (ViT-B) (86M), and ViT-Large (ViT-L) (307M) for the observation encoder. For predictor, it is designed as a lightweight (narrow) ViT. Specifically, the predictor has the same architecture as the corresponding observation encoder, differing only in embedding dimension and depth. For the ViT-S and ViT-B observation encoders, the predictor has a depth of 6 and embedding dimensions of 192 and 384, respectively. The ViT-L observation encoder uses a predictor with a depth of 12 and an embedding dimension of 384. Brain-JEPA is pretrained without [cls] token. For evaluation, we used the target encoder and average pooled its output to generate global fMRI representation. The main results in Section 4.3 along with the analysis Section 4.5 and 4.7 were all based on ViT-B pre-trained for 300 epochs. Refer to Appendix B for optimization and masking details.

---

Table I: Internal tasks of age and sex prediction on UKB 20% held-out. The mean (standard deviation) of Mean Squared Error (MSE). Pearson Correlation (p), and/or Accuracy (ACC), Fl score across 5 independent runs reported. t: the higher, the better; 4: the lower, the better. The best results are in bold, with * denoting significant improvement over previous approaches (p < 0.05).
Age
Sex
Methods
MSE 1
P
A
F1(%)) f
BrainNetCNN 6
(0,07))
0.225 (0.015)
77.86
(0.17)
BrainGNN S
0.931 (0.035
0.532 (00055
77.31 (0.33)
(0.33)
BNT 4
0363 (0/031)
0.447 (0.017)
(0.48(
(0.32)
TFS
(0.023)
0.487 (0.01))
82 (2.60:0..59)
83 (3.00)0.00)
BrainLM 14
(.612(0
0,632 (0.020)
96.17 (0.74)
(6.84
Brain-JEi JEPA
0.501*
0.718 (0.021)
88.17 (0.06)
88.58
Trained- Frmined-FFomm-eee Scratihhe
Table 2: External tasks of demographics and trait prediction on HCP-Aging.
Age
Sex
Neuroticism
Flanker
Methods
MSE.J.
pt
ACC (S)
(
MSB.
pt
MSB.
Pt
BrainNerCNNN 6
0.462 (.017)
(.23)
(2.25-0092)) (0.16
1.201 (187)
0.096
1.08 (096)
0:201 (018)
BrainGNN S
(015)
(.24)
72.7 (0:54)
74.09 (0.67)
1183(096)
(.00)(000)
(.982
0.309 ((06))
BNT A
(1085)
0.731 (.05))
(1.09)
(1.11)
1,199 ((991)
0.101 (0005)
(.937
0.307 (026)
BrainLM
14
0.331
(.028)
77.51 (1.55(
0,942
0.231 (.012)
0 (.054)
(.018
Brain-JEPP
0,298 (.017)
(.930)
81.52* (1.03)
84.26* (0.82)
0.897 (.055)
0,307 (.006)
(.038)
0.406* (.027)
Table 3: External tasks of brain disease diagnosis and prognosis on ADNI and MACC.
NC/MCI
Amyloid ve/ ve
NC/MII (Asian)
Methods
A
F
A
F
ACC(%) f
F1(%) 1
BrainNerCNN 6
(3.00
(3.18)
(2.00
59.43(L14)
(7.32(145))
(4.92
BrainGNN 5
71422 (2.95)
(7.00((400)
(3.48)
(2.35)
$5.69 (2.299
BNT 4
(4.12)
(3.54)
62.00
(0.53)
62.06 (3.58)
60.45 (4.52)
BruinLM 14
75.79 (1.05)
85.66 (1.27)
67.00
(8.82)
61.65 (1.35)
(3.26)
Brain-JEPA
(6.84(1.00)
86.32 (0.54)
71.00" (4.90)
75.97% (3.93)
65.98* (2.84)
64.67 (2.61)
4.3 Main results
Table T and 3 compare Brain-JEPA with the existing deep learning models for fMRI analysis and
foundation model BrainLM. We select the three deep learning baselines because they not only represent the previous state-of- ootte-of-hh-an- in fMRI analysis but also exemplify diverse model types: convolutional BrainGNN (5),
neural network (CNN)-based BrainNetCNN 6), graph neural network (GNN)-basee and transformer-based BNT 4 For a fair comparison, both Brain-JEPa and BrainLM MM utilized a ViT-B backbone and were fine-tuned for downstream tasks (Section 4.4] will discuss performance scaling with different model sizes, and Section 4.5 will examine lincar probing comparisons between the two models). BrainLM utilized [cls] token for downstream evaluation.
The results show that Brain-JEPA achieves oafe-of-thhearr performance in various downstream tasks on both the unseen data from the same pretrained cohort and other independent datasets. Brain-JEPA effectively captures fundamental demographic information such as age and sex, cognitive/personality variance (Neuroticism and Flanker), and dissase-relatee patterns for neurodegenerative diseases. Notably, Brain-JEPA demonstrates superior performance in classifying NC/MCI Asian ethnic groups one of the most challenging tasks for early diagnosis and prognosis of Alzheimer's Disease (AD) - even though was trained exclusively on the Caucasian cohort. Please refer to C.1 for additional results on more datasets and comparisons with more baselines.
7

---

Age Prediction (ICP-Aging)
Sex Prediction (HCP-Aging)
NC/MCI Classification (ADNI)
78.42%
Peassin Exerclatoon
84.55%
76.84%
K1 52%
79.0%%
0.758
66.79%
86M
307M
22M
861
507M
22M
86M
307M
Figure 3: Performance scaling of the model sizes.
Prediction (HCP-Aging)
Prediction (HCP-Aging)
Amyloid + (ADNI)
Parron Cieveiation
v13.01%
24.53%%
Accuraty
Accoraty
V1.11%
VIRS2%
V7.40%
28.9%%
-
BERinJJEA BrainMM
oT
probing
Figure 4: Fine-tuning linear probing.
4.4 Performance scaling
presents present the performance Brain-JEPA across various model sizes, using ViT-S, VTT-B, and ViT-L backbones. The results demonstrate that the larger model configuration consistently achieves better performance. Specifically, there is a clear trend of increasing accuracy/correlation with larger models, with Brain-JEPA using ViT-L consistently achieving the best performance. We also studies the scaling property with respect to dataset size, please refer to C.2 for additional results.
Figure
4.5 Linear probing
BrainLM initially showcases its performance improvements through fine-tuning, complemented by attached MLP [14]. However, effectively assess the representations learned during pretraining, off-the-shelf evaluations such as linear probing are essential. As depicted in Figure [ Brain-JEPA consistently outperforms BrainLM in linear probing and exhibits a smaller performance decline from fine-tuning to linear probing. This highlights the robustness and higher level of abstraction in the representations learned by Brain-JEPA.
4.6
Ablation study
Siu andi
Praia Goulle Reddiwing
We first compared Brain-JEPA with its ablated sions, employing sine and cosine functions 43 and anatomical locations [14] for ROI spatial positioning,
Farm Comstaaa
75/00/
as shown in Figure 5 Brain Gradient Positioning demonstrates superior performance over these two baseline methods. It indicates that Brain Gradient Po- sitioning facilitates natural and accurate placement of brain functional parcellations, enhancing the learning of brain dynamics. Next, we assessed the effectiveness of our proposed Spatiotemporal Masking by compar- ing Brain-JEPA, pretrained over various numbers epochs, to its ablated counterpart that utilizes standard multi-block sampling of targets (2I). This compari- son, illustrated in F highlights that not only does our proposed masking technique yield superior
22.15
(SCCAAl0))
(ICDNAppig)
LABING
Figure 5: Comparisons of spatial positional embedding (For the first task, refer to the left y axis for the Pearson's 's Correlation, with the right axis accuracy for the last two tasks).

---

Age Prediction ((gig)
Sex Prediction (ICPPAgingg)
NC/MCI classification (ADNT)
paarrooempr Muk-brock
Prinses Clrnciation
Varilla Mub-Blikk (ID0 epozt)
Accoracy
Accunayy
0.786
0.70
75.79%
74
0.5
12
73
0.99
Epochs (pretraining)
Epochs (pectraining)
Epocha (pretraining)
Figure 6: Comparisons of masking strategies.
performance, but also introduces a stronger inductive bias leading more efficient pretraining. Notably, Brain-JEPA achieves or even surpasses the peak performance of the ablated version, which was pretrained for 300 epochs, with significantly fewer epochs- only 100, 200, and respectively. For more ablation results regarding architectures and the number of gradient components, please refer to
4.7
Interpretation
NC/MC (Cancasian)
NC/MCI (Asian)
3.25
325
3.20
CN
Allenion
DAN
SMN
25
2.0%
D65
144
CAN
2PH
05
D40
C4N
W
Networks
Figure 7: Attention across different brain networks for NC/MCI classification.
With the Schaefer functional atlas 488, the brain network is categorized into seven distinct sub- networks: the control network (CN), the default mode network (DMN), the dorsal attention network (DAN), the limbic network (L.N), the salience ventral attention network (SAN), the somatomotor network (SMN), and the visual network (VN). To assess whether Brain-JEPA has captured the brain functional organization, we calculate the network-level attention for NC/MCI classification. For each ROI, we first average the self-attention across its 10 patches. Next, average the values of the ROIs within each aub-network and normalize them to obtain the network-level attention distribution. As shown in Figure 7 we found consistent patterns across both Caucasian and Asian ethnic groups, with the model highlighting the critical roles of the DMN, CN, SAN, and LN in cognitive impairment, consistent with previous literature [52 53 54
Conclusion
In this study, we developed Brain-JEPA, a brain dynamics foundation model based on the Joint- Embedding Predictive Architecture (JEPA). Brain-JEPA predicts abstract representations of sampled targets from observations during the pretraining stage. Utilizing Brain Gradient Positioning, Brain- JEPA encodes brain functional organization more naturally and accurately. With Spatiotemporal Masking, it effectively handles heterogeneous patches in fMRI time series. Brain-JEPA fosters generalizable and highly abstract representations of fMRI, achieving ffate-of-the-nd art performance across various tasks, including demographic prediction, trait prediction, and disease diagnosis and prognosis across different cohorts and ethnic groups. Our study provides new insights into applying large-scale self-supervised learning to brain activity modelling and contributes addressing key questions in AI for neuroscience.

---

Limitation and future work
We acknowledge several limitations in our study, which also serve as inspirations for future research: 1) Larger models: Due to limited computing resources, we have not tested larger models like ViT-H. We expect that larger models could further improve performance. 2) More diverse datasets: more diverse brain recording dataset for pretraining, including different ethnicity cohorts collected from various sites, scanning protocols, behavioral tasks, and disease groups, could enhance the generalizability and robustness of the representations learned by the model. 3) Fine- tine-grained interpretation: More thorough interpretation can be achieved through the attention mechanism, such as comparing cortical and subcortical regions, identifying salient ROIs and critical timesteps. This would enable more nuanced and complex spatiotemporal interpretations. 4) Multi-modal integration: Brain-JEPA sets potential foundation for integrating multimodal brain activity data such as MEG and EEG or even brain structure data like T1 -weeghtt MRI. The integration could enhance our understanding of brain structure, function, and their links to human behavior and mental disorders. Please refer to Appendix for the broader impact of Brain-JEPA.
7 Acknowledgements
This study was supported by the Singapore National Medical Research Council (NMRC/OFLCGI CCCGGMMM-OOOR -005, NMRC/CIRR//1455201. STA.:////ASSSSSS.CCCCAAAAA/// NMRRCCSA-SS0007/201. NMRC/MOH- 00707-01, NMRC/CG/435 M009/2017-NUHNNURT
CIRG2Inov-000 and HLCA23Feb-0004 RIE2020 AME Programmatic Fund from A*STAR, Singapore (No. A20G8b0102), Ministry of Education (MOE-T22 12EP40120-0007 & T2EP2-0223 022--2022202, MOE-T2EP20200-002 and Yong Loo Lin School of Medicine Research Core Funding, National University of Singapore, Singapore.
References
[1] Nikos K Logothetis. What we can do and what we cannot do with fmri. Nature, 453(7197): 869-878, 2008
[2] David J Heeger and David Ress. What does fmri tell about neuronal activity? Nature reviews neuroxcience, 3(2:1111511 2002.
2002.
[3] Nikos K Logothetis, Jon Pauls, Mark Augath, Torsten Trinath, and Axel Oeltermann. Neurophys- iological investigation of the basis of the fmri signal. nature, 412(6843): 150-157,
[4] Xuan Kan, Wei Dai, Hejie Cui, Zilong Zhang, Ying Guo, and Carl Yang. Brain network transformer. Advances in Neural Information Processing Systems, 35:25586-25599, 2022.
[5] Xiaoxiao Li, Yuan Zhou, Nicha Dvomek, Muhan Zhang. Siyuan Gao, Juntang Zhuang, Dustin Scheinost, Lawrence Staib, Pamela Ventola, and James S Duncan. Braingnn: Interpretable brain graph neural network for fmni analysis. Medical Image Analysis, 74:102233, 2021.
[6] Jeremy Kawahara, Colin J Brown, Steven P Miller, Brian G Booth, Vann Chau, Ruth E Grunau, Jill G Zwicker, and Ghassan Hamarneh. Brainnetenn: Convolutional neural networks for brain networks; towards predicting neurodevelopment. Neurolmage, 146: 10000-1844 2017.
[7] Hejie Cui, Wei Dai, Yanqiao Zhu, Xuan Kan, Antonio Aodong Chen Gu, Joshua Lukemire, Liang Zhan, Lifang Hc, Ying Guo, and Carl Yang. Braingb: a benchmark for brain network analysis with graph neural networks. JEEE transactions on medical imaging, 42(2):493-506, 2022.
[8] Andrew T Drysdale, Logan Grosenick, Jonathan Downar, Katharine Dunlop, Farrokh Mansouri, Yue Meng, Robert N Fetcho, Benjamin Zebley, Desmond J Oathes, Amit Etkin, et al. Resting state connectivity biomarkers define neurophysiological subtypes of depression. Nature medicine, 23 (1):28-38. 38, 2017.
[9] Tetsuya lidaka. Resting state functional magnetic resonance imaging and neural network classified autism and control. Cortex, 63:55-67, 2015.
[10] Zijian Dong, Yilei Wu, Zijiao Chen, Yichi Zhang, Yueming Jin, and Juan Helen Zhou. Prompt your brain: Scaffold prompt tuning for efficient adaptation of fmri pre-trained model. arXiv preprint
10

---

[11] Zijian Dong, Yilei Wu, Yu Xiao, Joanna Su Xian Chong, Yueming Jin, and Juan Helen Zhou. Beyond the snapshot: Brain tokenized graph transformer for longitudinal brain functional con- nectome embedding. In International Conference on Medical Image Computing and Computer- Assisted Intervention, pages
Springer, 2023.
[12] Yilei Wu, Zijian Dong, Chongyao Chen, Wangchunshu Zhou, and Juan Helen Zhou. Mixup your own pairs, arXiv preprint 211.:i:::3991663333 2023.
[13] Zijian Dong, Joanna Su Xian Chong, Bing Cai Kok, and Juan Helen Zhou, Coop-dhgnn: framework for joint classification and prediction of brain functional connectivity using sparse trajectory dataset with application to early dementia. In 2022 IEEE International Conference on Big Data (Big Data), pages 4972 4978, 2022. doi: 10. Bigt/////aaaaa6...22222 2022. 2022000000.222...0011 10021043.
[14] Josue Ortega Caro, Antonio Henrique de Oliveira Fonseca, Christopher Averill, Syed A Rizvi, Matteo Rosati, James L Cross, Prateek Mittal, Emanuele Zappala, Rahul Madhav Dhodapkar, Chadi Abdallah, et al. Brainlm: A foundation model for brain activity recordings. In The Twelfth International Conference on Learning Representations, 2023.
[15] Guangqi Wen, Peng Cao, Lingwen Liu, Jinzhu Yang, Xizhe Zhang, Fei Wang, and Osmar R Zaiane. Graph self- supervisdd learning with application brain networks analysis. IEEE Journal of Biomedical and Health Informatics, 2023.
[16] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.072 2021.
[17] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprini SrxXix:2303.08722 2023.
[18) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
[19] K. Hc, X. Chen, S. Xie, Y. Li, P. Dollar, and R. Girshick. Masked autoencoders are scalable vision learners. In 2022 IEEE/CVE Conference Computer Vision and Pattern Recognition (CVPR), pages 15979-15988, Los Alamitos, CA, USA, jun 2022. IEEE Computer Society, doi: 10.1 1./10/CP2RR2268.. 2035533.. https:tss:///o https///dd.....eeccmmuuurrcrccec org/ 10. N0.1109/CVPR52288.. 2022.01553
[20] César Sibbllleo-Gaudee and Richard C Reynolls. Methods for cleaning the bold fmri signal. Neuroimage, 154: 1444 2017.
[21] Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a joint- embedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16619-15229,00 2023.
[22] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 2020.
[23] Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2022-06-27. Open Review, 62(1), 2022.
[24] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[25] Stephen M Smith, Christian F Beckmann, Jesper Andersson, Edward Auerbach, Janine Bi- jsterbosch, Gwenaëlle Douaud, Eugene Duff, David A Feinberg, Ludovica Griffanti, Michael P Harms, et al. Resting-state fmri in the human connectome project. Neuroimage, 80: 148-1668 2013.
11

---

[26] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 245-255. leee, 2009.
[27] Peter Kim, Junbeom Kwon, Sunghwan Joo, Sangyoon Bae, Donggyu Lee, Yoonho Jung, Shinjae Yoo, Jiook Cha, and Taesup Moon. Swift: Swin 4d fimri transformer. In A. Oh, Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36, pages 42015-42033 Inc., 2023. URL/attpss hwwpiiinnnnnnuuuureeeeeee cc wwwpaperpppppeeeeceeeeeeeeeeeeee haapparrrellseslleaaaalaaaleeeee www.ffieealprrrr httppppappp/e///ip///pp 831 6G115190eee068886666677666.-6.60
Curran Associates,
Paper- Conference, pdd
[28] Ze Liu, Yutong Lin. Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[29] Yanwu Yang, Chenfei Ye, Guinan Su, Ziyao Zhang, Zhikai Chang, Hairui Chen, Piu Chan, Yue Yu, and Ting Ma. Brainmass: Advancing brain network analysis for diagnosis with large- scale sele-super learning. IEEE Transactions on Medical Imaging, pages 1-1, 2024. doi:
10.1 T0.11000/7M...2244444447
[30] Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence sequence learning. In International conference on machine learning, pages 1243- 1252. PMLR, 2017.
[31] M-Marsel Mesulam. From sensation cognition. Brain: a journal of neurology, 121(6):
1013-1052, 1998. ISSN 1460-2156.
[32] Daniel S. Margulies, Satrajit S. Ghosh, Alexandros Goulas, Marcel Falkiewicz, Julia M. Huntenburg. Georg Langs, Gleb Bezgin, Simon B. Eickhoff, F. Xavier Castellanos, Michael Petrides, Elizabeth Jefferies, and Jonathan Smallwood. Situating the default-mode network along principal gradient of macroscale cortical organization. Proceedings of the National Academy of Sciences, 1 113(4)):1574411277, 2016. doi: 10. 10. 1073/pass... 1eappaas..6688888 16088821 13. URL pnas org/ www.cooonttttsststeee 11////13/44///// 12574. 11122574..l. full. .pdf.
[33] Z. X. Zhou and X. N. Zuo. Editorial: Lifespan connectome gradients for a road to mental health. J Am Acad Child Adolese Psychiatry, 63(1):25-22 2024. ISSN 0890-8567. doi: 10.1016/) 1016/j.jaac 2023.08.00 123.08.006. 1527-5418 Zhou, Zi-Xuan Zuo, Xi-Nian Editorial United States 2023/09/02 Am doi: 10.1016/j.jaae 10..001.... 2023.08.006 Epub
Acad Child Adolese Psychiatry. 2024 Jan:63(1):22522)
2023 Aug 30.
[34] Sara Larivière, Reinder Vos de Wacl, Seok-Jun Hong, Casey Paquola, Shahin Tavakol, Alexander J. Lowe, Dewi V. Schrader, and Boris C. Bernhardt. Multiscale structure-function gradients in the neonatal connectome. Cerebral Cortex, 30( 10)455888 2020. ISSN 1047-3211. doi: 10. 0001093//errco//bbz00. https://////d org/ 10. wtt.//////o///////00000 cercor/
[35] Thuan Tinh Nguyen, Xing Qian, Eric Kwun Kei Ng, Mareus Qin Wen Ong, Zhen Ming Ngoh, Shayne S. P. Yeo, Jia Ming Lau, Ai Peng Tan, Birit F. P. Broekman, Evelyn C. Law, Peter D. Gluckman, Yap-Seng Chong. Samuele Cortese, Michael J. Mcaney, and Juan Helen Zhou. Variations in cortical functional gradients relate to dimensions of psychopathology in preschool children. Journal of the American Academy of Child & Adolescent Psychiatry, 2023, ISSN 0890- 8567. doi: https:///doi.ooo org/ pww.///aoooaa........aaaaaaaaaaa jaac 2023 05.029. URL http:/wwwwcciieee. sciencedirectt ccw/sccienc////cccccccccccccccce httm//ciennc//rtticlieee
[36] Hao-Ming Dong, Daniel S. Margulies, Xi-Nian Zuo, and Avram J. Holmes. Shifting gra- dients of macroscale cortical organization mark the transition from childhood to adoles- cence. Proceedings of the National Academy of Sciences, 118228)::2024488111 2021. doi: wttp:///....1000......2222002020 2024448118 doi:
10.1073/pnas 11073//pnas222244488188 2024448118 0000033/paa.2024488111
URL
[37] Richard A. I. Bethlehem, Casey Paquola, Jakob Seidlitz, Lisa Ronan, Boris Bernhardt, Cam-Cai Consortium, and Kamen A. Tsvetanov. Dispersion of functional gradients across the adult lifespan. Neurolmage, 222:1 17299-117299 2020. ISSN 1095-9572 1053-8111 doi: 10. j.neuroimage 2020. 1799.. 7/www. ncbi nlm. nih htw.::/w.............oomaaaaaaaa
nlm http:///www.aac....cc.aaaaa////n
12

---

[38] Xiuyi Wang, Daniel S. Margulies, Jonathan Smallwood. and Elizabeth Jefferies. A gradient from long-term memory to novel cognition: Transitions through default mode and executive cortex. Neurolmage, 220: 1et117074, 2020. ISSN 1053-8119. doi: https://doi.orr/100.100
neuroimage. 17704. URL https: htttps://wwes //tp:://wwwww.e www.:sseeiecedree..... com.s///ccccceene//trrii.l.
p11/$10538119203056007
[39] Seok-Jun Hong, Reinder Vos de Wael, Richard A. I. Bethlehem, Sara Lariviere, Casey Paquola, Sofie I.. Valk, Michael P. Milham, Adriana Di Martino, Daniel S. Margulies, Jonathan Small- wood, and Boris C. Bernhardt. Atypical functional connectome hierarchy in autism. Nature Communicationns 10(1): 2019. ISSN 2041-1723, doi: 10. 10011///41477701900944444 URL https://doi.. org/ 1038/ 0118///4111100009999----1111
[40] Yirong He, Qiongling Li, Zhenrong Fu, Debin Zeng, Ying Han, and Shuyu Li. Functional gradients reveal altered functional segregation in patients with amnestic mild cognitive impairment and alzheimer's s disease. Cerebral Cortex, 33(11):100366 2023. ISSN 1047-3211. doi: 10. 1...0///erccon/baad338....
[41] R. R. Coifman, S. Lafon, Geometric diffusions as a tool for harmonic analysis and structure definition of data: Dif- fusion maps. Proceedings of the National Academy of Sciences, 102(21):7426 7431, 2005. doi: 11000///ma.00000333444 10. 10.1073/pmas 050033340000
Lee, M. Maggioni, Nadler, F. Warner, and S. W. Zucker.
pnas PRLhhtps://dol/ org/ 10. 1073/ 10w.://as......oecmoom 0000//p03.0.0000000...00 0500334102 doi:
[42] Reinder Vos de Wael, Oualid Benkarim, Casey Paquola, Sara Lariviere, Jessica Royer, Shahin Tavakol, Ting Xu, Seok-Jun Hong, Georg Langs, Sofie Valk, al. Brainspace: toolbox for the analysis of macroscale gradients in neuroimaging and connectomics datasets. Communications biology, 3(1):103, 2020.
[43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URt/ neurips.....m cc/ cc/paper.p.... feraperfflles//eeeeee wwarr__flls__pppeee wtalls/paaeer/017/iiili pdf
[44]
L Miller, Fidel Alfaro-Almagro, Neal Bangerter, David L Thomas, Essa Yacoub, Junqian Xu, Andreas J Bartsch, Saad Jbabdi, Stamatios N Sotiropoulos, Jesper LR Andersson, et al. Multimodal population brain imaging in the uk biobank prospective epidemiological study. Nature neuroscience, 19(1 11:11233-1136, 2016.
[45] Maria A Di Biase, Robert E Smith, Andrew Zalesky, Caio Seguin, al. Connectomes for 40,000 uk biobank participants: a multi-modal, multi-scale brain network resource. Neuroimage, 283: 1 120407, 2023.
[46] Jennifer Stine Elam, Matthew F Glasser, Michael Harms, Stamatios Sotiropoulos, Jesper LR Andersson, Gregory C Burgess, Sandra W Curtiss, Robert Oostenveld, Linda J Larson-Prior, Jan Jan-Mathijs Schoffelen, et al. The human connectome project: a retrospective. Neurolmage, 244: 1 118543, 2021.
[47] Clifford R Jack Jr, Matt A Bernstein, Nick C Fox, Paul Thompson, Gene Alexander, Danielle Harvey, Bret Borowski, Paula J Britson, Jennifer L. Whitwell, Chadwick Ward, et al. The alzheimer's 's disease neuroimaging initiative (adni): Mri methods. Journal of Magnetic Resonance Imaging: An Official Journal of the International Society for Magnetic Resonance in Medicine, 27(4):685-691, 2008
[48] Alexander Schaefer, Ru Kong, Evan M Gordon, Timothy O Laumann, Xi-Nian Zuo, Avram Holmes, Simon Eickhoff, and BT Thomas Yeo. Local-global parcellation of the human cerebral cortex from intrinsic functional connectivity mri Cerebral cortex, 2809):3095-3114, 2018.
[49] Ye Tian, Daniel S Margulies, Michael Breakspear, and Andrew Zalesky. Topographic oroaaiiaiioogaanizaatiioo of the human subcortex unveiled with functional connectivity gradients. Nature neuroscience, 23 (11):1421-1433 2020.
13

---

[50] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory- exact attention with I In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[51] Tri Dao. F Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
[52] Puneet Talwar, Suman Kushwaha, Monali Chaturvedi, and Vidur Mahajan. Systematic review of different neuroimaging correlates in mild cognitive impairment and alzheimer's disease. Clinical neuroradiology, 31(4)9553-977 2021.
[53] Yvette Sheline and Marcus E Raichle. Resting state functional connectivity in preclinical alzheimer's disease. Biological psychiatry, 74(5):340-347, 2013.
[54] Matthew R Brier, Jewell B Thomas, and Beau M Ances. Network dysfunction in alzheimer's disease: refining the disconnection hypothesis. Brain connectiviry, (5)):999311,, 2014.
[55] Edmund R Thompson. Development and validation of an international english big big-five mini- markers. Personality and individual differences, 45(6):542-5433 2008.
[56] Barbara A Eriksen and Charles W Eriksen. Effects of noise letters upon the identification of target letter in nonsearch task. Perception & psychophysics, 16(1):143-149, 1974.
[57] Ronald Carl Petersen, Paul S Aisen, Laurel Beckett, Michael C Donohue, Anthony Collins Gamst, Danielle J Harvey, CR Jack Jr, William Jagust, Leslie M Shaw, Arthur W Toga, et al. Alzheimer's disease neuroimaging initiative (adni) clinical characterization. Neurology, 74(3): 201-209, 2010.
[58) Kevin Clark, Minh-Thang Luong, Quoe V. Le, and Christopher D. Manning. Electra: Pre-training text encoders as discriminators rather than generators. In International Conference on Learning Representations, 2020. htt.l/tppp////peenrreeeeeeee/eep
[59] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 448 456, Lille, France, 07-09 Jul 2015. PMLR. URL https://proceediing mlr (tteess//3//ioffe15hhho
[60] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL /tR/attpp:////errevvee http:////pppnerreeeeeet/trooomo net/forum?
id-Bkg6Ricqy7
[61] Boris Ginsburg, Igor Gitman, and Yang You. Large batch training of convolutional networks with layer-wise adaptive rate scaling, 2018. 8. URI.
rJ4uaX2aW
[62] Marc-Andre Schulz, B. T. Thomas Yco, Joshua T. Vogelstein, Janaina Mourao-Miranada: Jamaina Kather, Konrad Kording, Blake Richards, and Danilo Bzdok. Different scaling of linear models and deep learning in ukbiobank brain images versus machine-learning datasets. Nature Commu- nications, 41338,228, Aug 2020. ISSN 2041-1723. doi: 10. 10.003/44114---0--18877777 120-1777722- URL https:///oi... org/10. 10/0 11/1100.101///1 00008//44167--10 -20-41167-----1111 18037-2.
Jakob N.
[63] Armin Thomas, Christopher Ré, and Russell Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Infor- mation Processing Systems, volume 35, pages 21255 21269. Curran Associates, Inc..
2022. URL http:: https://proceeedingg neurips. cc/ http/pppr_fffee files/paper/2000 8600a9df1a0885 $5C0088FF.0778883500004444400000 Conference
14

---

A Task Details
A.1 Neuroticism
Neuroticism is a personality trait linked to negative emotions and is one the Big Five personality traits. People who score high in neuroticism tend to experience negative feelings more frequently than others [55]. The HCP uses the 60-question version of the NEO-FFI Short Form (ages 16+) questionnaire, which provides a quick, reliable, and accurate assessment of the Big Five personality traits: neuroticism, extraversion, openness, agreeableness, and conscientiousness. More detailed information be found in Lifessan HCP 2.0) Data Release Appendix 2: Details and References for Behavioral & Clinical Instruments.
A.2 Flanker
The Flanker task is designed to assess both attention and inhibitory control in participants [56]. It involves the participant focusing on a central stimulus while ignoring adjacent stimuli, which are either fish for ages 3-7 arrows for ages 8-85. Sometimes the central stimulus points in the same direction as the flanking stimuli (congruent) and sometimes in the opposite direction (incongruent). For participants aged 8-85, the task consists of twenty trials and takes about three minutes to complete. More detailed information can be found in the Lifespan HCP 2.0) Data Release Appendix 2: Details and References for Behavioral & Clinical Instruments.
A.3 NC/CCI
For the ADNT dataset [7]. the criteria for NC was as follows: 1) No subjective memory complaints, 2) preserved activities of daily living and cognitive function, 3) Mini-mental state examination (MMSE) score of between 24 to 30 inclusive, 4) Clinical Dementia Rating (CDR) score of O, and 5) education- adjusted score on delayed recall one paragraph from Wechsler Memory Scale Logical Memory Il of >=3 for 0-7 years of education, >= 5 for 8-15 years of education, and >= 9 for >=16 years of education. The criteria for MCI was as follows: 1) significant subjective memory complaints reported by the participant, clinician informant, 2) not significantly impaired in other cognitively domains, 3) essentially preserved activities daily living and does not meet criteria for diagnosis of dementia, 4) MMSE score of between 24 to 30 inclusive, 5) CDR score of 0.5, and delayed recall of one paragraph from Wechsler Memory Scale Logical Memory of 3-6 for 0-7 years of education, 5-9 for 15 years of education, and 9-11 for >=16 years of education [57].
eduucation-adjuster score
For the Asian disease cohort, all participants completed a locally validated neuropsychological test battery, which assessed seven domains: executive function, attention, language, visuomotor speed, verbal memory, and visual memory. Impairment in a particular domain was defined as failing at least half of the individual tests in domain, and failure in individual test determined using ddlucationadjjustee cut-offs of 1.5 standard deviations below established normal means. NC was defined as having impairment in all cognitive domains on the neuropsychological test battery, while MCI was defined as having an impairment in at least one cognitive domain of the neuropsychologica test battery. Detailed descriptions of the neuropsychological assessments and diagnostic criteria described in previous work which will be added upon acceptance.
A.4 Amyloid +/-
Participants from the ADNI cohort were also classified as amyloid positive or amyloid negative, using threshold of global [18F]-Florbetapir amyloid PET SUVR >= 1.11 to define amyloid positivity [57]
B Additional Implementation Details
Optimization for pre-training. The default settings detailed in Table We initialized all transformer blocks using the Xavier uniform method, as described in 19 The pre- pre-training process utilized four A100 GPUs, each with 40GB of memory.
Optimization for downstream tasks. The default settings for end-to-end fine-tuning and linear probing are detailed in Table For fine-tuning, following [19), applied layer-wise Ir decay (58).
15

---

Table 4: Pre-training settings. GAS: Gradient accumulation steps; BS: Batch size
config
value
optimizer
AdamW 660
optimizer momentum
B1,32 0.9.0.999
learning rate schedule
warmup cosine schedule 21
start learning rate
5x10-5
learning rate
1x10-3
final learning rate
1 1010-
weight decay schedule
cosine weight decay schedule
weight decay
0.04
final weight decay
0.4
EMA momentum schedule
linear 21
EMA start momentum
0.996
EMA final momentum
total batch size
GPUs GAS x BS BS
warmup epochs
40
patch size p
dimension of gradient vector
30
training epochs
300
Table 5: Settings of end-to-end fine-tuning and linear probe.
config
value of FT
value of LP
optimizer
Adam
LARS 61]
optimizer momentum
(0.9, 0.999)
0.9
learning rate schedule
cosine decay
19
cosine decay
base learning rate
0.001
0.01
weight decay
0.05
N.A.
layer-wise Ir decay batch size warmup epochs
0.75
N.A.
16
64
0
0
training epochs
50
50
Table 6: Hyperparameters for spatiotemporal masking.
Region
Mask ratio
observation block
Infe | 1.0), (0.84, 1.0))
target
(n (004)(00 0.6), (0.2, 0.6))
target 8
(1/g- (00115, 0.3), (0.0, 0.4))
target
In )= 0.3), (0.0, 0.4))
For linear probing, we incorporated an additional BatchNorm layer
59 before the linear head, per
19
The range ratios for obtaining the observation block and three target blocks introduced in Section 3.2 are presented in Tables 7/ denotes the mask range along the dimension for the block.
Masking,
We set K = which is the number of randomly sampled blocks in three target regions.
C Additional Results
C.1 Results on additional baselines and datasets
We incorpprated more baseline results for downstream tasks on external datasets in Tables 8 including commonly used SVM/SVR (62] and recent self-supervised learning methods. is observed that Brain- JEPA outperforms these models most tasks. We note that for the compared basclines, BrainMass
16

---

29 is a concurrent work. Additionally, CSM 63 and SwiFT [27 are not time series models; CSM utilizes text-like representations, while SwiFT operates on raw fMRI data.
Table 7: Results of additional baselines on HCP-aging.
Age
Sex
Methods
MSE L
Pf
ACC (%) f
Fl (%) f
SVM/SVR
0,586 (.019)
0.699(.022) (.022)
76.67 (1.88)
80.82 (1.15)
BrainMass
0.00())00))
0.831 (.014)
(3.87)
(3.37)
CSM
0.409 (:012)
0.733 (.023)
74.85 (1.11)
76.23 (0.37)
SwiFT
0,341 (.007)
0.755 (.063)
73.48 (2.20)
74.65 (2.32)
Brain-JEPA
0.298 (.01))
0.844 (.030)
81.52 (1.03)
84.26 (0.82)
Table 8: Results of additional baselines on ADNI.
NC/MCI
Amyloid aB ve/
Methods
1%))
FI (%) t
(%)
F1 (5) +
SVM/SVR
73.06 (5.16)
62.00
(5.44)
BrainMass
74.21 (5.10)
(3.36)
(7.40)7748)
69.29 (8.96)
CSM
68.42 (4.99)
(4.54)
63.00
(9.799
SwiFT
(4.46) (5.3))
(6.30)
(7.79.66.38)
Brain-JEPs
76.84 (1.05)
(0.54)
71.00 (4.90)
75.97 (3.93)
To further demonstrate the diversity of our downstream applications, we conducted additional experi- ments using two aging-related public datasets: (OASSS3 and CamCAN for AD conversion prediction in MCI participants and depression diagnosis, respectively. The results are shown in Table 9 By applying Brain-JEPA to five downstream datasets across eight distinct tasks totally, we have demon- strated its versatility in 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.
Table 9: AD conversion prediction and depression diagnosis on OASIS-3 and CamCAN datasets.
OASIS-3
CamCAN
Methods
AD Conversion
Depression
ACC t
F1 (
ACC (%)
FI (%)
SVM/SVR
(2.81)
(1.66)
(3.67))
(2.32)
BrainNetCNN
(2.45
(0.58)
62:73
(4.87)
BrainGNN
(2.00)
(4.34)
63.64 (4.98)
56.68 (3.266
BNT
(8.72)
64.73 (11.29)
65.45 (4.664)
(8.67)
BrainLM
(7.75)
62.67 (9.04)
(6.17)
(3.88
BrainMass
(6.00)
(6.55)
70.91 (2.23)
(3.56)
CSM
(4.90)
(5.49)
64.55 (4.445
56.08 (6.23)
SwiFT
(6.32)
(4.12)
69.09 (6.68)
61.78 (9.26)
Brain-JEPA
(7.35)
(7.92)
67.45 (2.87)
C.2 Scaling properties with respect to dataset size.
We compared the performance of Brain-JEPA trained with varying portions of the UKB pretraining dataset: 25%, 50%. 75% and 100% As shown Table TO, the performance improves as the dataset size increases, highlighting the scalability of Brain-JEPA in relation to the pretraining datasct size.
17

---

Table 10: Ablation on different dataset size for pretraining.
HCP-Aging
ADNI
Percentage
Age
Sex
NC/MCI
P
ACC (%)
ACC (56) 1
0.659 (.043)
68.03 (1.21)
67.89
50%
0.768 (.012)
74.24
(3.05
75%
0.813
77.42 (2.00)
74.74 (4.88)
100%
(.330))03)
(1.03)
76.84
C.3 Additional ablations
Stciitectures/rraewwwrkk. To thoroughly compare the performance between JEPA with anatomical locations (AL) and BrainLM (MAE-based), we extended our comparison to include all the tasks except for the three in the main content, well as two newly added datasets, OASIS-3 and CamCAN. The results shown Table 11 demonstrates that JEPA with AL outperforms BrainLM 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, will further investigate the possible interactions between the self-supervised learning framework and brain gradient positioning.
Table 11: Ablation position embedding.
UKB
HCP HCP-Aging Aging
Methods
Age
Sex
Neurotism
Flanker
p
(CC
pf
p
BrainIM
0.632 (0,020)
86.47 (0.74)
(.23)
0.318(,048)
Brain JEPA w AL.
0,686 (0.013)
84.11 (0.50)
1.267.000
(.274
ADNI
MACC
OASIS-3
CamCAN
Methods
Amy+/-
NC/MCI
Conversion
Depression
ACC (%))
ACC
(CC
AC)
BruinLM
67.00
61.65
65.00
70.00
Bruin-JEPA
65.00
(7.00
We further compared Brain-JEPA without JEPA architecture (i.e., BrainLM with contributions) to Brain-JEPA. As shown in Table 12 Brain-JEPA (JEPA framework) outperforms BrainLM (MAE
framework) with contributions consistently, demonstrating the superiority of JEPA framework.
Table 12: Comparisons of different frameworks.
HCP-Agine
ADNI
Methods
Age
Sex
Amy+/-
et
ACC (%) f
ACC (%) f
BrainLM
0.832 (.028)
74.39 (1.55)
67.00
BrainLM contributions
0.838 (.014)
76.36 (2.58)
70.00
JEPA w contributions
0.844 (003))
81.52 (1.03)
71.00 (4.90)
The number of gradient components. We compared the model performance between 3-dimensional (3-dim) and 30-dim brain gradient positioning, shown Table 13 The 30-dim model consistently outperformed the 3-dim model by large margin. This indicates that higher-dimensional brain gradients may encapsulate finer finer-grained information on brain network organization, which benefits the learning of brain dynamics.
18

---

Table 13: Comparison of different number of gradient components.
HCP-Aging
ADNI
Methods
Age
Sex
Amy+/-
PT
ACC (%)
ACC (%6
3-dim brain gradient
0.819 (.003)
76.96 (1.77)
(6.00)
30-dim brain gradient
0.844 (.030)
(1.52
71.00 (4.90)
D Broader Impact
The introduction of Brain-JEPA marks significant advancement in the interdisciplinaty field of AI and neuroscience, particularly in the brain activity analysis. An assessment of the broader impact of this model has across various dimensions:
D.1 Neuroscience and medical advancements
Brain-JEPA's capabilities in demographic prediction, disease diagnosis, and prognosis could revolution- ize how neurological disorders are diagnosed and treated. This may lead to earlier detection and more personalized therapeutic interventions, potentially improving outcomes for patients with conditions like AD, schizophrenia, or autism spectrum disorders. Furthermore, the model's innovative techniques, including Brain Gradient Positioning and Spatiotemporal Masking, offer new ways to understand the brain's 's functional organization. This could lead to breakthroughs in identifying how various cognitive processes are mapped the brain, aiding both basic science and clinical applications. On the other hand. by effectively predicting various traits, Brain-JEPA can aid in the study of the genetic and environmental influences on behavior and cognitive functions. This can enhance our understanding of the neural underpinnings of psychological traits and disorders.
D.2 Technological impact
Brain-JEPA sets a new standard Al's application to complex brain activity data with a novel brain functional coordinate system and masking strategy, which could spur further innovations and applications Al across different sub-fields of neuroscience. Furthermore, the model's success in performing well across different ethnic groups indicate potential for broad applications in diverse global settings, which is crucial for building inclusive and unbiased Al systems.
D.3 Ethical and social considerations
Ensuring the confidentiality and integrity of patient data while using such advanced Al systems is paramount. While Brain-JEPA has shown superior performance across different tasks, continuous monitoring for potential biases is essential, especially as the model is scaled and deployed in varied clinical settings. Besides, the deployment of advanced technologies like Brain-JEPA could exacerbate existing disparities in healthcare access unless carefully managed. Ensuring that these technologies benefit all segments of the population equally is critical.
19

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope?
Question:
Answer: [Yes]
Justification: Refer to Section I for our contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Refer to Section for the discussion of limitation and future work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full set assumptions and a complete (and correct) proof?
Question:
Answer: [NA]
20

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yss]
Justification: We have provided detailed information about the datasets in Section 4.1 and Appendix [A] The implementation details are presented in Section 442 Appendix The model checkpoints and anonymized code are provided in supplementary material with clear instructions in the readme file.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users). but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
21

---

Question: Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided model checkpoints and code in supplementary material with a readme file. The links to the publicly available datasets are provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https /nips cctpp/ccc/gguiiieeeeeeeee.e http:p//i...c...eeeeeeeeeeettaaa
ssionnolicy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have provided detailed information about the datasets in Section 4.1 and Appendix A We have also provided the implementation details and hyperparameters in
Section 4.2 and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All main results in Table 2 and Eonntain mean and standard deviation from 5 independent runs. The statistical significance of the experiments has been shown with * denoting significant improvement with p 0.05
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
22

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma enror bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We provide the computer resources in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclose whether full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The proposed research conforms with the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there a special considera- tion due laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Refer to Section potential positive socictal impacts and negative socictal impacts of Brain-JEPA.
Guidelines:
The answer NA means that there is no socictal impact of the work performed.
23

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic a algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The original papers that produced the code package dataset are all properly cited.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
24

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, pawpeets.ppprwwiiii. thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged to reach out the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yss]
Justification: The details of the dataset/code/mooel have been provided in 4.1 4.2 and Appendix B with the codes in supplementary materials.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codee/modll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, rememher to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and rescarch with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The datasets of human subjects used in this paper are either publicly available or from previous work which will cited upon acceptance.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The datasets of human subjects used in this paper either publicly available from previous work which will be cited upon acceptance.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
25

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "0v9l1yQxy7", "reviewer": "Reviewer_GJHn", "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.", "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. \n\nPositional 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. \n\nThe 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. \n\nI believe UK biobank uses three Siemens 3T scanners, not one. EDIT: manuscript was amended to reflect this. \n\nCould 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.", "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.\n\nIt 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.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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.\n\nThe 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.\n\nThe 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.\n\nSurprisingly, 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). \n\nMore 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. \n\nThe 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.\n\nYang et al., 2024 BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning \nFerrante 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\nThomas et al., 2023 Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data https://arxiv.org/abs/2206.11417\n\nFinally, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> *Code, data source, and preprocessing.*\n\nThank you for your valuable feedback.\n\nWe 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. \n\nWe utilized open-source preprocessing pipelines as outlined in [1] and [2]. We implemented robust scaling for normalization as practiced in BrainLM. \n\nBy sharing the complete codebase, subject IDs, and using open-source preprocessing pipelines, we believe that our results are highly reproducible. \n\n> *Other related work.*\n\nThank you for your important suggestions.\n\nWe 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.\n\nAlthough 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.\n\nIn 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.\n\n> *More downstream datasets.*\n\nThank you for the valuable suggestion.\n\nNote 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).\n\nTo 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.\n\n> *Temporal resolution.*\n\nWe 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.\n\nTo 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.\n\nFor future work, we plan to use learnable TR embeddings to enable the model to adapt to different TRs dynamically.\n\n> *Clarification on Fig 7.*\n\nBrain 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.\n\n> *Claim in abstract.*\n\nThank 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.\n\n> *UKB has three scanners.*\n\nIn revised version, we will revise \"one scanner\" to \"three scanners\" in the UKB introduction of our revised version.\n\n> *IRB approval.*\n\nWe have obtained all necessary IRB approvals for the datasets used for analysis.\n\nReferences:\n\n[1] Spatial topography of individual-specific cortical networks predicts human cognition, personality, and emotion. Cerebral Cortex 2019.\n\n[2] Global signal regression strengthens association between resting-state functional connectivity and behavior. Neuroimage 2019.\n\n[3] BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning. TMI 2024.\n\n[4] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.\n\n[5] SwiFT: Swin 4D fMRI Transformer. NeurIPS 2023.\n\n[6] Interpretable Graph Neural Networks for Connectome-Based Brain Disorder Analysis. MICCAI 2022.\n\n[7] Beyond the Snapshot: Brain Tokenized Graph Transformer for Longitudinal Brain Functional Connectome Embedding. MICCAI 2023.\n\n[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.\n\n[9] The organization of the human cerebral cortex estimated by intrinsic functional connectivity. J Neurophysiol 2011.\n\n[10] Correspondence of the brain's functional architecture during activation and rest. PNAS 2009.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4pZmDi7Iza", "reviewer": "Reviewer_M77o", "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.", "questions": "1. In lines 71 to 73, I wonder if the author made it italics on purpose or if it’s a formatting mistake.\n2. In line 154, what’s the shape/size of functional connectivity $c_i$ in ROI i?\n4. 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?\n3. 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.", "limitations": "The author has adequately discussed the limitations.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of using fMRI gradient information to guide the position encoding is very interesting and is proven to be effective.\n2. Overall the paper is nicely organized and written.\n3. 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:\n* 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.\n* 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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> *Ablation study on framework.*\n\nThank you for your feedback on our ablation study.\n\nTo 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.\n\n|         |      UKB      |      UKB     |   HCP-Aging  |   HCP-Aging  |     ADNI     |     Asian    |   OASIS-3    |    CamCAN    |\n|---------|:-------------|:------------|:------------|:------------|:------------|:------------|:-------------|:------------|\n|         |      Age      |      Sex     |   Neurotism  |    Flanker   |    Amy+/-    |    NC/MCI    | AD Conversion |  Depression  |\n|         |      $\\rho$ $\\uparrow$    |    ACC (%) $\\uparrow$  |     $\\rho$ $\\uparrow$    |     $\\rho$ $\\uparrow$    |    ACC (%) $\\uparrow$  |    ACC (%) $\\uparrow$  |ACC (%)  $\\uparrow$  |    ACC (%)  $\\uparrow$ |\n| BrainLM | 0.632 (0.020) | **86.47** (0.74) | 0.231 (.012) | 0.318 (.048) | **67.00** (7.48) | 61.65 (3.35) | 65.00 (7.75) | 70.00 (6.17) |\n|    Brain-JEPA *w* AL  | **0.686** (0.013) | 84.11 (0.50) | **0.267** (.003) | **0.374** (.022) | 65.00 (6.32) | **64.33** (1.80) | **67.00** (4.00) | **71.82** (6.03) |\n\nWe 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.\n\n> *Reconstruction of time series.*\n\nThank 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.\n\n> *Italic lines 71-73.*\n\nThank 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.\n\n> *Line 154, the value of $c$.*\n\nThank 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.\n\n> *Downsampling.*\n\nWe 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.\n\nTo 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.\n\nFor 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.\n\n> *Other self-supervised study on fMRI.*\n\nThank you for your suggestion on discussing other self-supervised studies on fMRI.\n\nWe 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.\n\nCSM 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.\n\nIn 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.\n\nReferences:\n\n[1] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.\n\n[2] SwiFT: Swin 4D fMRI Transformer. NeurIPS 2023.\n\n[3] BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning. TMI 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3ApDuKRBLO", "reviewer": "Reviewer_iLSi", "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:\n- they rely on a Joint-Embedding Predictive Architecture, and devise a specific masking strategy for brain data (referred to as spatio-temporal masking)\n- 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\n\nOverall, 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).", "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.\n\nThe 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?\nIn 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$?\n\nI 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).\nMoreover, 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?\n\n**Additional suggestions:**\n\nIn 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.\n\nThe 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.\n\nIt would be nice to indicate how much time the pre-training and fune-tuning phases required (lines 556-560).\n\nIn 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)\n\nThe 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:\n- 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.\n- 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.\n- 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).", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nMoreover, 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.\nHowever, some important pieces of information are missing at this stage, notably concerning how data was processed for downstream tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> *Full open-source code, preprocessing.*\n\nWe appreciate your important suggestions. We have supplemented the following materials/information:\n\n**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.\n\n**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. \n\nBy sharing the complete codebase, the list of subject IDs, and using open-source preprocessing pipelines, we believe that our results are highly reproducible.\n\n> *Data splitting of downstream datasets.*\n\nThank 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.\n\n> *Different TR for different datasets.*\n\nThank 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.\n\nTo 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.\n\nFor future work, we plan to use learnable TR embeddings to enable the model to adapt to different TRs dynamically.\n\n> *Dimension of brain gradient embedding.*\n\nThank 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. \n\nWe 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.\n\n> *Explain to neuroscience community.*\n\nWe 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.\n\n> *Fig7.*\n\nThank 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.\n\n> *Training time.*\n\nThank you for your inquiry regarding the training time.\n\nFor 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.\n\nFor 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.\n\n> *Fig1.*\n\nThank 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}$.\n\n> *Related work.*\n\nThank 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.\n\nReferences:\n\n[1] Spatial topography of individual-specific cortical networks predicts human cognition, personality, and emotion. Cerebral Cortex 2019.\n\n[2] Global signal regression strengthens association between resting-state functional connectivity and behavior. Neuroimage 2019.\n\n[3] Extracting Representations of Cognition across Neuroimaging Studies Improves Brain Decoding. PLoS Computational Biology \n\n[4] Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data. NeurIPS 2022.\n\n[5] Seeing Beyond the Brain: Conditional Diffusion Model with Sparse Masked Modeling for Vision Decoding. CVPR 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fgbz5LMxyR", "reviewer": "Reviewer_fytG", "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.\n\nThe 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.\n\nThe methods achieve strong empirical results, outperforming previous approaches such as BrainLM with a significant margin across evaluation protocols (fine-tuning, linear eval).", "questions": "See weaknesses.", "limitations": "Authors adequately addressed the limitations", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed approach demonstrates strong empirical results compared to previous work in the field.\n\n- 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.\n\n- 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?\n\n- 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?\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> *Pretraining scheme of baselines (dataset and computational budget).*\n\nThank 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).\n\n> *Scaling of dataset size.*\n\nThank you for your suggestion on exploring the impact of dataset size scaling.\n\nWe 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.\n\n|   |  HCP-Aging |  HCP-Aging |   ADNI   |\n|------|:------------:|:------------:|:------------:|\n|   |   Age   |   Sex   |     NC/MCI    |\n|   |   $\\rho$  |  ACC (%)  |     ACC (%)     |\n| 25% | 0.659 (.043) | 68.03 (1.21) | 67.89 (9.18) |\n| 50% | 0.768 (.012) | 74.24 (1.36) | 71.05 (3.86) |\n| 75% | 0.813 (.015) | 77.42 (2.00) | 74.74 (4.88) |\n| 100% | 0.844 (.030) | 81.52 (1.03) | 76.84 (1.05) |\n\n> *Combine contributions with BrainLM.*\n\nThank you for your insightful suggestion.\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and effort in reviewing our work. We appreciate the positive feedback and great interest in our work from all reviewers, along with their insightful questions and suggestions. We are pleased to see that the reviewers acknowledge and appreciate the following aspects:

1. The idea is interesting and effective. The model shows great generalizability, scaling property, with strong empirical results. (Reviewers **GJHn**, **M77o**, **iLSi**, **fytG**)

2. The work is very clear and well constructed. (Reviewers **GJHn**, **M77o**, **iLSi**)

3. Excellent contribution (Reviewer **GJHn**) and presentation (Reviewer **fytG**).

4. Great and clear ablation study. (Reviewers **GJHn**, **fytG**)

5. The code and the checkpoints of the models are shared. (Reviewer **GJHn**)

6. The paper is of interest to the computational neuroscience community. (Reviewer **iLSi**)

We have addressed each reviewer's questions and suggestions point-by-point. This includes adding experiments with more baselines, and evaluating on additional downstream datasets and tasks, as well as clarifying the temporal resolution. Additionally, we responded to questions related to ablation study about the framework's contribution, the gradient dimension, and the scaling of the pretraining dataset size.

We would like to note that for the added baselines, BrainMass [1] is a concurrent work published after our initial submission. Additionally, CSM [2] and SwiFT [3] are not time series models; CSM utilizes text-like representations, while SwiFT operates on raw fMRI data. We thank the reviewers for suggesting these works, and we believe that including comparisons with them strengthens our results. 

Additional results are shown below:

- **Table 1. Results of additional baselines on HCP-Aging.** 

|           |      Age     |      Age     |      Sex     |      Sex     |
|-----------|:------------|:------------|:------------|:------------|
|           |      MSE $\downarrow$     |     $\rho$ $\uparrow$    |    ACC (%) $\uparrow$  |    F1 (%) $\uparrow$   |
|  SVM/SVR  | 0.586 (.019) | 0.699 (.022) | 76.67 (1.88) | 80.82 (1.15) |
| BrainMass | 0.396 (.002) | 0.831 (.014) | 74.09 (3.87) | 75.78 (3.37) |
|    CSM    | 0.409 (.012) | 0.733 (.023) | 74.85 (1.11) | 76.23 (0.37) |
|   SwiFT   | 0.341 (.007) | 0.755 (.063) | 73.48 (2.20) | 74.65 (2.32) |
|     **Brain-JEPA**  | **0.298** (.017) | **0.844** (.030) | **81.52** (1.03) | **84.26** (0.82) |

- **Table 2. Results of additional baselines on ADNI.**

|           |    NC/MCI    |    NC/MCI    |    Amy+/-    |    Amy+/-    |
|-----------|:------------|:------------|:------------|:------------|
|           |    ACC (%) $\uparrow$  |    F1 (%) $\uparrow$   |    ACC (%) $\uparrow$  |    F1 (%) $\uparrow$   |
|  SVM/SVR  | 64.21 (5.16) | 73.06 (4.71) | 62.00 (4.00) | 63.84 (5.44) |
| BrainMass | 74.21 (5.10) | 81.36 (3.56) | 68.00 (7.48) | 69.29 (8.96) |
|    CSM    | 68.42 (4.99) | 76.74 (4.54) | 63.00 (9.80) | 65.89 (9.79) |
|   SwiFT   | 73.16 (5.31) | 80.46 (4.16) | 65.00 (6.32) | 67.79 (6.38) |
|    **Brain-JEPA**  | **76.84** (1.05) | **86.32** (0.54) | **71.00** (4.90) | **75.97** (3.93) |

- **Table 3. Additional tasks of AD conversion prediction and depression classification on OASIS-3 and CamCAN datasets.**

|             |    OASIS-3    |    OASIS-3    |    CamCAN    |    CamCAN    |
|:-----------:|:-------------|:-------------|:------------|:------------|
|             | AD Conversion | AD Conversion |  Depression  |  Depression  |
|             |    ACC (%) $\uparrow$   |     F1 (%) $\uparrow$   |    ACC (%) $\uparrow$  |    F1 (%) $\uparrow$   |
|   SVM/SVR   |  56.00 (2.81) |  52.05 (1.66) | 63.64 (3.07) | 56.79 (2.32) |
| BrainNetCNN |  62.00 (2.45) |  59.53 (0.58) | 62.73 (4.45) | 56.85 (4.47) |
|   BrainGNN  |  59.00 (2.00) |  56.53 (4.34) | 63.64 (4.98) | 56.68 (3.26) |
|     BNT     |  68.00 (8.72) | 64.73 (11.29) | 65.45 (4.64) | 55.32 (8.67) |
|   BrainLM   |  65.00 (7.75) |  62.67 (9.04) | 70.00 (6.17) | 64.18 (3.82) |
|  BrainMass  |  67.00 (6.00) |  66.53 (6.95) | 70.91 (2.23) | 63.56 (2.93) |
|     CSM     |  61.00 (4.90) |  61.97 (5.49) | 64.55 (4.45) | 56.08 (6.23) |
|    SwiFT    |  65.00 (6.32) |  66.80 (4.12) | 69.09 (6.68) | 61.78 (9.26) |
|     **Brain-JEPA**   |  **69.00** (7.35) |  **67.32** (7.92) | **72.73** (2.87) | **67.45** (1.57) |

- **Table 4. Comparisons of different frameworks.**

|             |   HCP-Aging  |   HCP-Aging  |     ADNI     |
|-------------|:------------|:------------|:------------|
|             |      Age     |      Sex     |    Amy+/-    |
|             |     $\rho$ $\uparrow$    |    ACC (%) $\uparrow$  |    ACC (%) $\uparrow$  |
| BrainLM     | 0.832 (.028) | 74.39 (1.55) | 67.00 (7.48) |
| BrainLM *w* contributions  | 0.838 (0.014) |  76.36 (2.58) | 70.00 (11.40) |
| JEPA *w* contributions    | **0.844** (.030) |  **81.52** (1.03) | **71.00** (4.90) |

- **Table 5. Comparison of different number of gradient components. ‘bg’ means brain gradient.**

|             |   HCP-Aging  |  HCP-Aging  |     ADNI     |
|-------------|:------------|:-----------|:------------|
|             |      Age     |     Sex     |    Amy+/-    |
|             |     $\rho$ $\uparrow$    |   ACC (%) $\uparrow$  |    ACC (%) $\uparrow$   |
| 3-dim bg  | 0.819 (0.003) | 76.96 (1.77) |  67.00 (6.00) |
| 30-dim bg | **0.844** (.030) | **81.52** (1.03) | **71.00** (4.90) |

The main results on additional baselines and datasets can also be found in the attached PDF.

If reviewers have any further questions or concerns, please let us know. We are happy to engage in further discussion.

References:

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

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

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

Meta-review: This submission generated much interest and discussion. The reviewers appreciated the contributed foundation model for fMRI time series, with interesting modeling developments and solid evaluation for several applications such as demographic prediction, disease diagnosis/prognosis, and trait prediction.
Decision: Accept (spotlight)