Multi-modal Graph Learning for Disease Prediction

Let $X=\left[x_{1},x_{2},\cdots,x_{N}\right]\in\mathbb{R}^{d_{in}\times N}$ denote the $d_{in}$-dimensional multi-modal features of $N$ patients and $Y=\left[y_{1},y_{2},\cdots,y_{N}\right]$ is the corresponding labels. Each patient is represented by $M$ modalities, so the multi-modal features also can denote as $X=\{X^{1},X^{2},\cdots,X^{M}\}$ and $X^{m}=\left[x^{m}_{1},x^{m}_{2},\cdots,x^{m}_{N}\right]\in\mathbb{R}^{d_{m}\times N}$, where $x^{m}_{i}$ represents the $d_{m}$-dimensional features of the $m$-th modality of $x_{i}$ and $d_{in}=\sum_{m=1}^{M}d_{m}$. Given the multi-modal features $X$, the issue we consider in this paper for multi-modal graph learning is to achieve an optimal latent graph structure inference based on multi-modal feature fusion considering inter-modal correlation, thus providing reliable graph support to GCNs for disease prediction or the other biomedical tasks.

As illustrated in Fig. 2, the overall framework of MMGL consists of three modules. In feature fusion phase, the multi-modal features $X$ is integrated into a fused single feature matrix $H=\left[h_{1},h_{2},\cdots,h_{N}\right]\in\mathbb{R}^{d\times N}$ by modal-attentional feature fusion module (MaFF). Then, in graph learning phase, the adjacency matrix $A\in\mathbb{R}^{N\times N}$ characterizing the latent graph structure is captured through a novel adaptive graph learning mechanism (AGL) based on the multi-modal modal-fused features. Finally, according to the learned adjacency matrix $A$ and fused feature matrix $H$, the disease prediction results can be obtained by GCN.

In a real diagnostic scenario, medical experts always need to analyze various multi-modal data of the patient to make a reliable decision, since the single-modal data lacks of providing enough information for an accurate diagnosis. Similarly, reliable Computer Aided Diagnosis (CADx) also requires to be capable of leveraging the complementarity of multi-modal data [15]. To capture inter-modal relevance and complementarity for efficient information integration, we propose a transformer-style [23] modal-attentional feature fusion module named MaFF, which fuses the features of each modality while preserving rich information.

Given the query vectors $Q^{m}$ of current modality and key-value pairs of the others (i.e., $K$-$V$), the inter-modal attention score can be calculated. In practice, considering the problems of space-efficiency and parallelization, the scaled dot-product is chosen as attention function [23]. Since each modal feature has different dimension, as shown in Fig. 1(c), we first use the projection matrices (i.e., $W_{Q}$, $W_{K}$, $W_{V}$) with fixed dimension $d_{f}$ to translate $X$ to $Q$, $K$, and $V$, which facilitates the computation of attention scores in the same dimensional space. For a specific patient $x$, we can obtain the corresponding query matrix $q=\{q^{m}\}^{M}_{m=1}$, key matrix $k=\{k^{m}\}^{M}_{m=1}$, and value matrix $v=\{v^{m}\}^{M}_{m=1}$. Then, the inter-modal attention score map $P$ for $x$ is computed as:

where $P_{ij}$ denotes the attention score between modality $i$ and $j$. $\tau$ is the scaling factor to control the hardness of attention, which is set to $\sqrt{d_{f}}$ like in  [23]. After that, we perform cross-modal aggregation of the value vectors of each modality based on the calculated inter-modal correlations $P$, thus the integrated feature $h$ of patient $x$ could be computed as:

where $W_{h}$ and $W_{m}$ are projection layers. Besides, the implementation of residual connection between $\sum_{j=1}^{M}P_{m,j}~{}v^{j}$ and the initial value vector $v^{m}$ can effectively avoid the gradient vanishing problem during training process.

Compared to the fusion module in [12, 4, 8], the modal-attentional fusion puts concerns on inter-modal correlation through a multi-modal attention mechanism. As a consequence, it tend to optimally combine the complementary information from different modalities. It’s worth noting that the multi-modal attention map $P$ is patient-specific, which is also applicable to inductive learning. And we could obtain global-level inter-modal correlation map by averaging $P$ over all patients. Besides,the proposed multi-modal attention mechanism can also be scalable to multi-head version easily.

