Re-Think and Re-Design Graph Neural Networks in Spaces of Continuous Graph Diffusion Functionals

Graph diffusion process. Given graph data $\mathcal{G}=(V,W)$ with $N$ nodes $V=\{v_{i}|i=1,…,N\}$, the adjacency matrix $W=\left[w_{ij}\right]_{i,j=1}^{N}\in\mathbb{R}^{N\times N}$ describes connectivity strength between any two nodes. For each node $v_{i}$, we have a graph embedding vector $x_{i}\in\mathcal{R}^{m}$. In the context of graph topology, the graph gradient $\left(\nabla_{\mathcal{G}}x\right)_{ij}=w_{ij}\left(x_{i}-x_{j}\right)$ indicates the feature difference between $v_{i}$ and $v_{j}$ weighted by the connectivity strength $w_{ij}$, where $\nabla_{\mathcal{G}}$ is a $\mathbb{R}^{N}\rightarrow\mathbb{R}^{N\times N}$ operator. Thus, the graph diffusion process can be formulated as $\frac{\partial x(t)}{\partial t}=div\left(\nabla_{\mathcal{G}}x(t)\right)$, where the evolution of embedding vectors $x=\left[x_{i}\right]_{i=1}^{N}$ is steered by the graph divergence operator.

Connecting GNN to graph diffusion. In the regime of GNN, the regularization in the loss function often measures the smoothness of embeddings $x$ over the graph by $x^{T}\Delta x$, where $\Delta=div(\nabla_{\mathcal{G}})$ is the graph Laplacian operator [27]. To that end, the graph smoothness penalty encourages two connected nodes to have similar embeddings by information exchange in each GNN layer. Specifically, the new graph embedding $x^{l}$ in the $l^{th}$ layer is essentially a linear combination of the graph embedding $x^{l-1}$ in the previous layer, i.e., $x^{l}=A_{W,\Theta}x^{l-1}$, where the matrix $A$ depends on graph adjacency matrix $W$ and trainable GNN parameter $\Theta$. After rewriting $x^{l}=Ax^{l-1}$ into $x^{l}-x^{l-1}=(A-I)x^{l-1}$, updating graph embeddings $x$ in GNN falls into a discrete graph diffusion process, where the time parameter $t$ acts as a continuous analog of the layers in the spirit of Neural ODEs [10]. It has been shown in [6] that running the graph diffusion process for multiple iterations is equivalent to applying a GNN layer multiple times.

GNN is a discrete model of Lagrangian mechanics via E-L equation. The diffusion process $\frac{\partial x(t)}{\partial t}=div\left(\nabla_{\mathcal{G}}x(t)\right)$ has been heavily studied in image processing in decades ago, which is the E-L equation of the functional $\mathop{\min}\limits_{x}\int_{\Omega}|\nabla x{|^{2}}dx$. By replacing the 1D gradient operator defined in the Euclidean space $\Omega$ with the graph gradient $(\nabla_{\mathcal{G}}x)_{ij}$, it is straightforward to find that the equation governing the graph diffusion process $\frac{\partial x(t)}{\partial t}=div\left(\nabla_{\mathcal{G}}x(t)\right)$ is the E-L equation of the functional $\mathop{\min}\limits_{x}\int_{\mathcal{G}}|\nabla_{\mathcal{G}}x{|^{2}}dx$ over the graph topology. Since the heat kernel diffusion is essentially the mathematical description of the inductive bias in current GNNs, we have established a mapping between the mechanics of GNN models and the functional of graph diffusion patterns in a continuous domain.

Tracing the smoking gun of over-smoothing in GNNs. In Fig. 1, we observed that the inductive bias of link-wise propagation is the major reason for excessive information exchange, which attributes to the over-smoothing problem in GNNs. An intuitive approach is to align the diffusion process with high-level properties associated with graph topology, such as network communities. After connecting the GNN inductive bias to the functional of graph diffusion process, we postulate that the root cause of over-smoothing is the isotropic regularization mechanism encoded by the ${\ell_{2}}$-norm. More importantly, connecting GNN to the calculus of variations offers a more principled way to design new deep models with mathematics guarantees and model mechanistic explainability.

The general roadmap for re-designing GNNs typically involves three major steps: (1) formulating inductive bias into the functional of graph diffusion patterns; (2) deriving the corresponding E-L equation; and then (3) developing a new deep model of GNN based on the finite difference solution of E-L equation. Since the graph diffusion functional is application-specific, we demonstrate the construction of new GNN models in the following two graph learning applications.

