[논문 리뷰] Brain Network Transformer

Xuan Kan, Wei Dai외 4명|arXiv (Cornell University)|2022. 10. 13.

Functional Brain Connectivity Studies인용 수 55

한 줄 요약

BrainNetTF는 연결 프로파일 노드 특징과 Orthonormal Clustering Readout을 사용하는 뇌 네트워크에 맞춘 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에서 표준화된 평가를 통해 모델 간 공정한 비교를 가능하게 합니다.

제안 방법

• 노드 특징을 에지 기반 위치 부호화 없이 향상된 노드 임베딩으로 변환하기 위해 2-layer Multi-Head Self-Attention (MHSA)을 사용합니다.
• OCRead를 통해 노드 임베딩을 소프트 클러스터링하여 클러스터 기반 풀링을 형성하고 그래프 수준 임베딩을 생성합니다.
• 직교 기초로 OCRead 클러스터 센터를 초기화하고 클러스터링 인식 읽기를 위한 소프트 클러스터 할당 P를 계산합니다.
• 직교 클러스터 센터가 소프트맥스 기반 읽기에서 분산을 최대화하고 풀링 신뢰성을 개선한다는 이론적 근거를 제시합니다.
• ABIDE에서 표준화된 평가 파이프라인을 제공하고 모델 간 공정한 비교를 위해 ABCD에서 테스트합니다.

실험 결과

연구 질문

• RQ1 RQ1: BrainNetTF가 ABIDE와 ABCD에서 최첨단 그래프 트랜스포머 및 뇌 네트워크 특화 모델과 비교하여 어떤 성능을 보이는가?
• RQ2 RQ2: OCRead 읽기가 다른 트랜스포머 아키텍처와 클러스터 초기화에서 성능에 어떠한 영향을 미치는가?
• 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에서 SAN 및 Graphormer보다 최대 6%p의 절대 AUROC 차이로 우수한 성능을 보입니다.
• OCRead는 트랜스포머 아키텍처 전반에서 예측력을 개선하고 다른 읽기 방법들보다 우수합니다.
• 직교 초기화는 특히 작은 클러스터 수에서 더 구별력 있는 클러스터 할당과 안정적인 성능을 제공합니다.
• BrainNetTF에서 학습된 어텐션 점수는 가능할 경우 기능 모듈과 일치하여 설명 가능성을 뒷받침합니다.
• BrainNetTF는 계산 면에서도 경쟁력이 유지되며 전체복잡도는 O(V^2)로 다른 뇌-네트워크 GNN과 비슷합니다.
• 모델은 BrainNetTF를 사용해 ABIDE에서 AUROC가 80.2±1.0, ABCD에서 96.2±0.3에 도달하며 두 데이터셋에서 정확도/민감도/특이도 개선을 보입니다.