XAIguiFormer: explainable artificial intelligence guided transformer for brain disorder identification

We sincerely appreciate your constructive suggestion to thoroughly evaluate the effectiveness of our XAI-guided attention mechanism. We have conducted additional experiments to quantitatively assess the relationship between attention entropy and model performance.

[W1] Figures are too far to the corresponding descriptive text, e.g., fig.2 in page 6 is described in page 9. I recommend moving Figure 2 closer to its description or adding a forward reference earlier in the text.

R1: We apologize for the inconvenience caused by the placement of Figure 2. We have relocated Figure 2 to page 7, positioning it closer to the corresponding descriptive text to improve readability. Additionally, we have added a forward reference earlier in the text to ensure that readers can easily locate the figure if needed.

[W2] The attention map by the proposed XAI has lower entropy, which is called an improvement in line 463. It confused readers to an accuracy improvement according to the proposed XAI attention map. Although authors were selling the story under such motivation, the relationship between entropy and accuracy are missed. Can you provide entropy-accuracy curves or other quantitative evidence demonstrating how the lower entropy attention maps directly contribute to improved model performance? I think this the key results to support the proposal in this paper, I will raise the score if these results are fit with the expectation.

R2: Thank you for highlighting this important concern regarding the relationship between attention entropy and model performance. To address this issue, we conducted additional experiments to quantitatively evaluate the relationship between attention entropy and Balanced Accuracy (BAC) under various scenarios. We analyzed the entropy-BAC pairs from (1) different warm-up XAI models and (2) different training epochs within the same XAIguiFormer model. The correlation coefficients between attention entropy and BAC in these scenarios were found to be -0.7 and -0.75, respectively. This negative correlation indicates that lower attention entropy (more concentrated attention) is consistently associated with improved performance. The detailed entropy-BAC curves are included in Appendix B of the revised manuscript.

On the other hand, a recently popular work [1], Differential Transformer, supports the hypothesis that sparse/concentrated (lower-entropy) attention patterns can improve performance. Differential Transformer argues that conventional attention mechanisms in Transformers tend to overallocate attention to irrelevant contexts. By introducing a differential approach to amplify attention toward relevant contexts while suppressing noise, the model encourages sparse and concentrated attention patterns, resulting in improved performance. Similarly, the XAI-guided attention in our method compels the Transformer to concentrate on sparse and relevant patterns through the XAI approach, thereby reducing attention entropy and filtering out irrelevant information.

Reference:

[1] Ye, Tianzhu, et al. Differential transformer. Submit to ICLR 2025.

[W2].a: Given results of the additional experiments shown in Fig.5, the proposed method is empirically proved. But I have to suggest the authors arrange Fig. 5 into the main text since there is enough white space in your manuscript, where detailed methodology steps can be kept in Appendix. Furthermore, it looks like a logistic regression (e.g., seaborn.regplot) can fit better than the linear regression authors currently used. I recommend that you improve the presentation accordingly.

[W2].b: Listing the evidence shown in previous works to support your motivation is effective in improving the soundness of your proposed methods. Instead of an ICLR submission, there are published peer-reviewed works you can review to do this, e.g., [1] has solid evidence to support sparse attention can improve accuracy in both empirical and theoretical aspects. I recommend you include it in the Introduction to enhance the motivation and proposal of this paper.

[1] NeuroPath: A Neural Pathway Transformer for Joining the Dots of Human Connectomes. NeurIPS, 2024, https://openreview.net/forum?id=AvBuK8Ezrg

Authors have revised their work with additional empirical evidence to support the proposal in this paper via [W2].a. They also have shown an aspiration of side evidence via [W2].b. My concern in [W3] cannot be resolved by quantitative evidence at this point. Conclusively, I'd like to raise my score if authors can arrange their responses into the manuscript.

We appreciate the opportunity to clarify and expand some experiments to improve our model. Please find below our point-by-point responses to your comments.

[W1] Authors should discuss further on why they chose to model EEG as connectome rather than time series, given some state-of-the-art EEG foundation models with time series as input, such as [2], especially it focuses on frequency domain as well. Experimental comparisons with these EEG models are also missed in the paper.