Dataset and benchmark methods. We evaluate the new GNN models derived from our proposed GNN framework in two different applications. First, we use three standard citation networks, namely Cora, Citeseer, and Pubmed [36] for node classification (the detailed data statistic is shown in Table 3 of Supplementary). We adopt the public fixed split [46] to separate these datasets into training, validation, and test sets. We follow the experimental setup of [9] for a fair comparison with six benchmark GNN models (vanilla GCN [27], GAT [39], GCNII [9], ResGCN [29], DenseGCN [29], GRAND [6]). Since our DC-layer can be seamlessly integrated into existing GNNs as a plug-in. The corresponding new GNN models (with DC-layer) are denoted GCN+, GAT+, GCNII+, ResGCN+, DenseGCN+, and GRAND+, respectively.

Second, we apply the GAN model in Section 2.2.2 to predict the concentration level of AD-related pathological burdens and their spreading pathways from longitudinal neuroimages. Currently, there is no in-vivo imaging technique that can directly measure the flow of information across brain regions. Here, our computational approach holds great clinical value to understand the pathophysiological mechanism involved in disease progression [26]. Specifically, we parcellate each brain into 148 cortical surface regions and 12 sub-cortical regions using Destrieux atlas [13]. The wiring topology of these 160 brain regions is measured by diffusion-weighted imaging [3] and tractography techniques [19]. The regional concentration levels AD pathology including amyloid, tau, and fluorodeoxyglucose (FDG) and cortical thickness (CoTh) are measured from PET (positron emission tomography) and MRI (magnetic resonance imaging) scans [23]. We use a total of $M=1,291$ subjects from ADNI [32], each having longitudinal imaging data (2-5 time points). The details of image statistics and pre-processing are shown in Sec. 6.1.2. Since we apply the flow prediction model to each modality separately, we differentiate them with X-FlowNet ($X$ stands for amyloid, tau, FGD, and CoTh).

Experimental setup. In the node classification task, we verify the effectiveness and generality of DC-layer in various number of layers ($L={2,4,8,16,32,64,128}$). All baselines use their own default parameter settings. Evaluation metrics include accuracy, precision and F1-score. To validate the performance of X-FlowNet, we examine (1) prediction accuracy (MAE) of follow-up concentration level, (2) prediction of the risk of developing AD using the baseline scan, and (3) understand the propagation mechanism in AD by revealing the node-to-node spreading flows of neuropathologies.

The main results of graph node classification and flow prediction are demonstrated in Section 3.2 and 3.3, respectively. Other supporting results such as ablation study and hyper-parameter setting are shown in Section 6 of the Supplementary document.

We postulate that by mitigating the over-smoothing issue, we can leverage the depth of GNN models to effectively capture complex feature representations in graph data. As shown in Table 1, we investigate the graph node classification accuracy as we increase the number of GNN layers by six benchmark GNN models and their corresponding plug-in models (indicated by ’+’ at the end of each GNN model name) with the DC-layer. The results demonstrate that: (1) the new GNN models generated from the GNN-PDE-COV framework have achieved SOTA in Cora (86.30% by GCNII+), Citeseer (75.65% by GRAND+), and Pubmed (80.10 % by GCNII+); (2) all of new GNN models outperforms their original counterparts with significant improvement in accuracy; (3) the new GNN models exhibit less sensitivity to the increase of model depth compared to current GNN models; (4) the new GNN models are also effective in resolving the gradient explosion problem [30] (e.g, the gradient explosion occurs when training GAT on all involved datasets with deeper than 16 layers, while our GAT+ can maintain reasonable learning performance even reaching 128 layers.)

It is important to note that due to the nature of the graph diffusion process, graph embeddings from all GNN models (including ours) will eventually become identical after a sufficiently large number of layers [11]. However, the selective diffusion mechanism (i.e., penalizing excessive diffusion across communities) provided by our GNN-PDE-COV framework allows us to control the diffusion patterns and optimize them for specific graph learning applications.

First, we evaluate the prediction accuracy between the ground truth and the estimated concentration values by our X-FlowNet and six benchmark GNN methods. The statistics of MAE (mean absolute error) by X-FlowNet, GCN, GAT, GRAND, ResGCN, DenseGCN and GCNII, at different noise levels on the observed concentration levels, are shown in Fig. 5 (a). It is clear that our X-FlowNet consistently outperforms the other GCN-based models in all imaging modalities.

Second, we have evaluated the potential of disease risk prediction and presented the results in Table 4 in Supplementary document, where our GNN-PDE-COV model not only achieved the highest diagnostic accuracy but also demonstrated a significant improvement (paired t-test $p<0.001$) in disease risk prediction compared to other methods. These results suggest that our approach holds great clinical value for disease early diagnosis.