Based on existing graph structures, GCNs in [17, 24, 6] learn node representations for downstream tasks through neighborhood aggregation or spectral convolution. However, it’s no trivial to obtain the available graph for some specific tasks in biomedical field. Therefore, the graph learning problem often needs to be considered for applying GCNs in biomedical tasks.

For graph learning, it is usually modeled as two kinds of forms: (i) learning a joint discrete probability distribution on the edges of the graph [7], (ii) learning a similarity metric of nodes. Since the former is non-differentiable and hard applicable to inductive learning, we take the graph learning problem in consideration from the perspective of similarity metric learning of nodes. Some previous methods have adopted radial basis function (RBF) kernel [22, 27], cosine similarity [10], or threshold-based metric(for discrete feature) [19, 13] as the similarity metric. However, these approaches still require careful manual tuning to construct a meaningful graph structure for downstream GCNs. Therefore, as illustrated in Fig 2, we propose a simple but effective learnable metric function, which could be jointly optimized with the downstream GCNs:

where $W_{A}$ is a learnable weight matrix and $\hat{A}_{ij}$ is computed as weighted cosine similarity between patient $i$ and $j$. Since there are few uni-directional effects between patients except for epidemics, thus the learned adjacency matrix $\hat{A}$ is symmetric that is also in accordance with the expectations of realistic population graph of patients.

Commonly, a realistic adjacency matrix is usually non-negative and sparse. However, since $\hat{A}$ is a fully connected graph that is computationally expensive and the element $\hat{A}_{ij}$ is ranging in [-1,1], we capture a non-negative sparse graph $A$ from $\hat{A}$ by applying the ReLU function, i.e., setting negative value elements in $\hat{A}$ to zero. Finally, the latent graph $A$ is obtained for downstream tasks.

Graph Regularization Due to the sensitivity of GCNs to the graph structure, graph learning has a significant impact on the performance of GCN in downstream tasks. Furthermore, the constraint on the sparsity, connectivity, and smoothness of the learned graph is also important for adaptive graph learning [3]. Here, the Dirichlet energy is used to measure the smoothness of a set of graph signals $\{h_{1},h_{2},\cdots,h_{N}\}$:

It can be seen from Eq. 4 that the smaller the distance between $h_{i}$ and $h_{j}$, the larger $A_{ij}$ will be. Hence, the smooth loss $\mathcal{L}_{smooth}$ is intended to making connections between similar nodes, which means to enforce smoothness of the graph signals on the learned graph $A$.Essentially, $L_{smooth}$ simultaneously serves to control the sparsity of A [11]. However, only utilizing $L_{smooth}$ may lead to the trivial solution (i.e., $A=0$). To avoid it, two additional regularization terms following [11] are imposed on $A$:

where the first term $\mathcal{L}_{con}$ uses logarithmic barrier to control the connectivity of $A$, and the second term $\mathcal{L}_{r}$ is a regularization term to avoid the excessive sparseness caused by $\mathcal{L}_{smooth}$. The total graph regularization loss is defined as $\mathcal{L}_{g}=\mathcal{L}_{smooth}+\alpha\mathcal{L}_{con}+\beta\mathcal{L}_{r}$.

TADPOLE [18]: As a subset of the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database (adni.loni.usc.edu), TADPOLE dataset contains features extracted from multi-modalites, which include MR, PET, cognitive tests, cerebro-spinal fluid (CSF) biomarkers, risk factors, clinical examinations and demographic information. For Alzheimer’s Disease prediction, we select 685 patients with 366-dimensional multi-modal features from TADPOLE, divided into 245 normal, 360 Mild Cognitive Impairment (MCI) and 80 Alzheimer’s Disease (AD) patients, respectively. Then, we divide the features according to the corresponding modalities and use the means of features for missing value filling.

To validate the effectiveness of modal-attentional feature fusion (MaFF) and adaptive graph learning mechanism (AGL), we replace MaFF with MLP and concatenation operation respectively, and replace AGL with kNN graph construction based on RBF kernel and construction method of popGCN respectively. Table. 2 shows the ablation study results on different modules in our models. Specifically, the performance of the constructed graph of popGCN is the worst, especially on the ABIDE dataset, indicating that hand-constructed graphs are indeed not a desirable choice. Besides, AGL achieves favorable performance despite the absence of MaFF, which again validates the effectiveness of adaptive graph learning. More importantly, it can be observed that the combination of MaFF and AGL achieves a performance that far exceeds other combinations.