Community-Aware Graph Transformer for Brain Disorder Identification

1 Introduction

Neurological disorders, including autism spectrum disorder (ASD), Parkinson's disease (PD), and major depressive disorder (MDD), impact the quality of life for hundreds of millions of people worldwide [Marx et al., 2023; Lord et al., 2020]. These diseases typically involve complex neurobehavioral manifestations and neurobiological mechanisms, making accurate diagnosis a significant challenge [Li et al., 2021b]. Traditional diagnostic methods primarily rely on behavioral assessments and clinical observations, which are subjective and prone to misdiagnosis [Fusar-Poli et al., 2022]. In recent years, resting-state functional magnetic resonance imaging (rs-fMRI) has demonstrated tremendous research and diagnostic potential, it can reveal functional connections among brain regions of interest (ROIs) by assessing changes in blood oxygenation level dependent (BOLD) signals [Luo et al., 2024; Liu et al., 2024a]. By analyzing rs-fMRI data, researchers can identify abnormal patterns in brain functional network, which are crucial for early diagnosis and formulation of treatment strategies [Kong et al., 2024].

∗Corresponding author

Besides the whole-brain functional network, brain function depends on collaborative interactions among various regions that are organized into distinct functional communities, namely sub-networks. Within each community, the regions are tightly interconnected and work together to perform specific cognitive and physiological tasks [Van Den Heuvel and Pol, 2010]. For example, the default mode network (DMN) is associated with self-reflection and internal thought, and the dorsal attention network (DAN) is critical for attention regulation. However, traditional deep learning methods for brain network analysis often fail to incorporate these community-specific associations. As illustrated in Figure 1, traditional methods rely on the brain region-level interaction model, which mainly focus on modeling direct interactions between pairs of ROIs, overlooking the intricate relationships between functional communities. In contrast, a more natural and biologically reasonable model is multi-level collaboration that not only considers information transmission at the ROI-level but also accounts for information interaction at the community-level, facilitating functional activities through intra- and inter-community interactions [Sporns and Betzel,

2016]. To address this issue, we propose to effectively aggregate information within and between communities, integrating ROI-level and community-level interactions. This approach captures the functional representation of brain regions by encompassing interactive information across local and global scales, extracting richer spatial representation, and improving model performance.

To further enhance the representation of complex brain graph data, recent studies suggest utilizing Graph Transformer, which provide a more flexible mechanism for capturing long-range dependencies compared to traditional graph neural network (GNN) [Kong et al., 2024; Yu et al., 2024]. However, existing approaches still face limitations. Previous Transformer-based studies often treat brain region nodes as sequentially arranged nodes without appropriate positional encoding, hindering the model's ability to accurately perceive the spatial structure of the brain network. Additionally, existing methods utilize the shortest path or adjacency matrix to help the attention mechanism capture structural information in graphs. However, they fail to fully exploit the available spatial structural prior knowledge, such as spatial distance and community information, despite extensive evidence demonstrating their significant impact on brain information transfer processes [Sporns and Betzel, 2016; Bassett and Bullmore, 2017]. To address these issues, we incorporate prior knowledge to optimize the Transformer architecture. Specifically, we design a novel positional encoding suitable for brain networks by combining Montreal Neurological Institute (MNI) coordinates with graph random walk strategies, effectively enhancing the model's spatial awareness. Furthermore, by integrating spatial distance and community structure information, the attention mechanism aligns more closely with the natural information transfer processes of the brain. This approach not only enhances the model's performance but also improves its biological interpretability.

In summary, we propose the community-aware graph Transformer (CAGT) method that effectively leverages prior knowledge of brain network to provide a novel deep learning architecture for neurological disease diagnosis. The main contributions are as follows.

• We design a dual-scale spatial feature fusion module that utilizes community structure to extract richer spatial representations from local to global scales.
• We design a prior guided graph Transformer that utilizes prior knowledge of community structure and topological information to construct positional encodings and attention weights, enhancing the spatial awareness of Transformer
• We identify key biomarkers for disease diagnosis, including connections between functional communities, this indicates that our method enhances diagnostic performance while preserving interpretability.

