GRAND: Graph Neural Diffusion
Abstract
We present Graph Neural Diffusion (GRAND), a model 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, achieving competitive results on many standard graph benchmarks.
Introduction
Graph Neural Networks (GNNs) are intimately connected to differential equations.
The seminal work of [12] was concerned with ﬁnding the ﬁxed points of differential equations using the Almeida-Pineda algorithm [1, 11].
The currently predominant message passing paradigm [6] can be modelled as a differential equation.
More recently, diffusion processes have been shown to be an effective preprocessing step for graph learning [7].
PDEs are among the most studied mathematical constructions, and historically, PDE-based methods have been used extensively in signal and image processing [10], computer graphics [13] and machine learning [5, 4].
Our goal is to show that the tools of PDEs can be used to understand existing GNN architectures and as a principled way to develop a broad class of new methods.
We focus on architectures that can be interpreted as information diffusion on graphs, modelled by the diffusion equation.
In doing so, we show that many popular GNNs can be derived from a single mathematical framework by different choices of the form of diffusion equation and discretisation schemes.
Standard GNNs are equivalent to the explicit single-step Euler scheme that is inefﬁcient and requires small step sizes.
We show that more advanced, adaptive multi-step schemes such as Runge-Kutta perform signiﬁcantly better and using implicit schemes, which are unconditionally stable, amounts to larger multi-hop diffusion operators.
Choosing different spatial discretisation amounts to graph rewiring , a technique recently used to improve the performance of GNNs [7, 2].
We show that appropriate choices within our framework allow the design of deep GNN architectures with tens of layers.
This is a feat hard to achieve otherwise due to feature oversmoothing [8, 9] and bottlenecks [2] – phenomena that are recognised as a common plight of most graph learning architectures.
Background
Diffusion equation Let x ( t ) denote a family of scalar-valued functions on Ω × [0 , ∞ ) representing the distribution of some property (which we will assume to be temperature for simplicity) on a domain 35th Conference on Neural Information Processing Systems (NeurIPS 2021), Sydney, Australia.
Ω at some time, and let x ( u, t ) be its value at point u ∈ Ω at time t .
According to Fourier’s law of heat conduction, the heat ﬂux h
= − g ∇ x, is proportional to the temperature gradient ∇ x , where g is the diffusivity describing the thermal conductance properties of Ω .
An idealized homogeneous setting assumes that g is a constant scalar throughout Ω .
More generally, the diffusivity is a inhomogeneous (position-dependent) function that can be scalar-valued (in which case it simply scales the temperature gradient and is isotropic ) or matrix-valued (in which case the diffusion is said to be anisotropic , or direction-dependent).
The continuity condition x t = − div( h ) (roughly meaning that the only change in the temperature is due to the heat ﬂux, as measured by the divergence operator, i.e., heat is not created or destroyed), leads to a PDE referred to as the (heat) diffusion equation ,
[Equation]
with the initial condition x ( u, 0) = x 0 ( t ) ; for simplicity, we assume no boundary conditions.
The choice of the diffusivity function determines if the diffusion is homogeneous ( g = c ), inhomogeneous ( g ( u, t ) ), or anisotropic ( A ( u, t ) ).
In the isotropic case, the diffusion equation can be expressed as
[Equation]
where ∆ x = div( ∇ x ) is the Laplacian operator.
Diffusion on manifolds
In our discussion so far we assumed some abstract domain Ω .
The structure of the domain is manifested in the deﬁnition of the spatial differential operators in the diffusion PDE.
In a general setting, we model Ω as a Riemannian manifold, and let X (Ω) and X ( T Ω) denote the spaces of scalar and (tangent) vector ﬁelds on it, respectively.
We denote by (cid:104) x, y (cid:105) and (cid:104)(cid:104) X , Y (cid:105)(cid:105) the respective inner products on X (Ω) and X ( T Ω) .
Furthermore, we denote by ∇ : X (Ω) →
X ( T Ω) and div = ∇ ∗ : X
( T Ω) → X (Ω) the gradient and divergence operators, which are adjoint w.r.t.
the above inner products: (cid:104)(cid:104)∇ x, X (cid:105)(cid:105) =
(cid:104) x, div( X ) (cid:105) .
Informally, the gradient ∇ x of a scalar ﬁeld x is a vector ﬁeld providing at each point u ∈ Ω the direction ∇ x ( u ) of the steepest change of x .
The divergence div( X ) of a vector ﬁeld X is a scalar ﬁeld providing, at each point, the ﬂow of X through an inﬁnitesimal volume.
The Laplacian ∆ x can be interpreted as the local difference between the value of a scalar ﬁeld x at a point and its inﬁnitesimal neighbourhood.
Diffusion equations on graphs
We now deﬁne diffusion equations on graphs , analogous to Section 2 and argue that formalizing
GNNs under the diffusion equation framework provides a principled and rigorous way to develop new architectures for graph learning.
3.1 Graph diffusion equation
Let G =
( V , E ) be an undirected graph with |V| = n nodes and |E| =
e edges, and let x and X denote features deﬁned on nodes and edges respectively.
The node and edge ﬁelds can be represented as n - and e -dimensional vectors assuming some arbitrary ordering of nodes.
We adopt the same notation for the respective inner products:
[Equation]
Here, w ij denotes the adjacency of G : w ij = w ji = 1 iff ( i, j ) ∈ E .
We tacitly assume edge ﬁelds to be alternating , so X ji = − X ij , and no self-edges, so ( i, i ) / ∈ E .
The gradient ( ∇ x ) ij = x
j − x i assigns the edge ( i,
j ) ∈
E the difference of its endpoint features and is alternating by deﬁnition.
Similarly, the divergence (div( X )) i assigns the node i the sum of the features of all edges it shares:
[Equation]
The two operators are adjoint, (cid:104)(cid:104)∇ x , X (cid:105)(cid:105)
= (cid:104) x , div( X ) (cid:105) .
We consider the following diffusion equation on the graph
[Equation]
with an initial condition x (0) .
Here we denote by G = diag( a ( x i ( t ) , x j ( t ) , t ))
an e × e diagonal matrix and a is some function determining the similarity between nodes i and j .
While in general a ( x i , x j , t ) can be time-dependent, we will assume a = a ( x i , x j ) for the sake of simplicity.
Plugging in the expressions of ∇ and div , we get
[Equation]
where A ( x ) = ( a ( x i , x j )) is the n × n attention matrix with the same structure as the adjacency of the graph (we assume a ij =
if ( i, j ) / ∈ E ).
Note that in the setting when A ( x ( t ))
=
A we get a linear diffusion equation that can be solved analytically as x ( t ) =
e ¯ A t x (0) .
3.2 Properties of the graph diffusion equation
Differential equation stability is closely related to the concept of robustness in machine learning; changes in model outputs should be small under small changes in inputs.
Formally, a solution x ( t ) of the PDE is said to be stable, if given any (cid:15) >
there exists δ > 0 such that for any solution ˆ x ( t ) , such that | x (0) − ˆ x (0) | ≤ δ , it is also the case that | x ( t ) − ˆ x ( t ) | ≤ (cid:15) for all t ≥ 0 .
In the linear case, it is sufﬁcient to show that the eigenvalues of ¯ A are non-positive (see the Supplementary Materials for proof).
For the general nonlinear case, we show
[Equation]
which follows from (i) the function
¯ A ( x ) x being continuous in x , (ii) the largest component of x ( t ) not increasing in time, and (iii) the smallest component is not decreasing in time.
Condition (i) holds as ¯ A is a composition of Lipschitz-continuous functions (cf.
equation (6)).
Deﬁning indices k = arg max i x i and l = arg min
i x i we have
[Equation]
since A is right stochastic, which proves (ii) and (iii).
Furthermore, the derivative ∂∂x A ( x ) is
Lipschitz-continuous (from the deﬁnition of the attention function we use).
Taken together with continuity in time, the requirements of Picard-Lindelöf are satisﬁed and the PDE is also well posed.
Graph Neural Diffusion
We now describe Graph Neural Diffusion (GRAND), a new class of GNN architectures derived from the graph diffusion formalism.
We assume a given graph G =
( V , E ) with n nodes and d -dimensional node-wise features represented as a matrix X in .
GRAND architectures implement the learnable encoder/decoder functions φ , ψ and a learnable graph diffusion process, to produce node embeddings
Y = ψ ( X
[Equation]
T )) ,
[Equation]
∂
X
[Equation]
is given by the graph diffusion equation (2).
Different GRAND architectures amount to the choice of the learnable diffusivity function G and spatial/temporal discretisations of equation (2).
The diffusivity is modelled with an attention function a ( ., . ) .
Empirically, scaled dot product attention [14] outperforms the Bahdanau [3] attention used in GAT [15].
The scaled dot product attention is given by
[Equation]
where W K and W Q are learned matrices, and d k is a hyperparameter determining the dimension of W k .
We use multi-head attention which is useful to stabilise the learning [15, 14] by taking the expectation, A ( X ) = 1 h (cid:80)
h
A h ( X ) .
The attention weight matrix A =
( a ( X i , X j )) is rightstochastic, allowing equation (8) to be written as
[Equation]
A broad range of discretisations are possible.
Temporal discretisations amount to the choice of numerical scheme, which can use either ﬁxed or adaptive step sizes and be either explicit or implicit.
Time forms a continuous analogy to the layer index, where each layer corresponds to an iteration of the solver.
When using adaptive time step solvers, the number of layers is not speciﬁed a-priori.
Explicit schemes invoke residual structures that are usually more complex than those employed in standard resnets and which follow directly from rigorous numerical stability results.
Implicit numerical schemes offer a natural way of trading off depth and width (spatial support of the diffusion kernel).
Spatial discretisation amounts to modifying the given graph, or building one in settings where no graph is given and the data can be assumed to lie in some feature space or on a continuous manifold.
When the input graph is given, we can rewire the given graph and use a different edge set in the diffusion equation.
While in general equation (7) is nonlinear due to the dependence of A on X , it becomes linear if the attention weights are ﬁxed inside the integral, ¯ A ( X ( t ))
=
¯ A (note that A is still parametric and learnable, but does not change throughout the diffusion process).
In this case, equation (7) can be solved analytically as X ( t ) =
e ¯ A t X (0) .
As ¯ A is a form of normalised Laplacian, all eigenvalues are non-positive and the steady state solution is given by the dominating eigenvector, which is the degree vector.
However, as ¯ A is learned, this limitation is not severe as the system can be (and in practice is) degenerate; the graph becomes (approximately) disconnected, with connected components permitted to have unique steady state solutions.
We call this model GRAND-l for linear to distinguish it from the more general GRAND-nl for non-linear.
The ﬁnal variant is GRAND-nl-rw (non-linear with rewiring), where rewiring is performed via a two step process: as a preprocessing step, the graph is densiﬁed using diffusion weights as in [7], and then at runtime the subset of edges to use is learned based on attention weights.
With explicit Euler, Equation (2) becomes:
[Equation]
where E (cid:48)
= { ( i, j ) : ( i, j ) ∈ E and a ij > ρ } with some threshold value ρ , is the ‘rewired’ edge set, which may now contain self-loops.
While a changes throughout the diffusion process, rewiring is only performed at the start of the epoch based on features at t = 0 .
GRAND shares parameters across layer/iteration and is thus more data-efﬁcient than conventional
GNNs.
[Table] [Table]
parameters.