Paper: BrainLM: A foundation model for brain activity recordings (ICLR-2024)

Reviewer comment (Evaluation):
In the section on related work you mention some models that were trained on fMRI to narrowly predict just one clinical variable, and, presumably, those models are state of the art in their respective predictions of that clinical variable. However, in Table 3 you seem to compare BrainLM against some other results. The question is whether the comparison presented in the table reflects the state of the art achieved by other models in predicting certain clinical variables, and not just state of the art using a certain method. The fact that an LSTM or an SVM performs poorly is not that interesting to see; what would be interesting is how much better or worse performance we get when compared to the state of the art method specifically trained to predict, for example, GAD-7. The same applies for other clinical scores, as a method does not have to predict all of them at once. (If, in fact, the methods you use in this comparison are the current state of the art, please clarify and this criticism will be withdrawn.)

Author response:
To clarify, the referenced MAE models trained on fMRI are only focused on stimuli-decoding based on visual areas. None of the models have been used for clinical variable prediction. In the literature, common approaches for clinical variable prediction involve the use of SVM and MLPs to correlation matrix input [1]. These are the models we have included in Table 2 (previous Table 3).

[1] Vieira, Sandra, Walter HL Pinaya, and Andrea Mechelli. “Using deep learning to investigate the neuroimaging correlates of psychiatric and neurological disorders: Methods and applications.” Neuroscience & Biobehavioral Reviews 74 (2017): 58-75.

Reviewer comment (Experiments):
For Table 3, are the "competitor" results obtained by you training these methods on the same data, or are these the numbers reported in the respective studies by the authors of those methods?

Author response:
We obtained the performance of competing methods ourselves. The methods designed for clinical variable prediction were trained on other datasets, preventing us from incorporating their performance directly. Furthermore, we requested access to their code for comparison but received no response. As a result, we constructed our own baselines based on the information available in their manuscript. We did extensive hyperparameter search for each method for fairness and presented the best performance of each method in our analysis.

Reviewer comment (Experiments):
In Table 1, it appears that allowing the model to see more context (less masking) yields significantly better performance. (a) Have you tried to train the model with even less masking (MR=0.1) and what were the results? (b) If the trend of "less masking = better performance" continues, then what would be the limit of this—at what point would using less masking start to hurt generalizability?

Author response:
Regarding generalization, the models with lower masking ratio (and thus more information) have an easier prediction task. As such, their in-distribution performance is higher. However, when tested with other masking ratios, their performance decreases significantly. As a result, for all subsequent comparisons, we opted to use the 75% masking model, which exhibits the best generalization among all models.

Reviewer comment (Experiments):
On page 5, last paragraph: When fine-tuning BrainLM to regress a metadata variable, did you predict all of the listed variables (age, neuroticism, PTSD, anxiety) at once as 4-head output, or was it a separate fine-tuning process per variable?

Author response:
All models were fine-tuned separately. We will clarify this in the manuscript.

Reviewer comment (Writing):
In Figure 4, the "CLS" abbreviation is not explained anywhere in the text.

Author response:
Thank you for pointing this out. The CLS token is a reference to a “CLASS” token, which is used for classification/regression after pretraining. This token effectively acts as an embedding which captures the whole dynamics. We have clarified this in the manuscript.

Reviewer comment (Experiments):
The network is very modestly sized for a foundation model, with 13M parameters—smaller than the MLP, LatentODE, and LSTM baselines, and much smaller than what would be considered even a small foundation model, such as HuBERT-Base with 90M parameters. While BrainLM is trained on a nontrivial amount of fMRI data using SSL, it is remarkably small for a foundation model. This would be acceptable if performance were good relative to baselines chosen based on prior results on fMRI decoding, but this does not appear to be the case.

Author response:
We appreciate the reviewer's perspective on the comparative scale of our model, especially when juxtaposed with the larger foundation models prevalent in other domains. Taking this constructive feedback into account, we have expanded our research scope and developed enhanced versions of our model, now featuring 111 million and 650 million parameters, respectively. This represents a substantial increase, up to 50 times the parameter count of our initial model.

The effectiveness of these larger models is evident in the revised Tables 1, 2 and Supp. Table 17, as well as in Figure 4 of our manuscript, which collectively demonstrate a significant boost in performance attributable to the increased parameter count. However, we were not able to finish training the 650M parameter model in Table 1 due to time constraints. In a future version of the manuscript, this model will be included.

In line with our commitment to collaborative research and open science, we plan to make all variants of these models accessible to the broader research community. By sharing these models, we aim to foster further exploration and innovation in the field, inviting other researchers to build upon, refine, or even challenge our findings.

Reviewer comment (Experiments):
The baselines are not a priori motivated in a way that would support the contributed model as truly SOTA. The NODE, LatentODE, and GCN citations point back to original papers without neuroscience applications, meaning the baselines themselves are novel, so it is unclear how impactful beating them should be.

Author response:
We acknowledge the reviewer's concern regarding the challenge of evaluating results in clinical variable prediction. We conducted a thorough literature review and found that the use of Support Vector Machines (SVM) or Multi-Layer Perceptrons (MLPs) trained on correlation matrices remains the predominant approach in the field (see references [1-3]). This prevalence influenced our decision to include these methods in our study.

Moreover, the absence of more recent methodologies in existing literature led us to develop our own baseline models for comparison. We ensured fairness in evaluation by allocating an equal amount of hyperparameter optimization effort to all models presented in our study.

