XAIguiFormer: explainable artificial intelligence guided transformer for brain disorder identification
Authors: Hanning Guo
Farah Abdellatif
Yu Fu
N. Jon Shah
Abigail Morrison
Jürgen Dammers
Venue: ICLR 2025 Poster (2025)
Keywords: EEG
Explainable Artificial Intelligence (XAI)
Explanation-Guided Learning
Transformer
Primary area: applications to neuroscience & cognitive science
TL;DR: The proposed XAIguiFormer is the first framework to employ XAI for enhancing transformer performance in neuroimaging data, extending its benefits beyond mere interpretability analysis.

Abstract: EEG-based connectomes offer a low-cost and portable method to identify brain disorders using deep learning. With the growing interest in model interpretability and transparency, explainable artificial intelligence (XAI) is widely applied to understand the decision of deep learning models. However, most research focuses solely on interpretability analysis based on the insights from XAI, overlooking XAI’s potential to improve model performance. To bridge this gap, we propose a dynamical-system-inspired architecture, XAI guided transformer (XAIguiFormer), where XAI not only provides explanations but also contributes to enhancing the transformer by refining the originally coarse information in self-attention mechanism to capture more relevant dependency relationships. In order not to damage the connectome’s topological structure, the connectome tokenizer treats the single-band graphs as atomic tokens to generate a sequence in the frequency domain. To address the limitations of conventional positional encoding in understanding the frequency and mitigating the individual differences, we integrate frequency and demographic information into tokens via a rotation matrix, resulting in a richly informative representation. Our experiment demonstrates that XAIguiFormer achieves superior performance over all baseline models. In addition, XAIguiFormer provides valuable interpretability through visualization of the frequency band importance. Our code is available at https://github.com/HanningGuo/XAIguiFormer.