2 RelatedWork

2.1 Local Information ProcessingMethods

The local information processing methods aggregate information through neighboring brain regions, in which GNN is a typical representative. Cui et al. proposed the brain graph neural network benchmark (BrainGB), which supports the combination of different node characteristics, message passing mechanisms, and pooling strategies, it is to standardize the brain network analysis process and provide a modular implementation framework [Cui et al., 2023]. Zhang et al. proposed BrainUSL, an unsupervised graph structure learning method that leverages graph generation, topology-aware encoding, and multi-view contrastive regularization to automatically learn discriminative network structures [Zhang et al., 2023]. Cui et al. proposed IBGNN, an interpretable framework that integrates an edge weight-aware message passing mechanism and learns a globally shared edge mask to identify brain regions and key connections associated with specific diseases [Cui et al., 2022]. Zheng et al. leverage the information bottleneck principle to construct disease-relevant subgraphs and evaluate their effectiveness of information retention by computing mutual information, improving diagnostic accuracy for psychiatric disorders and providing interpretable biomarkers [Zheng et al., 2024]. However, these methods overlook critical community information in the brain and do not break through the inherent limitations of GNN models. Because GNN is a local information processing method, due to its limited receptive field, long-range interaction between ROIs is achieved through increasing the number of layers, but multiple layers can lead to over smoothing.

2.2 Global Information ProcessingMethods

The global information processing methods can directly analyze the information interaction between any two brain regions, in which Transformer is a typical representative. Transformer uses self-attention mechanism to obtain global dependencies, enabling modeling of complex long-range interactions. Recent studies focus on improving the attention mechanism to better capture the intricate relationships in brain networks. For instances, Peng et al. enhanced feature discriminability by removing smaller singular value components from the attention matrix and applying geometric constraints, effectively learning discriminative graph representations of ROIs [Peng et al., 2024]. Cai et al. aimed to predict functional connections from structural connections, refining the attention mechanism to better model the relationship between brain structure and function [Cai et al., 2023]. Similarly, Kong et al. utilized graph topological information to guide the attention mechanism and extracted rich spatiotemporal information at multiple scales of brain networks [Kong et al., 2024]. However, these methods lack suitable positional encodings, which are proven to effectively enhance the spatial perception capabilities of Transformer in brain network analysis. To address this issue, Qu et al. introduced the gated graph Transformer (GGT), designing learnable structural positional encodings to optimize the learning of graph structural information, thereby enhancing cognitive ability prediction. [Qu et al., 2024]. Yu et al. designed a positional encoding based on random walk strategies, allowing the model to better understand long-range dependencies in brain networks [Yu et al., 2024]. However, these methods still do not incorporate community and spatial prior knowledge, which limits their effectiveness in enhancing the model's understanding of spatial relationships and information transmission mechanisms among ROIs.

3 Method

The overall framework of the proposed CAGT is presented in Figure 2. Firstly, the brain graph of each subject is constructed and partitioned into multiple community subgraphs. Then, a dual-scale spatial feature fusion module is designed to obtain representations of brain regions, which contain information exchange within and between community subgraphs. Finally, a prior guided graph Transformer is designed to realize brain disorder identification based on global information processing approach, in particular, the positional encoding is generated using a random walk strategy combined with MNI coordinates, the attention weights incorporate spatial distances and community structure information.

3.1 Brain Graph Construction

The brain is divided into N ROIs with a brain atlas, then the fMRI data of each subject can be represented as a matrix with the dimension of N × T , where N is the number of ROIs and T is the number of time points. By computing the Pearson correlation between each pair of ROIs, a symmetric functional connectivity (FC) matrix Xp ∈ RN×N is obtained, each row is the feature vector corresponding to a specific ROI. Xp records the whole-brain connectivity network of ROIs.

