A Generalized Neural Diffusion Framework on Graphs

Graphs, such as traffic networks, social networks, citation networks, and molecular networks, are ubiquitous in the real world. Recently, Graph Neural Networks (GNNs), which are able to effectively learn the node representations based on the message-passing manner, have shown great popularity in tackling graph analytics problems. So far, GNNs have significantly promoted the development of graph analysis towards real-world applications. e.g, node classifification (Abu-El-Haija et al. 2019; Lei et al. 2022; Wang et al. 2017), link prediction (Kipf and Welling 2016b; You, Ying, and Leskovec 2019), graph reconstruction (Ouyang et al. 2023), subgraph isomorphism counting (Yu et al. 2023), and graph classifification (Gao and Ji 2019; Zhang et al. 2018).

Some recent studies show that GNNs are in fact intimately connected to diffusion equations (Chamberlain et al. 2021; Wang et al. 2021; Thorpe et al. 2022), which can be considered as information diffusion on graphs. Diffusion equation interprets GNNs from a continuous perspective (Chamberlain et al. 2021) and provides new insights to understand existing GNN architectures, which motives some diffusion-based GNNs. For instance, (Wang et al. 2021) proposes continuous graph diffusion. (Thorpe et al. 2022) utilizes the diffusion process to handle oversmoothing issue. (Song et al. 2022) considers a graph as a discretization of a Riemannian manifold and studies the robustness of the information propagation process on graphs. Diffusion equation can also build a bridge between traditional GNNs and control theory (Zang and Wang 2020). Although the diffusion process and graph convolution are closely related, little effort has been made to answer: Is there a unified diffusion equation framework to formally unify the current GNN architectures? A well-informed answer can deepen our understanding of the learning mechanism of GNNs, and may inspire to design a broad new class of GNNs based on diffusion equation.

Actually, (Chamberlain et al. 2021) has explained GCN (Kipf and Welling 2016a) and GAT (Veličković et al. 2017) from diffusion equation. However, with more proposed GNN architectures, it is highly desired to formally revisit the relation between diffusion equation and GNNs. In this paper, we discover that many GNN architectures substantially can be unified with a general diffusion equation with the fidelity term, such as GCN/SGC [22], APPNP[14], GAT [19], AMP [16], DAGNN [15]. Basically, the diffusion equation describes that the change of a node representation depends on the movement of information on graphs from one node to its neighbors, and the fidelity term constraints that the change of a node representation depends on the difference with its initial feature. Furthermore, we show that the unified diffusion framework can also be derived from an energy function, which explains the whole framework as an energy minimization process in a global view. Compared with other unified frameworks (Zhu et al. 2021; Ma et al. 2021), our framework is from the diffusion perspective, which has many advantages. For example, diffusion-based methods are able to address the common plights of graph learning models such as oversmoothing (Chamberlain et al. 2021; Wang et al. 2021). What’s more, the diffusion equation can be seen as partial differential equations (PDEs) (Chamberlain et al. 2021), thus introducing many schemes to solve the graph diffusion equation such as explicit scheme, implicit scheme, and multi-step scheme, some of which are more stable and converge faster.

Based on the above findings, we can see that the diffusion process employed by most current GNNs just considers the first-order diffusion equation, which only diffuses messages among 1-hop neighbors. That is, the first-order diffusion has the underlying homophily assumption among 1-hop neighbors. While we empirically discover that the labels of 2-hop neighborhoods actually appear monophily property (Altenburger and Ugander 2018), i.e., nodes may have extreme preferences for a particular attribute which are unrelated to their own attribute and 1-hop neighbors’ attribute, but are more likely to be similar with the attribute of their 2-hop neighbors. Simply put, monophily can induce a similarity among 2-hop neighbors without requiring similarity among 1-hop neighbors. So when the 1-hop neighbors are heterophily-dominant or have noise, the 2-hop neighbors will provide more relevant context. Therefore, a more practical diffusion process should take both the first-order and second-order neighbors into account. How can we design a new type of graph diffusion networks satisfying the above requirement based on our framework?

In this paper, we design a new high-order neighbor-aware diffusion equation in our proposed diffusion framework, and then derive a High-order Graph Diffusion Network (HiD-Net). Specifically, our model simultaneously combines the first-order and second-order diffusion process, then we regularize the diffusion equation by minimizing the discrepancy between the estimated and the original graph features. The whole diffusion equation is finally integrated into the APPNP architecture. With second-order diffusion equation, HiD-Net is more robust against the attacks and more general on both homophily and heterophily graphs. We theoretically prove that HiD-Net is essentially related with the second-order random walk. We also provide the convergence guarantee that HiD-Net will converge to this random walk’s limit distribution as the number of layers increases, and meanwhile, the learned representations do not converge to the same vector over all nodes. The contributions of this paper are summarized as follows:

• •

  We propose a novel generalized diffusion graph framework, consisting of diffusion equation and fidelity term. This framework, formally establishing the relation between diffusion process with a wide variety of GNNs, describes a broad new class of GNNs based on the discretized diffusion equations on graphs and provides new insight to the current graph diffusion/neural networks.

• •

  We discover the monophily property of labels, and based on our diffusion framework, we propose a high-order graph diffusion network, HiD-Net, which is more general and robust. We theoretically build the relation between HiD-Net and second-order random walk, together with the convergence guarantee.

• •

  Our extensive experiments on both the homophily and heterophily graphs clearly show that HiD-Net outperforms the popular GNNs based on diffusion equation.

Graph convolutional networks. Recently, graph convolutional network (GCN) models (Bruna et al. 2013; Defferrard, Bresson, and Vandergheynst 2016; Kipf and Welling 2016a; Veličković et al. 2017; Hamilton, Ying, and Leskovec 2017) have been widely studied. Based on the spectrum of graph Laplacian, (Bruna et al. 2013) generalizes CNNs to graph signal. Then (Defferrard, Bresson, and Vandergheynst 2016) further improves the efficiency by employing the Chebyshev expansion of the graph Laplacian. (Kipf and Welling 2016a) proposes to only aggregate the node features from the one-hop neighbors and simplifies the convolution operation. (Veličković et al. 2017) introduces the attention mechanisms to learn aggregation weights adaptively. (Hamilton, Ying, and Leskovec 2017) uses various ways of pooling for aggregation. More works on GNNs can be found in surveys (Wu et al. 2020b; Zhou et al. 2020).

Diffusion equation on graphs. Graph Heat Equation (GHE) (Chung and Graham 1997), which is a well-known generalization of the diffusion equation on graph data, models graph dynamics with applications in spectral graph theory. GRAND (Chamberlain et al. 2021) studies the discretized diffusion PDE on graphs and applies different numerical schemes for their solution. GRAND++ (Thorpe et al. 2022) mitigates the oversmoothing issue of graph neural networks by adding a source term. DGC (Wang et al. 2021) decouples the terminal time and propagation steps of linear GCNs from a perspective of graph diffusion equation, and analyzes why linear GCNs fail to benefit from deep layers. ADC (Zhao et al. 2021) strategies to automatically learn the optimal diffusion time from the data. However, these works focus on specific graph diffusion network, thus there is not a framework to formally unify the GNNs.

The unified GNN framework. (Zhu et al. 2021) establishes a connection between different propagation mechanisms with a unified optimization problem, and finds out that the proposed propagation mechanisms are the optimal solution for optimizing a feature-fitting function over a wide class of graph kernels with a graph regularization term. (Ma et al. 2021) establishes the connections between the introduced GNN models and a graph signal denoising problem with Laplacian regularization. It essentially is still an optimization solving framework from the perspective of signal denoising. However, our framework is based on the diffusion equation, where the advantages are two fold: one is that diffusion-based methods are able to address the oversmoothing problem (Chamberlain et al. 2021; Wang et al. 2021). The other is that the diffusion equation can be seen as partial differential equations (PDEs) (Chamberlain et al. 2021) and thus can introduce many schemes that have many good properties such as fast convergence rate and high stability.