Third, we examine the spreading flows of tau aggregates in CN (cognitively normal) and AD groups. As the inward and outward flows shown in Fig. 5(b), it is evident that there are significantly larger amount of tau spreading between sub-cortical regions and entorhinal cortex in CN (early sign of AD onset) while the volume of subcortical-entorhinal tau spreading is greatly reduced in the late stage of AD. This is consistent with current clinical findings that tau pathology starts from sub-cortical regions and then switches to cortical-cortical propagation as disease progresses [28]. However, our Tau-FlowNet offers a fine-granularity brain mapping of region-to-region spreading flows over time, which provides a new window to understand the tau propagation mechanism in AD etiology [14].

By substituting Eq. 2 into Eq. 1 in the main manuscript, we can obtain the objective function of subscript $z$ (we temporarily drop $i$ for clarity):

Next, we convert Eq. LABEL:function2 into a minimization problem as follows:

By letting the derivative with respect to $z_{i}$ to zero, we have the following equation

Since $z$ might be in high dimensional space, solving such a large system of linear equations under the constraint $|z|\leq 1$ is oftentimes computationally challenging. In order to find a practical solution for $z$ that satisfies the constrained minimization problem in Eq. 10, we resort to the majorization-minimization (MM) method [16]. First, we define:

By setting $z^{l}$ as point of coincidence, we can find a separable majorizer of $M(z)$ by adding the non-negative function

to $M(z)$, where $\beta$ is greater than or equal to the maximum eigenvalue of $\nabla_{\cal G}x\nabla_{\cal G}{x^{\intercal}}$. Note, to unify the format, we use the matrix transpose property in Eq. 13. Therefore, a majorizer of $M(z)$ is given by:

And, using the MM approach, we can obtain the update equation for $z$ as follows:

where $b={z^{l}}+\frac{1}{\beta}{\nabla_{\cal G}}(\frac{2}{\lambda}{x^{0}}-{\nabla_{\cal G}}x{z^{l}})$.

Then, the next step is to find $z\in{\cal R}^{N}$ that minimizes $z^{\intercal}z-2bz$ subject to the constraint $|z|\leq 1$. Let’s first consider the simplest case where $z$ is a scalar:

The minimum of ${{z^{2}}-2bz}$ is at $z=b$. If $b\leq 1$, then the solution is $z=b$. If $|b|\geq 1$, then the solution is $z=sign(b)$. We can define the clipping function as:

as illustrated in the middle of Fig. 3 in the main text, then we can write the solution to Eq. 16 as $z=clip(b,1)$.

Note that the vector case Eq. 15 is separable - the elements of $z$ are uncoupled so the constrained minimization can be performed element-wise. Therefore, an update equation for $z$ is given by:

where $l$ denotes the index of the network layer, the representation of $(l+1)^{th}$ is given by Eq. (1) in the main manuscript. Because the optimization problem is convex, the iteration will converge from any initialization. We may choose, say $z^{0}=0$. We call this the iterative diffusion-clip (DC) algorithm.

This algorithm can also be written as

By scaling $z$ with a factor of $\lambda/2$, we have the following equivalent formulations:

We summarize the process of the diffusion-clip (DC) layer in Algorithm 1 (it is similar to the iterative shrinkage threshold algorithm [12]):

According to the introduction of Secction 2.2.2 (Eq. 4 and Eq. 5) in the main manuscript, we summarize the following equations,

Since the PDE in Eq. 5 in the main manuscript is equivalent to the E-L equation of the quadratic functional $\mathcal{J}(u)=\mathop{\min}\limits_{u}\int_{\mathcal{G}}\alpha\otimes|\nabla_{\mathcal{G}}u{|^{2}}du$ (after taking $\phi$ away), we propose to replace the ${\ell_{2}}$-norm integral functional $\mathcal{J}(u)$ with TV-based counterpart

We then introduce an auxiliary matrix $f$ to lift the undifferentiable barrier, and reformulate the TV-based functional as a dual min-max functional

where we maximize $f$ such that $\mathcal{J}_{TV}(u,f)$ is close enough to $\mathcal{J}_{TV}(u)$. Using Gâteaux variations, we assume $u\to u+\varepsilon a$, $f\to f+\varepsilon b$, and the directional derivatives in the directions $a$ and $b$ defined as ${\left.{\frac{{d{\cal J}}}{{d\varepsilon}}(u+\varepsilon a)}\right|_{\varepsilon\to 0}}$ and ${\left.{\frac{{d{\cal J}}}{{d\varepsilon}}(f+\varepsilon b)}\right|_{\varepsilon\to 0}}$. Given a functional $\mathcal{J}_{TV}(u,f)$, its Gâteaux variations is formulated by:

Since we assume either $u$ is given at the boundary (Dirichlet boundary condition), the boundary term $\alpha\otimes f\cdot a$ can be dropped. After that, the derivative of $\mathcal{J}_{TV}(u,f)$ becomes: