Paper: Learning transferrable and interpretable representation for brain network (ICLR-2024)

Reviewer comment (Experiments):
The proposed BrainMAE is based on self-supervised masked auto-encoding, so it needs to be compared with the existing self-supervised learning-based models, including the following:  
[1] Shi, Chenwei, et al. "Self-supervised pretraining improves the performance of classification of task functional magnetic resonance imaging." Frontiers in Neuroscience 17 (2023).  
[2] Malkiel, Itzik, et al. "Self-supervised transformers for fMRI representation." International Conference on Medical Imaging with Deep Learning. PMLR, 2022.  
[3] Thomas, Armin, Christopher Ré, and Russell Poldrack. "Self-supervised learning of brain dynamics from broad neuroimaging data." Advances in Neural Information Processing Systems 35 (2022): 21255-21269.

Author response:
Thank you for raising this concern. Following your suggestion, we added a downstream experiment in $\textbf{Appendix F}$ of our revised submission using the HCP Task Dataset, same to the setup in [1]. The models are trained to make classification between 20 mental states using fMRI signals, with each state lasting tens of seconds — a notable difference from the downstream tasks in section 3.3. Despite this, $\textbf{Tables 10 and 11}$ show our models outperforming other state-of-the-art, self-supervised learning models designed for transient state modeling. Further details are available in $\textbf{Appendix F}$.
    
Additionally, we added $\textbf{Appendix G.3}$ to compare our model's performance in gender prediction with [2], as shown in $\textbf{Table 15}$. Our model outperforms TTF by a significant margin.
    
We haven't included a comparison with [3] due to the unavailability of their code. However, we would be happy to compare with them and following their exact experiment setup for our model evaluation if you have any additional concerns.

Reviewer comment (Experiments):
Since the word embedding includes information about its meaning, the word embeddings in NLP can be used in all sentence positions. However, the position of ROIs is never changed in training sessions, which could limit the model's ability to learn ROI traits but could learn only absolute position information. It can be considered as learnable positional encoding. Since the graph attention mechanism is permutation-invariant, the changes in ROI order should not affect the performance if the ROI embeddings have their characteristics. Even though the paper shows the evaluation of pretrained ROI embeddings in Figure 3, it still needs additional evidence that these embeddings do not incorporate absolute position information.

