B RAIN N ETWORK T RANSFORMER
Abstract
Human brains are commonly modeled as networks of Regions of Interest (ROIs) and their connections for the understanding of brain functions and mental disorders.
Recently, Transformer-based models have been studied over different types of data, including graphs, shown to bring performance gains widely.
In this work, we study Transformer-based models for brain network analysis.
Driven by the unique properties of data, we model brain networks as graphs with nodes of fixed size and order, which allows us to (1) use connection profiles as node features to provide natural and low-cost positional information and (2) learn pairwise connection strengths among ROIs with efficient attention weights across individuals that are predictive towards downstream analysis tasks.
Moreover, we propose an O RTHONORMAL C LUSTERING R EADOUT operation based on selfsupervised soft clustering and orthonormal projection.
This design accounts for the underlying functional modules that determine similar behaviors among groups of ROIs, leading to distinguishable cluster-aware node embeddings and informative graph embeddings.
Finally, we re-standardize the evaluation pipeline on the only one publicly available large-scale brain network dataset of ABIDE, to enable meaningful comparison of different models.
Experiment results show clear improvements of our proposed B RAIN N ETWORK T RANSFORMER on both the public ABIDE and our restricted ABCD datasets.
The implementation is available at https://github.com/Wayfear/BrainNetworkTransformer .
Introduction
Brain network analysis has been an intriguing pursuit for neuroscientists to understand human brain organizations and predict clinical outcomes [50, 59, 58, 5, 18, 27, 52, 29, 58, 28, 41, 44, 31].
Among various neuroimaging modalities, functional Magnetic Resonance Imaging (fMRI) is one of the most commonly used for brain network construction, where the nodes are defined as Regions of Interest (ROIs) given an atlas, and the edges are calculated as pairwise correlations between the blood-oxygenlevel-dependent (BOLD) signal series extracted from each region [54, 53, 59, 16].
Researchers observe that some regions can co-activate or co-deactivate simultaneously when performing cognitiverelated tasks such as action, language, and vision.
Based on this pattern, brain regions can be classified into diverse functional modules to analyze diseases towards their diagnosis, progress understanding and treatment.
Nowadays Transformer-based models have led a tremendous success in various downstream tasks across fields including natural language processing [56, 17] and computer vision [20, 10, 55].
Recent efforts have also emerged to apply Transformer-based designs to graph representation learning.
GAT [57] firstly adapts the attention mechanism to graph neural networks (GNNs) but only considers the local structures of neighboring nodes.
Graph Transformer [21] injects edge information into the attention mechanism and leverages the eigenvectors of each node as positional embeddings.
SAN [40] further enhances the positional embeddings by considering both eigenvalues and eigenvectors and improves the attention mechanism by extending the attention from local to global structures.
Graphomer [64], which achieves the first place on the quantum prediction track of OGB Large-Scale
Challenge [30], designs unique mechanisms for molecule graphs such as centrality encoding to enhance node features and spatial/edge encoding to adapt attention scores.
However, brain networks have several unique traits that make directly applying existing graph
Transformer models impractical.
First, one of the simplest and most frequently used methods to construct a brain network in the neuroimaging community is via pairwise correlations between
BOLD time courses from two ROIs [43, 35, 13, 63, 69].
This impedes the designs like centrality, spatial, and edge encoding because each node in the brain network has the same degree and connects to every other node by a single hop.
Second, in previous graph transformer models, eigenvalues and eigenvectors are commonly used as positional embeddings because they can provide identity and positional information for each node [15, 26].
Nevertheless, in brain networks, the connection profile, which is defined as each node’s corresponding row in the brain network adjacency matrix, is recognized as the most effective node feature [13].
This node feature naturally encodes both structural and positional information, making the aforementioned positional embedding design based on eigenvalues and eigenvectors redundant.
The third challenge is scalability.
Typically, the numbers of nodes and edges in molecule graphs are less than 50 and 2500, respectively.
However, for brain networks, the node number is generally around 100 to 400, while the edge number can be up to 160,000.
Therefore, operations like the generation of all edge features in existing graph transformer models can be time-consuming, if not infeasible.
[Figure]
In this work, we propose to develop B RAIN N ETWORK T RANSFORMER (B RAIN N ET TF), which leverages the unique properties of brain network data to fully unleash the power of Transformer-based models for brain network analysis.
Specifically, motivated by previous findings on effective GNN designs for brain networks [13], we propose to use the effective initial node features of connection profiles.
Empirical analysis shows that connection profiles naturally provide positional features for Transformer-based models and avoid the costly computations of eigenvalues or eigenvectors.
Moreover, recent work demonstrates that GNNs trained on learnable graph structures can achieve superior effectiveness and explainability [35].
Inspired by this insight, we propose to learn fully pairwise attention weights with Transformer-based models, which resembles the process of learning predictive brain network structures towards downstream tasks.
One step further, when GNNs are used for brain network analysis, a graph-level embedding needs to be generated through a readout function based on the learned node embeddings [37, 43, 13].
As is shown in Figure 1(a), a property of brain networks is that brain regions (nodes) belonging to the same functional modules often share similar behaviors regarding activations and deactivations in response to various stimulations [7].
Unfortunately, the current labeling of functional modules is rather empirical and far from accurate.
For example, [3] provides more than 100 different functional module organizations based on hierarchical clustering.
In order to leverage the natural functions of brain regions without the limitation of inaccurate functional module labels, we design a new global pooling operator, O RTHONORMAL C LUSTERING R EADOUT , where the graph-level embeddings are pooled from clusters of functionally similar nodes through soft clustering with orthonormal projection.
Specifically, we first devise a self-supervised mechanism based on [60] to jointly assign soft clusters to brain regions while learning their individual embeddings.
To further facilitate the learning of clusters and embeddings, we design an orthonormal projection and theoretically prove its effectiveness in distinguishing embeddings across clusters, thus obtaining expressive graph-level embeddings after the global pooling, as illustrated in Figure 1(b).
Finally, the lack of open-access datasets has been a non-negligible challenge for brain network analysis.
The strict access restrictions and complicated extraction/preprocessing of brain networks from fMRI data limit the development of machine learning models for brain network analysis.
Specifically, among all the large-scale publicly available fMRI datasets in literature, ABIDE [6] is the only one provided with extracted brain networks fully accessible without permission requirements.
However, ABIDE is aggregated from 17 international sites with different scanners and acquisition parameters.
This inter-site variability conceals inter-group differences that are really meaningful, which is reflected in the unstable training performance and the significant gap between validation and testing performance in practice.
To address these limitations, we propose to apply a stratified sampling method in the dataset splitting process and standardize a fair evaluation pipeline for meaningful model comparison on the ABIDE dataset.
Our extensive experiments on this public ABIDE dataset and a restricted ABCD dataset [8] show significant improvements brought by our proposed B RAIN
N ETWORK T RANSFORMER .
Background and Related Work 2.1 GNNs for Brain Network Analysis
Recently, emerging attention has been devoted to the generalization of GNN-based models to brain network analysis [42, 2].
GroupINN [62] utilizes a grouping-based layer to provide explainability and reduce the model size.
BrainGNN [43] designs the ROI-aware GNNs to leverage the functional information in brain networks and uses a special pooling operator to select these crucial nodes.
IBGNN [14] proposes an interpretable framework to analyze disorder-specific ROIs and prominent connections.
In addition, FBNetGen [35] considers the learnable generation of brain networks and explores the explainability of the generated brain networks towards downstream tasks.
Another benchmark paper [13] systematically studies the effectiveness of various GNN designs over brain network data.
Different from other work focusing on static brain networks, STAGIN [39] utilizes
GNNs with spatio-temporal attention to model dynamic brain networks extracted from fMRI data.
2.2 Graph Transformer
Graph Transformer raises many researchers’ interest currently due to its outstanding performance in graph representation learning.
Graph Transformer [21] firstly injects edge information into the attention mechanism and leverages the eigenvectors as positional embeddings.
SAN [40] enhances the positional embeddings and improves the attention mechanism by emphasizing neighbor nodes while incorporating the global information.
Graphomer [64] designs unique mechanisms for molecule graphs and achieves the SOTA performance.
Besides, a fine-grained attention mechanism is developed for node classification [68].
Also, the Transformer is extended to larger-scale heterogeneous graphs with a particular sampling algorithm in HGT [32].
EGT [33] further employs edge augmentation to assist global self-attention.
In addition, LSPE [22] leverages the learnable structural and positional encoding to improve GNNs’ representation power, and GRPE [49] enhances the design of encoding node relative position information in Transformer.
B RAIN N ETWORK T RANSFORMER 3.1 Problem Definition
In brain network analysis, given a brain network X ∈ R V × V , where V is the number of nodes (ROIs), the model aims to make a prediction indicating biological sex, presence of a disease or other properties of the brain subject.
The overall framework of our proposed B RAIN N ETWORK
T RANSFORMER is shown in Figure 2, which is mainly composed of two components, an L -layer attention module MHSA and a graph pooling operator OCR EAD .
Specifically, in the first component of MHSA , the model learns attention-enhanced node features Z L through a non-linear mapping
X → Z L ∈ R V × V .
Then the second component of OCR EAD compresses the enhanced node embeddings Z L to graph-level embeddings Z G ∈ R K × V , where K is a hyperparameter representing the number of clusters.
Z G is then flattened and passed to a multi-layer perceptron for graph-level predictions.
The whole training process is supervised with the cross-entropy loss.
[Figure]
Linear 𝑀
Project
Flatten
Biological sex,Autism spectrum disorder…
T i m e fMRI
ROI#
ROI#
ROI#3 ROI# 𝐸 !
×𝐿 Layers 𝒁 !
𝑷
FCN
OCRead MHSA
[Figure]
3.2 Multi-Head Self-Attention Module (MHSA)
To develop a powerful Transformer-based model suitable for brain networks, two fundamental designs, the positional embedding and attention mechanism, need to be reconsidered to fit the natural properties of brain network data.
In existing graph transformer models, the positional information is usually encoded via eigendecomposition, while the attention mechanism often combines node positions with existing edges to calculate the attention scores.
However, for the dense (often fully connected) graphs of brain networks, eigendecomposition is rather costly, and the existence of edges is hardly informative.
ROI node features on brain networks naturally contain sufficient positional information, making the positional embeddings based on eigendecomposition redundant.
Previous work on brain network analysis has shown that the connection profile
X i · for node i , defined as the corresponding row for each node in the edge weight matrix X , always achieves superior performance over others such as node identities, degrees or eigenvector-based embeddings [43, 35, 13].
With this node feature initialization, the self-connection weight x ii on the diagonal is always equal to one, which encodes sufficient information to determine the position of each node in a fully connected graph based on the given brain atlas.
To verify this insight, we also empirically compare the performance of the original connection profile with two variants concatenated with additional positional information, i.e. , connection profile w/ identity feature and connection profile w/ eigen feature.
The results indeed show no benefit brought by the additional computations (c.f. Appendix B).
As for the attention mechanism, previous work [13] has empirically demonstrated that integrating edge weights into the attention score calculation can significantly degrade the effectiveness of attention on complete graphs, while the generation of edge-wise embedding can be unaffordable given a large number of edges in brain networks.
On the other hand, the existence of edges provides no useful information for the computation of attention scores as well because all edges simply exist in complete graphs.
Based on the observations above, we design the basic B RAIN N ETWORK T RANSFORMER by (1) adopting the connection profile as initial node features and eliminating any extra positional embeddings and (2) adopting the vanilla pair-wise attention mechanism without using edge weights or relative position information to learn a singular attention score for each edge in the complete graph.
Formally, we leverage a L -layer non-linear mapping module, namely Multi-Head Self-Attention ( MHSA ), to generate more expressive node features Z L = MHSA( X ) ∈ R V × V .
For each layer l , the output Z l is obtained by
[Equation]
where Z 0 = X , ∥ is the concatenation operator, M is the number of heads, l is the layer index,
W l O , W l,m Q , W l,m K , W l,m V are learnable model parameters, and d l,m K is the first dimension of W l,m K .
3. O RTHONORMAL C LUSTERING R EADOUT (OCR EAD )
The readout function is an essential component to learn the graph-level representations for brain network analysis ( e.g. , classification), which maps a set of learned node-level embeddings to a graphlevel embedding.
Mean( · ) , Sum( · ) and Max( · ) are the most commonly used readout functions for
GNNs.
Xu et al.
[61] show that GNNs equipped with Sum( · ) readout have the same discriminative power as the Weisfeiler-Lehman Test.
Zhang et al.
[66] propose a sort pooling to generate the graph-level representation by sorting the final node representations.
Ju et al.
[34] present a layer-wise readout by extending the node information aggregated from the last layer of GNNs to all layers.
However, none of the existing readout functions leverages the properties of brain networks that nodes in the same functional modules tend to have similar behaviors and clustered representations, as shown in Figure 1(a).
To address this deficiency, we design a novel readout function to take advantage of the modular-level similarities between ROIs in brain networks, where nodes are assigned softly to well-chosen clusters with an unsupervised process.
Formally, given K cluster centers, each center has V dimensions, E ∈ R K × V , a Softmax projection operator is used as the function to calculate the probability P ik of assigning node i to cluster k ,
P ik =
[Equation]
where ⟨· , ·⟩ denotes the inner product and Z L is the learned set of node embeddings from the last layer of MHSA module.
With this computed soft assignment P ∈ R V
×
K , the original learned node representation Z L can be aggregated under the guidance of the soft cluster information, where the graph-level embedding Z G is obtained by Z G
= P ⊤ Z L .
However, jointly learning node embeddings and clusters without ground-truth cluster labels is difficult.
To obtain representative soft assignment P , the initialization of K cluster centers E is critical and should be designed delicately.
To this end, we leverage the observation illustrated in Figure 1(b), where orthonormal embeddings can improve the clustering of nodes in brain networks w.r.t.
the functional modules underlying brain regions.
Orthonormal Initialization .
To initialize a group of orthonormal bases as cluster centers, we first adopt the Xavier uniform initialization [25] to initialize K random centers and each center contains
V dimensions C ∈ R K × V .
Then, we apply the Gram-Schmidt process to obtain the orthonormal bases E , where
[Equation]
In the next section, we theoretically prove the advantage of this orthonormal initialization.
3.3.1 Theoretical Justifications
In OCR EAD , proper cluster centers can generate higher-quality soft assignments and enlarge the difference between P from different classes.
[51, 46] showed the advantages of orthogonal initialization in DNN model parameters.
However, none of them proves whether it is an ideal strategy to obtain the cluster centers.
We propose two methods from the perspective of statistics as follows.
Firstly, to discern features of different nodes, we would expect a larger discrepancy among their similarity probabilities indicated from the readout.
One way to measure the discrepancy is using the variance of P for each feature.
Let ¯ P ≡ 1 /K denote the mean of any discrete probabilities with K values.
Variance of P measures the difference between P and ¯ P .
We average over the feature vector space: if the result is small, then there is a large tendency that different P approaches ¯ P and hence cannot be discerned easily.
Specifically, the following theorem holds for our function Eq.
Theorem 3.1.
For arbitrary r > 0 , let B r = {Z ∈ R V ; ∥Z∥ ≤ r } denote the round ball centered at origin of radius r with Z being fracture vectors.
Let V r be the volume of B r .
The variance of
Softmax projection averaged over
[Equation]
Despite the concise form, it is unclear whether the above integral has an elementary antiderivative.
Even though, we can circumvent this problem and a rigorous proof is given in Appendix C.
The second statistical method shows that for general readout functions without a known analytical form, initializing with orthonormal cluster centers has a larger probability of gaining better performance.
To set up the proper statistical scenario, we assume that the unknown readout is obtained by a regression of some samples ( ˆ Z ( s ) , ˆ E ( t ) , ˆ P ( st ) ) .
This formally converts the exact functional relationship between Z i · , E k · and P ik to a statistical relationship :
[Equation]
with P T being the probability truly reflecting similarities between nodes and clusters and ϵ i denoting the stochastic error.
It is almost impossible to find P T , but by computing the so-called variation inflation factor [47], we show that regression in orthonormal case has a higher accuracy than that in non-orthonormal case.
Combining with a hypothesis testing, we obtain the following
Theorem 3.2.
The significance level α E k · which reveals the probability of rejecting a well-estimated pooling is lower when sampling from orthonormal centers than that from non-orthonormal centers.
More details can be seen in Appendix C.
3.4 Generalizing OCR EAD to Other Graph Tasks and Domains
In this work, we tested the proposed OCR EAD on functional connectivity (FC) based brain networks.
Other popular modalities of brain networks include structural connectivities (SC), which describe the anatomical organization of the brain by measuring the fiber tracts between brain regions [4].
In
SC-based brain networks, ROIs that are positionally close to each other on the structural connectivity networks tend to share similar connection profiles.
This means the idea of OCR EAD is also naturally applicable to SC networks, where the orthonormal clustering is based on the physical distances instead of the functional modules on FC.
At a higher level, the idea of our proposed OCR EAD is not confined to graph-level prediction tasks on brain networks but can also be generalized to other graph learning tasks and domains.
Precisely, there is a growing tendency in node/edge level prediction tasks to enhance the node/edge representation learning by utilizing the subgraph embeddings around each target node/edge [67, 65].
In this process, substructure learning needs to be performed on the subgraphs, where our proposed OCR EAD can be adapted for compressing a set of node embeddings to subgraph embeddings.
Besides, OCR EAD is also potentially useful for other types of graphs in the biomedical domains.
For example, for proteinprotein interaction networks, proteins can be implicitly grouped by families that share common evolutionary origins [48], whereas for gene expression networks, genes can be grouped based on the latent pathway information [36].
Both of them are potential directions for the future application of OCR EAD , among many others driven by biological or other types of prior knowledge regarding underlying node/edge groups.
Experiments
This section evaluates the effectiveness of our proposed B RAIN N ETWORK T RANSFORMER (B RAIN N ET TF) with extensive experiments.
We aim to address the following research questions:
RQ1.
How does B RAIN N ET TF perform compared with state-of-the-art models of various types?
RQ2.
How does our proposed OCR EAD module perform with different model choices?
RQ3.
Does the learned model of B RAIN N ET TF exhibit consistency with existing neuroscience knowledge and suggest reasonable explainability?
4.1 Experimental Settings
Datasets.
We conduct experiments on two real-world fMRI datasets.
(a) Autism Brain Imaging
Data Exchange (ABIDE) :
This dataset collects resting-state functional magnetic resonance imaging (rs-fMRI) data from 17 international sites, and all data are anonymous [6].
The used dataset contains brain networks from 1009 subjects, with 516 (51.14%) being Autism spectrum disorder
(ASD) patients (positives).
The region definition is based on Craddock 200 atlas [12].
As the most convenient open-source large-scale dataset, it provides generated brain networks and can be downloaded directly without permission request.
Despite the ease of acquisition, the heterogeneity of the data collection process hinders its use.
Since multi-site data are collected from different scanners with different acquisition parameters, non-neural inter-site variability may mask inter-group differences.
In practice, we find the training unstable, and there is a significant gap between validation and testing performances.
However, we discover that most models can achieve a stable performance if we follow an appropriate stratified sampling strategy by considering collection sites during the training-validation-testing splitting process for ABIDE.
Training curves in Appendix A also show how different models achieve a stabler performance on our designed new splitting settings than the random splitting.
Therefore, we use ABIDE as one of the benchmark datasets in this work, and we share our re-standardized data splitting to provide a fair evaluation pipeline for various future methods.
(b) Adolescent Brain Cognitive Development Study (ABCD) :
This is one of the largest publicly available fMRI datasets with restricted access (a strict data requesting process needs to be followed to obtain the data) [8].
The data we use in the experiments are fully anonymized brain networks with only biological sex labels.
After the quality control process, 7901 subjects are included in the analysis, with 3961 (50.1%) among them being female.
The region definition is based on the
HCP 360 ROI atlas [24].
Metrics.
The diagnosis of ASD is the prediction target on ABIDE, while biological sex prediction is used as the evaluation task for ABCD.
Both prediction tasks are binary classification problems, and both datasets are balanced between classes.
Hence, AUROC is a proper performance metric adopted for fair comparison at various threshold settings, and accuracy is applied to reflect the prediction performance when the threshold is 0.5.
Besides, since the model is mainly for medical applications, we add two critical metrics for diagnostic tests, Sensitivity and Specificity, which respectively refer to true positive rate and true negative rate.
All reported performances are the average of 5 random runs on the test set with the standard deviation.
Implementation details.
For experiments, we use a two-layer Multi-Head Self-Attention Module and set the number of heads M to 4 for each layer.
We randomly split 70% of the datasets for training, 10% for validation, and the remaining are utilized as the test set.
In the training process of
B RAIN N ET TF, we use an Adam optimizer with an initial learning rate of 10 − 4 and a weight decay of 10 − 4 .
The batch size is set as 64.
All models are trained for 200 epochs, and the epoch with the highest AUROC performance on the validation set is used for performance comparison on the test set.
The model is trained on an NVIDIA Quadro RTX 8000.
Please refer to the repository and
Appendix G for the full implementation of B RAIN N ET TF.
Computation complexity.
In B RAIN N ET TF, the computation complexity of Multi-Head SelfAttention Module and OCR EAD are O ( LMV 2 ) and O ( KV ) respectively, where L is the layer number of Multi-Head Self-Attention Module, V is the number of nodes, M is the number of heads, and K is the number of clusters in OCR EAD .
The overall computation complexity of B RAIN N ET TF is thus O ( V 2 ) , which is on the same scale as common GNNs on brain networks such as BrainGNN [43] and BrainGB [13].
4.2 Performance Analysis (RQ1)
We compare FBNETGEN with baselines of three types.
The details about how to tune hyperparameters of various baselines can be found in Appendix F. Besides, Appendix E shows the comparison of the number of parameters between our model and other baseline models, which shows that the parameter size of B RAIN N ET TF is larger than GNN and CNN models but smaller than other transformer models.
(a) B RAIN N ET TF vs. other graph transformers.
We compare B RAIN N ET TF with two popular graph Transformers, SAN [40] and Graphormer [64].
In addition, we also include a basic version of B RAIN N ET TF without OCR EAD , composed of a Transformer with a 2-layer
Multi-Head Self-Attention and a CONCAT-based readout named VanillaTF.
Our B RAIN N ET TF outperforms SAN and Graphormer by significant margins, with up to 6% absolute improvements on both datasets.
VanillaTF also surpasses SAN and Graphormer.
We believe this downgraded performance of existing graph transformers results from their design flaws facing the natures of brain networks.
Specifically, both the preprocessing and the training stages of the Graphormer model accepts only discrete, categorical data.
A bin operator has to be applied on the adjacency matrix, coarsening the node feature from connection profiles and dramatically hurting the performance.
Furthermore, since brain networks are complete graphs, key designs like centrality encoding and spatial encoding of Graphormer cannot be appropriately applied.
Similarly, for SAN, experiments in
Appendix B show that adding eigen node features to connection profiles cannot improve the model’s performance.
Besides, the benchmark paper [13] reveals that injecting edge weights into the attention mechanism can significantly reduce the prediction power.
Furthermore, Appendix D shows our
B RAIN N ET TF is much faster than other graph transformers due to special optimizations towards brain networks.
(b) B RAIN N ET TF vs. neural network models on fixed brain networks.
We further introduce another three neural network baselines on fixed brain networks.
BrainGNN [43] designs ROI-aware GNNs for brain network analysis.
BrainGB [13] is a systematic study of how to design effective GNNs for brain network analysis.
We adopt their best design as the BrainGB baseline.
BrainnetCNN [37] represents state-of-the-art of specialized GNNs for brain network analysis, which models the adjacency matrix of a brain network similarly as a 2D image.
As is shown in Table 1,
B RAIN N ET TF consistently outperforms BrainGNN, BrainGB and BrainnetCNN.
(c) B RAIN N ET TF vs. neural network models on learnable brain networks.
Unlike classical GNNs, FBNETGEN [35],
DGM [38] and BrainNetGNN [45] hold a similar idea, which is to apply GNNs based on a learnable graph.
FBNETGEN achieves SOTA performance on the ABCD dataset for biological sex prediction, and the learnable graphs can be seen as a type of attention score.
Experiment results show that our proposed B RAIN N ET TF beats all three of them on both datasets.
[Table]
4.3 Ablation Studies on the OCR EAD Module (RQ2) 4.3. OCR EAD with varying readout functions
We vary the readout function for various Transformer architectures, including SAN, Graphormer and
VanillaTF, to observe the performance of each ablated model variant.
The results shown in Tabledemonstrate that our OCR EAD is the most effective readout function for brain networks and improves the prediction power across various Transformer architectures
[Table]
[Table] [Table]
OCR EAD with varying cluster initializations
To further demonstrate how the design of OCR EAD influences the performance of B RAIN N ET TF
[Table]
we investigate two key model selections, the initialization method for cluster centers and the cluster number K .
For the initialization, three different kinds of initialization procedures are compared,
namely (a) Random : the Xavier uniform [25] is leveraged to randomly generate a group of centers, which are then normalized into unit vectors; (b) Learnable : the same initial process as Random, but the generated centers are further updated with gradient descent; (c) Orthonormal : our proposed process as described in Eq.
Specifically, we test each initialization method with the cluster number K equals to 2, 3, 4, 5, 10, 50, 100.
The results of adjusting these two hyper-parameters on ABIDE and ABCD datasets are shown in Figure 3(a).
We observe that: (1) When cluster centers are orthonormal, the model’s performance increases with the number of clusters ranging from 2 to 10, and then drops with the cluster number rising from 10 to 100, suggesting the optimal cluster number to be relatively small, which leads to less computation and is consistent with the fact that the typical number of functional modules are smaller than 25; (2) With a sufficiently large cluster number, all three initialization methods, Random,
Learnable and Orthonormal, tend to reach similar performance, but orthonormal performs stably better when the number of clusters is smaller; (3) It is also notable that our OCR EAD consistently achieves the best performance over other initialization methods regarding smaller standard deviations.
[Figure]
4.4 In-depth Analysis of Attention Scores and Cluster Assignments (RQ3)
Figure 3(b) displays the self-attention score from the first layer of Multi-Head Self-Attention.
The attention scores are the average across all subjects in the ABCD test set.
This figure shows that the learned attention scores well match the divisions of functional modules based on available labels, demonstrating the effectiveness and explainability of our Transformer model.
Note that since there exists no available functional module labels for the atlas of the ABIDE dataset, we cannot visualize the correlations between attention scores and functional modules.
Figure 4 shows the cluster soft assignment results P on nodes in OCR EAD with two initialization methods.
The cluster number K is set to 4.
The visualized numerical values are the average
P of all subjects in each dataset’s test set.
From the visualization, we observe that (a) Base on
Appendix H, orthonormal initialization produces more discriminative P between classes than random initialization; (b) Within each class, orthonormal initialization encourages the nodes to form groups.
These observations demonstrate that our OCR EAD with orthonormal initialization can leverage potential clusters underlying node embeddings, thus automatically grouping brain regions into potential functional modules.
Discussion and Conclusion
Neuroimaging technologies, including functional magnetic resonance imaging (fMRI) are powerful noninvasive tools for examining the brain functioning.
There is an emerging nation-wide interest in conducting neuroimaging studies for investigating the connection between the biology of the brain, and demographic variables and clinical outcomes such as mental disorders.
Such studies provide an unprecedented opportunity for cross-cutting investigations that may offer new insights to the differences in brain function and organization across subpopulations in the society (such as biological sex and age groups) as well as reveal neurophysiological mechanisms underlying brain disorders (such as psychiatric illnesses and neurodegenerative diseases).
These studies have a tremendous impact in social studies and biomedical sciences.
For example, mental disorders are the leading cause of disability in the USA and roughly 1 in 17 have a seriously debilitating mental illness.
To
[Figure]
address this burden, national institutions such as NIH have included brain-behavior research as one of their strategic objectives and stated that sound efforts must be made to redefine mental disorders into dimensions or components of observable behaviors that are more closely aligned with the biology of the brain.
Using brain imaging data to predict diagnosis has great potential to result in mechanisms that target for more effective preemption and treatment.
In this paper, we presents B RAIN N ETWORK T RANSFORMER , a specialized graph Transformer model with O RTHONORMAL C LUSTERING R EADOUT for brain network analysis.
Extensive experiments on two large-scale brain network datasets demonstrate that our B RAIN N ET TF achieves superior performance over SOTA baselines of various types.
Specifically, to model the potential node feature similarity in brain networks, we design OCR EAD and prove its effectiveness both theoretically and empirically.
Lastly, the re-standardized dataset split for ABIDE can provide a fair evaluation for new methods in the community.
For future work, B RAIN N ET TF can be improved with explicit explanation modules and used as the backbone for further brain network analysis, such as digging essential neural circuits for mental disorders and understanding cognitive development in adolescents.
Acknowledgments
This research was supported in part by the University Research Committee of Emory University, and the internal funding and GPU servers provided by the Computer Science Department of Emory University.
The authors gratefully acknowledge support from NIH under award number R01MH and R01MH118771. The content is solely the responsibility of the authors and does not necessarily represent the official views of the National Institutes of Health.
Data used in the preparation of this article were obtained from the Adolescent Brain Cognitive
DevelopmentSM (ABCD) Study ( https://abcdstudy.org ), held in the NIMH Data Archive (NDA).
This is a multisite, longitudinal study designed to recruit more than 10,000 children age 9-10 and follow them over 10 years into early adulthood.
The ABCD Study ® is supported by the
National Institutes of Health and additional federal partners under award numbers U01DA041048,
U01DA050989, U01DA051016, U01DA041022, U01DA051018, U01DA051037, U01DA050987,
U01DA041174, U01DA041106, U01DA041117, U01DA041028, U01DA041134, U01DA050988,
U01DA051039, U01DA041156, U01DA041025, U01DA041120, U01DA051038, U01DA041148,
U01DA041093, U01DA041089, U24DA041123, U24DA041147.
A full list of supporters is available at https://abcdstudy.org/federal-partners.html .
A listing of participating sites and a complete listing of the study investigators can be found at https://abcdstudy.org/consortium_ members/ .
ABCD consortium investigators designed and implemented the study and/or provided data but did not necessarily participate in the analysis or writing of this report.
This manuscript reflects the views of the authors and may not reflect the opinions or views of the NIH or ABCD consortium investigators.
The ABCD data repository grows and changes over time.
The ABCD data used in this report came from NIMH Data Archive Release 4.0 (DOI 10.15154/1523041).
DOIs can be found at https://nda.nih.gov/abcd .
A Training Curves of Different Models with or without StratifiedSampling
In Figure 5, we demonstrate the training curves of different models with or without stratified sampling based on site information from ABIDE.
The curves of different variants display similar patterns across three model architectures in a single run.
We remove Graphormer since its performance is much worse than others.
Specifically, it is shown that (a) with stratified sampling, the performance gap between validation and test on ABIDE is much smaller than the one without stratified sampling; (b) stratified sampling can stabilize the training process on ABIDE, especially for VanillaTF and
B RAIN N ET TF.
B Transformer Performance with Different Node Features
We compare the performance of Transformer model equipped with different node features.
The results are shown in Table 3, where connection profile represents the corresponding row for each node in the adjacency matrix, identity feature initializes a unique one-hot vector for each node, and eigen feature generates a k -dimensional feature vector for each node from the k eigenvectors based on the eigendecomposition on the adjacency matrix.
Empirical observations demonstrate that adding identity or eigen node features to connection profiles cannot improve the model’s performance.
[Table]
C Statistical Proof of the Goodness with Orthonormal Cluster Centers
We propose two statistical methods to prove the goodness in orthonormal case since it is impractical to directly compare the performance of the orthonormal and non-orthonormal initializations.
C.1 Proof of Theorem 3.
We state Theorem 3.1 here and show the proof details.
Theorem C.1.
For arbitrary r > 0 , let B r = {Z ∈ R V ; ∥Z∥ ≤ r } denote the round ball centered at origin of radius r with Z being fracture vectors.
Let V r be the volume of B r .
The variance of
Softmax projection averaged over B r
[Equation]
attains maximum when E is orthonormal.
Proof.
For simplicity, we first consider the two-dimensional case with two cluster centers E 1 , E 2 .
Since we integrate over the round ball B r , spherical symmetry allows us to set E 1 = (1 , 0) and
E 1 = (cos( ϕ ) , sin( ϕ )) with ϕ ∈ [0 , π 21 ] being the angle between E 1 and E 2 under polar coordinates.
Then the Softmax readout Eq.
(2) can be rewritten as:
[Equation]
where θ is the angle between Z and E 1 and ρ is the norm of Z .
Hence, the integral is
[Equation]
Our aim is to show that the integral F ( ϕ ) attains its maximum when E 1 , E 2 are orthogonal.
It is unclear whether the above integral has an elementary antiderivative.
Thus, instead of evaluating the integral directly, we firstly prove two symmetric properties of the integrand f
( ρ, θ, ϕ ) : (a) It is straightforward to show that f ( ρ, θ + kπ, ϕ ) =
f ( ρ, θ, ϕ ) for k ∈ N .
That is, f is periodic for π on the first argument θ .
(b) We have
[Equation]
which means f is symmetric with respect to θ
= ϕ 2 + π 2 + kπ .
As the integrand f ( ρ, θ, ϕ ) is periodic, we are allowed to compare F ( ϕ 1 ) , F ( ϕ 2 ) via
[Equation]
The integral domain [ ϕ 2 , ϕ 2 + 2 π ] is taken according to the second symmetry property of f and can be significant for the following trick: we take the directional derivative of f along v = (1 , 2) tangent to the straight line θ = ϕ 2 :
[Equation]
[Equation]
It is easy to check that in the above integral domain and for any ρ > 0 , Df ( v ) is always non-negative.
Hence,
[Equation] [Equation]
when ϕ 1 ≤ ϕ 2 .
After taking integral, F ( ϕ 1 ) ≤ F ( ϕ 2 ) and thus it attains maximum in the orthonormal case ( ϕ = π 2 ).
Comparing F ( ϕ 1 ) , F ( ϕ 2 ) without adjusting the integral domain as above cannot give a clear result because the simple partial derivative ∂f/∂ϕ oscillates around zero.
Higher dimensional cases follow similarly by employing spherical and hyperspherical coordinates.
C.2 Proof of Theorem 3.
Theorem 3.2 deals with a more general case: comparing the performance of an arbitrary readout P defined by orthonormal cluster centers with non-orthonormal ones.
We regard P as an estimated similarity probability between nodes and clusters and solve this problem from the perspective of statistics.
The estimation is considered as a regression of samples ( ˆ Z ( s ) , ˆ E ( t ) , ˆ P ( st ) ) from node features, cluster centers and similarity probabilities.
We then judge the estimation relative to true similarity probability P T .
Although it is almost impossible to find an analytic formula for P T , we can indirectly judge the quality of estimation.
To clarify the idea, we introduce some basic concepts from statistics and prove our results on a statistical basis.
C.2.1 Background Knowledge of Regression Analysis
We first consider process samples by logistic regression with cluster centers as categorical variables .
Intuitively, non-orthonormal centers correlate with each other, which means there is an overlap among categorical variables and makes it hard to identify the decision boundary that leads to a failed classification.
However, as far as we know, it is unclear how to compare overlaps between orthonormal and non-orthonormal variables rigorously.
Thus, we simply process samples by a general nonlinear regression.
The regression process is linearized by the Gauss-Newton algorithm to facilitate the analysis.
We judge the goodness-of-fit describing the degree to which the regression function fits its observed value, and then conduct a hypothesis test.
The goodness-of-fit is measured by coefficient of determinate R 2 [47]:
Definition C.2.
We consider a regression with r independent main variables:
[Equation]
Let ˆ x p =
(ˆ x p 1 , ..., ˆ x ps ) ⊤ and ˆ y =
(ˆ y 1 , ..., ˆ y s )
⊤ be data sets (samples) associated with fitted values ˇ y =
(ˇ y 1 , ..., ˇ y s ) .
Each difference e q
=
ˆ y q − ˇ y q is called a residue .
We denote the mean of ˆ x p and ˆ y by ¯ x p , ¯ y .
The variability of data set can be measured by the total sum of squares (SST), the sum of squares of residuals (SSR) and the explained sum of squares (SSE) defined as (where p = 1 , 2 , ..., r q = 1 , 2 , ..., s ):
[Equation]
In linear regression, SSR + SSE = SST and the coefficient of determination R 2 is defined as:
[Equation]
Conceptually, SSE is the error cost by regression of main variables.
Thus by definition, R 2 reveals the percentage of errors that main variables can explain in the total error SST.
The value of R 2 is bounded by 1 .
A large value of R 2 indicates a better fitting.
However, it should be noted that an extremely-large R 2 could indicate overfitting.
In our problem, since our regression is nonlinear, the sum of SSR and SSE is less than SST [1].
Therefore, measuring goodness-of-fit by R 2 in nonlinear regression is inaccurate.
A common strategy to remedy this problem is approximating nonlinear functions by polynomials via Gauss-Newton algorithm .
We provide a brief introduction here, and more details can be found in [1]: for a nonlinear model f k with parameter δ , in a small neighborhood of δ T -the true value of δ , we have the linear expansion:
[Equation]
Or briefly, we write it by vector notation :
[Equation]
where F ( δ − δ T ) stands for the dot product of derivatives and differences of parameters from Eq.
Suppose δ ( γ ) is an approximation to the least-squares estimation δ of our model, for δ close to δ ( γ ) , we rewrite the expansion as:
[Equation]
where ˇ P denotes a fitted value of P and F ( γ ) ( δ − δ ( γ ) ) again means a dot product.
Applying this to the residual vector e ( δ ) , we have:
[Equation]
Thus, the norm
[Equation]
The right-hand side is minimized with respect to δ when
[Equation]
[Equation]
This suggests that given a current approximation δ ( γ ) , the next approximation should be:
[Equation]
Expanding the nonlinear function f as polynomials and modifying the parameter δ as above, we can use R 2 to measure the goodness-of-fit .
To acquire higher accuracy in a general nonlinear regression, one can make a elaborated goodness-of-fit test for specific fitting functions e.g., [9, 11].
We do not discuss this sophisticated method as it is out of the scope of this paper.
C.2.2 Comparing R 2 by Variance Inflation Factor
The proof of Theorem 3.2 consists of two steps: (a) we first prove that the regression accuracy , the accuracy when regressing P is higher when sampling from orthonormal cluster centers (Theorem
C.4), and consequently (b) higher regression accuracy increases appraisal accuracy , the accuracy when appraising an estimated value in hypothesis testing (Theorem C.6).
In this subsection, we compare regression accuracy.
we fix Z i when regressing P via the fitted value ˇ P ( E k ) .
Statistically, the expectation E( P ) of all readouts is identified as the true similarly probability P T .
In regression analysis, the Ordinary Least Squares (OLS) guarantees asymptotically unbiased estimations.
That is, when the sample size s is large enough, it can be regarded as an unbiased estimation [47]:
[Equation]
Therefore, the better the goodness-of-fit reflected by R 2 , the smaller the variance of estimation.
To compare this, we use the concept of variance inflation factor which reflects the inflation of weights of variables in regression:
Definition C.3.
The variance inflation factor (VIF) p is defined as:
[Equation]
[Equation]
where R 2 p is the coefficient of multiple determination when X p is regressed by the r-1 other variables in the model from Eq.
Remark.
We discuss more details about VIF in the following context [47].
For simplicity, we denote the following collection of samples and regression coefficients:
[Equation]
In the regression model Eq.
(13), the estimation ˇ β p of regression coefficients β p are obtained by
Ordinary Least Squares (OLS):
[Equation]
[Equation]
We standardize the regression equation by covariance matrices σ y of ˇ y and the variance σ q of ˆ x p as
[Equation]
and
[Equation]
Similarly to Eq.
(25), standardized estimation of regression coefficients are equal to
[Equation]
On the other hand, the covariance matrix of the estimated regression coefficients is
[Equation]
where σ 2 is the error term variance for X (cf.
Definition C.2).
After standardization, it is noted that
X ∗⊤ X ∗ is just the correlation matrix r XX of X ∗ .
Hence, by Eq.
(29) we obtain:
[Equation]
Let (VIF) p be the p -th diagonal element of the matrix r − 1 XX .
The variance of β ∗ p is equal to:
[Equation]
The diagonal element (VIF) p is just the variance inflation factor for ˇ β ∗ p .
The variance of β ∗ p can also be written as [47]
[Equation]
With the previous discussion, we conclude that
[Equation]
where R 2 p is defined in C.3.
Theorem C.4.
Let
[Equation]
where r denotes the number of variables in Eq.
Then VIF ≥ 1 with equality holds if and only if the variables are orthogonal.
Proof.
To prove this, we need to generalize the definition of R 2 .
By definition,
[Equation]
Substituting Eq.
(27) into the above identity, we have
[Equation]
and by Eq.
, we conclude that
[Equation]
As the finial step, we compute R 2 p from Definition C.3 by Eq.
It should be noted that according to Definition C.3, R 2 p is the goodness-of-fit when X p is regressed by the r-1 other variables.
These variables are uncorrelated in orthonormal case.
Hence r XY = 0 , R 2 p = 0 and VIF = 1 .
Remark.
In statistics, when a variable’s VIF is greater than 1, or equivalently R 2 p ̸ = 0 , the influence of this variable on the whole estimation is inflated.
It breaks the so-called absence of multicollinearity , a fundamental principle in multiple regression analysis, and hence causes more error.
Since SSE is a constant value, the error generated by the inflation would be counted into SSR, which leads to a decrease in R 2 by Definition C.2 (see [47, 1] for more details).
C.2.3 Statistical Hypothesis Testing
The previous discussion verifies that regressing with orthonormal samples attains a higher goodnessof-fit .
In other words, it achieves a higher regression accuracy.
Tools from hypothesis testing are borrowed here to determine the appraisal accuracy mentioned at the beginning of Section C.2.2.
We first introduce mean squared error (MSE) commonly used in statistics [19]:
Definition C.5.
Recall that the residue e q =
(ˆ y q − ˇ y q ) from Definition C.2.
Then,
[Equation]
As mentioned in C.2.1, a small coefficient of determination R 2 indicates a large SSR and hence leads to a large MSE.
As a result of Theorem C.4, MSE is minimized in the orthonormal case.
We now assume a domain centered at the true value P T of radius d , and treat the outside space W as the rejection region .
Statistically, if the distance between ˇ P and P T is less than a small enough d , we can regard them as the same.
Intuitively, if fitted values ˇ P are largely scattered from the true value P T , that is, when MSE is large, it can interfere with our judgment of whether P can be identified with P T .
Rigorously, we make a hypothesis testing and analyze the probability of rejecting a well-estimated readout function.
We prove in the following that when sampling from orthonormal cluster centers, a higher regression accuracy (Theorem C.4) guarantees a lower MSE and therefore increases the appraisal accuracy.
Theorem C.6.
The significance level α E k · reveals that the probability of rejecting a well-estimated readout is lower when sampling from orthonormal centers than sampling from non-orthonormal centers.
Proof.
Let P be a readout function such that ∥ P T − P ∥ ≤
d for small enough d .
Statistically, we can treat them as the same and simply write ˇ P = P T .
In hypothesis testing , we define null hypothesis
H 0 and alternative hypothesis H 1 by
[Equation]
in which H 1 means that we reject a well-estimated readout with H 0 having the opposite meaning.
The rejection region for this test is thus given as W = { ˇ P ̸ = P T } .
As a conventional procedure in hypothesis testing , we take a suitable test statistic T E k ( Z i ) whose distribution f is known [19].
It is used to compute the probability that ˇ P is in the rejection region.
The corresponding probability distribution is called potential function g ( θ ) for W in this setting:
[Equation]
where the significance level α E k is the upper bound of the probability of making mistakes (formally called type I error )
By Theorem C.4 and Remark C.5, MSE is minimized in the orthonormal case.
It can be treated as a variance of distribution f .
Then by Vysochanskij–Petunin inequality , a refinement of Chebyshev inequality, the integration over W with orthonormal cluster centers E k is smaller than that with non-orthonormal cluster centers E ′ k :
[Equation]
As the result holds true for any well-chosen T E k ( Z i ) , α E k ≤ α E ′ K , this finishes the proof.
D Running Time
Table 4 shows that state-of-the-art models of Graphormer and SAN are much slower than our
B RAIN N ET TF and VanillaTF, mainly because their implementations are not optimized toward the unique properties of brain networks.
Specifically, let e be the number of edges and v be the number of nodes.
The calculation of Graphormer and SAN optimizes the case where e ≪ v 2 .
However, brain networks usually have a small number of nodes but dense connections, i.e., e ≃ v 2 .
Therefore the optimized sparse graph operations in PyTorch Geometric [23] do not work properly.
On the other hand, since the number of nodes in brain networks is usually relatively small (less than 500), we can directly speed up the calculation using matrix multiplication, which is what we did in B RAIN N ET TF and VanillaTF.
Besides, the edge feature generation operator in Graphormer further increases the burden on its computing time.
[Table]
E Number of Parameters
[Table]
F Parameter Tuning
For BrainGB, BrainGNN, FBNetGen, we use the authors’ open-source codes.
For SAN and
Graphormer, we folk their repositories and modified them for the brain network dataset.
For
BrainNetCNN and VanillaTF, we implement them by ourselves.
We use the grid search for some important hyper-parameters for these baselines based on the provided best setting.
To be specific, for
BrainGB, we search different readout functions {mean, max, concat} with different message-passing functions {Edge weighted, Node edge concat, Node concat}.
For BrainGNN, we search different learning rates {0.01, 0.005, 0.001} with different feature dimensions {100, 200}.
For FBNetGen, we search different encoders {1D-CNN, GRU} with different hidden dimensions {8, 12, 16}.
For
BrainNetCNN, we search different dropout rates {0.3, 0.5, 0.7}.
For VanillaTF, we search the number of transformer layers {1, 2, 3} with the number of headers {2, 4, 6}.
For SAN, we test LPE hidden dimensions {4, 8, 16}, the number of LPE and GT transformer layers {1, 2} and the number of headers {2, 4} with 50 epochs training.
For Graphormer, we test encoder layers {1, 2} and embed dimensions {256, 512}.
Furthermore, since the rebuttal time is pretty short, we do not have enough time to dig two new baselines, BrainnetGNN and DGM, which may be why their performance is worse than others.
G Software Version
[Table]
H The Difference between Various Initialization Methods
To show orthonormal initialization can produce more discriminative P between classes than random initialization, we calculate the difference score d based on the formula
[Equation]
where V is the number of nodes and K is the number of clusters.
After running the t-test, we found the margins between random and orthonormal on both ABIDE and ABCD are significant, which is consistent with our conclusion.
[Table]