BrainMoE: Cognition Joint Embedding via Mixture-of-Expert Towards Robust Brain Foundation Model

W1

Thank you for the comment. Our evaluations are exactly cross-subject. Each fMRI in the five clinical datasets corresponds to one unique subject, and the evaluations are subject-wise (L.191-L.206). Since two non-clinical datasets are for cognitive state recognition, there are multiple cognitive states per subject, but the train/test split in our multi-fold cross-validation is still subject-level (L.207).

W2

The running time of BrainMoE has fulfilled the real-time requirement for fMRI applications, which takes a much longer time (10-15mins) than BrainMoE (133-287ms) for data acquisition using a common 3T scanner [1]. Additionally, BrainMoE has been implemented with the pre-computation of cognitive embeddings from complex experts to amortize the high computational costs from multiple brain experts. For instance, BrainJEPA experts have high FLOPs (5.32T) and compute cognitive embeddings prior to BrainMoE finetuning, while BrainMass experts with much lower FLOPs (26.96M) have the cognitive embeddings computed on-the-fly. This trick enhanced efficiency as listed in Table 5, where BrainJEPA with the trillion-level FLOPs costs less time than BrainMass with the million level FLOPs. We will include the discussion of this possible amortized efficiency pruning in the “Limitation” section of the final version.

[1] Kumar, V. A., et al. "Recommended resting-state fMRI acquisition and preprocessing steps for preoperative mapping of language and motor and visual areas in adult and pediatric patients with brain tumors and epilepsy." American Journal of Neuroradiology 45.2 (2024): 139-148.

Q1

Parameters are increased $N$ times compared to the previous single expert model, where $N$ is the number of experts. However, increasing the parameter amount by fusing multiple experts naively, e.g., a late fusion MoE [2], sometimes cannot boost the robustness of brain modeling (80.70 -> 73.33 for AD prediction). The table below shows the seven downstream applications using baseline (BrainMass), the late fusion MoE, and our BrainMoE with the parameter amount of each model. BrainMoE, in contrast, consistently has the best performance.

Note that late fusion is a strategy for handling multimodal data where the final output is the weighted sum of predictions from each independent expert selected by the router. Unlike BrainMoE has the cognition adapter to implicitly utilize multiple expert embeddings via Transformer, the late fusion directly combining expert task embeddings suffers from a data unbalance issue for different domain experts (Table A for Reviewer ZhRL). Therefore, the performance boost of BrainMoE is not only advanced by more parameters but also by our new ML component, the cognition adapter.

[2] Han, Xing, et al. "Fusemoe: Mixture-of-experts transformers for fleximodal fusion." Advances in Neural Information Processing Systems 37 (2024): 67850-67900.

Q2

BrainMoE approach uses cognitive embeddings so that the cognition adapter can learn from the latent space, which is the information bottleneck between data space and task space (Fig. 1). The latent space is a lower-dimensional representation of the input data, is designed to capture the most essential and relevant information from the original data while discarding less important details [3]. The task embeddings, on the other hand, capture the characteristics of a specific downstream application instead of various cognition-related features in the latent space by domain experts, while either clinical or non-clinical brain modeling is highly relevant to cognitive function or dysfunction. The less robust performance of late fusion MoE further provides empirical evidence to support that cognitive embeddings are better than task embeddings.

[3] Wang, Shiye, et al. "Self-supervised information bottleneck for deep multi-view subspace clustering." IEEE Transactions on Image Processing 32 (2023): 1555-1567.

Q3

The scalability of MLP has been analyzed in Appendix Figure 8. Curves of the train loss at the 200th epoch by MLP are increasing with more parameters, using more layers and hidden channels in 3 datasets, demonstrating the bad scalability in brain modeling. In contrast, our Transformer-based adapter maintains a decreasing train loss at the 200th epoch when there are more parameters. Note that lightweight convolutional heads are specific for image data, while FC is structured data, and BOLD is a timeseries data, so that convolutional heads cannot be adapted to our cases.

Q4

No, there are no observable identical or conflicting cognitive embeddings from different brain experts in the observation of the correlation between cognitive embeddings. The table below is the mean value of Pearson’s correlation between BrainMass experts specified for 12 different cognitive states, where the maximum correlation is 0.51 [Language, Working memory] and the minimum is -0.47 [Vismotor, Rest]. The absolute value of correlation is mostly less than 0.5, indicating experts are not in conflict or redundant to each other.

Q5

Yes, it could. We selected a multimodal (fMRI+EEG) visual stimuli dataset [4] for the evaluation. There are 388 fMRI-EEG pairs from 22 healthy subjects under 8 cognitive states (age$\in$[23,51], F:M=1:1). The EEG expert is implemented with the pre-trained CBraMod [5]. Since the preprocessed fMRI in this dataset was using the Shaefer atlas and only Cog. Classif. experts in our work have been pretrained with Shaefer400, BrainMoE was evaluated on this FMRIEEG dataset with one CBraMod expert and 12 Cog. Classif. experts. This experiment is a subject-level 5-fold cross-validation. Results of cognitive state classification, sex classification, and age regression are shown in the tables below. Better performance by multimodal BrainMoE further indicates the value of our approach.

The table below shows the cognitive state classification performance of BrainMoE, where 'Prec' is precision, and 'Rec' is recall.

cognitive state

The table below shows the sex prediction performance of BrainMoE.

sex

The table below shows the age regression performance of BrainMoE, where 'Abs. Corr.' is the absolute Pearson's correlation between ground truths and predictions, and 'MSE' is the mean squared error.

age

[4] Telesford, Qawi K., et al. "An open-access dataset of naturalistic viewing using simultaneous EEG-fMRI." Scientific Data 10.1 (2023): 554.

[5] J. Wang et al., ‘CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding’, in The Thirteenth International Conference on Learning Representations, 2025.

Paper format 1

We appreciate your suggestion, and we are committed to thoroughly revising the manuscript. Specifically, we will highlight the innovation of BrainMoE and provide more comprehensive supporting results as shown in this rebuttal.

Paper format 2

We will revise the font color to dark green.

Thank you for the detailed responses. I agree with most points and as in my original review I think this paper should be accepted. Thanks for clarifying the cross-subject evaluation also.

Brief

Thank you for your insightful comments and questions. The main concern of your reviews is derived from the misunderstanding of our methodological novelty and the lack of evaluations other than classification. We added results from a generic MoE to show the novelty of our BrainMoE approach, as well as regressive and generative applications.

W1

Along with the agreement of our novelty by reviewer Yxzq, one key novelty is innovating the BrainMoE approach that can significantly enhance the robustness of brain foundation model by pre-training with tasking-state fMRI, which was not fully explored in previous studies. Due to various sample sizes of pre-training datasets (Table A for Reviewer ZhRL), the data unbalance issues of MoE are amplified in the brain modeling area, as indicated by the performance of a generic MoE framework, the late fusion MoE [4], in Table A. To address this limitation, our BrainMoE presents innovative learning strategies for our new cognitive MoE adapter, to tackle the data unbalance issue.

Effectiveness of BrainMoE: We can see late fusion MoE can sometimes outperform the baseline, but it is consistently worse than BrainMoE except for the smallest dataset Taowu (n=40). Unlike BrainMoE has the cognition adapter to implicitly utilize expert embeddings, the late fusion explicitly combines expert predictions, therefore cannot handle the unbalance issue.

Table A

Regarding top expert combinations (top 1 if the probability > 0.9, otherwise top 2) chosen by the router, it favors the top activated experts rather than the domain experts across all datasets.

As listed in Table B, the combination from late fusion consistently contains Rest (n=29,971), Hariri (n=24,672), and Relational (n=1,609) experts. In contrast, BrainMoE has diverse and similar top combinations for heterogeneous and homogeneous applications, respectively (Table C). The contrast of selected combinatory patterns between routers trained by BrainMoE approach and the late fusion provides neuroscientific insight and supports our methodological novelty.

Table B

Table C

W2

The challenge of heterogeneity from downstream tasks among our 7 datasets is as critical as previous brain foundation models (Appendix Table 6). The number of different downstream datasets is the highest, i.e., 5 cohorts, compared to the most recent brain foundation models. In addition, downstream tasks are also more variable than others since there are four types of brain disorders (AD, ASD, PD, SZ) and two types of non-clinical tasks (sex, cognitive state). We will add the variations of our downstream tasks to Table 6.

W3, Lim.2

We appreciate your suggestion, and we are committed to thoroughly revising the manuscript. Specifically, we will highlight the innovation of BrainMoE and provide more comprehensive supporting results as shown in this rebuttal.

Q1

BOLD signal is framed in a spatial-temporal data structure. One novelty of BrainMoE is modeling the brain with arbitrary input type (see L.68 point #2). Input of brain experts is denoted by X (L.106), which has variable temporal dimensions. Arbitrary brain expert maps X to Z (L.137) to have the same dimension $C_{hid}$ for BrainMoE modules. Although in practice, input shapes $C_{in}=M$ for BrainMass experts and $C_{in}=T$ for BrainJEPA experts are different, their cognitive embeddings both have the same shape $C_{hid}$. Therefore, BrainMoE can work with various input shapes. We will revise the method section with some examples of input shape to smooth the description of BrainMoE handling variable input shape.

Q2

The motivation behind this is considering the variations of brain expert architecture to the extent of supervision type. The utilization of binary supervision complements the gap between standard brain classification models, e.g., transformer-based [3], and the brain foundation models. Empirical evidence in our work shows that binary supervision can benefit the performance of BrainMoE on clinical (Cog. Classif. in Table 3) and non-clinical applications (All-in-one in Table 4), indicating the potential of standard brain classification architectures for brain foundation modeling. We will clarify this motivation in the Introduction section and discuss the results in the Conclusion section.

Q3

We illustrated the overall framework of BrainMoE and the learning strategy in Fig 3, including pre-training and fine-tuning. Fig 4 illustrates the novel blocks in our framework, particularly the brain expert and the cognition adapter. Fig 4a shows the pre-training of brain experts corresponding to Fig 3a instead of b, the downstream tasks are therefore not involved in this step. On the other hand, Fig 4b shows the fine-tuning for downstream, and it corresponds to Fig 3b. We will combine Fig 3 and Fig 4 as one figure to make the method description clear.

Lim.1

Age regression for 6 datasets has been evaluated for the best baseline, 68k version of BrainMass, and BrainMoE. We can see the performance of BrainMoE is the best. The improvement is especially significant for ABIDE (n=1,025, 6-58 yrs) with MSE 36.77 -> 4.86. This empirical evidence further supports the robustness of BrainMoE.

Table D

Visual decoding task for a new dataset NSD [1] has also been evaluated for MindEye2 [2] as the baseline and BrainMoE. We pre-trained two MindEye2s as the specific experts for long-term (novel trials) and short-term memory (easy/hard trials), respectively. Since visual decoding is a generative task, output contains much higher dimensions (256x1664 vs. class number 2-7) than downstream tasks focused on our study. Therefore, we skipped our cognition adapter by weighted summing the diffusion prior of two experts with the BrainMoE routing probabilities. Both baseline and BrainMoE are pretrained with subjects 2-7 and finetuned with subject 1 on the entire 40 sessions. The final train and test losses, cosine similarity, and Mean Squared Error (MSE) during finetuning are listed in Table E. Given the evidence that the performance of BrainMoE is better than the single expert MindEye2, there is potential for BrainMoE to expand to visual decoding.

Table E

[1] 10.1038/s41593-021-00962-x

[2] 10.48550/arXiv.2403.11207

[3] 10.48550/arXiv.2210.06681

[4] 10.48550/arXiv.2402.03226

W1

We appreciate your suggestion, and we are committed to thoroughly revising the method section. Specifically, we will highlight the innovation of BrainMoE and connect Figs 3 and 4 with a smooth description as in this rebuttal: We illustrated the overall framework of BrainMoE in Fig. 3, including pre-training and fine-tuning, respectively. Fig. 4 illustrates the architecture of novel blocks shown in Fig. 3, particularly the brain expert and the cognition adapter. Figure 4 (a) shows the pre-training of brain experts corresponding to the individual expert in Figure 3 (a). On the other hand, Fig. 4 (b) shows the fine-tuning for downstream, and it corresponds to Fig. 3 (b), where the cognition adapter is the main block in the fine-tuning stage. We will combine Fig. 3 and Fig. 4 as one figure to clarify the method description.

W2

The motivation behind this is considering the variations of brain expert architecture to the extent of supervision type. The utilization of classification loss complements the gap between standard brain classification models, e.g., transformer-based [3][4], and the brain foundation models. Empirical evidence in our work shows that classification loss derived from standard brain modeling can benefit the performance of BrainMoE on clinical (Cog. Classif. in Table 3) and non-clinical applications (All-in-one in Table 4), indicating the potential of standard brain classification architectures for brain foundation modeling. We will clarify this motivation in the Introduction section and discuss the results in the Conclusion section.

W3

Here, we collect the expert routing probabilities of BrainMoE on clinical datasets and count the combination of top experts (top 1 if the probability > 0.9, otherwise top 2). We can observe the most frequent combination (Table C for reviewer PTHj):

• AD: [Working memory] 31.1%, [Rest, Motor] 20%
• ASD: [Language, Relational (relationships between stimuli)] 14.9%, [Rest, Relational] 8.3%
• PD: [Language, Relational] 20%, [Vismotor (color change and button press), Relational] 9.1%
• SZ: [Gambling, Relational] 18.5%, [Facename (face and name recall)] 14.3%

The relatively more frequent activation of these expert combinations, e.g., [Language, Relational] for ASD and PD, suggests that BrainMoE captures complementary neural patterns that individual experts might miss.

On the other hand, combinations for different applications align with neuroscientific knowledge:

• For AD, the dominance of [Working memory] aligns with that AD is commonly attributed to central executive impairment and frontal lobe dysfunction [1].
• For ASD, [Language, Relational] mirrors impairments in relational memory processing, where individuals with ASD show difficulties in forming relations between items and encoding relational but not item information [2], along with the language network showing atypical functional connectivity [3].
• For PD, [Vismotor, Relational] reflects the visuomotor processing deficits, including impaired decision-making cascades and altered cortical sensorimotor processing [4].
• For SZ, those combinations reflect deficits in reward learning and decision-making processes, as well as impaired facial processing and recognition abilities that are characteristic of the disorder [5].

[1] 10.1016/j.cortex.2010.12.002

[2] 10.1002/aur.1493

[3] 10.1002/aur.2171

[4] 10.3390/brainsci13081173

[5] 10.3389/fneur.2019.00990

Q1

Sorry for the confusing description of the expert number. The number of experts (N) is derived from the number of distinct cognitive states of arbitrary input type (FC or BOLD) required by an arbitrary expert model architecture. Expert number is not strictly equal to the number of cognitive tasks but $N_{arch}\times N_{cog}$, where $N_{arch}$ is the number of available model architecture and $N_{cog}$ is the number of distinct cognitive states. In practice, there are three types of expert architecture (BrainMass, BrainJEPA, and Cog. Classif.) implemented in our BrainMoE All-in-one, so that $3\times 12=36$. We will revise the description in L.113-L.120 to clarify the expert number.

Q2

HCPYA has an age range 22 to 35 including young adults. But yes, given ABIDE has a younger population (age < 22), the performance of BrainMoE on ABIDE (ASD) is impacted as the worst among binary classifications (Table 1).

Q3

No. UKB has 2: Resting and task-Hariri. HCPA has 4: Resting, task-Carit, task-Facename, and task-Vismotor. HCPYA has 7: Motor, gambling, language, social, relational, emotion, and working memory. In total, there are 12 different cognitive states.

Q4

It is a fair comparison on HCPA and HCPYA datasets since the same cohorts of data were used for different models in the pre-training stage. Baselines and BrainMoE both have the same 68k pre-training data with the same number of training iterations. Although there are multiple experts in BrainMoE, each expert only sees a specific cognitive state, and each sample only has one cognitive state. Thus, BrainMoE has no additional pre-training epochs or iterations than baseline models, resulting in a fair comparison.

Lim.1

We appreciate your suggestion, and we are committed to thoroughly revising the ‘Limitations’ section with the discussion of confounding effects from study/scanner difference.

Brief

Thank you for your insightful comments and questions. The main concern of your reviews stands on the lack of evaluation per individual expert and late fusion. We have demonstrated all evaluations in our rebuttal. However, the evaluation per individual expert is lengthy and uncommon in previous brain modeling and MoE studies. We will keep them brief in the final version.

W1

Since resting-state fMRI was mainly used in previous brain foundation models (see L.40 in the manuscript), we included only resting-state experts as baselines. Here, we evaluated the standard baselines, a single expert per tasking-state, and reran the experiments of Tables 1 and 2 in the manuscript. Tables A and B show results for BrainMass-MLP baseline.

Baseline results show that there are consistently existing task-specific experts (e.g., language for AD, working memory for PD) outperforming the resting-state baselines (30k version in the manuscript), confirming that task-state fMRI contains valuable information for brain modeling. Conclusively, BrainMoE still holds the best performance compared to all datasets and baselines. This supports that the utilization of cognitive embeddings from BrainMoE leads to more robustness of brain modeling than naively training a single task of fMRI. Furthermore, the expert combinations learned by BrainMoE and their strong relevance to downstream applications, shown in W2 response, such as Language for ASD, and Working memory for PD, offer valuable neuroscientific insights that current baseline models fail to capture.

The late fusion of independent experts is also tested across 7 datasets. Please refer the results to Table A for Reviewer PTHj. The superior performance of BrainMoE and the contrast of selected expert routing patterns between BrainMoE and the late fusion provide neuroscientific insight and support our methodological novelty.

Table A

Table B

W2, Q3

Here, we collect the expert routing probabilities of BrainMoE on clinical datasets and count the combination of top experts (top 1 if the probability > 0.9, otherwise top 2). We can observe the most frequent combination (Table C for reviewer PTHj):

• AD: [Working memory] 31.1%, [Rest, Motor] 20%
• ASD: [Language, Relational (relationships between stimuli)] 14.9%, [Rest, Relational] 8.3%
• PD: [Language, Relational] 20%, [Vismotor (color change and button press), Relational] 9.1%
• SZ: [Gambling, Relational] 18.5%, [Facename (face and name recall)] 14.3%

The relatively more frequent activation of these expert combinations, e.g., [Language, Relational] for ASD and PD, suggests that BrainMoE captures complementary neural patterns that individual experts might miss.

On the other hand, combinations for different applications align with neuroscientific knowledge:

• For AD, the dominance of [Working memory] aligns with that AD is commonly attributed to central executive impairment and frontal lobe dysfunction [1].
• For ASD, [Language, Relational] mirrors impairments in relational memory processing, where individuals with ASD show difficulties in forming relations between items and encoding relational but not item information [2], along with the language network showing atypical functional connectivity [3].
• For PD, [Vismotor, Relational] reflects the visuomotor processing deficits, including impaired decision-making cascades and altered cortical sensorimotor processing [4].
• For SZ, those combinations reflect deficits in reward learning and decision-making processes, as well as impaired facial processing and recognition abilities that are characteristic of the disorder [5].

[1] 10.1016/j.cortex.2010.12.002

[2] 10.1002/aur.1493

[3] 10.1002/aur.2171

[4] 10.3390/brainsci13081173

[5] 10.3389/fneur.2019.00990

W3

BrainMoE outperformed all baselines from the largest (n=4,863) to the smallest dataset (n=40), demonstrating the superior robustness. On the other hand, the running time of BrainMoE has fulfilled the real-time requirement for fMRI applications, which takes a much longer time (10-15mins) than BrainMoE (133-287ms) for data acquisition using a common 3T scanner [6].

The possible amortized efficiency gains can be the pre-computation of cognitive embeddings. In fact, our experiments for BrainMass experts (FLOPs: 26.96M) have the cognitive embeddings computed on-the-fly, while BrainJEPA (FLOPs: 5.32T) computes them prior to the finetuning. Since the cognitive embeddings are frozen and fixed per data, pre-computation of the complex experts, BrainJEPA in our case, can save the time of iterative computation during BrainMoE finetuning. This trick enhanced efficiency as listed in Table 5, where BrainJEPA with the trillion-level FLOPs costs less time than BrainMass with the million level FLOPs. We will include the discussion of this possible amortized efficiency pruning in the “Limitation” section of the final version.

[6] 10.3174/ajnr.A8067

Q1

The number of experts (N) equals $N_{arch}\times N_{cog}$, where $N_{arch}$ is the number of available model architectures and $N_{cog}$ is the number of distinct cognitive states. E.g., there are three types of architecture (BrainMass, BrainJEPA, and Cog. Classif.) implemented in our BrainMoE All-in-one, so that $3\times 12=36$. We will revise the description in L.113-L.120 to clarify the expert number.

Q2

Yes, the design of BrainMoE cognition adapter can address the unbalanced expert.

As shown in Fig. 4 (b) and Eq. (3), the final prediction of BrainMoE is not the explicit fusion of expert embeddings but the transformation of input FC or BOLD weighted by the cross-attention between expert embeddings. This implicit usage of experts avoids their overfitted or underfitted cognition embeddings to directly bias the output of BrainMoE. Thus, it addressed unbalanced brain experts supported by empirical evidence:

The pattern comparison (Table B vs. Table C for Reviewer PTHj) between the router trained by our BrainMoE approach and the generic late fusion MoE. Late fusion favors the top-activated experts rather than the domain experts across all datasets. As in Table B for Reviewer PTHj, the top combination among 7 datasets consistently contains Rest (n=29,971), Hariri (n=24,672), and Relational (n=1,609) experts. In contrast, as in Table C for Reviewer PTHj, BrainMoE has diverse top combinations for heterogeneous downstream tasks and similar top combinations for homogeneous applications, aligning with neuroscientific knowledge

Thank you for your response and recognizing our work! Regarding to the last concern about the ease of stratification of data, the standard BIDS format [1] of neuroimaging data strictly requires the filename to have the cognitive state during imaging, i.e., BIDS format filename is like: {subject-name}_{session-name}_{task-name}_{run-number}_{more-info}.nii.gz. This makes the data stratification in our approach can be easily extended to other standard neuroimaing cohorts with BIDS format.

[1] Gorgolewski, Krzysztof J., et al. "The brain imaging data structure, a format for organizing and describing outputs of neuroimaging experiments." Scientific data 3.1 (2016): 1-9.