Adaptive Graph Diffusion Networks

This is a pytorch implementation of the paper Adaptive Graph Diffusion Networks.

Environment

We conduct all experiments on a single Tesla V100 (16Gb) card. The maximum memory cost on certain datasets may be up to ~32Gb. We implement our methods with Deep Graph Library (DGL). Check requirements.txt for more associated python packages.

Reproduce

We offer associated shell scripts for reproducing our results. Remember to modify your own dataset root path in each directory. We also offer other common models and possible sampling methods in some directories (not fully evaluated).

Feel free to utilize and modify this repository, but remember to briefly introduce and cite this work:D

Method

(There exist some rendering mistakes when using \sum_{xxx}^{xxx} in latex scripts)

We propose Adaptive Graph Diffusion Networks (AGDNs) to extend receptive fields of common Message Passing Nueral Networks (MPNNs), without extra layers (feature transformations) or decoupling model architecture (resulting graph convolution restricted in the same space). Following the historical path of GNNs, that spectral GNNs have elegant analyzability but usually have poor scalability and performance, we generalize the graph diffusion to be more spatial. In detail, for an MPNN model (usually GAT in this repository), we replace the graph convolution operator in each layer with a generalized graph diffusion operator. The generalized graph diffusion operator is defined as follows:

$$\tilde{\boldsymbol H}^{(l,0)} = \boldsymbol H^{(l-1)}\boldsymbol W^{(l)},$$

$$\tilde{\boldsymbol H}^{(l,k)} = \overline{\boldsymbol A}\tilde{\boldsymbol H}^{(l,k-1)},$$

$${\boldsymbol H}^{(l)}=\sum^{K}_{k=0}{\boldsymbol \Theta}^{(k)}\otimes\tilde{\boldsymbol H}^{(l,k)}+\boldsymbol{H}^{(l-1)}\boldsymbol{W}^{(l),r},$$

where $\otimes$ denotes the element-wise matrix multiplication. We describe the above proceedures in a node viewpoint:

$$\tilde{\boldsymbol h}^{(l,0)}_i=\boldsymbol h^{(l-1)}_i\boldsymbol W^{(l)},$$

$$\tilde{\boldsymbol h}^{(l,k)}i=\sum{j\in \mathcal{N}i}\left(\overline{A}{ij}\tilde{\boldsymbol h}^{(l,k-1)}_j\right),$$

$$h^{(l)}{ic}=\left[\sum{k=0}^{K}\left(\theta_{ikc}\tilde{h}^{(l,k)}{ic}\right)+\sum{c'=1}^{d^{(l-1)}}\left(h^{(l-1)}{ic'}W^{(l),r}{c'c}\right)\right],$$

To obtain the possibly node-wise or channel-wise weighting coefficients $\theta_{ikc}$, we propose two mechanisms: Hop-wise Attention (HA) and Hop-wise Convolution (HC).

Hop-wise Attention (HA) is a GAT-like attention mechanism and utilizes a $2\times d$ query vector $\boldsymbol a_{hw}$ to induce node-wise weighting coefficients $\boldsymbol \Theta^{HA} \in \mathbb R^{N\times (K+1)}$:

$$\omega_{ik} = \left[\tilde{\boldsymbol h}^{(l,0)}{i}\left\lvert\right\rvert\tilde{\boldsymbol h}^{(l,k)}{i}\right]\cdot \boldsymbol a_{hw},$$

$${\theta}{ik}^{HA}=\frac {{\rm exp}\left({\sigma}\left(\omega{ik}\right)\right)}{\sum {k=0}^{K} {{\rm exp}\left({\sigma}\left(\omega{ik}\right)\right)}},$$

$$\theta^{HA}{ikc}=\theta^{HA}{ik},\forall i,k,c$$

Hop-wise Convolution (HC) directly defines a learnable channel-wise convolution kernel $\boldsymbol \Theta^{HC}\in \mathbb R^{(K+1)\times d}$:

$$\theta^{HC}{ikc}=\theta^{HC}{kc},\forall i,k,c$$

HA and HC are just examples, we expect more mechanisms from the community.

In addition, we utilize hop-wise Positional Embedding (PE) on some datasets to enhance hop-wise position information (PE may increase 1~2 Gb memory cost):

$$\tilde{\boldsymbol h}^{(l,k)} = \tilde{\boldsymbol h}^{(l,k)} + \boldsymbol p^{(l,k)}$$

Performance

ogbn-arxiv:

ogbn-proteins:

ogbn-products:

Heterophily-prone datasets: Chameleon & Squirrel & Actor (Test Accuracy (%)):

ogbl-ppa:

ogbl-ddi:

ogbl-citation2:

Runtime and Parameters

ogbn-proteins (The inference runtime on another RTX 6000 (48Gb) card of RevGNN is not reported in its paper):

ogbl-ppa, ogbl-ddi, ogbl-citation2:

Extra tricks

For ogbn-arxiv: BoT, self-KD and GIANT-XRT.

For ogbn-ddi: AUC loss from PLNLP.

Reference

1. BoT (Repository, Paper)
2. Self-KD(Repository)
3. GIANT-XRT(Repository, Paper)
4. PLNLP(Repository, Paper)
5. GraphSAINT(Repository, Paper)
6. DeeperGCN (Repository, Paper)
7. RevGNNs (Repository, Paper)
8. UniMP (Repository, Paper)
9. SEAL (Repository, Paper)