Author response:
Apologies for the confusion. To clarify, in our implementation of embedding-informed graph attention, we use a transformer block, where the ROI embeddings are linearly transformed using different learnable weight matrices for the key and query. This design makes the adjacency matrix inherently asymmetric, enabling it to capture the asymmetric information flow between ROIs. The introduction of these learnable weight matrices means that the attention mechanism $\textbf{isn't permutation invariant}$. We have thoroughly revised Section 2.1 in our revised submission to better explain this aspect.

Reviewer comment (Evaluation):
One of the main reasons for learning and exploiting the ROI embeddings was to mitigate and circumvent the inherent noise in the fMRI signal. However, from the experimental results, the authors concluded that the lower performance of DG-BrainMAE than that of SG-BrainMAE was due to such inherent noise. These are conflicting.

Author response:
1. We politely disagree "the author concluded that the lower performance of DG-BrainMAE than that of SG-BrainMAE was due to such inherent noise" because we do not make such conclusion. This misunderstanding might stem from section 3.3.3 discussing the performance of baseline dynamic FC models versus fixed FC models.
    
2. Our ablation study in section 3.3.4 demonstrates the effectiveness of ROI embedding-informed attention. Specifically, we found that DG-BrainMAE outperforms FDG-BrainMAE in the downstream tasks.
    
3. The DG-BrainMAE, introduced as a complement to SG-BrainMAE, focuses on the dynamic reconfiguration of ROI connectivity. Our experiments show that DG-BrainMAE excels SG-BrainMAE in identifying transient mental state fluctuations (Appendix F of our revised submission).

Reviewer comment (Experiments):
The paper doesn't use "traditional" ML models (eg, SVM, random forests) directly on the flatten upper-triangle of the FC matrix as baselines. Based on my experience when running new deep learning models on FC connectivity, what I've found is that a simple traditional ML model (with a reasonable simple hyperparameter search) achieves comparable, if not better, metrics than many DL-based models on some datasets. Therefore, to evaluate the utility of this model, it would be important to see how BrainMAE compares to these more "traditional" ML models.

Author response:
Thank you for raising this point. To address this concern, we analyzed the performance on the three downstream tasks with three classification models (SVM, Logistic Regression and Random Forrests), and three regression models (ordinary linear regression, ridge regression, Elastic Net) with details in $\textbf{Appendix G.4}$ and $\textbf{Tables 16-19}$ of our latest revised submission. 

Indeed, traditional ML models demonstrate performance levels comparable to those of baseline deep learning methods in Gender classification tasks. This observation can be attributed to the well-established understanding that FC matrix harbors significant information of human traits. However, for more complex regression tasks, such as task performance prediction on NSD dataset, which necessitate inferring intricate brain states from the dynamic of fMRI signals, simple ML models are less effective. In such scenarios, deep learning methods, endowed with the robust capability for representation learning, are able to achieve superior results. Nevertheless, our models continuously outperform these baseline methods by significant margin.

Reviewer comment (Evaluation):
The paper is missing important metrics (e.g., sensitivity and specificity) in the evaluation section of the (binary) gender prediction. Given there's an appendix, and well-known limitations with accuracy and AUROC, I don't think it is a subjective comment to request these two other metrics in a paper that has a clear connection with the medical domain.

Author response:
Thank you for pointing this out. We acknowledge the importance of these metrics and have added them (sensitivity, specificity, and F1-score) to $\textbf{Tables 16-19}$ in the appendix of our revised submission. With these additional metrics, our method continues to significantly outperform other baseline models.

Reviewer comment (Presentation):
Page 5 includes a part that is confusing in terms of readability. We have figure 2 mentioning SG-, DG-, and FDG- models, but right below in "BrainMAE Variants" only two are mentioned. Technically, they are not connected in the flow of the text, but visually they are together in the paper.

Author response:
Thank you for bringing this to our attention. To address this confusion, we will rename the FDG-BrainMAE to 'vanilla-BrainMAE' in our final revision. This reflects its use of TSE without embedding-informed graph attention, aligning with the typical use of 'vanilla' as model for comparison, which should improve clarity.

Reviewer comment (Novelty):
Was there a particular reason for the authors to choose masked AEs in favour of other AEs? From the Introduction section it doesn't seem clear to me why this was the case. If the reason is just that they haven't been explored before in this field, I think it's a reasonable motivation (and results show it might have been a good choice). However, it would be good to say it why this was preferred to, for example, (V)AEs or other encoding variants as motivation for this new method.

Author response:
Our motivation stems from the idea that masking a large set of ROIs and predicting their signals from unmasked ROIs can enhance the model's understanding of ROI relationships and improve ROI embedding learning. This approach is a key component of our method and has shown robust ROI representation learning. Additionally, in $\textbf{section 3.3.4's Ablation Study}$, we compared FDG-BrainMAE and FDG-BrainAE. Despite having the same architecture, they differ in pretraining methods: masked autoencoding for FDG-BrainMAE and standard autoencoding for FDG-BrainAE. Our results show that FDG-BrainMAE consistently surpasses FDG-BrainAE in all three downstream tasks by large margin, highlighting its effectiveness.

Reviewer comment (Presentation):
The paper does a very good job explaining the different components of the model with regards to how it learns/creates the unsupervised brain representation. However, the paper gets a bit confusing (and even a bit unexpected) when in section 2.1 it is said in the "Autoencoder" component that the decoder is only used in pre-training phase (why is there such a division?); then, on section 3.1 when it is said that behaviour/demographic measurements are used (for what?); then, section 3.2.3 once again mentions that somehow we have a pretrained model for all datasets (why so much computational complexity?). Only in section 3.3.1 it is explained that the authors appended a task-specific linear head based on previous studies, and tables 1/2/etc also make it obvious that there are classification and regression tasks. For all these reasons, I hope it's clear why figure 1 (and consequently section 2) needs a quick mention to this task-specific head for better clarity/readability.

Author response:
Sorry for the confusion. To clarify, the decoder referenced in the 'Autoencoder' section is solely employed for reconstructing fMRI signals during the pretraining phase. It is not utilized in the finetuning of downstream tasks. As detailed in Appendix A.2 and shown in Figure 5, we integrate a task-specific linear head atop the CLS token from the encoder for downstream task finetuning.
> Q9: Where are behavior / demographic measurements used?

For clarification, the behavior and demographic measurements are exclusively used in the downstream fine-tuning tasks on the HCP-3T dataset. These measurements are not incorporated into the pretraining phase of the models.

> Q10: Why to have pre-trained models for all datasets?

Thank you for your questions. To clarify, in Section 3.2.3, SG-BrainMAE is pretrained individually on each dataset, rather than collectively on all datasets. This approach was chosen to assess if the ROI embeddings, when pretrained separately on each dataset, would lead to convergent representations. As demonstrated in Figure 3 and detailed in Appendix D, we indeed observed convergence in the ROI embeddings.

Reviewer comment (Experiments):
Why haven't the authors included the task-specific fMRI data from HCP, but decided to include the task-specific fMRI data from the NSD dataset?

Author response:
The NSD dataset contains direct measurement of memory task accuracy and behavioral arousal measurement such as response time. These metrics are ideal for testing whether the model is able to capture essential brain states from fluctuation of fMRI signals. 
    
In $\textbf{Appendix F}$ of our revised submission, we newly added a downstream task using the HCP task dataset, where the models are required to classify 20 mental states, with each representing distinct type of task and with duration of tens-of-seconds. Our model surpasses current state-of-the-art methods that are specifically designed for transient fMRI modeling[1].

Reviewer comment (Writing):
Can the authors clarify what is the difference between the DG- and FDG-BrainMAE models? The FDG-BrainMAE is introduced as the model but without the "static graph mechanism". Isn't this basically the same as the DG-BrainMAE? My guess is that what the authors mean is that instead of the static graph mechanism in *the first layer* of DG-TSE, the FDG-BrainMAE's TSE module used only the dynamic graph representation across all blocks from the very beginning?

Author response:
Your interpretation is absolutely right! The key distinction between DG-BrainMAE and FDG-BrainMAE lies in the use of the embedding-informed attention module. DG-BrainMAE incorporates this module in its first layer, whereas in FDG-BrainMAE, all layers within the TSE utilize self-attention.

Reviewer comment (Experiments):
How are the folds created in the 5-fold cross-validation procedure? Is it in a stratified fashion independently for each downstream task? Are the authors careful to include subjects never seen before in training completely separated in the test sets?

Author response:
Our 5-fold cross-validation follows the standard approach, randomly assigning samples for testing in each fold, with the remaining divided into training and validation. In the three downstream tasks, we didn’t use stratified cross-validation nor separating traning and testing set based on subject identity.
    
However, we recognize this concern and in our newly added HCP task for mental state classification (detailed in $\textbf{Appendix F}$), we conducted cross-validation based on subjects, ensuring no overlap between training and testing groups. Our proposed models consistently outperformed baseline methods.

Reviewer comment (Evaluation):
Do the authors have any hypothesis on why the initial and final task blocks were particularly important for the model to make the final predictions?