Published as a conference paper at ICLR 2025
XAIGUIFORMER:
EXPLAINABLE ARTIFICIAL INTELLIGENCE GUIDED TRANSFORMER FOR BRAIN DISORDER IDENTIFICATION
Hanning Guol-?, Farah Abdellatif Yu Fu', Jon. N. Shah 4.6,,6 Abigail Morrison
Jürgen Dammers
Institute of Neuroscience and Medicine (INM-4), Forschungszentrum Jülich, Germany
2 Department of Computer Science 3 Software Engineering, RWTH Aachen University, Germany
3 School of Information Science and Engineering, Lanzhou University, China
4 Department of Neurology, RWTH Aachen University, Germany
5 SERRA-BRIN--AE -ranslitaanlliionall Medicine, Germany
Institute of Neuroscience and Medicine (INM-11), JARA, Forschungszentrum Jülich, Germany
7 Institute of Advanced Simulation (IAS-6), Forschungszentruum Jülich, Germany
8 Department of Psychiatry, Psychotherapy and Psychosomatics, RWTH Aachen University, Germany
(h. guo, j d.damerss J fz- ffzzjuell fafffjjjliccheeee ich.de
ABSTRACT
EEG-based connectomes offer a low-cost and portable method to identify brain disorders using deep learning. With the growing interest model interpretabil- ity and transparency, explainable artificial intelligence (XAI) is widely applied understand the decision of deep learning models. However, most research fo- cuses solely on interpretability analysis based on the insights from XAI, over- looking XAI's potential improve model performance. To bridge this gap. we propose dlynamica-systte-innspie (XAlguiFormer), where XAI not only provides explanations but also contributes enhancing the transformer by refining the originally coarse information in self- attention mechanism to capture more relevant dependency relationships. der not to damage the connectome's topological structure, the connectome to- kenizer treats the single-band graphs as atomic tokens to generate a sequence in the frequency domain. To address the limitations of conventional positional encoding in understanding the frequency and mitigating the individual differ- ences, integrate frequency and demographic information into tokens via a rotation matrix, resulting in a richly informative representation. Our experi- ment demonstrates that XAIguiFormer achieves superior performance over all baseline models. In addition, X XAlguiFormer provides valuable interpretability through visualization of the frequency band importance. Our code is available http://jihuu..oo///aannnnno/o///
architecture, XAI guided transformer
or-
INTRODUCTION
Scalp electroencephaloggapyy (EEG), an objective physiological method for detecting the human brain's neural activity, has been widely used in various applications, including disease diagnosis (Yang al., 2024), seizure detection (Tang et al., 2022), action and motor imagery recognition Venkatachalam et al.) 2020), and emotional analysis Yi et al. 2024). These non-invasive signals are collected by electrodes placed on the curved surface of the scalp, resulting in inherent non- Euclidean structures. The so-called connectome is widely used metric characterize the complex patterns and connectivity of the brain this non-standard geometry (Pang et al.) 2022). For instance, NeuroPath (Wei et al. (2024) is graph transformer designed to model human connectomes and learn the relationship between neural pathways and brain functions. Additionally, EEG contains rich frequency domain information and often analyzed over specific frequency bands. Consequently, multi-frequency band graphs are a powerful way to identify brain disorders using EEG signals.
"Corresponding author

---

Published as a conference paper at ICLR 2025
In the past few years, transformer- models have achieved great success in versatile applications, such as natural language processing (Vaswani et al. 2020), and brain signal processing (Yang al. 20024) Across these diverse modalities, it is essential for transformers to design specific token generation methods based on the characteristics of the data. Existing approaches typically treat words, image patches, or time series segments as the basic units for embedding tokens. However, this segmentation strategy for tokenizing multi-band graphs disrupts their topological structure. Therefore, it is necessary develop a novel tokenizer that divides multi-band graphs into atomic tokens while preserving sufficient graph structure.
2017), computer vision (Dosovitskiy et al.
Second, since transformers lack inherent mechanisms to capture the position of the input, explicitly injecting positional encoding is a common solution perceive such latent information. However, is this necessary for multi- band connectomes deriving from EEG data? The primary purpose of positional encoding is explicitly add information that transformers struggle to learn inherently, and this type of information is crucial for characterizing specific modalities. Accordingly, explicitly adding information to the input embeddings should take into account the intrinsic properties of the data. In the case of the multi-frrquencc band connectomes, each connectome inherently implies a specific frequency range and demographic information (i.c., age and gender) that also leads to significant anatomical, functional, and biochemical differences in the brain (Zaidi 2010). However, conventional positional encoding methods fail to encode these intrinsic characteristics in EEG data. Therefore, another challenge is to inject implicit information, including frequency and individual differences due to demographics, into the transformer in a manner similar to positional encoding.
Third, the widespread application of transformers in medical diagnosis might be hindered by their black-box nature, as clinicians are concerned not only with the performance of transformers but also with the clinical rationale. To open the black box of deep learning models, explainable artificial intelligence (XAI) provides a promising solution for bringing transparency complex and opaque deep learning models. Some interpretable models have emerged to offer various insights into a model's decision (Zhao et al., 2022. Li al. 202T). However, almost all research in neuroimaging stops at the visualization purpose of explanation, neglecting to use this information for improving the performance of the transformer. On the other hand, transformers, especially large models, always benefit from transfer learning, knowledge distillation and convolutional operators. The common idea behind these techniques is to introduce additional knowledge from external datasets, teacher models or prior knowledge, respectively. Interestingly, XAI is also able to provide valuable insights to the transformer from the perspective of "self-refinmmet" without the demand for additional datasets and models (Stammer al. 2023). In other words, XAI-based transformer improvement is perfectly suited for the limited amount of EEG data. Therefore, a key challenge is to take full advantage of XAI to improve the performance of the transformer.
To address the mentronned challenges, we propose XAI guided transformer (XAlguiFormer), novel architecture for EEG- based brain disorder classificatton. XAlguiFormer takes multi-frequent band connectomes as input, constructed by two complementary methods (i.e. coherence (COH) and weighted phase lag index (wPLI), see Appendix [A.3 for details) from preprocessed EEG signals. XAIguiFormer consists of three key components: connectome tokenizer, rotary frequency encoding with demographics (dRoFE), and XAI guided self-attention. The connectome tokenizer treats each single-band graph as an atomic token, forming a graph-wise frequency sequence. Compared to fragmenting the graph, our method can capture more comprehensive brain connectivity patterns in the frequency domain. In contrast to conventional positional encoding that adds position information to tokens, explicitly integrate frequency and demographic information into tokens by the rotation matrix, thereby creating more informative representation. To introduce valuable explanations to improve the transformer, XAlguiFormer employs XAI to refine the originally coarse query and key vectors within the self-attention mechanism, resulting in a more concentrated attention.
To the best of our knowledge, this is the first study employing XAI to enhance transformer perfor- mance in neuroimaging data, extending its benefits beyond mere interpretability analysis. The main contributions of this paper are as follows:
we propose a novel approach to tokenize the connectome across frequency bands without destroying the topological structure of single-band graphs. It can be easily extended to other multi-graph scenarios, such as dynamic functional connectivity.
we propose a fusion method to explicitly inject the frequency and demographic information into tokens, improving the model's 's understanding of frequency and mitigating the negative effects of individual differences.
2

---

Published as a conference paper at ICLR 2025
we propose to use XAI to direetly enhance transformer performance rather than focusing only on analysing the visual interpretability. Inspired by the dynamical system, we propose an XAI guided self- self-attention mechanism along with an XAI guided loss function, directing the model's focus towards more relevant dependency relationships.
2 RELATED WORK
Generally, there are two primary types of positional encoding for trans- formers. One is the absolute positional encoding (Vaswani al., 2017). which injects the absolute position of tokens to the input embeddings through sinusoidal or learnable embeddings. The other is the relative positional encoding (Shaw et al., 2018), which encodes relative position information into the self-attention mechanism. Rotary Positional Embedding (RoPE) (Su et al. 2024) is spe- cial relative position embedding that uses the absolute position to represent the relative position via a rotation matrix. In addition to these general-purge cenerra-purpose methods, several task-specific positional encodings have been developed to exploit the unique properties of specific types of data. For in- stance, Informer (Zhou et al., 2 considers the intrinsic characteristics of time series, such as week, month, year and holidays, and proposes hierarchical (local-global) time stamps as positional encoding forecast long sequence time series. Brain-JEPA (Dong et 2024) also considers the inherent functional relationships among brain regions, thereby developing brain gradient positioning to inject temporal, spatial and functional information into tokens. In the EEG field, MMM (Yi al,) 2024) introduces munti-diimessoonll position encoding to capture intrinsic geometric layout of elec- trodes. To preserve the spatiotemporal features of EEG data, BIOT (Yang et al., 2024) adds channel embedding and relative position embedding to each token. However, those techniques fail to inte- grate frequency bands and individual differences into tokens, causing missing essential information in multi-band connectomes.
Positional Encoding.
Explainable Artificial Intelligence. The increasing demand for model transparency and trust- worthiness in recent years has driven the development of various explanation methods for black-box deep learning models. A prominent category is the saliency method, which identifies the most im- portant input features contributing to the model's decision, such as LRP (Bach et al. 2015) and SHAP (Lundberg & Lee| 2017). In contrast, some inherently interpretable deep learning models have been proposed provide explanations during training, such as CRATE (Yu al. 2024 and BrainGNN et. al. 2021). Although these visualizations of explanations can help researchers iden- tify flawed model reasoning, they often lack the ability directly guide improvements in the model's internal parameters or architecture during training. To this end. Weber et al. (2023) conducted com- prehensive overview of XAI techniques for enhancing various properties of deep learning models, demonstrating how explanations can be leveraged to improve components in deep learning models, including data, feature representations, loss functions, gradients, and trained models. However, few studies have fully utilized XAI to improve transformers.
3 PRELIMINARY KNOWLEDGE
Rotary Positional Embedding. The RoPE is a novel position embedding method originally de- signed for natural language processing to encode the relative positions among tokens in sequence, which is widely applied in large language models. The core idea behind RoPE is to rotate the query and the key & by certain angles to characterize the position In the complex domain, the rotary matrix is defined as
R ime
(1)
where d is the embedding dimension, i is the imaginary number, and θ is pre-defined parameter given by
θ = {0, = 10000 2td | {0, 1, ... .......d.2..... d/2-11
The rotary matrix is then applied to q and & for injecting the position information by
/(q.:)(ar.. W ((q.k)) = n²._m W ((..))
(3)
where om is the m-th token in a sequence and W Wgg..: represents the weights specific to the q and k. Finally, the attention value with RoPE is calculated by
Am, = Re < ( fa (2, Wx)
where Re < denotes the real part of the complex number. In this formulation, RoPE naturally incorporates relative position information through the inner product of the rotation matrix, which exhibits remarkable performance.
3

---

Published as a conference paper at ICLR 2025
Connectome Tokenizer
Concel
Pepaenc requence
Global Average Pooking
Graph Naund
Scale SottMox
(i)
Aggregate
ReaSlorm
Explanation
GRoFE
Varila MuB: Hoad Attorion
XAI Guided Muts-
Curpt
Kurpt
kada'e
Head Altentioonn
W
W.
W.
Shamed welgre
Connectome Takenber
Commetooe Tokerize
E5 rallo
Explanation
The architecture of XAlguiFormer. XAlguiFormer forward process can be described as follows: (i) construct multi-frequency band connectomes and generate a sequence in the frequency domain by connectome tokenizer, (ii) forward pass the frequency sequence by vanilla transformer, (iii) obtain refined features by explaining the vanilla transformer, (iv) feedforward refined features and the frequency sequence through XAI guided transformer. Subsequently, the MLP is employed as the classification head to predict the label of the brain disorder.
Figure 1:
Understanding Self-attention Mechanism via Dynamical System Perspective. Huang et al. 2023) explain the self-attention mechanism from the mathematical perspective of dynamical sys- tem. In such a framework, the transformer with residual block can be written
Lt-1 Tt (zt; Wo) (a1; Wg), fu fe(;;;;
(4)
Feature mop
Attention value
where at is the input sequence of the linear layer f(: W (g,k,o)) with the learnable parameters W Wig in the 1-th block, & denotes the Hadamard product and ø the self-attention module. Further, equation can be reformulated
ordinary differential equation (ODE)
Lt-1 2t
(a W.)
(5)
AL
Here, At can be interpreted as the attention value generated by the self-attention module o(fg. fa), which serves as an adaptive step size in ODE. Assuming f(ar,)) are all invertible, the attention value function can be transformed to
(1 Sa Ws)) = &Ss So (v), Sa 1o
(b)) 3(v) = H (a W.))
(6)
Coarse Info
Here, (, 2/ (a:+11 - 2x²) is kind of coarse stiffness information in feature trajectory at Taking into account all these equations, it can be concluded that the attention value function & takes coarse information fi (æz; W.) as input and refines to obtain finer estimation of infor- mation, further generating adaptive attention values. This implies the importance of establishing better input for self-attention and the resulting attention values to capture refined information in the transformer.
4 METHODOLOGY
In this study, we introduce the XAI guided transformer, novel architecture that not only interprets the model's decisions but also utilizes explanations enhance the performance of the transformer in an end-to-end fashion. Instead of only considering the position in the sequence, our model integrates the frequency and demographic information into tokens by specially designed rotary matrix. As illustrated Figure [ the model architecture consists of 4 steps: (i) generate token for each fre- quency band connectome using the graph neural network, (ii) pass the frequency sequence through

---

Published as a conference paper at ICLR 2025
the vanilla transformer encoder while inject the frequency and demographic information to tokens by dRoFE, (iii) explain the vanilla transformer to obtain refined features, (iv) run the feedforward path again by feeding refined query and key to XAI guided transformer. Finally, the XAI guided loss function is employed to train the entire model.
4.1 PROBLEM FORMULATION
represent the brain connectomes across various fre- quency bands Freq = {6, low a. high a. low B, mid 3. high B, low 7, 0/8/3 ratio). For each frequency band connetome 9j (V.E,) EEG channels are defined graph nodes y (11, v,,-..,, where e is the number of EEG channels and f E Freq. The node features at fre- quency band f are denoted by H(Co [ h .... h 1 where h! is the coherence-based connectivity patterns of node v, at frequency band ʃ (see Appendix [A.3] for details). The edge features at frequency band f are indicated by E = g where <i,, is the wPL.-based connec- tivity between nodes v; and Uy at frequency band f. The adjacency matrix of the single-band graph Gy is represented A Rexe. In addition, let Demog = [age, gender] be the demographic information. The main objective of this study is to establish the mapping from (G, Demog) to brain disorder diagnosis Y by learning function F (G, Demog) Y.
Let the multi-graph 9 = {9o,
4.2 CONNECTOME TOKENIZER
A fundamental operation for transformers is the design of tokenizers tailored to specific data modali- ties. For images, ViT splits the image into sequence of flattened patches (Dosovitskiy al. 2020) Although the multi-frequengy band connectome G can be regarded as a multi-channel "image" due its grid-like structure, partitioning the connectome into patches risks destroying its connectivity and results in the loss of topological information. Inspired by the tokenization method of DETR (Carion et al. 2020), we adopt strategy where the entire single-band connectome 9y serves as the basic unit to generate tokens. To this end, the GINEConv operator (Hu et al. 2020 is employed to update node representations at (-th layer by
(1+1) eerf
MLP((1 +
Rh 1 ogd + +c(,))
DyEN(es)
where N(w;) represents the set of nodes adjacent to v; and is a trainable parameter. Notably, lever- aging the GNN also introduces a beneficial inductive bias our model, particularly on relatively small datasets. Subsequently, single-bane band token 2 is obtained by averaging the node embeddings the last layer of GNN using
f MEAN((1 1 E V})
(8)
Finally, those token embeddings 2f are flattened into frequency sequence (mss 2g...., Faya! E [f which are then fed into the XAI-guided transformer encoder.
X
4.3 ROTARY FREQUENCY ENCODING WITH DEMOGRAPHICS
Conventional RoPE only considers intrinsic positional properties of the sequential data in text and image. However, positional information is less valuable in the frequency sequence X because its order can be arbitrary, and treating X as a simple sequence risks losing the intrinsic frequency infor- mation. Hence, we encode the frequency information by rotating each frequency band embedding with angles determined by frequency bounds
R2 et0 R4. 12
where fi, fu the lower and upper bound of the frequency band f, and i is the imaginary number. To achieve smoother frequency encoding in the short sequence X compared to the exponentially scaled one in the language sequence, the 0 at the 1-th dimension of tokens is defined as
0: = 4nt/d, E {o, d/4 d/4-1 1}
(10)
Moreover, age and gender are known to influence differences in individual connectomes Mijalkov al. 2023). To alleviate the individual differences, propose to integrate demographics into
5

---

Published as a conference paper at ICLR 2025
rotary frequency encoding (dRoFE) by encoding demographie information the magnitude of the vector. Taking together, the rotation matrix of dRoFE is defined as
Rt, Demog...22 age e gender, R "ff.Demo,.2+4 2t-11 age e 0z
gender
(11)
Finally, dRoFE is applied on the query q and key le to explicitly inject frequency and demographic information by
dRoF E (q,k) f. Demog) = j. W (q.6)
(12)
In practical implementation, dRoFE introduces the multiplication of Euler's formula to q and k, which can produce the same attention value qk" = Re[g.k]] but reduces computational burden (Heo et al. (024).. Here, Re[-] denotes the real part of the complex number and qi k; refer to the complex vectors.
4.4 XAI GUIDED TRANSFORMER
XAT Guided Self- self-Attention. According to the insights from the section preliminary knowledge, the ability properly generate refined input and attention value is essential for the performance of the self-attention mechanism. To enhance the transformer's ability to capture refined information, propose theoretic-inspired self-attention mechanism, named XAI guided selffattention. The value of XAI lies in its capability to identify important features and filter out irrelevant information for self-attention input. Consequently, the XAl-generated importance maps for Q and K are more refined than the original ones. Further, the attention value is calculated based on these underlying important features identified by XAI, effectively removing noise to generate a more precise attention distribution that concentrates on important dependency relationships by
so fomax Qeepp KT Expl
Attn Qezpi, K Kempl V)
(13)
where Qeapi and K Keept are the feature importance scores generated by DeepLift (Shrikumar et al. 2017). To access those valuable explanations, we first process the frequency sequence X via the vanilla transformer encoder. We then employ DeepLift to pinpoint which parts are most influential for the vanilla transformer. In general, the explanation algorithm explanation method from XAL. By making use of XAI, we can improve originally coarse informa- tion capture more focused attention value, ultimately enhancing the representative ability of the transformer.
be realized with any local
Finally, the XAlguiFormer model comprises a stack of several encoder blocks. As illustrated in Figure [ each XAIguiFormer block consists of XAI guided multi-head self-attention module, followed by 2-layer MLP with GeGLU activation function. A residual connection is applied after XAI guided multi-head attention and the feed forward network, followed by an RMSNorm layer for normalization.
XAI Guided Loss Function. To boost the performance of model under the supervision of XAI, we propose joint loss function optimize the parameters
L -11 (1-o) LOF Mers y) + aLLF (Heept y)
(14)
Coarse Prediction
Refined Prediction
where LoE refers to cross-entropy loss, Dere and Jeups denote the labels predicted by vanilla and XAI guided self-attention mechanisms, respectively. The parameter controls the magnitude of
XAI guidance. In our experiments, XAlguiFormer is not very sensitive to a as long as it is larger than 0.3.
EXPERIMENTS
5.1 DATASETS
We evaluated the performance of our XAlguiFormer model two publicly available datasets:
The TUH Abnormal BEG Corpus (TUAB) is one of the largest open-source EEG datasets, which collects data from 1385 healthy controls (normal) and 998 subjects having brain disorder (abnor- mal), with total of 2993 sessions (Obeid & Picone, 2016). All electrophysiology recordings contain 19 common EEG channels and 2 reference channels across subjects. We employ the XAlguiFomer to perform binary classification for identifying brain disorders this dataset.

---

Published as a conference paper at ICLR 2025
Table I: Comparison of model performance with baseline methods the TUAB dataset.
TUAB
Methods
Model Size
FLOPs
BAC
AUC-PR
AUROC
FFCL 2022b
2.4M
0.83G
0.7848 + 0.0038
0.8448 0.0065
0,8569 0.0051
SPaRCNet 2023
0.79M
0.26G
0.7896 + 0.0018
0.8414 0.0018
0.8676 = 0.0012
BIOT (2024)
3.2M
1.9G
0.7959 1 0.0057
0.8792 + 0.0023
0.8815 + 0.0043
S3T (202T)
3,5M
0.22G
0.7966 + 0.0023
0.8521 + 0.0026
0.8707 + 0.0019
LaBraM-Base 2024
5.8M
2.7G
0.8140 = 0.0019
0.8965 = 0.0016
0.9022 0.0009
Corr-DCRNN 2022
0.69M
0.21G
0.7672 = 0.0050
0.8586 0.0044
0.8564 0.0040
IGGNNet (2024)
1.4M
0.64G
0.7711 + 0.0030
0.0010 0.8594
0.8574 = 0.0011
XAlguiFormer (Ours)
3.5M
1.6G
0.8205 0.0027
0.8965 0.0079
0,9000 0.0046
Table 2: Comparison of model performance with baseline methods on the TDBRAIN dataset.
TDBRAIN
Methods
Model Size
FLOPs
BAC
AUC-PR
AUROC
FFCL 2022b
2.4M
0.83G
0.5394 1 0.0400
0.4258 + 0.0278
0.5583 1 0.0490
SPaRCNet 2023
0.79M
0.26G
0.4786 1 0.0153
0.3448 + 0.0194
0,4580 + 0,0224
BIOT (2024)
3.2M
1.9G
0.6320 + 0.0082
0.6007 0.0205
0.7175 0.0176
S3T (2021)
3.5M
0.22G
0.4937 + 0.0214
0.3500 0.0174
0.4739 = 0.0195
I.aBraM-BR Base 2024
5.8M
2.7G
0.6456 + 0.0089
0.5438 = 0.0058
0.7147 - 0.0145
Cor-DDCRNN (2022
0.69M
0.21G
0.6045 : 0.0058
0.4951 + 0.0131
0.6586 + 0.0158
LGGNet (2024)
1.4M
0.64G
0.6152 = 0.0098
0.5058 = 0.0161
0.6929 = 0.0174
XAlguiFormer (Ours)
3.5M
1.6G
0.6635 1: 0.0080
0.5961 1 0.0136
0.7814 1 0.0156
Research Archive for Insights in Neurophysiology (TDBRAIN) is a clinical lifespan EEG biobank comprising diverse patients with major depressive disorder, at- tention deficit hyperactivity disorder and obsessive-compulsive disorder (Van Dijk et al. 2022). We utilized eyes-closed condition with 26 channels of EEG data to carry out multi-class classification distinguish brain disorders.
The Two
For further details the demographic distribution and preprocessing steps of those datasets, as well as the construction method of XAlguiFormer's input connectomes, please refer to Appendix
Additionally, details on the baseline models and evaluation metrics can be found in Appendix
5.2 IMPLEMENTATION DETAILS
Our experiments are implemented using Python 3.10.12, PyTorch 2.0.I+CUDAIL. and PyTorch- Geometric 2.3.1
single NVIDIA A100 GPU. The connectome input dimensions are 19 for TUAB and 26 for TDBRAIN, both of which are mapped into 128 as hidden dimensions. The model consists of l-yerr GNN and 12-layer XAIguiFormer with 4 attention heads, training by AdamW optimizer with 3 = (0.9, 0.99) and 6 = 1e-8. For the TUAB dataset, the learning rate is set to 4e-5 with a weight decay of le-4, while for the TDBRAIN dataset, the learning rate is set to 5e-5 with a weight decay of le-5. The loss weight for XAI-guided training is set to 0.7. Due to the different sizes of the TUAB and TDBRAIN datasets, the batch sizes are set 512 and 64, respectively. More details about hyperparameters can be found in Appendix G
5.3 COMPARISON WITH BASELINE METHODS
Table I and 2 present the performance comparison between XAlguiFormer and baseline models on the TUAB and TDBRAIN datasets, where we report subject-wise performance by averaging the predictions of all 30-second samples belonging to the same subject. The results demonstrate that XAlguiFormer outperforms all baselines on various evaluation metrics for both datasets, except for the AUC-PR in the TDBRAIN dataset. In particular, on the relatively small TDBRAIN dataset, the XAlguiforner achieved significant improvement in performance, suggesting that incorporating intrinsic and explanation information can help the model better understand latent patterns that are challenging to learn from the small dataset. Despite being trained from scratch, our model exhibits superior performance compared to the pretrained BIOT and self-supervised Corr-DCRN model, which implies model can effectively capture more valuable information through the dRoFE and
7

---

Published as a conference paper at ICLR 2025
Table 3: Ablation studies dRoFE algorithm and XAI guided transformer.
Datasets
Metric
XAlguiFormer
w/o dRoFE
w/o XAI guided
BAC
0,8205 + 0.0027
0.8018 = 0.0048
0.8070 + 0.0066
TUAB
AUC-PR
0.8965 0.0079
0.8791 = 0.0046
0.8971 + 0.0047
AUROC
0.9000 = 0.0046
0.8846 = 0.0047
0.8997 + 0.0032
BAC
0.6635 = 0.0080
0.6249 = 0.0097
0.6357 = 0.0060
TDBRAIN
AUC-PR
0.5961 + 0.0136
0.5135 = 0.0092
0.5347 = 0.0164
AUROC
0.7814 0.0156
0.6983 0.0096
0.7101 + 0.0131
XAI modules. Furthermore, while the majority of baseline methods are the black box, the XAlgu- iFormer can elucidate its decision-making process, thereby enhancing interpretability alongside its performance.
5.4 ABLATION STUDIES
To analyze the effectiveness of proposed dRoFE and XAI guided transformer, conducted ablation experiments on both datasets as illustrated in Table B In general, it is observed that both algorithms contribute to the overall performance improvement.
Effectiveness of dRoFE. To rigorously verify the effectiveness of dRoFE, we replace dRoFE with RoPE for our model instead simply removing it. This setting ensures a fair comparison, as the transformer without any positional information would inherently lack sense of position, skewing the results.
that dRoFE leads to a significant enhancement in performance, achieving +1.87% BAC the TUAB and +3.86% BAC on the TDBRAIN. This suggests that incorporating frequency and demographic information more beneficial for the frequency sequence than simple relative position. As result, dRoFE can effectively capture more comprehensive information and mitigate the influence of individual differences.
Effeetiveness of XAI Guided Transformer. this experiment, we ablate the XAI guided self- attention mechanism and exclude the XAI guided item from the entire loss function. As shown in Table 3. the inclusion XAI guidance results in a BAC increase of +1.35% on the TUAB and +2.78% on the TDBRAIN, demonstrating that XAI does contribute to enhancing the performance of the model by introducing additional knowledge. Interestingly, the TDBRAIN dataset benefits more from XAI guidance. This is likely because the model trained small dataset relies more on supplementary information to learn complex representations. In contrast, the model trained on a larger dataset, like TUAB, can learn relatively general knowledge and thus reduce its dependence on the additional knowledge.
Warm Up XAI to Guide Transformer on TUAB
83.0
62.5
Balanced Accuracy (%)
81.0
60.5
80.0
79.5
0
10
20
30
40
Epoch
Figure 2: Warmup start XAI to guide transformer, which shows average balanced accuracy and standard deviation across five random seeds on the TUAB. Each point on the line, from left to right, corresponds to different model configuration where XAI is activated at 0% (normal XAlgu- iFormer). 10%, 20%, 30%, and 100% (without XAI guidance) of the total epochs.
8

---

Published as a conference paper at ICLR 2025
(A)
(B)
XAI Guided Attention Maps on TUAB
XAI Guided Attention Maps TDBRAIN
Head o
Hood
Head 2
Head
Head o
Hoad
Head 2
Head
Entrmpy
Entrapy
Entroppy 0.91
Entrrayy 0:27
Entropy: 1.00
Entrapy
Enttropyy 0.79
Entropy
Vanilla Attention Maps on TUAB
Vanilla Attention Maps on TDBRAIN
Head
Head
Head
Entropy 0.51
1.13
Figure 3: Comparison attention maps between XAI guided multi-head self-attention and vanilla multi-head self-attention. As shown in the figure, XAI guided attention exhibits tern, suggesting a concentration on more crucial dependency relationships and the elimination of distractions from irrelevant tokens. Attention entropy, detailed in Appendix [ is used quan- titative measure to assess the concentration of attention maps. Smaller attention entropy indicates that the attention is more concentrated. To facilitate visual comparison, an identical colorbar used for each attention map. In each attention map, the rows from top to bottom and columns from left to right represent the frequency sequence low a, high a, low 3, mid B, high 3, 1.low y, 0/ ratio.
refined pat-
Warmup Start XAI Guidance. We hypothesize that XAlguiFormer benefits from explanations derived from relatively good source model. Since the source model would improve during training, we explore warmup start strategy where XAI guidance activated after 10%, 20%, and 30% of the total training epochs on the TUAB dataset. As illustrated in Figure 2 while warm-started XAI guidance consistently outperforms the vanilla transformer, it achieves lower performance compared to starting XAI from the beginning of training. This is because activating XAI guidance during training alters the parameter optimization space of the model, as reflected by observed jitter in the loss function when starting XAI guidance. Despite this, it indirectly verifies the effectiveness of XAI in enhancing model performance.
(A)
(B)
Entropy-BAC Relationship Curve During the Training Process on TUAB
Entropy-BAC Relationship Curve of Warm Up XAI on TUAB
1.00
Attonston Entrosy
0.50
Attintionn
0.25
3.00
o
-0.25
0.5
70
75
80
80.75
81.00
81.25
$1.50
81.75
82.00
Be Ancuracy (%)
Belanced Accuracy
Figure 4: The attention entropy and BAC relationship curve on the TUAB dataset. (A) The attention entropy and BAC relationship under different epochs. The pair of attention entropy and BAC is obtained during the training process. Each point the scatter plot corresponds to BBAC pair at a specific training epoch. (B) The attention entropy and BAC relationship under different warmup XAI models. The points from left to right correspond to the activation of XAI 100% (without XAI), 10%, 30%, 20%, 0% (normal XAlguiFormer) of the total training epoch.
5.5 CONCENTRATED ATTENTION PATTERN
Figure 3 visually compares the attention maps generated by the XAI guided and vanilla multi-head self-attention mechanisms. is evident that XAI guided attention maps focus more on specific dependency relationships, whereas the vanilla attention maps appear more scattered. This improve- ment is because XAlguiFormer can calculate a more precise attention score based on the refined query and key. Consequently, XAlguiFormer captures dependency relationships and generates

---

Published as a conference paper at ICLR 2025
adaptive values from the perspective of the dynamical system, leading to more effective representa- tions.
To further explore the impact of concentrated attention on the model performance, we conduct ex- periments quantitatively evaluate the relationship between attention entropy and BAC under dif- ferent scenarios. Specifically, we analyzed the entropy-Bek -BAC pairs from (1) different training epochs within the same XAIguiFormer model (Figure 4 and (2) different warm-up XAI models (Figure 4B). These results show that the correaltion coefficients between attention entropy and BAC are -0.7 and -0.75, respectively. This negative correlation suggests that more concentrated attention, characterized by lower entropy, is associated with improved performance.
5.6 VISUALIZATION OF INTERPRETABILITY
Figure represents the importance of different frequency bands in model for the TUAB and TDBRAIN datasets. The frequency band importance can be directly extracted from the explainer used in XAlguiFormer. As depicted, the theta/beta ratio significantly contributes
model. In psychiatry, this ratio serves as biomarker for attentional control (Putman et al. 2014), stress assessment (Wen & Aris 2020) and anxious inhibition (Putman al., 22100, Consequently, it is able to capture functional aspects of normal and abnormal subject behavior (Schutter & Kenemans, 2022]. The low and high alpha bands also play a crucial role our model, though their functions differ slightly (Debnath et al. (021]. The connectome of high alpha is negatively correlated with elevated mood, whereas the amplitude in low alpha is positively correlated (Petsche et al. 1997). This distinction is why split the alpha band into subbands. Furthermore, the analysis of nartow EEG bands has proven to be more useful for assessing EEG alterations (Ponomareva et al. 2014).
Overall, identifying these frequency bands significantly enhances the interpretability of model.
(A)
Frequency Band Importance on TUAB
Frequency Band Importance on TDBRAIN
thetalbeta ratio
2,28
thetafoete racio
1.79
high alpha
sipha
law alpha
1.05
high beta
0.85
Requency fledd
gamme
0.56
0.77
Frequency
0.79
delta
0.73
theta
0.77
mid
high alpha
theta
0.72
deta
0.71
0.0
0.5
1.0
1.5
2.0
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
Contribution
Contribution
Figure 5: Frequency band importance the TUAB and TDBRAIN datasets
CONCLUSION
In this paper, we propose theoretically inspired architecture, XAI guided transformer (XAlgu- iFormer) that not only provides trustworthy explanations but also leverages XAI insights to improve the performance the transformer. XAlguiFormer takes advantage of multi-band connectomes constructed from EEG signals as input. By treating the single-band graph as an atomic token, the connectome tokenizer preserves the interaction between brain regions. To overcome the model's difficulty in learning the intrinsic information of EEG data, we introduce dRoFE to integrate the intrinsic frequency and demographic information into tokens, thereby mitigating the negative im- pact of individual differencee. Our experiment demonstrates that XAlguiFormer outperforms all baseline methods in the brain disorder classification task. Future efforts can explore self self-supervised pretraining of larger models and expand to additional EEG tasks. Finally, hope this research can draw attention to employing XAI to enhance EEG-based deep learning models, extending its benefits beyond visualization.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors gratefully acknowledge the Gauss Centre for Supercomputing e.V. (www.gauss centre.eu) for funding this project by providing computing time through the John von Neumann Institute for Computing (NIC) on the GCS Supercomputer JUWELS at Jülich Supercomputing Cen- tre (JSC).
REFERENCES
Angelos Angelidis, Willem van der Does, Lemmy Schakel, and Peter Putman. Frontal eeg theta/beta ratio as an electrophysiological marker for attentional control and its test-retest reliability. Bio- logical Psychology, 121:49-52, 2016. ISSN 0301-0511.
Stefan Appelhoff. Austin J. Hurst, Aamna Lawrence, Adam Li, Yorguin José Mantilla Ramos, Chris- tian O' O'Reilly, Liang Xiang, and Jonte Dancker. PyPREP: A Python implementation of the pre- processing pipeline (PREP) for EEG data., March 2022.
Sebastian Bach, Alexander Binder, Grégoire Montavon, Frederick Klauschen, Klaus- Müller, and Wojeiech Samek. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PloS one, 10(7):e0130140, 2015.
Nima Sagdely-Shamlo Tim Mullen, Christian Kothe, Kyung-Min Moin--Min&i Su, and Kay Robbins. The prep pipeline: standardized preprocessing for large-scale eeg analysis. Frontiers in neuroinformaties, 9:16, 2015.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pp. 213-229 Springer, 2020.
Ranjan Debnath, Natalie Viola Miller, Santiago Morales, Kaylee R. Seddio, and Nathan A. Fox. Investigating brain electrical activity and functional connectivity in adolescents with clinically elevated levels of adhd symptoms in alpha frequency band. Brain Research, 1750:147142, 2021.
ISSN 0006-8992
Yi Ding, Neethu Robinson, Chengxuan Tong, Qiuhao Zeng, and Cuntai Guan. Lggnet: Learn- ing from globbll -aaph representations for brain brain- computer interface. IEEE Transactions on Neural Networks and Learning Systems, 35(7):9773-977 9786, 2024.
Zijian Dong, Li Ruilin, Yilei Wu, Thuan Tinh Nguyen, Joanna Su Xian Chong, Fang Ji, Nathanael Ren Jie Tong, Christopher Li Hsian Chen, and Juan Helen Zhou. Brain-jepa: Brain dynamics foundation model with gradient positioning and spatiotemporal masking. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An im- age is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Representations, 2020.
Alexandre Gramfort, Martin Luessi, Eric Larson, Denis A. Engemann, Daniel Strohmeier, Christian Brodbeck, Roman Goj, Mainak Jas, Teon Brooks, Lauri Parkkonen, and Matti S. Hämäläinen. MEG and EEG data analysis with MNE-Python. Frontiers in Neuroscience, 7(267):1-13, 2013.
Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. G-mixup: Graph data augmentation for graph classification. In International Conference on Machine Learning, pp. 2022.
PMLR,
Byeongho Heo, Song Park, Dongyoon Han, and Sangdoo Yun. Rotary position embedding for vision transformer, 2024.
W Hu, B Liu, Gomes, Zimik,, Liang, V Pande, and Leskovec. Strategies for pre-training graph neural networks. In International Conference Learning Representations (ICLR), 2020.
11

---

Published as a conference paper at ICLR 2025
Zhongzhan Huang, Mingfu Liang, Jinghui Qin, Shanshan Zhong, and Liang Lin. Understanding self-attention mechanism via dynamical system perspective. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 2412-1422, 2023.
Weibang Jiang, Liming Zhao, and Bao-liang Lu. Large brain model for learning generic represen- tations with tremendous eeg data in bei. In The Twelfth International Conference on Learning Representations, 2024.
Jin Jing, Wendong Ge, Shenda Hong, Marta Bento Fernandes, Zhen Lin, Chaoqi Yang, Sungtae An, Aaron Struck, Aline Herlopian, Ioannis Karakis, et al. Development of expert-level classifica- tion of seizures and rhythmic and periodic patterns during eeg interpretation. Neurology, 100(17): el c1750-e1762. 2023.
Adam Li, Jacob Feitelberg, Anand Prakash Saini, Richard Höchenherger, and Mathieu Scheltienne. Mne-icalabel: Automatically annotating ica components with iclabel in python. Journal of Open Source Software, 7(76))4484, 2022a.
Hongli Li, Man Ding, Ronghua Zhang, and Chunbo Xiu. Motor imagery eeg classification algorithm based on enn-lstm feature fusion network. Biomedical signal processing and control, 72:103342, 2022b.
Xiaoxiao Li, Yuan Zhou, Nicha Dvomek, Muhan Zhang, Siyuan Gao, Juntang Zhuang, Dustin Scheinost, Lawrence H Staib, Pamela Ventola, and James S Duncan. Braingnn: Interpretable brain graph neural network for finr analysis. Medical Image Analysis, 74:102233, 2021.
Scott M Lundberg and Su-In Lee. A unified approach to interpreting model predictions. Advances in neural information processing systems, 30, 2017.
Mite Mijalkov, Dániel Veréb, Oveis Jamialahmadi, Anna Canal-Gareia, Emiliano Gómez-Ruiz, Di- dac Vidal-Piñeiroo Stefano Romeo, Giovanni Volpe, and Joana B Pereira. Sex differences in mul- tilayer functional network topology over the course of aging 37543 biobank participants. Network Neuroscience, 7(1):351-376, 2023.
Iyad Obeid and Joseph Picone. The temple university hospital eeg data corpus. Frontiers in neuro- science, 10: 10:195498, 2016.
Jiaoyan Pang, Hanning Guo, Xiaochen Tang, Yu Fu, Zhengwu Yang, Yongchao Li, Na An, Jing Luo, Zhijun Yao, and Bin Hu. Uncovering the global madulaed brain network in chunk decomposition with chinese characters. Neuroimage, 147:118826 2022.
Petsche, Steven Kaplan, A Von Stein, and Filz. The possible meaning of the upper and lower al- pha frequency ranges for cognitive and creative tasks. International journal of psychophysiology. 26(1-3):77-97, 1997.
Natalya Ponomareva, Sergey Klyushnikov, Natalya Abramycheva, Daria Malina, Nadejda Sche- glova, Vitaly Fokin, Irina Vvanova-Smolenskaia and Sergey Illarioshkin. Alpha-theta border eeg abnormalities in preclinical huntington's 's disease. Journal of the neurological sciences, 344(1-2): 114-120, 2014.
Peter Putman, Jacobien van Peer, Ioulia Maimari, and Steven van der Werff. Eeg theta/beta ratio in relation to modulated response-inhibition attentional control, and affective traits. Biological psychology, 83(2):73-78, 2010.
Peter Putman, Bart Verkuil, Elsa Arias-Garcia, Ioanna Pantazi, and Charlotte van Schie. Eeg theta/beta ratio as a potential biomarker for attentional control and resilience against deleteri- ous effects of stress attention. Cognitive, Affective, & Behavioral Neuroscience, 14:782-791 2014.
Dennis J. L. G. Schutter and J. Leon Kenemans (eds.). Theta-Bett Power Ratio: An Electrophysio- logical Signature of Motivation, Attention and Cognitive Control. Oxford University Press, The Oxford Handbook of EEG Frequency, 08 2022.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions, 2018
12

---

Published as a conference paper at ICLR 2025
Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje. Learning important features through propagating activation differences. In International conference on machine learning, pp. 3145- 3153. PMLR. 2017.
Yonghao Song, Xueyu Jia, Lie Yang, and Longhan Xie. Transformer-based spatial-temporal feature learning for eeg decoding, 2021.
Wolfgang Stammer, Felix Friedrich, David Steinmann, Hikaru Shindo, and Kristian Kersting, Learn- ing by self-explaining, 2023.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 168:127063, 2024.
Siyi Tang, Jared Dunnmon, Khaled Kamal Saab, Xuan Zhang, Qianying Huang, Florian Dubost, Daniel Rubin, and Christopher Lee-Messer. Self-supervised graph neural networks for improved electreencephalorrapiic seizure analysis. In International Conference Learning Representa- tions, 2022.
Hanneke Van Dijk, Guido Van Wingen, Damiaan Denys, Sebastian Olbrich, Rosalinde Van Ruth, and Martijn Arns. The two decades brainclinies research archive for insights neurophysiology (idbrain) database. Scientific data, 91):333, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention you need. Advances in neural informa- fion processing systems, 30, 2017.
K Venkatachalam, A Devipriya, J Maniraj, M Sivaram, A Ambikapathy, and S Amiri Iraj. A novel method of motor imagery classification using eeg signal. Artificial intelligence in medicine, 103: 101787, 2020.
Leander Weber, Sebastian Lapuschkin, Alexander Binder, and Wojeiech Samek. Beyond explaining: Opportunities and challenges of xai-based model improvement. Information Fusion, 92: 12:114-176, 2023. ISSN 1566-2535.
Ziquan Wei, Tingting Dan, Jiaqi Ding, and Guorong Wu. Neuropath: A neural pathway transformer for joining the dots of human connectomes. arXiv preprint arXiv:2409. 2024.
Tee Yi Wen and SA Mohd Aris. Electroencephalloram (ceg) stress analysis alpha/beta ratio and theta/beta ratio. Indonesian Journal of Electrical Engineering and Computer Science, 17(1): 175-182, 2020.
Chaoqi Yang. M Westover, and Jimeng Sun. Biot: Biosignal transformer for cross-data learning in the wild. Advances in Neural Information Processing Systems, 36, 2024.
Ke Yi, Yansen Wang, Kan Ren, and Dongsheng Li. Learning agpooocy-ggooticc eeg representations with geometry-aware modeling. Advances Neural Information Processing Systems, 36, 2024.
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin Ha- effele, and Yi Ma. White-box transformers via sparse rate reduction. Advances in Neural Infor- mation Processing Systems, 36, 2024.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy train strong classifiers with localizable features. In Proceed- ings of the IEEE/CVF international conference on computer vision, pp. 6023-6032, 2019.
Zeenat F Zaidi. Gender differences in human brain: a review. The open anatomy journal, 2(1), 2010.
Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Kanhao Zhao, Boris Duka, Hua Xie, Desmond J Oathes, Vince Calhoun, and Yu Zhang. A dynamic graph convolutional neural network framework reveals new insights into connectome dysfunc- tions in adhd. Neurolmage, 246:1 18774, 2022.
13

---

Published as a conference paper at ICLR 2025
Haoyi Zhou, Shanghang Zhang. Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. Proceedings of the AAAI Conference on Artificial Intelligence, (5(12)11106-11 11115,, May 2021.
14

---

Published as a conference paper at ICLR 2025
A DEMOGRAPHICS, PREPROCESSING AND CONNECTOME CONSTRUCTION
A.I DEMOGRAPHIC DISTRIBUTION
As demographic information plays crucial role in our proposed dRoFE, it is essential to ensure unbiased dataset splits on demographic distribution. For TUAB dataset, we follow the official train and eval splits. We further split the official train randomly into 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. Tables S and 6 illustrate the demographic distributions (e.g., age, gender, and brain disorder categories) across the train/val/test splits for TUAB and TDBRAIN. Overall, the distributions of age, gender, and brain disorder status are balanced across the splits, minimizing bias and ensuring fair model evaluation.
Table 4: Gender statisties of the TUAB dataset.
Splits
Abnormal
Normal
train
602 (M) vs. 609 (F)
544 (M) vs. 690(F)
val
65 (M) vs. 70 (F)
59 (M) vs. 78 (F)
test
63 (M) vs. (F)
65 (M) vs. 85 (F)
Table 5: Gender statistics of the TDBRAIN dataset.
Splits
ADHD
MDD
OCD
train
89 (M) 42(F)
117 (M) 110(F)
21 (M) 17(F)
val
13 (M) vs. 15 (F)
13 (M) vs. 36 (F)
6 (M) vs. 2(F)
test
22 (M) vs. 6 (F)
25 (M) vs. 24 (F)
5 (M) 3(F)
Table 6: Age distributions of the TUAB and TDBRAIN datasets.
TUAB
TDBRAIN
Age distributions
train
val
test
train
val
test
0-10
0
34
5
10-20
59
3
52
9
8
20-30
360
32
34
63
15
18
30-40
321
41
49
69
21
18
40-50
490
45
52
81
16
20
50-60
511
59
51
59
12
11
60-70
381
51
32
27
4
4
>00
316
40
52
3
A.2 PREPROCESSING
The raw HEG signals are preprocessed with MNE-Python 1.2.1 (Gramfort et al., 2013). follow- ing these steps: Initially, bad channels are automatically detected and removed using pyprep 0.4.2 Bigdely-Shambl et al. 2015 Appelhoff et al., (022). Subsequently, the signals are filtered with 1-45 Hz band-pass zero sero-phase Hamming windowed sine FIR filter (-6dB). The filtered signals are resampled to 250 Hz and segmented into 30-second samples. Next, these samples subjected to independent component analysis (ICA) using the Infomax algorithm and then artifactual compo- nents identified by the pretrained classifier MNE-ICLabel 0.4 (Li al., (202aa) rejected. The previously detected bad channels are then repaired using the spherical spline interpolation. Finally, artifact-free samples recomputed against the average reference.
A.3 CONNECTOME CONSTRUCTION
The multi-frequency band connectomes are constructed using MNNE-Connectivine 0.5.0. First, the 30-second preprocessed signals are further segmented into ten 3-second epochs. For these 3-second
15

---

Published as a conference paper at ICLR 2025
epochs, two complementary
and weighted phase lag index (wPLI) are ployed to estimate channel-level functional connectivity on the following 9 frequency bands: delta (2-4Hz), theta (4-8Hz), low alpha (8-10Hz). high alpha (10-12Hz), low beta (12- mid beta(18- 21Hz), high beta (21-30Hz), low gamma (30-45Hz) as well as the theta/beta ratio serving as a reliable biomarker of attention deficit hyperactivity disorder (Angelidis al.] 2016). Next, the func- tional connectivity from both methods in each frequency band are averaged over the ten epochs. The motivation for employing two complementary construction methods lies in their ability capture distinct brain activity patterns (Li et al., [021), thereby providing multi-view information. Finally, we aggregate two types of functional connectivity by using coherence to build node features and wPLI to build edge connections, generating the multi-frequency band connectomes.
B BASELINES AND METRICS
To comprehensively compare our method with relevant baseline methods, consider the following models: (i) FFCL (Li et al. 2022b) fuses multilevel spatial-temporal features based on CNN and LSTM architecture, (ii) SPaRCNet (Jing al. 2023) is ID-CNN based model with dense residual connections, (iii) BIOT Yang et al. (024)) is a biosignal transformer that learns embeddings from EEG signals of various formats, (iv) S3T (Sog et al. 2021 EEG transformer perceive the spatial and temporal features for decoding EEG signals, (v) LaBraM (Jiang et al. (024) is pre- trained foundation model that can handle diverse EEG signals with varying channels and lengths, (vi) Cor-DDCRN Taag et al.) 2022) is a self-supervised graph neural network that captures dy- namic brain connectivity, (vii) LGGNet (Ding et al., 2024) a neurologically inspired graph neural network to learn local-global-gen To evaluate the performance, we use the following metrics: balanced accuracy (BAC), area under precision-recal curve (AUC-PR) and area under the receiver operating characteristic curve (AU- ROC) for the TUAB and TDBRAIN datasets. We report the average performance and standard deviation across five different random seeds on the TUAB and TDBRAIN datasets. In all scenarios, the best models are trained on the training set, selected from the validation set, and evaluated on the test set. The BAC is employed as the monitor scores to select the best model for the TUAB and TDBRAIN datasets.
representations from EEG.
C ATTENTION ENTROPY
Although Figure clearly shows that XAI guided self-attention is more concentrated, introduce the attention entropy to quantitatively measure the concentration of attention maps. The attention entropy for each token (row) is defined
[Fregl E(A,) = Σ
Ai. log(Ai,)
(15)
where represents the attention map. The overall attention entropy for each head averaged over all token-wise attention entropy. A small attention entropy suggests that the attention tends to be concentrated while a large attention entropy indicates that the attention is more distributed across various tokens.
As illustrated in Figure and the attention entropy, XAI guided self-attention maps are generally more concentrated than vanilla attention maps. This suggests that XAI guided self-attention is better at focusing attention on the most relevant dependency relationships, as the attention score is calcu- lated based underlying important features. In contrast, the lack of concentration in the vanilla self-attention can lead to the failure to extract the most relevant information.
D DATA AUGMENTATION
Mixup Zhang et al. 2018) is popular data augmentation technique in computer vision and natural language processing, where random pairs data samples and their corresponding labels are linearly interpolates to generate synthetic data. However, directly applying Mixup to graph data remains challenging problem due to the characteristics of graphs: they have irregular node, not well-aligned node order and unique typology in non-Fuclideaan space. Although the multi-band connectomes
16

---

Published as a conference paper at ICLR 2025
in this study have regular node numbers and aligned node order, it is essential to ensure that the generated connectomes preserve the key topological features of the source connectomes during data augmentation. Therefore, G-Mixup (Han et al., 2022) is employed to augment connectomes, which linearly interpolates the estimated graphons from graphs and then generate synthetic graphs from the mixed graphons. Formally, can be described as
Graphon Estimation:
G1 Wg,, 92 W3,
(16)
Graphon Mixup:
Wz = AWg, + (1 g)
(17)
Graph Generation:
(I1 12, ... , In] G(k, Wz)
(18)
Label Mixup: yr Aye, + (1 A)ye.
(19)
where We, and Wg, are the estimated graphons, controls the contribution for both estimated graphons, I (2,, J J is the set of synthetic graphs generated from mixed graphon Wz, G(k, Wz) is a random graph with nodes, and the yo, and y/g, are the ground-truth labels.
Moreover, inspired by G-Mixup and CutMix (Yun et al., 2019), we extend this strategy into G- CutMix. Instead of applying CutMix directly to non-Euclidean graph data, we replace the removed structure in the graphon We, with the corresponding structure from Wg.. ensuring that the aug- mented connectomes retain critical topological features from the source connectomes. Finally, both G-Mixup and G-CutMix are employed as the data augmentation strategies in the XAlguiFormer.
E ROBUSTNESS ANALYSIS OF DIFFERENT EXPLAINERS
To assess the robustness of different XAI algorithms in the XAI guided attention mechanism, we employ two additional popular XAI methods, GradCAM and Integrated Gradients, as explainers. Tables and 8 present the performance achieved by these XAI methods on the TUAB and TD- BRAIN datasets. The results indicate that substituting the original DeepLift does not result in significant decline in performance.
fact, a slight improvement in the overall performance is ob- served when DeepLift is replaced with GradCAM. Consequently, XAlguiFormer is able to achieve stable performance without being over reliant on specifie XAI algorithm. In addition, the time complexity of different XAI methods impacts the computational efficiency of XAlguiFormer. Con- sidering both performance and efficiency, GradCAM and DeepLift are recommended as preferred explainers for the XAlguiFormer.
Table 7: Robustness analysis of different explainers on the TUAB dataset.
TUAB
Explainers
FLOPs
BAC
AUC-PR
AUROC
DeepLift
1.6G
0.8205 0.0027
0.8965 + 0.0079
0.9000 + 0.0046
GradCAM
0.95G
0.8240 + 0.0082
0.8963 + 0,0039
0.9010 + 0.0030
Integrated Gradients
35.7G
0.8210 0.0047
0.8923 + 0.0069
0.8964 + 0.0051
Table 8: Robustness analysis of different explainers on the TDBRAIN dataset.
TDBRAIN
Explainers
FLOPs
BAC
AUC-PR
AUROC
DeepLift
1.6G
0.6635 0.0080
0.5961 + 0.0136
0.7814 + 0.0156
GradCAM
0.95G
0.6553 + 0.0163
0.6149 + 0.0155
0.7996 + 0.0143
Integrated Gradients 35.7G
0.6569 + 0.0138
0.5979 + 0.0231
0.7874 + 0.0194
Furthermore, we evaluate the robustness of the calculated frequency band importance by changing the XAI algorithms. Figures (6] and illustrate the frequency band importance generated by Grad- CAM and Integrated Gradients, respectively. The rankings of the theta/beta ratio, high and low on TUAB remain consistent across different XAI methods. Similarly, on the TDBRAIN dataset, the rankings of the theta/beta ratio, low and high are largely consistent. These results demonstrate that the ranking of the most important frequency bands remains stable, indicating strong robustness across different XAI algorithms.
17

---

Published as a conference paper at ICLR 2025
(A)
Frequency Band Importance on TUAB
Frequency Band Importance on TDBRAIN
thetalbete retio
thetafbeta rate
0.14
high
low alpha
0.14
low alpha
0.10
Freuueicy
mid beta
0.12
Frequency
high alpha
0.09
gamma
theta
0.11
delta
0.09
mid beta
High acta
0.10
0.20
Contribution
Contributio
Figure 6: Frequency band importance generated by GradCAM on the TUAB and TDBRAIN datasets (A)
(B)
Frequency Band Importance TUAB
Frequency Band Importance TDBRAIN
theta/beta
2.24
thetw/heta
7.03
high alpha
4.86
high beta
low alpha
4.33
low alpha
3.89
Prequency
detta
4.00
Frequency
gamma
3.57
low beta
5:46
low beta
3.02
theta
high beta
high slpha
3.36
deta
Contribution
Figure 7: Frequency band importance generated by Integrated Gradients on the TUAB and TD- BRAIN datasets
F ABLATION STUDIES ON THE CONTROLLING PARAMETER &
To investigate the impact of the controlling parameter in the XAI guided loss function, we con- duct series of experiments with different settings ranging from 0.1 0.9 in increments of 0.1. Figure Saure&&presse presents the mean and standard deviation of the BAC across five different random seeds on the TUAB and TDBRAIN datasets, where GradCAM employed as the explainer within XAIgu- iFormer. The results indicate that the optimal value is 0.8 for TUAB and 0.9 for TDBRAIN, However, when alpha reaches 0.9 on TUAB, the BAC begins to decline. Therefore, we conclude that alpha = 0.7 0.8 serves experimental results indicate that when a is lower than 0.3, the performance declines significantly. This is likely because, with low a, the explainer's guidance becomes insufficient to meaningfully influence the learning process, leading to degraded performance.
generally effective setting for both datasets. In addition, these
(A)
Ablation Studies of the Controlling Parameter on TUAB
Ablation Studies of the Controlling Parameter on TDBRAIN
68
Ralancvel Acrurary (50)
Balanced Accuaacy
82
60
58
0.2
0.4
0.5
0.8
0.2
0.4
aloha
Figure 8: Ablation studies of the controlling parameter a on the TUAB and TDBRAIN datasets
18

---

Published as a conference paper at ICLR 2025
G
HYPERPARAMETER CONFIGURATIONS
Table Hyperparameter configurations of XAlguiFormer on the TUAB and TDBRAIN datasets.
Hyperparameters
TUAB
TDBRAIN
Input dimensions
19
26
Connectome tokenizer
Hidden dimensions
128
128
Layer number
4
4
Head number
MLP ratio
4
4
XArguiFormer
Layer scale init
le-3
1e-3
Layer number
12
12
Optimizer
AdamW
AdamW
Adam 8
(0.9,0.99)
(0.9,0.99)
Peak learning
4e-5
5e-5
Warmup learning rate
le-7
le-7
Learning rate scheduler
Cosine
Cosine
Training settings
Batch size
512
64
Weight decay
le-4
1e-5
Warmup epochs
5
5
Training epochs
50
100
Label smoothing
0.1
0.1
XAI guided loss weight
0.7
0.7
G-Mixup probability
1.0
1.0
Data augmentation
G-CutMix probability
0.5
0.5
H LIMITATIONS AND OUTLOOK
First, although we have evaluated XAlguiFormer on two large-scale clinical datasets (TUAB and TDBRAIN), they may not fully capture the diversity of real-world clinical scenarios. Our work represents an initial step toward exploring the positive impact of explanation-guidee learning on enhancing performance in the EEG field. Second, it should be noted that XAlguiFormer currently provides explanations at the frequency band level, rather than at the level of functional connectivity. This limitation arises because XAlguiFormer focuses token importance (query, key, and value vectors) to calculate refined attention values, which restricts explanations to the token level. As a result, the explanations are confined to the token (frequency band) level and do not offer detailed insights into the functional connectivity patterns. In the future, we plan to explore additional clinical datasets further verify the effectiveness of our model. Additionally. aim explore pretraining strategies for XAIguiFormer to further improve its performance.
19


Reviews:
[{"review_id": "qavu42DZuo", "reviewer": "Reviewer_A5ot", "summary": "This work proposed a dynamical-system-inspired architecture, XAI guided transformer (XAIguiFormer), where XAI not only provides explanations but also contributes to enhancing the transformer by refining the originally coarse information in the self-attention mechanism to capture more relevant dependency relationships.", "questions": "see above. The questions are about the description of warmup start XAI guidance and the robustness of the explainer module in the proposed method.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This method proposed a fusion strategy to explicitly inject the frequency and demographic information into tokens, improving the model’s understanding of the frequency and mitigating the negative effects of individual differences.\n2. This method proposed to use XAI to directly enhance transformer performance rather than focusing only on analyzing the visual interpretability.", "weaknesses": "1. Robustness is unclear. How about the performance after changing to other explainers in the module in addition to using DeepLift? How about the robustness?\n2. 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?\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "n0vDvzzU0I", "reviewer": "Reviewer_D4N1", "summary": "The paper introduces XAIguiFormer, a transformer model guided by explainable AI (XAI) techniques to enhance model performance. The authors integrate both frequency and demographic information into the model’s tokenization process, demonstrating that these features are essential for the observed performance improvements. XAIguiFormer achieves superior results compared to baseline models", "questions": "3. 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 dataset. 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? \n4. 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?\n5. Figure 4 illustrates the importance of difference frequency bands, could the author further clarity how the frequency information can be extracted from XAIguiFormer?\n6. Could the authors include some discussion on the limitations of the presented method?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors clearly define the specific issue they aim to address and effectively outline their solution approach. The introduction is clear and acknowledges the relevant literature and highlighting the limitations of current methods. The authors also provide a thorough comparison with a range of other methods and evaluate their approach on two public datasets. Another strength is that they intend to share their code upon acceptance, promoting transparency and reproducibility. The ablation studies also highlights the importance of both proposed suggestion (positional embedding and models guided by interpretability)", "weaknesses": "2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "TMwflx040X", "reviewer": "Reviewer_LoKT", "summary": "The paper introduces XAIguiFormer, a transformer model that uses explainable AI to both interpret and improve brain disorder detection from EEG data. The model features novel tokenization and encoding methods to preserve brain network patterns, achieving better performance than existing approaches while providing insights into which brain wave frequencies are most important for diagnosis.", "questions": "The proposed method assumes explanations from an imperfectly trained model can improve performance, would that propagate early training errors?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper uses XAI to actively improve model performance rather than just for interpretation.\n\nThe connectome tokennization approach and demographic-aware frequency encoding are also novel.\n\nThere are comprehensive experiments and ablation studies.", "weaknesses": "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 [1], especially it focuses on frequency domain as well. Experimental comparisons with these EEG models are also missed in the paper.\n\nIt seems like only temporal frequency is considered in positional encoding, how about spatial frequency for different brain regions? Some work like [2] developed spatial connectome based positioning, which should at least be discussed in the paper.\n\nLimited dataset scope: Only tested on two datasets (TUAB and TDBRAIN) which may not fully represent real-world clinical diversity.\n\nDiscussion of training time, computational efficiency compared to the baselines would be appreciated.\n\nThe model relies heavily on DeepLift's accuracy without exploring alternative explanation methods ot validating explanation quality.\n\n[1] Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. ICLR 2024.\n\n[2] Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking. NeurIPS 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "AjHYI3iKBM", "reviewer": "Reviewer_vgdo", "summary": "Authors proposed XAIguiFormer, the first framework to employ XAI for enhancing transformer performance in neuroimaging data. The explanability and accuracy are both improved as shown in experiments, while the improvement of accuracy might be not caused by the higher explanability.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Clear motivations.\n\n2. Good paper writing.\n\n3. Idea of explanability guiding self-attention is novel.", "weaknesses": "1. 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.\n\n2. 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.\n\n3. The proposed method leads to a new attention map that is more determined and hence easier to be explained. Even there is an qualitative assessment given neuroscience literatures, more details are needed on how the results in fig.4 were computed. Can you include a quantitative evaluation of how well the attention maps align with established neuroscience knowledge?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)