To detect the community structure in the brain network, we adopt the approach described in [Zhu et al., 2022] to rearrange the rows and columns of Xp. As shown in Figure 2 (a), the N ROIs are grouped into the C community modules, and the ROIs belonging to the same community are arranged together. This produces a rearranged FC matrix X ∈ RN×N and a community index Cidx that records the starting positions of each community module in the matrix. This rearrangement only changes the order of ROIs in the FC matrix, preserving the original connectivity information for each ROI. By retaining the top-k strongest functional connections for each ROI, we refine X to construct the topology of brain graph. This sparsification process extracts the most significant functional connections and eliminates redundant information.

As a result, we construct a brain graph G(V,E), where the adjacency matrix is denoted as A ∈ RN×N , and the node features are represented by X ∈ RN×N . This brain graph can be partitioned into C community subgraphs using the community index Cidx. For the c-th community subgraph Gc(Vc,Ec), the adjacency matrix Ac and the node features Xc are defined as follows:

where Ac ∈ RNc×Nc , Xc ∈ RNc×N , Nc is the number of nodes in the c-th community. The total number of nodes in all communities satisfies the equation N = N1 +N2 +N3 + · · · + NC. By partitioning the brain graph into C community subgraphs, it can facilitate the extraction of both intracommunity and inter-community aggregate information for further analysis.

3.2 Dual-Scale Spatial Feature Fusion

To capture information exchange within and between communities under the multi-level collaboration model, we extract features at both the node scale and the community scale, and further integrate these two levels to achieve a comprehensive local-to-global spatial representation of the brain network.

Local Information Extraction Within Communities

For each community subgraph, we employ the Graph Isomorphism Network with Edge Features (GINE) [Hu et al., 2019] to aggregate local information. The message-passing process is defined as:

where MLP(·) represents a multi-layer perceptron, ϵ is a learnable parameter, N(v) denotes the set of neighbors of node v, and eu,v is the edge feature between nodes u and v. Therefore, the community-internal information is updated

where Xc ∈ RNc×d and Hc ∈ RNc×d represent the initial and updated node feature within the community, respectively, and Ac is the adjacency matrix of the community subgraphs. Using convolution operations on each subgraph, it can effectively capture the features and relationships of ROIs within the community. The local feature maps are then concatenated as:

This process better reflects the functional connectivity among ROIs within each community.

Global Information Extraction Among Communities

To further extract global information, we perform a pooling operation on the internal features Hc of each community to obtain the community representation, the global feature map based on communities is obtained by concatenating them:

where P(·) denotes a combination of average pooling and max pooling, retaining essential features while smoothing the feature map. Next, we construct a connectivity matrix M ∈ RC×C based on the number of edges between communities, the element M(i, j) represents the number of functional connections between the i-th and j-th communities, and it is calculated as follows:

where Vi and Vj denote the node sets of the i-th and j-th communities, respectively. By normalizingM, we obtain the adjacency matrix of the global community graph:

This results in a community-level graph Gm(Vm,Em), where each node represents a community, and the edge weights reflect the strength of the functional connections between communities. Stronger connections correspond to larger edge weights. On the community-level graph, we also employ GINE to aggregate inter-community information, effectively capturing features and relationships between com-

Fusion of Local-to-Global Information

Inspired by the concept of residual connections, we integrate the local and global information with original feature X to obtain the final brain region feature map:

where the lin(·) represents a linear layer and operation expand(·) replicates Hg back to the pre-pooling dimensions based on the community indices:

and repeat(·) denotes a matrix replication operation.

3.3 Positional Encoding for Brain Networks

For non-Euclidean brain graph data, traditional sinusoidal positional encodings cannot effectively express the spatial relationships among ROIs. To better adapt to the unique characteristics of brain functional networks, we combine Random Walk Positional Encoding (RWPE) [Dwivedi et al., 2021] with MNI coordinates to create a unique positional encoding specifically tailored for brain functional networks. It can capture the centrality and importance of brain region nodes and their intricate spatial and topological relationships, significantly enhancing the Transformer's spatial awareness capabilities. Specifically, RWPE utilizes a random walk diffusion strategy. The probability of node i returning to itself after S steps of random walks is calculated as:

where D is the degree matrix, and A is the adjacency matrix of the graph. Based on this probability, the positional encoding after S steps of random walks is computed as:

It reflects the degree centrality of ROIs and global structure information of the brain graph.

The MNI coordinates directly provide the threedimensional spatial location of ROIs. By incorporating these coordinates into the positional encoding, the model is better to comprehend the relative positions of ROIs in threedimensional space. Specifically, suppose the coordinate of the i-th ROI is pi = [xi, yi, zi], the brain's spatial coordinate matrix PMNI is:

Then, the final positional encoding can be defined as:

where norm(·) normalizes the MNI coordinates, and ∥ denotes matrix concatenation. The positional encoding is added to the model input as follows:

where lin(·) maps the PE to a d-dimensional space. The final input ˜X ∈ RN×d integrates the local-to-global features Xlg with the positional encoding, allowing the model to capture both spatial and graph structural information.

3.4 Prior-GuidedMulti-head Attention

The attention mechanism in Transformer updates brain region information based on feature similarity between ROIs. However, neuroscience research reveals that brain information transmission is much more complex. It is strongly influenced by spatial distance and community structure, where physically adjacent regions tend to exhibit stronger functional connections and more frequent information exchanges. Moreover, connections within the same brain community are typically tighter, allowing efficient information transmission through functional synergy [Sporns and Betzel, 2016; Betzel and Bassett, 2017; Bassett and Bullmore, 2017]. Building on these findings, we integrate these prior knowledge into the multi-head self-attention mechanism, enabling the model to more accurately simulate brain information transfer process, enhancing interpretability and predictive accuracy.

Spatial Distance Weighting

The distance matrix D is computed using the Euclidean distance between nodes based on their MNI coordinates:

where (xi, yi, zi) and (xj , yj , zj ) denote the MNI coordinates of brain region nodes vi and vj, respectively. Subsequently, the distance weight matrix Dw is derived using a Gaussian kernel function:

where the exponential decay function increase the weight of short-distance connections, and α is a learnable parameter that is initially set to 1.

Community-based Weighting

We further design a community weight matrix that assigns higher weights to information transfer within the same community. For nodes within the same community, we set a learnable weight θc (initialized to 1) to form the community weight matrix Cw, the matrix is defined as follows:

Prior-Guided Attention Weight Matrix

By combining the distance weight matrix Dw and the community weight matrix Cw, we construct a prior-guided attention weight matrix that better reflects the brain's information transmission processes and activity mechanisms:

where F0 = ˜X, ∥ represents the concatenation operator, H is the number of attention heads, l is the layer index, and WlO , Wl,h Q ,Wl,h K ,Wl,h V are learnable model parameters. dl,h is the dimensionality of the query and key projections, σ represents softmax operator and λ1, λ2 are both learnable parameters initialized to 0.2 (via grid search). Ultimately, after L layers of the Transformer framework that integrates community spatial structure information, we can obtain feature map FL ∈ RN×d .

The final readout step involves aggregating global-level node embeddings to obtain an advanced representation of the brain graph. We employ the OCRead layer [Kan et al., 2022] to aggregate the learned node embeddings. The OCRead model begins by mapping the node data to a hidden space using an encoder. Then it computes standard orthogonal clustering centers E ∈ RK×N and employs a soft assignment mechanism to allocate nodes to these centers. The graph-level embedding ZG is calculated using the formula ZG = S · F L, where S ∈ RK×N is the learnable assignment matrix generated by OCRead. This process effectively captures the intricate patterns of brain activity. Finally, the graph-level embedding ZG undergoes dimension reduction and flattening before being fed into a multi-layer perceptron to produce accurate classification predictions. The entire process is optimized using supervised learning with cross-entropy loss.

4 Experiments and Results

4.1 Materials

We extensively evaluate the proposed method on three datasets, namely ABIDE I for autism [Craddock et al., 2013], REST-MDD for major depressive disorder [Yan et al., 2019] and Taowu and Neurocon for Parkinson's disease [Badea et al., 2017], where ABIDE I and REST-meta-MDD are large datasets with thousands of subjects, TaoWu and Neurocon is a relatively small dataset with about a hundred subjects.

ABIDE I is an open-source ASD diseases database comprising data from 17 international sites, available at https: //fcon 1000.projects.nitrc.org/indi/abide. The Preprocessed Connectomes Project (PCP) has preprocessed the fMRI data for each subject, resulting in rs-fMRI data for 1,035 subjects, including 505 ASD patients and 530 normal controls (NC).

REST-MDD is a publicly available MDD diseases database comprising 25 study cohorts, accessible at http:

//rfmri.org/REST-meta-MDD. The data were preprocessed using the Data Processing Assistant for Resting-State fMRI (DPARSF) toolbox [Yan et al., 2016]. Following the official recommendation to exclude overlapping data from site S4, the refined dataset consists of 2,380 rs-fMRI scans, including 1,276 MDD patients and 1,104 NC.

TaoWu and Neurocon datasets are among the earliest image datasets available for Parkinson's research. We followed the preprocessing steps outlined in [Liu et al., 2024b] using the DPARSF toolbox. The TaoWu dataset ultimately includes 39 subjects, with 20 PD patients and 19 NC. The Neurocon dataset includes 43 subjects, with 27 PD patients and 16 NC.

For each subject, the entire brain is parcellated into 200 regions using the CC200 atlas [Craddock et al., 2012]. These regions are further clustered into eight sub-networks based on the previous research [Yeo et al., 2011; Lawrence et al., 2021], namely Cerebellum and Subcortical Structures (CB & SC), Visual Network (VN), Somatomotor Network (SMN), DAN, Ventral Attention Network (VAN), Limbic Network (LN), Frontoparietal Network (FPN), and DMN.

4.2 Experimental Setting

We evaluate the model's performance using ten-fold crossvalidation. The entire network is trained end-to-end with the Adam optimizer. Key parameters included a batch size of 64, a total of 70 epochs, an initial learning rate of 1 × 10-4 , and a weight decay of 1 × 10-6 . The top-k connections for each node are retained as edges, with their values fixed at 30. The RWPE dimension S is set to 30. The Transformer architecture is configured with 2 layers and 8 heads. A dropout rate of 0.2 is applied to the GINE, Transformer, and the final fully connected classification layer. All codes are implemented using the PyTorch and PyG libraries. Experiments are conducted on a Windows server equipped with an Intel® Core™ i7-10700 CPU (2.90 GHz), a GeForce GTX 3080 Ti GPU and 32 GB of RAM. The code is available at https://github.com/null-cks/CAGT.

4.3 Comparison with ExistingMethods

We compare our CAGT model with two types of models: 1) GNN-based models, including GINE [Hu et al., 2019] , BrainGNN [Li et al., 2021a] , IBGNN [Cui et al., 2022] , BrainUSL [Zhang et al., 2023], and BrainIB [Zheng et al., 2024]; 2) Transformer-based models designed for graph data, including vanilla-Transformer (vanillaTF), BrainNetTransformer (BrainNetTF) [Kan et al., 2022], Com-BrainTF [Bannadabhavi et al., 2023], ALTER [Yu et al., 2024], GBT [Peng et al., 2024], Contrasformer [Xu et al., 2024] and GraphGPS [Ramp´aˇsek et al., 2022]. For these models, we used the open-source code from the original papers. We modified only the validation method to ten-fold cross-validation to align with our experimental setup, while maintaining the original model architectures.

Table 1 presents the results of three classification tasks in three public datasets, using accuracy (ACC), sensitivity (SEN), and specificity (SPE) as primary metrics. As can be seen, Graph Transformer-based methods generally outperform GNN-based methods on large datasets like ABIDE I and REST-MDD, but struggle with small datasets such as Tawowu and Neurocon due to their higher parameter count, which increases the risk of overfitting. Large datasets such as REST-MDD exhibit more consistent data distributions across folds, leading to smaller standard deviations, while small datasets such as Taowu and Neurocon show greater variability. Notably, our CAGT model consistently achieves superior performance in all tasks, with accuracy improvements of 2.72% for ASD, 1.19% for MDD, and 8.46% for PD. This demonstrates the advantage of our approach in effectively extracting comprehensive representations of brain connectomes by leveraging sub-network community characteristics and brain spatial topological information.

4.4 Ablation Studies

To evaluate the effectiveness of the CAGT model's components, we conduct ablation studies across three tasks, focusing on the dual-scale spatial feature fusion module (DF), positional encoding (PE), and prior-guided multi-head attention (PA). The results are summarized in Table 2. The DF module enhances spatial representation by fusing local features of sub-networks with global features of the whole brain, achieving performance improvements of 2.92%, 2.96%, and 13.07% across the three tasks. This demonstrates the effectiveness of leveraging community-related information. The PA module further strengthens long-range interactions between brain regions and incorporates prior knowledge of brain information transmission mechanisms, contributing significantly to performance gains of 5.87%, 3.38%, and 9.7% in the respective tasks. Additionally, the inclusion of PE improves the model's spatial awareness of the brain network, increasing accuracy by 1.54%, 1.07%, and 3.26%. These results confirm the positive impact of each component on the model's diagnostic capabilities.

4.5 Biomarker Detection

CAGT can also detect key functional communities and brain regions associated with brain disorder diagnosis. Taking ABIDE I as an example, we separately compute the average attention scores Sasd and Snc for correctly classified ASD patients and NC subjects obtained in our model. Then, the inter-community relationships is quantified by averaging the attention scores among their corresponding ROIs, this results in two inter-community connectivity matrix ¯Sasd ∈ R8×8 and ¯Snc ∈ R8×8 . By comparing ¯Sasd and ¯Snc, we can detect abnormal connectivity patterns in ASD patients. Specifically, the DMN shows the highest attention weight and demonstrates stronger connectivity with both VAN and DAN in ASD group, aligning with findings reported in [Padmanabhan et al., 2017]. Conversely, the SMN exhibits weaker connectivity with other sub-networks in ASD group, reflecting its impact on sensory and motor processing due to diminished connectivity [Marco et al., 2011; Mostofsky and Ewen, 2011]. The inter-community relationships of ASD is shown in Figure 3(a), visualizing the distinctive connectivity patterns. Furthermore, Figure 3(b) and 3(c) present the top thirty most significant functional connections and the top ten critical brain regions, identified by sparsifying the attention score matrix Sasd. Key brain regions including the frontal gyrus, precentral gyrus, putamen and insula are known to play crucial roles in sensory, motor, and cognitive processing. These results align with the existing neuroscience literature [Sussman et al., 2015; Plitt et al., 2015; Sapey-Triomphe et al., 2023].

5 Conclusion

In this work, we propose a novel framework that integrates sub-network community information and brain graph topology into the Transformer architecture to enhance brain disorder identification. The dual-scale spatial feature fusion module effectively combines local community-level features with global whole-brain features, enhances the model's ability to capture hierarchical brain interactions. The prior-guided graph Transformer incorporates brain-specific positional encoding and adds prior knowledge of brain information transmission mechanisms to the Transformer, aligning with brain biological characteristics and improving model interpretability. Our method not only achieves superior performance compared to existing methods but also reliably identifies key biomarkers associated with neurological disorders.

Acknowledgments

This work was supported in part by the Natural Science Research Project of Anhui Educational Committee under Grant 2023AH050081; in part by the Collaborative Innovation Project of Key Laboratory of Philosophy and Social Sciences in Anhui Province under Grant HZ2302; in part by the Scientific Research Project of Colleges and Universities in Anhui Province under Grant 2024AH040115; and in part by the Anhui Province Science and Technology Innovation and Tackling Key Problems Project under Grant 202423k09020041.