We are open to suggestions from the reviewer regarding the inclusion of additional methods. If there are specific, newer approaches that the reviewer believes would enrich our comparative analysis, we are more than willing to incorporate these into our study. Our commitment is to provide a comprehensive and robust benchmarking that reflects the current state of the art in the field.

[1] Woo, Choong-Wan, Luke J. Chang, Martin A. Lindquist, and Tor D. Wager. “Building better biomarkers: brain models in translational neuroimaging.” Nature Neuroscience 20, no. 3 (2017): 365-377.
[2] Vieira, Sandra, Walter HL Pinaya, and Andrea Mechelli. “Using deep learning to investigate the neuroimaging correlates of psychiatric and neurological disorders: Methods and applications.” Neuroscience & Biobehavioral Reviews 74 (2017): 58-75.
[3] Orru, Graziella, William Pettersson-Yeo, Andre F. Marquand, Giuseppe Sartori, and Andrea Mechelli. “Using support vector machine to identify imaging biomarkers of neurological and psychiatric disease: a critical review.” Neuroscience & Biobehavioral Reviews 36, no. 4 (2012): 1140-1152.

Reviewer comment (Presentation):
The legend for Fig. 5 lists LatentODE, LSTM, and NODE, whereas the text talks about ODENets and non-pretrained BrainLM. The non-pretrained BrainLM results should be added or the text corrected. Additionally, the text could make a clearer mapping from legend to citations (e.g., "NODE (Chen et al.), LatentODE (Rubanova et al.)"). It is also surprising that there is an unpretrained BrainLM baseline in Table 4 but not Table 3.

Author response:
We thank the reviewer for pointing out the inconsistency between the legend and the text. We have addressed the issue. We have also added the citations to the methods used in the table. Regarding the absence of a non-pretrained BrainLM model, due to the time constraints associated with this deadline, our efforts were primarily concentrated on experiments pertaining to model scaling. Consequently, we were unable to train the BrainLM 13M model from scratch for clinical variable regression within the stipulated time frame. However, we recognize the importance of this aspect and plan to include it in a future iteration of the manuscript. This addition will undoubtedly provide a more comprehensive understanding of the model's capabilities and limitations.

Reviewer comment (Experiments):
Something strange seems to be going on with NODE training—it performs much worse than other methods and produces a large negative $R^2$ value in Table 4, likely indicating severe model mismatch and raising concern that training may have failed entirely for this model.

Author response:
We agree with the reviewer's observation regarding the problematic nature of the results when employing the Neural ODE (NODE) model on the fMRI dataset. During the training phase, we shared similar concerns about NODE's performance with this specific type of data. In response, we undertook extensive experimentation to enhance its effectiveness.

Despite our rigorous attempts, we concluded that NODE's inherent structural design might not be optimally configured for the unique properties of our fMRI data. We speculate that the substantial dimensionality coupled with the intrinsic noise prevalent in fMRI datasets poses formidable challenges for NODE, limiting its capability to accurately represent and analyze such intricate data. 
We remain open to suggestions and would eagerly incorporate any modifications or alternative approaches recommended by the reviewer to enhance our methodological framework.

Reviewer comment (Evaluation):
Was Figure 3 randomly sampled or cherry-picked?

Author response:
The example showcased in Figure 3 was selected through a random sampling process. This approach was adopted to ensure the representativeness of the data presented. The robustness of our model is reflected in the high R2 values obtained in the prediction tasks for both UK Biobank (UKB) and Human Connectome Project (HCP) datasets. These values are indicative of a strong fit and suggest that a similarly high level of accuracy would be observed across other randomly sampled recordings.

Reviewer comment (Experiments):
Why are the 1% and 100% training results of interest? Should we not expect that more data improves performance? If it is surprising, why not show a trend over more than just two data sizes?

Author response:
The new Figure 4 and tables in your manuscript effectively demonstrate the scalability of the BrainLM models. By including models trained on varying data sizes (0.1%, 1%, 10%, 50%, and 100%) and showcasing models with different parameter sizes, we provide a comprehensive view of how model performance scales with data size and complexity. This approach not only reinforces the notion that larger models yield better performance but also offers valuable insights into the efficiency and effectiveness of our BrainLM models across different training scenarios.

Reviewer comment (Evaluation):
For statistics on Fig. 5, how were they done? It is hard to imagine BrainLM beating LSTM and LatentODE at later timepoints given the visualized error bars. Was it something like BrainLM vs. the average of the rest, such that the effect of NODE dominated?

Author response:
For each comparison between BrainLM and the models LSTM, LatentODE, and NODE, statistical significance was calculated. We've adopted the convention of marking an asterisk in our figures whenever BrainLM's performance significantly surpasses all three models in these pairwise comparisons. This methodology, along with a detailed explanation, will be included in the figure's description to ensure a comprehensive understanding of the comparative performance metrics.

Reviewer comment (Experiments):
What motivated the specific network architectures chosen? They vary substantially in number of parameters.

Author response:
Our objective was to encompass a broad spectrum of potential models, with a specific focus on those that have demonstrated efficacy in modeling spatiotemporal dynamics, including LatentODE, GCN, NODE, and LSTM. Given the novelty of applying these methods to fMRI data, as we previously noted, it was imperative for us to establish our own benchmarks for baseline comparison. This approach not only facilitates a comprehensive understanding of each model's capabilities in this new context but also lays the groundwork for future explorations in the field.