R1: Thank you for your insightful comment. There are two primary reasons for employing the connectome as input. First, compared to time series data, functional connectivity offers a superior advantage in modeling the interactions between channels or brain regions, as the brain is a complex communication and information processing system [1]. Additionally, we constructed the connectome by aggregating two distinct types of functional connectivity that provide complementary information, thereby providing multi-view information. Second, when developing XAIguiFormer, we carefully considered the additional computational burden introduced by the XAI method. Due to the nature of EEG multi-channel data, patching multi-variable time series into tokens results in a relatively long sequence length, which increases the computational demands of the XAI module. By utilizing the connectome as input and implementing a specially designed connectome tokenizer, we were able to reduce the sequence length, thereby enhancing computational efficiency while preserving essential information.

In our comparative experiments, we evaluated our approach against the S3T model and the pretrained BIOT model, both of which utilize EEG time series data. Following the reviewer’s suggestion, we also included the time series model LaBraM [2] as a baseline and conducted the comparison experiment.

[W2] It seems like only temporal frequency is considered in positional encoding, how about spatial frequency for different brain regions? Some work like [3] developed spatial connectome based positioning, which should at least be discussed in the paper.

R2: Thank you for highlighting this important aspect and referencing the relevant work. The Brain Gradient Positioning method proposed by Brain-JEPA incorporates both temporal and spatial information into tokens by introducing functional connectivity gradients. This approach captures the functional relationships among brain regions and integrates them with temporal information from fMRI time series segments to encode positional information. The input to the transformer in Brain-JEPA consists of time series segments from multiple brain regions, thereby creating a “two-dimensional” structure. One dimension represents the spatial and functional relationships among regions, while the other dimension is the temporal information of the time series segments. In this context, both spatial and temporal information are essential intrinsic characteristics of fMRI data.

In contrast, the input tokens of the transformer in XAIguiFormer are generated from the connectome constructed in the frequency domain. These tokens encapsulate the spatial relationships among the channels through the connectome structure and connectome tokenizer. Unlike multi-channel or multi-region time-series tokens, the frequency band connectome tokens in our method do not inherently contain spatial relationships. Consequently, spatial information is less critical in our framework compared to frequency and demographic information, which are explicitly prioritized to enhance model performance. Additionally, we include and discuss the Brain-JEPA as related work in our paper.

Reference:

[1] Seguin C, Sporns O, Zalesky A. Brain network communication: concepts, models and applications[J]. Nature reviews neuroscience, 2023, 24(9): 557-574.

[2] Jiang W, Zhao L, Lu B. Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. ICLR 2024.

[3] Dong Z, Ruilin L, Wu Y, et al. Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking. NeurIPS 2024.

We appreciate your suggestion and believe that this additional information enhances the clarity and comprehensiveness of XAIguiFormer.

[W1] Section 5.1 provides a brief description of the datasets; however, given the importance of demographic information to the proposed method, the authors should expand on this aspect in that section. Additionally, it would be helpful to know if demographics were considered when creating the train/test/validation split. How was the data split? Are the different diseases balanced on both datasets? Furthermore, what is the distribution of patients with brain disorders versus healthy individuals across splits?

R1: Thank you for pointing out this important consideration, as demographic information plays a crucial role in our proposed dRoFE. For the TUAB dataset, we followed the official train and eval splits. We further split the official training set randomly into our train and val sets, while using the official eval set as our test set. For TDBRAIN, we randomly split the entire dataset into train/val/test sets. As shown in these tables, the distributions of age, gender, and brain disorder status are balanced across the splits, thereby minimizing bias and ensuring fair model evaluation.

[Q1] Line 369: The authors mention that the BAC, AUC-PR are the average performance and standard deviation across five different random seeds on the TUAB and TDBRAIN datasets. I am assuming that the models were re-trained 5 times using different train/val splits while test data was kept constant, could the authors clarify if this is case? I am guessing that the reason why the models were re-trained 5 times is due to computational limitations, could the authors elaborate on the computational costs of their method compared to the other baseline methods?

R2: Unlike k-fold cross-validation, we employed a hold-out strategy to split the datasets into train, val, and test sets, where all splits remain constant across experiments. The reason for training the model five times is to assess the stability of the model under different random seeds. This stability reflects the model's robustness to variations in random weight initialization and data batch ordering, both of which are influenced by the random seed.