Notations. Consider an undirected graph as $G=(V,E)$ with adjacency matrix $\mathbf{A}\in\mathbb{R}^{n\times n}$, where $V$ contains $n$ nodes $\{v_{1},\dots,v_{n}\}$ and $E$ is the set of edges. The initial node feature matrix is denoted as $\mathbf{X}^{(0)}\in\mathbb{R}^{n\times q}$, where $q$ is the dimension of node feature. We denote the neighbors of node $i$ at exactly $k$ hops/steps away as $N_{k}(i)$. For example, $N_{1}(i)=\{j:(i,j)\in E\}$ are the immediate neighbors of $i$.

Diffusion is a physical process that equilibrates concentration differences without creating or destroying mass. This physical observation can be easily cast in the diffusion equation, which is a parabolic partial differential equation. Fick’s law of diffusion describes the equilibration property (Weickert 1998):

where $J$ is the diffusion flux, which measures the amount of substance that flows through a unit area during a unit time interval. $f$ is the diffusivity coefficient, which can be constant or depend on time and position. $\nabla u$ is the concentration gradient. This equation states that a concentration gradient $\nabla u$ causes a flux $J$ which aims to compensate for this gradient. The observation that a change in concentration in any part of the system is due to the influx and outflux of substance into and out of that part of the system can be expressed by the continuity equation:

where $t$ denotes the time. Plugging in Fick’s law  (1) into the continuity equation we end up with the diffusion equation:

As $\operatorname{div}$ is the sum of the second derivatives in all directions, please note that normal first order derivatives and second order derivatives are on continuous space and can not be generalized directly to graph which is on discrete space. As (Chamberlain et al. 2021) defined, the first derivative is the difference between the feature of a node and its neighbor. And the second derivative can be considered as the difference between the first derivatives of the node itself and its neighbors.

For better illustration, we provide an example. Consider a chain graph in Figure  1, where $i$, $i+1$, and $i-1$ are the indexes of the nodes. The feature of node $i$ is denoted as $x_{i}$.

The first order derivatives on node $i$ is defined as $x_{i+1}-x_{i}$ and $x_{i}-x_{i-1}$. The diffusion flux from node $j$ to node $i$ at time $t$ on a graph is:

The second order derivative is the difference of first order derivatives: $(x_{i+1}-x_{i})-(x_{i}-x_{i-1})=(x_{i+1}-x_{i})+(x_{i-1}-x_{i})$. We notice that the chain graph only has one dimension, so the divergence of node $i$ on a chain graph is equal to its second order derivative: $\operatorname{div}(\nabla x_{i})=(x_{i+1}-x_{i})+(x_{i-1}-x_{i})$.

Thus, on a normal graph, we have the generalized form: $\operatorname{div}(\nabla x_{i})=\sum_{j\in N_{1}(i)}(x_{j}-x_{i})$.

Here we normalize the diffusion process utilizing the degree of the nodes to down-weight the high-degree neighbors, and we have $\operatorname{div}(\nabla x_{i})=\sum_{j\in N_{1}(i)}\dfrac{\tilde{A}_{ij}}{\sqrt{\tilde{d_{i}}}\sqrt{\tilde{d_{j}}}}(x_{j}-x_{i})$, where $\tilde{A}_{ij}$ is the element of $\tilde{\mathbf{A}}=\mathbf{A}+\mathbf{I}$, and $d_{i}=\sum_{j}\tilde{A}_{ij}$.

So the diffusion equation on node $i$ can be defined as:

The diffusion equation models the change of representation $x_{i}^{(t)}$ with respect to $t$, which depends on the difference between the nearby nodes, implying that the greater the difference between a node and its neighbors, the faster it changes.

