Multi-Modal Graph Neural Network with Transformer-Guided Adaptive Diffusion for Preclinical Alzheimer Classification

Technical Summary: Multi-Modal Graph Neural Network with Transformer-Guided Adaptive Diffusion for Preclinical Alzheimer Classification

Problem Statement
The graphical representation of the brain, constructed from regions of interest (ROIs) and their connectome features, offers critical insights for diagnosing neurodegenerative diseases like Alzheimer's Disease (AD). However, existing Graph Neural Networks (GNNs) face inherent limitations in interpreting these networks. Traditional convolutional approaches struggle to aggregate information from distant neighborhoods, while attention-based methods often fail to capture sufficient node-centric information, particularly regarding critical characteristics from pivotal nodes. These shortcomings hinder the identification of disease-specific variations when utilizing diverse features from multiple imaging modalities (e.g., MRI, PET). Furthermore, the heterogeneity introduced by multiple biomarker magnitudes complicates the analysis of local versus global graph properties.

Methodology: GTAD Framework
The authors propose GTAD (GNN with Transformer-guided Adaptive Diffusion), an integrated framework designed to aggregate both short-range (local) and long-range (global) graph properties. The architecture consists of three primary components:

1. Modality-wise Adaptive Convolution Block:

   • The model treats the brain network as an undirected graph $G$ with a normalized Laplacian $\hat{L}$.
   • For each imaging modality $m$, an encoder processes node features $x^m$ using a heat-kernel convolution. The convolution operation is defined as $H^m_z = \sigma_z(e^{-s^m \hat{L}} H^m_{z-1} W^m_z)$, where $s^m$ represents a trainable, node-wise scale vector specific to each modality.
   • Unlike fixed kernels, these scales are learned to adaptively capture the local characteristics of individual nodes for different modalities. This allows the model to determine the optimal neighborhood range for each ROI based on the specific imaging data.

2. Modality-wise Self-Attention Block:

   • The embeddings $\{H^m_Z\}$ from the convolution block are fed into a Transformer-based self-attention mechanism.
   • Unlike standard Transformers that often treat all tokens uniformly, GTAD assigns a specific attention head to each modality. This allows the model to integrate long-range information across the graph that is not captured by the local convolution.
   • The multi-head self-attention computes global characteristics by jointly attending to information from different modalities across various ROIs. The block includes residual connections and layer normalization to stabilize learning.

3. Transformer-Guided Scale Update:

   • The global representation $B_P$ from the attention block is passed to a downstream classifier to predict the disease label.
   • The loss function (cross-entropy) is backpropagated not only to update weights but also to update the modality-specific scales $s^m$.
   • An $\ell_1$ norm regularization is applied to ensure positive scales for the heat-kernel. This mechanism allows the Transformer to guide the adaptive diffusion process, effectively tuning the local receptive fields based on the global classification task.

Key Contributions

1. Novel Framework: The paper proposes a framework that aggregates short-range properties via adaptive diffusion kernels and long-range properties via multi-head attention, addressing the limitations of purely convolutional or purely attention-based GNNs.
2. Superior Performance: The model demonstrates state-of-the-art performance in graph classification tasks for preclinical AD, outperforming various baselines including convolution-based GNNs (GCN, GAT), diffusion-based GNNs (GraphHeat, GDC, ADC, LSAP), and Graph Transformers (NodeFormer, DIFFormer, SGFormer).
3. Interpretability: The framework provides interpretability by learning node-wise scales and attention scores. This allows for the identification of key ROIs associated with preclinical AD stages, offering insights into which anatomical regions are most critical for specific imaging modalities.

Experimental Results
Experiments were conducted on the Alzheimer's Disease Neuroimaging Initiative (ADNI) dataset, comprising 919 preclinical subjects classified into Control (CN), Significant Memory Concern (SMC), and Early Mild Cognitive Impairment (EMCI). The study utilized structural brain networks from Diffusion Tensor Imaging (DTI) and ROI measures from FDG-PET, Amyloid-PET, and MRI (cortical thickness).

• Classification Accuracy: GTAD achieved the highest accuracy across all experimental combinations of modalities. Notably, it reached 96.3% accuracy when using all imaging features, outperforming the next best Transformer-based method (SGFormer at 95.1%).
• Robustness: The model maintained high performance even in difficult conditions, such as predicting early AD stages using cortical thickness and $\beta$-Amyloid, where accuracy reached 94.5%.
• Ablation Study: Removing the adaptive convolution layer or the multi-modal attention mechanism resulted in significant performance drops, confirming that the joint training of local adaptive scales and global attention is essential for the model's effectiveness.

Interpretability Findings
The study analyzed the learned scales and attention scores to identify disease-specific variations:

• Adaptive Scales: The model identified that ROIs with small trained scales rely heavily on neighboring information, while those with large scales require distant features. The scales varied significantly across modalities for the same brain region, providing a multi-dimensional understanding of subnetworks.
• Key ROIs: The model highlighted specific regions with high importance, including the lingual gyrus (detected with the highest importance rate across all modalities), hippocampus (high importance in FDG), and putamen (high in Cortical Thickness and $\beta$-Amyloid). These regions are known to be among the first affected in AD progression.

Significance
The authors claim that GTAD offers a practical solution for preclinical AD classification and interpretation. By effectively combining local and global graph properties through a Transformer-guided diffusion process, the model not only improves diagnostic accuracy but also identifies key anatomical regions associated with the disease. This dual capability supports early diagnosis and the potential for disease prevention strategies by elucidating the specific biomarker variations linked to preclinical stages of Alzheimer's.