When developing XAIguiFormer, we carefully considered the additional computational burden caused by the introduction of the XAI method. Therefore, we employed the connectome as input and designed a specialized connectome tokenizer to reduce the sequence length, thereby improving computational efficiency while preserving essential information. XAIguiFormer demonstrates better efficiency (lower FLOPs) than transformer-based baselines, such as LaBraM and BIOT, despite the inclusion of the additional computation cost from the XAI algorithm.

[Q2] In the text line 335 the authors mention: “XAIguiFormer is not sensitive to \alpha as long as it is larger than 0.3”. Could the authors report those results and discuss why they believe that there is a lower threshold but not an upper threshold?

R3: Thank you for raising this point. We have included a detailed relationship curve between $\alpha$ and BAC in Appendix F, where $\alpha$ varies from 0.1 to 0.9 in increments of 0.1. The experimental results indicate that when $\alpha$ is lower than 0.3, the performance declines significantly. This is likely because, with a low $\alpha$, the explainer’s guidance becomes insufficient to meaningfully influence the learning process, leading to degraded performance. On the other hand, our results show no upper threshold led to a significant performance drop. This could be because larger $\alpha$ values allow the explainer's guidance to dominate the attention mechanism without overwhelming the model's ability to learn effectively.

We sincerely appreciate your thorough review and valuable comments, which have significantly contributed to enhancing the robustness of XAIguiFormer and clarifying our description of warmup XAI. Please find below our point-by-point responses to your comments.

[W1] Robustness is unclear. How about the performance after changing to other explainers in the module in addition to using DeepLift? How about the robustness?

R1: Thank you for your insightful comment regarding the robustness of our method with respect to different explainers. To address this, we conducted additional experiments by replacing the original DeepLift explainer with two widely used alternatives: GradCAM and Integrated Gradients. Our results show that the replacement of DeepLift does not result in a significant performance degradation. Interestingly, we observed a slight improvement in the overall performance when GradCAM was utilized instead of DeepLift. These findings highlight that XAIguiFormer is robust and performs stably across various explainers, indicating that its effectiveness is not overly dependent on a specific XAI algorithm.

[W2] How to calculate the frequency band importance? Is the result from the explainer inside the proposed model? If so, can it get the same conclusion after changing the explainer?

R2: The importance of the frequency band is derived from the explainer within the XAIguiFormer. To assess the robustness of the calculated frequency band importance, we replaced the original explainer (DeepLift) with GradCAM and Integrated Gradients in XAIguiFormer and extracted the frequency band importance for comparison. Figures 6 and 7 in Appendix E illustrate the frequency band importance generated by GradCAM and Integrated Gradients, respectively. The rankings of the theta/beta ratio, high and low $\alpha$ on TUAB remain consistent across different explainers. Similarly, on the TDBRAIN dataset, the rankings of the theta/beta ratio, low $\alpha$ and high $\beta$ are largely consistent, where low $\alpha$ ranks fourth in importance identified by Integrated Gradients. These results demonstrate that the ranking of the most important frequency bands remains stable, indicating strong robustness across different XAI algorithms.

[W3] In Fig 3, why does the accuracy of warm-started XAI keep decreasing after training? The accuracy at the beginning is the best one, so how does it prove the effectiveness of the proposed method? It becomes worse after training the proposed one.

R3: We apologize for any confusion and are pleased to provide clarification. Figure 3 is not a training curve for a single model utilizing warm-started XAI. Instead, it represents a line chart demonstrating the relationship between different warm-started XAI models and their balanced accuracy (BAC) after training. Each point on the line, from left to right, corresponds to a different model configuration, where XAI is activated at 0% (normal XAIguiFormer), 10%, 20%, 30%, and 100% (without XAI) of the total epochs during the course of training process. The y-axis represents the final BAC for each model upon completion of training. Since Figure 3 presents independent final results for each configuration, it does not imply that the accuracy of a single model decreases over the course of training. The results depicted in Figure 3 demonstrate that models with warm-started XAI guidance consistently outperform the vanilla transformer, thereby validating the effectiveness of XAI in enhancing model performance.

[Q1] see above. The questions are about the description of warmup start XAI guidance and the robustness of the explainer module in the proposed method.

R4: We hope our responses above have addressed reviewer’s concerns on the robustness of the explainer module and clarified the concepts of warmup start XAI guidance.