[论文解读] 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.

研究动机与目标

• 将脑网络动机化并建模为以 ROI 为节点、相关性为边的密集图。
• 利用连接特征作为自然、成本较低的节点特征，而不是基于特征值的位置编码。
• 在完整脑图上学习全对全注意力，以预测二元脑相关的结果。
• 引入 OCRead，以获得反映脑网络模块化组织的聚类感知、具有表达力的图嵌入。
• 在 ABIDE 上标准化评估，以实现跨模型的公平比较。

提出的方法

• 使用两层多头自注意力（MHSA）将节点特征转换为增强的节点嵌入，不使用基于边的位置信息编码。
• 通过 OCRead 生成图级嵌入，该方法对节点嵌入进行软聚类，使用正交初始化以形成基于簇的池化。
• 用正交基初始化 OCRead 的簇中心，并计算软簇分配 P 以进行聚类感知读出。
• 给出理论依据，证明正交簇中心在基于 softmax 的读出中最大化方差并提高池化的可靠性。
• 提供在 ABIDE 上的标准化评估流程并在 ABCD 上进行测试，以确保模型之间的公平比较。

实验结果

研究问题

• RQ1 RQ1：在 ABIDE 和 ABCD 上，BrainNetTF 相对于最先进的图变换模型和脑网络特定模型的表现如何？
• RQ2 RQ2：OCRead 读出在不同 Transformer 架构和簇初始化下对性能的影响如何？
• RQ3 RQ3：注意力模式和 OCRead 的簇分配是否与已知的功能模块对齐并提供可解释性？
• RQ4 RQ4：正交初始化和簇数量对 OCRead 的性能与稳定性有何影响？
• RQ5 RQ5：OCRead 是否可以推广到其他脑网络模态（如结构连接组）或其他图领域？

主要发现

类型	方法	ABIDE AUROC	ABIDE 准确率	ABIDE 灵敏度	ABIDE 特异性	ABCD AUROC	ABCD 准确率	ABCD 灵敏度	ABCD 特异性
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 在 ABIDE 和 ABCD 上的绝对 AUROC 最多领先 SAN 与 Graphormer 高出 6%。
• OCRead 在不同 Transformer 架构上提升预测能力并优于其他读出方法。
• 正交初始化带来更具判别性的簇分配和更稳定的性能，特别是在较小的簇数时。
• 注意力分数在可用的情况下与功能模块对齐，支持可解释性。
• BrainNetTF 在计算方面保持竞争力，总体复杂度为 O(V^2)，与其他脑网络 GNNs 相当。
• 模型在 BrainNetTF 下实现 AUROC 为 80.2±1.0（ABIDE）和 96.2±0.3（ABCD），并在两个数据集上实现相应的准确率、特异性、灵敏度的提升。