GRAND: Graph Neural Diffusion

In plain words

A graph neural network (GNN) updates each node by mixing in information from its neighbours, and stacks such updates in layers. GRAND looks at that stacking as spreading: information flows from each node to its neighbours, the way heat spreads through a metal plate. Each layer becomes a small step of time in a continuous spreading process.

The spreading is controlled by weights on the edges. Those weights come from attention, so the model learns how much each neighbour should contribute. Choosing how the time steps are computed changes the model. Simple steps behave like ordinary GNN layers, while more careful solvers keep the spreading stable even with large steps.

The paper presents this as a principled way to design deep graph models. The same framework also covers rewiring the graph, meaning the spreading can use edges other than the ones in the input graph.

TL;DR

• GRAND models node features as a diffusion process on a graph and treats GNN layers as discretisations of it.
• Attention weights define the diffusion operator. Linear (frozen attention) and nonlinear (recomputed attention) versions are given.
• Explicit Euler is stable for step sizes below 1; implicit schemes are stated to be stable for any step size.
• Author-reported accuracy is competitive with GCN, GAT and ODE-based GNNs on Cora, Citeseer, PubMed and several larger benchmarks tab. 1 tab. 2.
• The released code differs from the paper in several places (learned scaling, self-loops, attention normalisation, solver settings). See "Paper vs code".

Problem

Standard GNN layers behave like the simplest explicit time step of a diffusion equation, which needs small steps to stay accurate §1. Deep GNNs suffer from oversmoothing, where node features converge to uninformative values, and from bottlenecks, where a fixed-size node state must summarise an exponentially growing neighbourhood §1. The paper aims to address both with a framework in which depth, step size and the graph structure are separate design choices §1 §4.

Method

Setting. Node features are written as a vector per node. Gradient and divergence operators on the graph are defined so that they are adjoint to each other, which gives a graph version of the diffusion equation §3.1. The continuous diffusion equation is

\frac{\partial \mathbf{x}(t)}{\partial t} = \mathrm{div}\left[\mathbf{G}(\mathbf{x}(t),t)\,\nabla \mathbf{x}(t)\right] eq. 1

with edge-wise diffusivity G built from a similarity function a(x_i, x_j) §3.1. When the attention matrix is written as A(x), the same equation reads

\frac{\partial \mathbf{x}}{\partial t} = (\mathbf{A}(\mathbf{x}) - \mathbf{I})\,\mathbf{x} = \bar{\mathbf{A}}(\mathbf{x})\,\mathbf{x} eq. 2

where A has the adjacency structure of the graph and is right-stochastic (each row sums to 1) §3.1 §4.

Method steps (GRAND).

1. Encode. A learned linear map turns the input features into the initial state, X(0) = φ(X_in) §4 app. C.
2. Diffuse. Integrate the ODE in §4 from time 0 to time T. The state at T is the node embedding §4.
3. Attention. Each edge weight is a scaled dot-product attention score, normalised with a softmax over neighbours, and averaged over h heads eq. 10 §4. A GAT-style (Bahdanau) attention was tried and performed worse on all datasets app. B.
4. Decode. A learned linear map sends X(T) to class scores, and training uses cross-entropy over labelled nodes §4 app. C.
5. Variants. GRAND-l keeps the attention fixed during integration, which gives a linear ODE. GRAND-nl recomputes attention at every solver step. GRAND-nl-rw also rewires the graph (see step 7) §4 §6.1.
6. Solve. Explicit Euler (eq. 6) and implicit Euler (eq. 8) are the basic schemes. Multi-step Adams methods (explicit Adams–Bashforth and implicit Adams–Moulton with a predictor–corrector) are used with fourth order (s = 4) §3.3. The paper's experiments use Dormand–Prince (dopri5) as the main adaptive solver and a Runge–Kutta method on ogbn-arxiv §6.1 app. G.
7. Rewire. Attention weights are thresholded, keeping edges with weight above a threshold ρ, and the surviving edge set is fixed during the integration. Densification uses diffusion weights from Klicpera et al. §4 §6.1.

Explicit forward Euler is

\mathbf{x}^{(k+1)} = \left(\mathbf{I} + \tau\,\bar{\mathbf{A}}(\mathbf{x}^{(k)})\right)\mathbf{x}^{(k)} eq. 7

with q_ii = 1 − τ Σ_l a_il and q_ij = τ a_ij on edges §3.3. Implicit (semi-implicit) Euler is

\left(\mathbf{I} - \tau\,\bar{\mathbf{A}}(\mathbf{x}^{(k)})\right)\mathbf{x}^{(k+1)} = \mathbf{x}^{(k)} eq. 8

which needs a linear solve per step §3.3.

Attention is

a(\vec{X}_i,\vec{X}_j) = \operatorname{softmax}\left(\frac{(\mathbf{W}_K \vec{X}_i)^\top \mathbf{W}_Q \vec{X}_j}{d_k}\right) eq. 10

with learned W_K and W_Q, and d_k a hyperparameter §4.

Pseudocode. The following follows the paper's equations. The released code differs in several details (see "Paper vs code").

# GRAND forward pass, nonlinear variant  [§4] [app. C]
# Input:  features (N, d_in); edges (E, 2) with (i, j) pairs; time T
# Output: class scores (N, C)

def grand_forward(features, edges, T, tau):
    # Encode each node with a learned linear map  [§4] [app. C]
    state = encoder(features)                         # (N, d)
    # Number of explicit steps covering [0, T]  [§3.3]
    num_steps = round(T / tau)
    for step in range(num_steps):
        # Recompute attention from the current state  [eq. 10] [§4]
        attention = attention_weights(state, edges)   # (E,)
        # Explicit Euler update, stable for 0 < tau < 1  [eq. 7] [thm. 1]
        state = state + tau * (spread(attention, state) - state)
    # Map the terminal state to class scores  [§4] [app. C]
    return decoder(state)                             # (N, C)

# Edge attention, averaged over heads  [eq. 10] [§3.2] [§4]
# Input:  state (N, d); edges (E, 2) as (i, j) pairs; d_k head size
# Output: weight per edge (E,), summing to 1 over the neighbours of i

def attention_weights(state, edges):
    # Key and query projections of every node  [eq. 10]
    keys = state @ W_K                                # (N, d_k) per head
    queries = state @ W_Q                             # (N, d_k) per head
    # Scaled dot product on each edge  [eq. 10]
    scores = []
    for (i, j) in edges:
        scores.append(dot(keys[i], queries[j]) / d_k)
    # Softmax over the neighbours j of each node i  [eq. 10] [§3.2]
    weights = row_softmax(scores, edges)
    # Average over heads to one weight per edge  [§4]
    return mean_over_heads(weights)

The paper's Dormand–Prince and Runge–Kutta schemes are standard adaptive and fixed-step methods; the paper gives the RK4 update in app. F.3.

Training. Cross-entropy on training nodes app. C. Gradients are obtained by backpropagating through the solver, or by the adjoint method (Pontryagin's maximum principle) when memory is limited §4. The paper says the adjoint method is used for all datasets §6.1 Complexity.

Stability. Under right-stochastic A and continuity of the attention function, the maximum and minimum of the state are non-increasing and non-decreasing respectively:

\max_i x_i(0) \ge x_i(t) \ge \min_i x_i(0) \quad \forall t \ge 0 eq. 3

This is the basis for the stability claims §3.2. Theorem 1 says explicit Euler is stable for 0 < τ < 1, and Theorem 2 says implicit Euler is unconditionally stable for τ > 0 §3.3 [thm. 1] [thm. 2]. The proofs are in app. D and app. E.

Hyperparameters

Items marked "not stated" do not appear in the paper text, tables or appendices. Code defaults and search ranges are listed separately in Notes, as they come from the repository, not the paper.

Results

Author-reported. Setup: 20 random initialisations with Planetoid splits tab. 1 or 100 random splits tab. 2. Baselines are the paper's own runs, with results from Shchur et al. used where available §6.1. Hardware and optimiser details: see Hyperparameters.

Planetoid splits (test accuracy in %, mean ± std over 20 initialisations) tab. 1

• GDE and GODE are the best of their model variants (three and six models). † Results from running the authors' code with their stated hyperparameters and PyTorch Geometric data readers tab. 1.

Random splits (test accuracy in %, mean ± std over 20 initialisations and 100 random splits) tab. 2

* Uses node labels. † Uses 1.5M parameters.

The paper's text says the GRAND variants take first place on all but one dataset, where they are second. On ogbn-arxiv, GAT is slightly better and, per the text, uses about 20 times as many parameters §6.1 Performance. Note the large standard deviations for some baselines in tab. 2.

Depth §6.2 fig. 2. With fixed RK4 steps of 1.0 and T varied, GRAND keeps its accuracy as depth grows, while GCN and GCN with residual connections are said to lose about 50% after 4 layers. This is a figure-based claim fig. 2.

Solver choice on Cora §6.3 fig. 3, left. Explicit Adams–Bashforth is unstable for every step size except τ = 0.005. Implicit Adams–Moulton is stable at all step sizes tried and, for large steps, converges faster than dopri5. The paper notes that the implicit solver's equations become harder to solve as the step grows §6.3.

Rewiring on Cora §6.3 fig. 3, right. Keeping the largest K coefficients per node trades sparsity for speed. Accuracy holds until K is below 8 §6.3.

MNIST superpixels §6.4 fig. 4. A qualitative illustration only: nonlinear GRAND weights diffusion within class boundaries and preserves the image, while Laplacian diffusion blurs it. No numbers are reported.

Paper vs code

Code: official repository twitter-research/graph-neural-pde, linked from the paper's experiments section §6. License: Apache-2.0 (LICENSE file in the repository). Commit: af9047e0860e0ce6dc6333adcb4366769c795014. The repository also contains the later BLEND code, and the defaults in src/best_params.py come from later tuning README.md. Paths below are relative to the repository root.

• Attention scale. Paper divides scores by d_k eq. 10. Code divides by √d_k function_transformer_attention.py.
• Attention direction. Paper puts W_K on the source node and W_Q on the target eq. 10. Code puts the query on the source and the key on the target function_transformer_attention.py. This is a naming swap, and it only matters if weights are copied between implementations.
• Learned scale on the diffusion term. Code multiplies the diffusion term by α = sigmoid(α_train), with α_train initialised to 0, so α starts at 0.5 function_laplacian_diffusion.py base_classes.py. The paper has no such factor eq. 2 eq. 11. It can be disabled with --no_alpha_sigmoid run_GNN.py.
• Source term. Code can add β·x(0) to the dynamics (add_source) function_transformer_attention.py. The paper has no such term. The shipped defaults enable it for Cora, Citeseer and PubMed best_params.py.
• Softmax replacement. square_plus replaces the softmax with a squareplus normalisation function_transformer_attention.py utils.py. The paper uses the softmax eq. 10. The shipped defaults enable it for Cora, Citeseer, PubMed and CoauthorCS best_params.py.
• Self-loops. The paper assumes no self-edges in the input graph §3.1. Code adds self-loops with weight self_loop_weight (default 1.0) to the attention graph run_GNN.py block_transformer_attention.py.
• Attention normalisation. Paper's stability argument needs right-stochastic A §3.2. The code default normalises over the source index (rows) run_GNN.py. The shipped defaults for Cora, Citeseer and CoauthorCS set attention_norm_idx to 1, which normalises over the target index instead best_params.py best_params.py.
• Frozen versus recomputed attention. All shipped defaults set function: laplacian best_params.py. With the attention block, attention is computed once per forward pass and held fixed during integration, which is the linear GRAND-l behaviour block_transformer_attention.py function_laplacian_diffusion.py. The nonlinear GRAND-nl behaviour needs function: transformer, which recomputes attention at every function evaluation function_transformer_attention.py. The rewiring block also computes attention once per forward pass block_transformer_rewiring.py.
• Decoder nonlinearity. Code applies ReLU before the linear decoder GNN.py GNN.py. Appendix C writes the decoder as affine app. C.
• Dropout placement. Code applies dropout to the input and before the decoder, outside the ODE GNN.py GNN.py. The paper does not state dropout values. Appendix B says dropout does poorly inside adaptive solvers.
• Adjoint. Paper uses the adjoint method for all datasets §6.1 Complexity. Code default is off run_GNN.py. The shipped defaults turn it on for PubMed, CoauthorCS, Computers, Photo and ogbn-arxiv, and leave it off for Cora and Citeseer best_params.py.
• Solver on ogbn-arxiv. Paper uses a Runge–Kutta method §6.1. The shipped default uses method: dopri5 with adjoint_method: rk4 best_params.py. The example command in ray_tune.py uses rk4 for both ray_tune.py.
• Tolerances. Appendix G gives atol = ts × 10^-12 and then, in the same sentence, atol = ts × 10^-6; the second is presumably rtol app. G. Code sets atol = tol_scale × 1e-7 and rtol = tol_scale × 1e-9 base_classes.py.
• Early stopping. Code runs validation-based early stopping inside the solver. It evaluates validation accuracy at intermediate solver steps, up to earlystopxT × T (default 3), and records the test accuracy at the best validation point early_stop_solver.py run_GNN.py run_GNN.py. The paper does not describe this. --no_early turns it off.
• Augmentation. Paper uses ANODE augmentation for Cora and Citeseer §6.1. Shipped defaults set augment to False for both best_params.py.
• Regularisation. Paper uses kinetic energy and Jacobian regularisation for larger datasets §6.1. The shipped defaults set kinetic_energy, jacobian_norm2 and regularise off for all datasets best_params.py. The kinetic-energy term is 0.5 times the mean squared dynamics regularized_ODE_function.py.
• Search protocol. Paper uses 1000 trials and a grace period of 10 epochs §6.1. Code defaults are 20 samples and grace period 5 ray_tune.py. The shipped defaults use grace periods of 20 or 25 best_params.py.
• Split protocol. Paper uses 100 random splits and 20 initialisations §6.1. Code defaults are num_splits = 1 run_GNN.py and num_init = 1 ray_tune.py. The shipped defaults use 1 to 8 splits and 1 or 2 initialisations best_params.py.
• Rewiring method. Paper densifies with diffusion weights from Klicpera et al., then keeps edges with attention above a fixed ρ §4 §6.1. The rewiring block densifies with random new edges by default and thresholds at a per-epoch quantile rather than a fixed ρ block_transformer_rewiring.py block_transformer_rewiring.py. The GDC preprocessing is separate, in apply_gdc graph_rewiring.py.
• Reproduction entry point. The README's example is cd src; python run_GNN.py --dataset Cora, which loads the shipped defaults README.md. I found no per-table configuration file for the paper's exact numbers. The paper's own hyperparameters are not stated, so the card cannot say which configuration produced tab. 1 or tab. 2.
• Adams experiments. The implicit Adams–Moulton run is in src/run_explicit_implicit_exp.py, which sets implicit_adams. The explicit Adams option is commented out run_explicit_implicit_exp.py.

Pitfalls

• Explicit step size. Theorem 1 gives 0 < τ < 1 as a sufficient condition for explicit Euler stability. The paper does not state what happens above it [thm. 1]. Its Cora experiment finds the explicit Adams method unstable for every step except τ = 0.005 §6.3.
• Implicit cost. Each implicit step solves a linear system. The paper replaces exact inversion with a few solver iterations but gives no count §3.3. Larger steps need more iterations §6.3. The "implicit is faster" result is one experiment on Cora with the paper's settings §6.3.
• Stability needs right-stochastic attention. The maximum–minimum bound depends on the rows of A summing to 1 §3.2 eq. 3. Column-normalised attention, which the shipped Cora defaults use, does not satisfy this argument (see "Paper vs code").
• Stability covers a fixed graph. Theorems 1 and 2 are for a fixed right-stochastic operator. The paper gives no stability result for the rewired variant §3.2 §4.
• Attention is recomputed at every evaluation in GRAND-nl. Attention over the edge set needs O(|E'| d) memory, and runtime scales with the number of function evaluations §6.1 Complexity. Stiff dynamics can make the solver hang. The code raises an exception after max_nfe evaluations function_laplacian_diffusion.py, and its help text warns that stiff ODEs hang without a cap run_GNN.py. The paper does not mention a cap.
• Tolerance scale is a tuned knob. The paper says loosening the tolerances by 10× to 1000× improves both accuracy and runtime app. G. The tolerance scale is a searched hyperparameter in code ray_tune.py. Its value is not stated for any reported result.
• Model selection matters. Reported accuracy depends on which epoch or solver step is chosen by validation. The paper describes only "highest validation accuracy" §6.1. The code's in-solver early stopping can select a time point beyond T, which changes the numbers (see "Paper vs code").
• Unstated training details. Optimiser, learning rate, weight decay, epochs, dropout, T, number of heads, d_k, hidden size, ρ, K, attention initialisation value and the predictor–corrector threshold are all "not stated" in the paper. The code's search ranges are the only guide ray_tune.py.
• Variance. Some baselines in tab. 2 have large standard deviations (for example GAT on Computer, 78.0 ± 19.0). Several gaps between GRAND and the baselines are within one standard deviation. The paper reports no significance tests.
• Parameter comparisons. The text gives 70K parameters for GRAND and 1.63M for GAT on ogbn-arxiv §6.1 Number of parameters. The table footnote gives 1.5M for GAT tab. 2. The two GAT figures disagree.
• Appendix typo. The tolerance formula in app. G lists atol twice, with 10^-12 and 10^-6. Read the second as rtol, and check against the code (see "Paper vs code").
• Equation and figure numbers. Numbers in this card were counted from the LaTeX source, which was not compiled. Check them against the PDF before citing them externally.
• No rerun. The reported accuracies were not reproduced in this compilation.

Lineage

• Builds on: Neural ODEs (Chen et al., 2018), for the continuous-depth view and the adjoint method §1 §3.4 §6.1. Graph diffusion for densification (Klicpera et al., 2019) §4. ANODE augmentation (Dupont et al., 2019) §6.1.
• Contrasts with: GAT (Veličković et al., 2018), whose update is the explicit single-step form the paper generalises. GRAND uses scaled dot-product attention instead §3.4 §4 app. B.
• Compared against: GCN, GAT, MoNet, GraphSAGE, CGNN, GDE, GODE and LanczosNet §6.1.
• The same repository also holds the later BLEND code, which extends this framework with positional encodings (README).