[論文レビュー] Brain Network Transformer

Xuan Kan, Wei Dai+4 more|arXiv (Cornell University)|Oct 13, 2022

Functional Brain Connectivity Studies被引用数 55

ひとこと要約

BrainNetTF は脳ネットワーク向けに調整された Transformer ベースのモデルで、接続プロファイルのノード特徴と正交クラスタリング読み出しを使用して、ABIDE および ABCD データセット上で優れたグラフレベルの表現と予測を実現します。

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 pair-wise connection strengths among ROIs with efficient attention weights across individuals that are predictive towards downstream analysis tasks. Moreover, we propose an Orthonormal Clustering Readout operation based on self-supervised 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 Brain Network Transformer on both the public ABIDE and our restricted ABCD datasets. The implementation is available at https://github.com/Wayfear/BrainNetworkTransformer.

研究の動機と目的

• Motivate and model brain networks as dense graphs with ROIs as nodes and correlations as edges.
• Leverage connection profiles as natural, low-cost node features instead of eigen-based positional encodings.
• Learn fully pairwise attention on complete brain graphs to predict binary brain-related outcomes.
• Introduce OCRead to obtain cluster-aware, expressive graph embeddings reflecting modular brain organization.
• Standardize evaluation on ABIDE to enable fair cross-model comparisons.

提案手法

• Use a two-layer Multi-Head Self-Attention (MHSA) to transform node features into enhanced node embeddings without edge-based positional encodings.
• Produce graph-level embeddings via OCRead, which soft-clusters node embeddings with orthonormal initialization to form cluster-based pooling.
• Initialize OCRead cluster centers with orthonormal bases and compute soft cluster assignments P for clustering-aware readout.
• Formulate a theoretical justification showing that orthonormal cluster centers maximize variance in softmax-based readouts and improve pooling reliability.
• Provide a standardized evaluation pipeline on ABIDE and test on ABCD to ensure fair comparisons across models.

実験結果

リサーチクエスチョン

• RQ1 RQ1: How does BrainNetTF perform versus state-of-the-art graph transformers and brain-network-specific models on ABIDE and ABCD?
• RQ2 RQ2: How does the OCRead readout influence performance across different transformer architectures and cluster initializations?
• RQ3 RQ3: Do attention patterns and OCRead cluster assignments align with known functional modules and provide explainability?
• RQ4 RQ4: What is the impact of orthonormal initialization and cluster count on OCRead’s performance and stability?
• RQ5 RQ5: Can OCRead be generalized to other brain network modalities (e.g., structural connectomes) or other graph domains?

主な発見

タイプ	方法	ABIDE AUROC	ABIDE Accuracy	ABIDE Sensitivity	ABIDE Specificity	ABCD AUROC	ABCD Accuracy	ABCD Sensitivity	ABCD Specificity
Graph Transformer	SAN	71.3±2.1	65.3±2.9	55.4±9.2	68.3±7.5	90.1±1.2	81.0±1.3	84.9±3.5	77.5±4.1
Graph Transformer	Graphormer	63.5±3.7	60.8±2.7	78.7±22.3	36.7±23.5	89.0±1.4	80.2±1.3	81.8±11.6	82.4±7.4
Transformer	VanillaTF	76.4±1.2	65.2±1.2	66.4±11.4	71.1±12.0	94.3±0.7	85.9±1.4	87.7±2.4	82.6±3.9
Fixed Network	BrainGNN	62.4±3.5	59.4±2.3	36.7±24.0	70.7±19.3	91.9±0.3	83.1±0.5	84.6±4.3	81.5±3.9
Fixed Network	BrainGB	69.7±3.3	63.6±1.9	63.7±8.3	60.4±10.1	91.9±0.3	83.1±0.5	84.6±4.3	81.5±3.9
Fixed Network	BrainNetCNN	74.9±2.4	67.8±2.7	63.8±9.7	71.0±10.2	93.5±0.3	85.7±0.8	87.9±3.4	83.0±4.4
Learnable Network	FBNETGNN	75.6±1.2	68.0±1.4	64.7±8.7	62.4±9.2	94.5±0.7	87.2±1.2	87.0±2.5	86.7±2.8
Learnable Network	BrainNetGNN	55.3±1.9	51.2±5.4	67.7±37.5	33.9±34.2	75.3±5.2	67.5±4.7	67.7±5.7	68.0±6.5
Learnable Network	DGM	52.7±3.8	60.7±12.6	53.8±41.2	51.1±40.9	76.8±19.0	68.6±8.1	40.5±29.7	95.6±4.2
Ours	BrainNetTF	80.2±1.0	71.0±1.2	72.5±5.2	69.3±6.5	96.2±0.3	88.4±0.4	89.4±2.6	88.4±1.5

• BrainNetTF outperforms SAN and Graphormer on ABIDE and ABCD by up to 6% absolute AUROC.
• OCRead improves prediction power across Transformer architectures and outperforms other readouts.
• Orthonormal initialization yields more discriminative cluster assignments and stable performance, especially with smaller cluster counts.
• Attention scores learned by BrainNetTF align with functional modules where available, supporting explainability.
• BrainNetTF remains competitive in computation, with overall complexity O(V^2) comparable to other brain-network GNNs.
• The model achieves AUROC of 80.2±1.0 (ABIDE) and 96.2±0.3 (ABCD) with BrainNetTF, and corresponding accuracy/specificity/sensitivity improvements on both datasets.