Author response:
Thank you for raising this question. We hypothesize that the attention score might mirror specific brain states, such as transient arousal states critical for performing memory tasks. In $\textbf{section 3.3.5}$ of our revised paper, we explored this by analyzing the relationship between averaged behavioral brain arousal measurements (inverse response time) and averaged attention scores. We noticed that both metrics display similar patterns within and across task blocks, with high values at the start and end of these blocks. This pattern could indicate that our model inherently captures changes in transient arousal states, which may be vital for predicting memory task accuracy. Furthermore, we examined the correlation between fMRI representations and a well-defined brain arousal index, finding a strong correlation (details in $\textbf{Appendix E.2.2}$).

Reviewer comment (Novelty):
While the BrainMAE model is presented as a fresh approach to fMRI data, it's hard to overlook the fact that all Transformer Layers, Masked Autoencoders (MAE) and Dynamic Graph Neural Networks are somewhat antiquated techniques. These methods are not only dated in the broader machine learning landscape but have also lost their novelty in computational neuroscience applications [1,2,3]. Frankly, the proposed method comes across as a 'engineering model'.

Author response:
Thank you for raising this concern. ML methods are indeed not new in computational neuroscience, but there are inherent challenges with fMRI, a critical tool for brain research. Unlike other signals such as EEG, fMRI's unique temporal resolution and spatial coverage require distinct modeling approaches.

Many existing ML models for fMRI, including the two papers you cited, primarily compute functional connectivity matrices and often overlook the importance of dynamic brain state changes. These dynamics carry rich information, proved to be a major contribution to the FC, and relates to a variety of brain function, showing promise for understanding how the brain actually works[1].
    
On the other hand, deep learning methods that consider fMRI's dynamic nature frequently underperform compared to fixed-FC approaches, partly due to fMRI's inherent noise. Moreover, several dynamic-FC approaches employ extended time windows (30-60 seconds) to manage noise and spurious connections, but this compromises their sensitivity to brief, transient brain state changes that typically happen within 10-20 seconds[1][2].
    
Our model addresses these challenges by introducing a novel ROI embedding-informed attention mechanism and a uniquely designed MAE training scheme. This approach not only overcomes the limitations seen in previous dynamic-FC models but also consistently outperforms state-of-the-art methods. The significance of these new components is evident in our ablation study, where their removal results in a substantial performance decrease. 

Because of these specifically designed features, our model holds promise for use in neuroscience and medical research. It's particularly suited for unraveling how transient fMRI fluctuations contribute to overall brain representations in various circumstances, aiding in the study of brain traits, states, and diseases.

Reviewer comment (Theory):
No theoretical guarantees of the proposed method. The BrainMAE seems to be too heuristic.

Author response:
We acknowledge that, at this stage, our approach is more heuristic and grounded in empirical insights from neuroscience and machine learning. The absence of formal theoretical guarantees is a valid point, and we view it as an important area for our future research. Our current focus has been on leveraging intuitive understandings of brain functionality and the nature of fMRI signals to develop a model that shows promising empirical results. We aim to complement these findings with more theoretical analysis in our subsequent work, and we believe this will strengthen the method's scientific contribution.

Reviewer comment (Theory):
You mention your method is 'interpretable' in the title. However, I don't figure out the true interpretability of the method since BrainMAE has no scientific inductive bias incorporated. This is just a stacking of deep-learning-based techniques, which make the model highly black-boxed. Please refer to the Questions Section for my further concerns.

Author response:
Thank you for highlighting this point. We realize our initial submission lacked clarity on this aspect, and we've addressed it in our revised submission with detailed explanations of the scientific inductive biases we've incorporated:
    
1. ROI embedding (section 2.1): Each ROI has its unique inherent characteristics and we resemble it as the words in the NLP context by assigning each ROI with learnable vector. Section 3.2.3 demonstrates that these ROI embeddings are informative and consistently converge across different datasets.
    
2. ROI embedding-informed attention module (section 2.1): 
    
* As brain ROI activities are intrinsically governed by both structural and functional networks, we consider building adjacency matrix based on ROI embeddings to represent such functional networks. 
    