However, how fast $x_{i}^{(t)}$ changes should not only depend on the representation difference between node $i$ and its neighbors, otherwise, it will cause oversmoothing issue, i.e., as the diffusion process goes by, the nodes are not distinguishable. Based on this phenomenon, we think that the representation change of $x_{i}^{(t)}$ should be also related with the node feature $x_{i}^{(0)}$ itself, i.e., if the difference between $x_{i}^{(t)}$ and $x_{i}^{(0)}$ is small, the change of $x_{i}^{(t)}$ should also be small. Then we add another fidelity term and obtain our general graph diffusion framework as follows:

where $\alpha$, $\beta$ are coefficients.

Remark 1.  (6) can be derived from the energy function:

where $\theta$ represents the position of the nodes, and $\Omega$ represents the entire graph domain. The corresponding Euler–Lagrange equation, which gives the necessary condition for an extremum of  (7), is given by:

(8) can also be regarded as the steady-state equation of  (6). Based on the energy function, we can see that  (6) constrains space variation and time variation of the diffusion process, indicating that the representations of the graph nodes will not change too much between nearby nodes, as well as not change too much from the initial features.

Remark 2. The framework (6) is closely related to many GNNs, such as GCN/SGC (Wu et al. 2019), APPNP (Klicpera, Bojchevski, and Günnemann 2018), GAT (Veličković et al. 2017), AMP (Liu et al. 2021), DAGNN (Liu, Gao, and Ji 2020), as demonstrated by the following propositions. We provide the proofs of all the subsequent propositions in Appendix.

Proposition 1. With $\alpha=0$, $\beta=1$, $\Delta t=1$ and $f=1$ in  (6), the diffusion process in SGC/GCN is:

Proposition 2. Introducing $\eta$ as coefficient, with $\alpha=1$, $\beta=1-\dfrac{1}{\eta}$, $\Delta t=1$ and $f=1$ in  (6), the diffusion process in APPNP is:

Proposition 3. With $\alpha=0$, $\beta=1$, $\Delta t=1$ and the learned similarity coefficient $f_{ij}^{(t)}$ between nodes $i$ and $j$ at time $t$ in  (6), the diffusion process in GAT is:

Proposition 4. With stepsize $\epsilon$ and coefficient $\lambda$, $\beta_{i}^{(t)}=\max\left(1-\frac{\epsilon\lambda}{\left\|(1-2\epsilon(1-\lambda))\mathbf{X}_{i}^{(t)}+2\epsilon(1-\lambda)\tilde{\mathbf{A}}\mathbf{X}_{i}^{(t)}-\left(\mathbf{X}^{(0)}\right)_{i}\right\|_{2}},0\right)$. Let $\alpha=1-\beta_{i}^{(t)}$, $\beta=2\epsilon(1-\lambda)\beta_{i}^{(t)}$, $\Delta t=1$ and $f=1$ in  (6), the diffusion process in AMP is:

Proposition 5. With $\alpha+\beta=1$, $\Delta t=1$ and the learned coeffiecient $f(t)$ at time $t$ satisfying $\sum_{t=0}^{T}(\beta f(t))^{t}=\frac{1}{\alpha}$ in  (6), the diffusion process in DAGNN is:

In this paper, we propose a generalized diffusion graph framework, which establishes the relation between diffusion equation with different GNNs. Our framework reveals that current graph diffusion networks mainly consider the first-order diffusion equation, then based on our finding of the monophily property of labels, we derive a novel high-order diffusion graph network (HiD-Net). HiD-Net is more robust and general on both homophily and heterophily graphs. Extensive experimental results verify the effectiveness of HiD-Net. One potential issue is that our model utilizes a constant diffusivity coefficient, and a future direction is to explore a learnable diffusivity coefficient depending on time and space. Our work formally points out the relation between diffusion equation with a wide variety of GNNs. Considering that previous GNNs are designed mainly based on spatial or spectral strategies, this new framework may open a new path to understanding and deriving novel GNNs. We believe that more insights from the research community on the diffusion process will hold great potential for the GNN community in the future.

This work is supported in part by the National Natural Science Foundation of China (No. U20B2045, 62192784, U22B2038, 62002029, 62172052, 62322203).