BrainIB: Interpretable Brain Network-based Psychiatric Diagnosis with Graph Information Bottleneck

TL;DR

BrainIB, a new graph neural network (GNN) framework to analyze functional magnetic resonance images (fMRI) by leveraging the famed information bottleneck (IB) principle, is proposed and is able to identify the most informative edges in the brain (i.e., subgraph) and generalizes well to unseen data.

Abstract

Developing a new diagnostic models based on the underlying biological mechanisms rather than subjective symptoms for psychiatric disorders is an emerging consensus. Recently, machine learning-based classifiers using functional connectivity (FC) for psychiatric disorders and healthy controls are developed to identify brain markers. However, existing machine learning-based diagnostic models are prone to over-fitting (due to insufficient training samples) and perform poorly in new test environment. Furthermore, it is difficult to obtain explainable and reliable brain biomarkers elucidating the underlying diagnostic decisions. These issues hinder their possible clinical applications. In this work, we propose BrainIB, a new graph neural network (GNN) framework to analyze functional magnetic resonance images (fMRI), by leveraging the famed Information Bottleneck (IB) principle. BrainIB is able to identify the most informative edges in the brain (i.e., subgraph) and generalizes well to unseen data. We evaluate the performance of BrainIB against 3 baselines and 7 state-of-the-art brain network classification methods on three psychiatric datasets and observe that our BrainIB always achieves the highest diagnosis accuracy. It also discovers the subgraph biomarkers which are consistent to clinical and neuroimaging findings. The source code and implementation details of BrainIB are freely available at GitHub repository (https://github.com/SJYuCNEL/brain-and-Information-Bottleneck/).

Figures (8)

• Figure 1: Pipeline from rs-fMRI raw data to the functional graph. The resting-state fMRI raw data are preprocessed and then parcellated into regions of interest (ROIs) according to the automated anatomical labelling (AAL) atlas. The functional connectivity (FC) matrices are calculated using Pearson correlation between ROIs. From the FC we construct the brain functional graph $\mathcal{G}=\left(A,X \right)$, where $A$ is the graph adjacency matrix characterizing the graph structure ($A\in\left\{ 0,1\right\}^{n\times n}$) and $X$ is node feature matrix. Specifically, $A$ is a binarized FC matrix, where only the top 20-percentile absolute values of the correlations of the matrix are transformed into ones, while the rest are transformed into zeros. For node feature $X$, $X_{k}$ for node $k$ can be defined as $X_{k}=\left [ \rho_{k1},\dots, \rho_{kn}\right ] ^{\text{T}}$, where $\rho_{kl}$ is the Pearson’s correlation coefficient for node $k$ and node $l$.
• Figure 2: Architecture of our proposed BrainIB. The framework mainly consists of three modules: subgraph generator module, graph encoder module and mutual information estimator module. Given an input graph $\mathcal{G}$ and corresponding graph label $Y$, the subgraph generator is used to sample subgraph $\mathcal{G}_{\text{sub}}$ from input graph $\mathcal{G}$. The graph encoder is used to learn graph embeddings $Z$ or $Z_{\text{sub}}$ from either $\mathcal{G}$ or $\mathcal{G}_{\text{sub}}$, where $Z_{\text{sub}}$ be used for graph classification by computing cross-entropy loss $\mathcal{L}_{\text{CE}}\left ( \mathcal{G}_{\text{sub}},Y \right )$. The mutual information estimation module evaluates the mutual information $I\left ( \mathcal{G}_{\text{sub}},\mathcal{G} \right )$ between $\mathcal{G}$ and $\mathcal{G}_{\text{sub}}$.
• Figure 3: Subgraph Generator. IB-Subgraph is generated from the input graph with the edge selection strategy. Given an input graph $\mathcal{G}=\left(A,X \right)$, $X$ is sent to a linear layer followed by a sigmoid function to obtain edge attention mask $P$ and ensure $p_{ij} \in\left [ 0,1 \right ]$, where $p_{ij}$ represents edge selection probability between node $i$ and node $j$. Subsequently, we binarize $P$ to obtain edge assignment $\mathcal{S}$. In order to make sure the gradient, w.r.t., $p_{ij}$ is computable, we leverage the gumbel-softmax reparameterization trick to update edge assignment $\mathcal{S}$. To ensure that a sufficient number of edges are retained, $P$ is reshaped to $K$-dimensional matrix and then is binarized to generate edge assignment $\mathcal{S}$ with Gumbel-Softmax approach. Next, $\mathcal{S}$ is transform to $n \times n$ matrix. Finally, the IB-subgraph $\mathcal{G}_{\text{sub}}$ can be extracted through $A_{\text{sub}}=A\odot \mathcal{S}$, where $\odot$ is entry-wise product and $A$, $A_{\text{sub}}$ represent adjacency matrix of the input graph $\mathcal{G}$, IB-subgraph $\mathcal{G}_{\text{sub}}$, respectively.
• Figure 4: The influence of hyperparameter $\beta$ on performance using the ABIDE and Rest-meta-MDD datasets.
• Figure 5: The influence of different graph pooling methods on performance using the ABIDE and Rest-meta-MDD datasets. SUM, AVE, MAX and SOPOOL represent sum-pooling, average-pooling, max-pooling and bilinear mapping second-order pooling, respectively.

Theorems & Definitions (2)

• Definition 4.1
• Definition 4.2