* Considering the asymmetric flow of information between ROIs, we apply two separate learnable linear transformations to the embeddings before attention computation, introducing asymmetry in the adjacency matrix.
    
* Based on the synchronized nature of fMRI signals among functionally relevant ROIs, we eliminate self-loops in the attention mechanism. This design helps prevent bias towards individual nodes and enhances feature learning by reducing susceptibility to input noise, similar to a group voting process.
    
3. Masked Autoencoding (section 2.2)
    
* We segment the fMRI using 15-second time window aligning with the typical duration of transient events in fMRI data.
    
* As functionally relevant brain regions show synchronization in performing certain brain functions, we mask large portion of fMRI ROI signals and optimize the model to reconstruct the masked signal from unmasked signal, which is useful to discover underlying functional relevance between ROIs thus improving the learning of ROI embeddings.

Reviewer comment (Evaluation):
For your Subsection 3.3.5, Representation and Interpretation Analysis. I hold the point that analyzing the representations of fMRI using principal component analysis (PCA) and self-attention scores are too post hoc and lacks scientific insights. In fact, the representations extracted by BrainMAE and attention layers are still black-box.

Author response:
Thanks for pointing out. To address this concern, we performed three additional representation analyses.

1. In Appendix E.1, we apply the pretrained SG-BrainMAE on the never-seen NSD dataset and extract the representation for each fMRI run. Surprisingly, in tSNE analysis, we found the subject identity is clearly separated from each other, suggesting the pretrained model is able to discern unseen subject identity purely from their fMRI signals.

2. In Appendix E.2, same as the above setup, within each subject, in tSNE analysis, the fMRI type (task/rest) can be discernable for most of the subjects, suggesting the pretrained model's ability to recognize different brain states without specifically trained to.

3. In Appendix D.3, we categorized subjects from the HCP-Aging dataset into three age groups: young, middle, and old. We pre-train SG-BrainMAE individually for each group and study the age effects on the learned ROI embeddings. The modular structure of the network constructed based on the embeddings, quantified by modularity, indicates a reduction with age increasing, aligning with the established findings in neuroscience research. 

Additionally, regarding the self-attention analysis in section 3.3.5 (Figure 4), we conducted further assessments.  Interestingly, we found the attention score correlates to the behavioral arousal measurements, the inverse reaction time, and across- and within-task block, suggesting the model is aware of the change of transient brain state.

Reviewer comment (Writing):
The writing is generally good, but the flow of the paper can be improved as certain parts can be hard to follow. For example, the notations in Section 2 are not listed clearly enough: "node features" are described in equation (3) before the actual input fMRI segment X is formally introduced in Section 2.2.

Author response:
Thank you for pointing this out. In our revised submission, for the definition of the node features, we added sentence as follows: "Within this graph, each ROI is considered a node, with its node feature represented by the ROI's time-series signal of length $\tau$ or $x\in R^{\tau}$".

Reviewer comment (Experiments):
Some of the statements may not be very well supported, mostly regarding the dynamic graph attention.

Author response:
Sorry for the confusion and thank you for sharing your concern. We recognize the potential misunderstanding about 'dynamic graph attention,' which is simply the self-attention used in transformer models and not a major novelty of our work. To improve readability, we've omitted this description in our revised submission. Additionally, we've extensively revised Section 2.1 to more clearly convey our motivation, the novelty, and the inductive biases considered in our model's development. Feel free to suggest more comments about them and we would be happy to address them further.

Reviewer comment (Presentation):
What exactly is shown in Figure 3A? Is this a representative t-SNE plot for one subject? Or did the authors take the mean ROI embeddings across all subjects and plot t-SNE afterwards? Or something else?

Author response:
Apologies for the confusion. The ROI embeddings are not derived from a single subject or an average across subjects. They are, in fact, model weights learned globally using data from all subjects in the dataset, akin to word embeddings in NLP. These pretrained ROI embeddings represent the rich ROI-related information "summarized" from all the training samples.