GRAND: Graph Neural Diffusion

Benjamin Paul Chamberlain; James Rowbottom; Maria Gorinova; Stefan Webb; Emanuele Rossi; Michael M. Bronstein

Source-linked AI summary

GRAND: Graph Neural Diffusion

Benjamin Paul Chamberlain, James Rowbottom, Maria Gorinova, Stefan Webb, Emanuele Rossi, Michael M. Bronstein

arXiv:2106.10934v2Jun 21, 2021cs.LGstat.ML

TL;DR

Graph learning needs architectures that can become deep while addressing oversmoothing, bottlenecks, and stability. GRAND treats GNNs as discretisations of graph diffusion PDEs, using temporal and spatial discretisation choices to construct linear and nonlinear models. The resulting variants perform competitively across standard benchmarks, while the framework’s shared parameters reduce model size; the approach currently constrains hidden-state width and parameter sharing across layers.

Problem
Graph learning architectures commonly face depth, oversmoothing, bottleneck, and stability challenges, motivating a principled framework for developing new GNNs.
Method
GRAND models GNNs as discretisations of diffusion PDEs on graphs, using temporal schemes, spatial graph structure, and learnable encoder/decoder functions.
Results
GRAND variants perform among the best methods, achieving first place on all but one dataset and second place on that remaining dataset.
Takeaways & Limitations
Appropriate discretisation and numerical schemes allow GRAND to train very deep GNNs and achieve superior performance on popular benchmarks.
Takeaways & Limitations
The current formulation requires a constant hidden-state size and the same parameters θ across all layers.

Abstract

from arXiv · show

We present Graph Neural Diffusion (GRAND) that approaches deep learning on graphs as a continuous diffusion process and treats Graph Neural Networks (GNNs) as discretisations of an underlying PDE. In our model, the layer structure and topology correspond to the discretisation choices of temporal and spatial operators. Our approach allows a principled development of a broad new class of GNNs that are able to address the common plights of graph learning models such as depth, oversmoothing, and bottlenecks. Key to the success of our models are stability with respect to perturbations in the data and this is addressed for both implicit and explicit discretisation schemes. We develop linear and nonlinear versions of GRAND, which achieve competitive results on many standard graph benchmarks.

1. Introduction

GRAND uses PDE diffusion as a framework for understanding GNNs and designing new architectures. Its discretisation choices address depth, stability, oversmoothing, and bottlenecks while yielding competitive benchmark performance.

PDE tools provide a principled framework for understanding existing GNNs and developing a broad class of new methods.
Many popular GNN architectures arise from different diffusion equations and discretisation schemes within one mathematical framework.
Standard GNNs correspond to explicit single-step Euler discretisations, whereas Runge-Kutta and implicit schemes offer adaptive or unconditionally stable alternatives.
Appropriate discretisation choices enable GRAND architectures with tens of layers despite oversmoothing and bottlenecks.
GRAND includes linear and nonlinear architectures, stability conditions, and ablations, achieving competitive performance on popular benchmark datasets.

2. Background

The paper develops diffusion from continuous domains and manifolds, where gradients, divergences, and diffusivity determine how quantities spread. These concepts motivate diffusion-based processing beyond simple Euclidean spaces.

Diffusion models movement from higher to lower concentration until the distribution becomes equalized.
The diffusion equation combines temporal change with spatial operators such as gradient, divergence, and the Laplacian.
Diffusivity determines whether diffusion is homogeneous, inhomogeneous, isotropic, or anisotropic.
Diffusion on manifolds: On Riemannian manifolds, gradient and divergence operate between scalar and tangent vector fields and are adjoint under their inner products.
Applications of diffusion equations: Diffusion equations have been applied to nonlinear image filtering, including noise removal while preserving sharp edges.

3. Diffusion equations on graphs

Graph diffusion discretises continuous diffusion on graph nodes and edges, with temporal schemes corresponding to GNN layers and spatial choices corresponding to graph structure. Stability analysis distinguishes explicit and implicit updates, while GRAND leverages these choices for deeper and more efficient architectures.

Graph diffusion equation: Graph node and edge features are represented as vectors, with graph gradient and divergence operators acting on them as adjoint discrete operators.
Graph diffusion equation: The graph diffusion equation uses an attention-based diffusivity, yielding a nonlinear operator that becomes analytically solvable when attention is fixed.
Solving the graph diffusion equation: Explicit updates apply the diffusion operator directly, whereas implicit updates require solving a linear system and can use dense multi-hop operators.
Stability: The explicit scheme is stable for 0 < τ < 1, while the implicit scheme is unconditionally stable for any τ > 0.
Temporal discretisation: A GNN layer corresponds to one discrete time step, so repeated diffusion iterations correspond to repeatedly applying a GNN layer.
Connection to existing architectures: Most GNNs are explicit single-step schemes, including a residual GAT formulation under learnable attention and no intermediate nonlinearity.
Connection to existing architectures: Spatial discretisation can rewire the graph, changing diffusion support for scalability, denoising, bottleneck avoidance, or numerical convenience.

4. Graph Neural Diffusion

GRAND derives graph neural architectures from a learnable diffusion process, with temporal and spatial discretisation choices determining the model structure. It includes linear, nonlinear, and rewired variants, with parameter sharing and numerical-solver-based training.

GRAND maps node features through learnable encoder, graph diffusion, and decoder functions to produce node embeddings.
Different GRAND architectures select the learnable diffusivity function together with spatial and temporal discretisations of the graph diffusion equation.
The diffusivity uses an attention function, with scaled dot-product attention empirically outperforming the Bahdanau attention used in GAT.
Time can use fixed or adaptive, explicit or implicit numerical schemes, while graph rewiring changes the edge set used by the diffusion equation.
GRAND-l fixes attention weights during integration, whereas GRAND-nl updates them at each numerical-integration step; GRAND-nl-rw additionally rewires edges.
Parameter sharing across layers or iterations makes GRAND more data-efficient than conventional GNNs and supports backpropagation or adjoint-based training.

5. Related work

Related work connects GNNs with diffusion processes, neural ODEs, and PDEs. GRAND differs from neural-PDE approaches that solve a given PDE by using diffusion as the basis for constructing GNNs.

Diffusion-based image-processing and graphics methods extended beyond flat Euclidean spaces to embedded manifolds and non-Euclidean geometries.
Earlier graph methods used diffusion operators, Laplacian filters, eigenvectors, and learned kernels for graph convolution or graph reweighting.
Neural ODE research introduced continuous-depth models and extensions involving augmentation, regularization, stochastic equations, and graph neural networks.
Neural-PDE approaches generally use neural networks to solve a specified PDE, whereas GRAND uses diffusion concepts to develop graph neural architectures.

6. Results

GRAND variants perform competitively across node-classification benchmarks and retain performance at greater depth. Experiments also examine numerical-scheme stability, graph sparsity, and learned diffusion behavior.

6.1. Node classification benchmarks: Experiments compare linear, nonlinear, and rewired GRAND variants across citation, coauthor, Amazon, and other graph datasets using repeated splits or initializations.
6.1. Node classification benchmarks: GRAND uses 70K parameters on ogb-arxiv, compared with 143K for GCN, 219K for SAGE, and 1.63M for GAT.
6.1. Node classification benchmarks: GRAND variants rank among the best methods on common node-classification benchmarks, placing first on all but one dataset and second on the remaining dataset.
6.1. Node classification benchmarks: On ogb-arxiv, GRAND is slightly below the best GAT result, although that GAT uses 20 times as many parameters.
6.2. Depth: As integration time increases, GRAND maintains performance while GCN and residual GCN degrade by 50% after four layers.
6.3. Choice of discretisation scheme: The explicit Adams method is unstable except at step size τ = 0.005, whereas the implicit Adams method remains stable for all tested step sizes.
6.3. Choice of discretisation scheme: Larger implicit-method step sizes can reduce solver steps but make the implicit equations harder to solve and require more iterations.
6.4. Graph rewiring: Graph sparsification makes all methods faster, while accuracy remains similar until sparsity impedes information flow at K < 8.

7. Conclusion

GRAND frames GNNs as discretisations of graph diffusion PDEs, using numerical schemes and stability analysis to support deep architectures and competitive benchmark performance. The framework remains constrained by a deliberately simplified diffusion equation and shared parameters across layers.

GRAND defines a new class of GNNs by discretising diffusion PDEs on graphs.
The framework includes stability analysis and supports principled choices for diffusion temporal and spatial operators.
Appropriate discretisation and numerical schemes enable very deep graph neural networks with superior performance on popular benchmarks.
The considered diffusion equation is intentionally chosen for mathematical tractability, limiting the model to functions of the form ∂x/∂t = f(x(t), t, θ).
The model requires a constant hidden-state size and the same parameter set θ across all layers.

A. Datasets

The paper evaluates attention-based diffusion on graph benchmarks and reports that scaled dot product attention outperforms the Bahdanau attention used in GAT across datasets.

GRAND can use any right stochastic attention matrix, with experiments using multiheaded Bahdanau attention previously applied to graphs.
Scaled dot product attention performed better than Bahdanau attention for all datasets.The paper suggests dropout behavior inside adaptive timestep numerical ODE solvers may explain this difference.

C. Full Training Objective

The training objective encodes node labels by decoding the terminal value of an evolutionary PDE, while the learned dynamics are integrated from an encoded input state. Linear GRAND admits an eigenspace-based stability analysis requiring non-positive real parts of the relevant eigenvalues.

Training optimises cross-entropy loss between one-hot node labels and predictions decoded from the terminal PDE state.
The PDE starts from an input encoder φ: R^din → R^d and produces predictions through a linear decoder ψ: R^d → R^dclass.
The nonlinear GRAND dynamics are represented by a learned system ∂x/∂t = f(x(t), t, θ).
For linear GRAND, the solution is expressed through a matrix exponential and analyzed after eigenvalue decomposition.
Convergence as t →∞ requires Re(λ̄_i) ≤0 for every eigenvalue, which holds because A is right stochastic and λ̄_i are eigenvalues of A−I.

E.1. Proof of theorem 1: Stability of explicit Euler

The explicit Euler scheme is stable only when the time step satisfies 0 < τ < 1, while the implicit scheme is stable for all choices of τ.

Stability requires the amplification factor ||Q(t)|| < 1, established by showing Q(t) is right stochastic with spectral radius at most 1.
The implicit scheme forms a Markov matrix Q whose spectral radius is bounded by unity, making it stable for every τ.

F. General Multistep Methods

General multistep methods combine explicit or implicit updates, with higher order reducing approximation error but increasing computational cost; RK4 provides the reported best trade-off.

Order: A method of order p has single-step error proportional to τ^(p+1) and interval-wide error proportional to τ^p.
Order: Function evaluations grow super-linearly with order after 4, motivating RK4 as a speed–accuracy trade-off.
Order: Forward Euler uses α_1 = −1 and β_1 = 1, with all other multistep coefficients equal to zero.
Runge-Kutta 4: RK4 or its adaptive-step variants outperform lower-order methods in all reported experiments.
Adaptive step size: Adaptive solvers compare interwoven order-p and order-(p−1) methods to adjust the step size with little extra computational cost.

G. Adaptive step size implementation details

GRAND’s adaptive ODE implementation uses relative and absolute tolerances to control solver steps, and larger-than-default tolerances improved accuracy while reducing runtime.

Adaptive step size implementation details: Adaptive ODE solvers use relative tolerance rtol and absolute tolerance atol to assess the new step size.
Adaptive step size implementation details: Using rtol and etol values 10–1000 times larger than defaults improved prediction accuracy and reduced runtime.
Adaptive step size implementation details: The adjoint method uses separate forward and backward tolerance scales because they control ODE-solution and gradient accuracy, respectively.