Beltrami Flow and Neural Diffusion on Graphs

1 Introduction . The majority of graph neural networks ( GNNs ) are based on the message passing paradigm [ 30 ] , wherein node features are learned by means of a non-linear propagation on the graph . Multiple recent works have pointed to the limitations of the message passing approach . These include ; limited expressive power [ 80 , 95 , 9 , 7 ] , the related oversmoothing problem [ 60 , 62 ] and the bottleneck phenomena [ 1 , 93 ] , which render such approaches inefficient , especially in deep GNNs . Multiple alternatives have been proposed , among which are higher-order methods [ 54 , 7 ] and decoupling the propagation and input graphs by modifying the topology , often referred to as graph rewiring . Topological modifications can take different forms such as graph sampling [ 32 ] , kNN [ 43 ] , using the complete graph [ 86 , 1 ] , latent graph learning [ 89 , 36 ] , or multi-hop filters [ 92 , 73 ] . However , there is no agreement in the literature on when and how to modify the graph , and a single principled framework for doing so . A somewhat underappreciated fact is that GNNs are intimately related to diffusion equations [ 16 ] , a connection that was exploited in the early work of Scarselli et al . [ 76 ] . Diffusion PDEs have been historically important in computer graphics [ 83 , 11 , 51 , 64 ] , computer vision [ 13 , 18 , 6 ] , and image processing [ 65 , 82 , 90 , 85 , 26 , 12 ] , where they created an entire trend of variational and PDE-based approaches . In machine learning and data science , diffusion equations underpin such popular manifold learning methods as eigenmaps [ 5 ] and diffusion maps [ 20 ] , as well as the family of PageRank algorithms [ 63 , 14 ] . In deep learning , differential equations are used as models of neural networks [ 16 , 19 , 25 , 94 , 71 , 98 ] and for physics-informed learning [ 72 , 22 , 75 , 21 , 81 , 47 ] . Main contributions In this paper , we propose a novel class of GNNs based on the discretised non-Euclidean diffusion PDE in joint positional and feature space , inspired by the Beltrami flow [ 82 ] used two decades ago in the image processing literature for edge-preserving image denoising . We show that the discretisation of the spatial component of the Beltrami flow offers a principled view on positional encoding and graph rewiring , whereas the discretisation of the temporal component can replace GNN layers with more flexible adaptive numerical schemes . Based on this model , we introduce Beltrami Neural Diffusion ( BLEND ) that generalises a broad range of GNN architectures and shows state-of-the-art performance on many popular benchmarks . In a broader perspective , our 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . approach explores new tools from PDEs and differential geometry that are less well known in the graph ML community . 2 Background . Beltrami flow Kimmel et al . [ 40 , 82 , 39 ] considered images as 2-manifolds ( parametric surfaces ) ( Σ , g ) embedded in some larger ambient space as z ( u ) = ( u , αx ( u ) ) ⊆ Rd+2 where α ≥ 0 is a scaling factor , u = ( u1 , u2 ) are the 2D positional coordinates of the pixels , and x are the ddimensional colour or feature coordinates ( with d = 1 or 3 for grayscale or RGB images , or d = k2 when using k × k patches as features [ 12 ] ) . In these works , the image is evolved along the gradient flow of a functional S [ z , g ] called the Polyakov action [ 68 ] , which roughly measures the smoothness of the embedding1 . For images embedded in Euclidean space with the functional S minimised with respect to both the embedding z and the metric g , one obtains the following PDE : ∂z ( u , t ) ∂t = ∆Gz ( u , t ) , z ( u , 0 ) = z ( u ) , t ≥ 0 , ( 1 ) and boundary conditions as appropriate . Here ∆G is the Laplace-Beltrami operator , the Laplacian operator induced on Σ by the Euclidean space we embed the image into . Namely , the embedding of the manifold allows us to pull-back the Euclidean distance structure on the image : the distance between two nearby points u and u+ du is given by d ` 2 = du > G ( u ) du = du21 + du 2 2 + α 2 d∑ i=1 dx2i , ( 2 ) where G = I + α2 ( ∇ux ( u ) ) > ∇ux ( u ) is a 2 × 2 matrix called the Riemannian metric . The fact that the distance is a combination of the positional component ( distance between pixels in the plane , ‖u− u′‖ ) and the colour component ( distance between the colours of the pixels , ‖x ( u ) − x ( u′ ) ‖ ) is crucial as it allows edge-preserving image diffusion . When dealing with images , the evolution of the first two components of ( z1 , z2 ) = u is a nuisance amounting to the reparametrisation of the manifold and can be ignored . For grayscale images ( the case when d = 1 and z = ( u1 , u2 , x ) ) , this is done by projection along the dimension z3 , in which case the Beltrami flow takes the form of an inhomogeneous diffusion equation of x , ∂x ( u , t ) ∂t = 1√ detG ( u , t ) div ( ∇x ( u , t ) √ detG ( u , t ) ) t ≥ 0 . ( 3 ) The diffusivity a = 1√ detG = 1√ 1 + α2‖∇x‖2 ( 4 ) determining the speed of diffusion at each point , can be interpreted as an edge indicator : diffusion is weak across edges where ‖∇x‖ 1 . The result is an adaptive diffusion [ 65 ] popular in image processing due to its ability to denoise images while preserving their edges . For cases with d > 1 ( multiple colour channels ) , equation ( 3 ) is applied to each channel separately ; however , the metric G couples the channels , which results in their gradients becoming aligned [ 38 ] . Special cases In the limit case α = 0 , equation ( 3 ) becomes the simple homogeneous isotropic diffusion ∂ ∂tx = div ( ∇x ) = ∆x , where ∆ = ∂2 ∂u21 + ∂ 2 ∂u22 is the standard Euclidean Laplacian operator . The solution is 1Explicitly , by minimising the functional with respect to the embedding one finds Euler-Lagrange ( EL ) equations that can be used to dictate the evolution process of the embedding itself . 2 given in closed form as the convolution of the initial image and a Gaussian kernel with time-dependent variance , x ( u , t ) = x ( u , 0 ) ? 1 ( 4πt ) d/2 e−‖u‖ 2/4t ( 5 ) and can be considered a simple linear low-pass filtering . In the limit t → ∞ , the image becomes constant and equal to the average colour.2 Another interpretation of the Beltrami flow is passing a Gaussian on the manifold ( see Figure 1 , bottom ) , which can locally be expressed as non-linear filtering with the bilateral kernel [ 85 ] dependent on the joint positional and colour distance ( Figure 1 , top ) , x ( u , t ) = 1 ( 4πt ) d/2 ∫ R2 x ( v , 0 ) e−‖u−v‖ 2/4te−α 2‖x ( u,0 ) −x ( v,0 ) ‖2/4tdv . ( 6 ) For α = 0 , the bilateral filter ( 6 ) reduces to a simple convolution with a time-dependent Gaussian . 3 Discrete Beltrami flow on graphs . We now develop the analogy of Beltrami flow for graphs . We consider a graph to be a discretisation of a continuous structure ( manifold ) , and show that the evolution of the feature coordinates in time amounts to message passing layers in GNNs , whereas the evolution of the positional coordinates amounts to graph rewiring , which is used in some GNN architectures . 3.1 Graph Beltrami flow . Let G = ( V = { 1 , . . . , n } , E ) be an undirected graph , where V and E denote node and edge sets , respectively . We further assume node-wise d-dimensional features xi ∈ Rd for i = 1 , . . . , n. Denote by zi = ( ui , αxi ) the embedding of the graph in a joint space C × Rd , where C is a d′-dimensional space with a metric dC representing the node coordinates ( for simplicity , we will assume C = Rd ′ unless otherwise stated ) . We refer to ui xi and zi as the positional , feature and joint coordinates of node i , respectively , and arrange them into the matrices U , X , and Z , of sizes n × d′ , n × d , and n× ( d′ + d ) . For images , Beltrami flow amounts to evolving the embedding z along div ( a ( z ) ∇z ) , with a a diffusivity map.3 Thus , we consider the graph Beltrami flow to be the discrete diffusion equation ∂zi ( t ) ∂t = ∑ j : ( i , j ) ∈E′ a ( zi ( t ) , zj ( t ) ) ( zj ( t ) − zi ( t ) ) zi ( 0 ) = zi ; i = 1 , . . . , n ; t ≥ 0 . ( 7 ) We motivate our definition as follows : gij = zj − zi and di = ∑ j : ( i , j ) ∈E gij are the discrete analogies of the gradient ∇z and divergence div ( g ) , both with respect to a graph ( V , E ′ ) that can be interpreted as the numerical stencil for the discretisation of the continuous Laplace-Beltrami operator in ( 3 ) . Note that E ′ can be different from the input E ( referred to as ‘ rewiring ’ ) . As discussed in Section 3.3 , most GNNs use E ′ = E ( input graph is used for diffusion , no rewiring ) . Alternatively , the positional coordinates of the nodes can be used to define a new graph topology either with E ( U ) = { ( i , j ) : dC ( ui , uj ) < r } , for some radius r > 0 , or using k nearest neighbours . This new rewiring is precomputed using the input positional coordinates ( i.e. , E ′ = E ( U ( 0 ) ) ) or updated throughout the diffusion ( i.e. , E ′ ( t ) = E ( U ( t ) ) ) . Therefore , ( 7 ) can be compactly rewritten as ∂zi ( t ) ∂t = div ( a ( z ( t ) ) ∇zi ( t ) ) . The function a is the diffusivity controlling the diffusion strength between nodes i and j and is assumed to be normalised : ∑ j : ( i , j ) ∈E′ a ( zi , zj ) = 1 . The dependence of the diffusivity on the 2Assuming appropriate boundary conditions . 3Note that in equation ( 3 ) the diffusivity function a coincides with ( det ( G ( u , t ) ) ) − 1 2 . Similarly to [ 82 , Section 4.2 ] we have neglected the extra term 1/a appearing in ( 3 ) . 3 embedding z matches the smooth PDE analysed in e.g . [ 82 , Section 4.2 ] and is consistent with the form of attention mechanism used in e.g . [ 88 , 86 ] . In matrix-form , we can also rewrite ( 7 ) as ( ∂ ∂tU ( t ) , ∂ ∂tX ( t ) ) = ( A ( U ( t ) , X ( t ) ) − I ) ( U ( t ) , X ( t ) ) ( 8 ) U ( 0 ) = U ; X ( 0 ) = αX ; t ≥ 0 , where we emphasise the evolution of both the positional and feature components , coupled through the matrix-valued function A aij ( t ) = { a ( ( ui ( t ) , xi ( t ) ) , ( uj ( t ) , xj ( t ) ) ) ( i , j ) ∈ E ( U ( t ) ) 0 else . representing the diffusivity . The graph Beltrami flow produces an evolution of the joint positional and feature coordinates , Z ( t ) = ( U ( t ) , X ( t ) ) . In Section 3.3 we will show how the evolution of the feature coordinates X ( t ) results in feature diffusion or message passing on the graph , the core of GNNs . As noted in Section 2 , in the smooth case the Beltrami flow is obtained as gradient flow of an energy functional when minimised with respect to both the embedding and the metric on the surface ( an image ) . When the embedding takes values in the Euclidean space , this leads to equations of the form ( 3 ) with no channel-mixing and an exact form of the diffusivity determined by the pull-back G of the Euclidean metric . To further motivate our approach , it is tempting to investigate whether a similar conclusion can be attained here . Although in the discrete case the operation of pull-back is not well-defined , we are able to derive that the gradient flow of a modified graph Dirichlet energy gives rise to an equation of the form ( 7 ) . We note though that the gradient flow does not recover the exact form of the diffusivity implemented in this paper . This is not a limitation of the theory and should be expected : by requiring the gradient flow to avoid channel-mixing and imitate the image analogy in [ 82 ] and by inducing a discrete pull-back condition , we are imposing constraints on the problem . We leave the theoretical implications for future work and refer to the Supplementary Materials for a more thorough discussion , including definitions and proofs . Theorem 1 . Under structural assumptions on the diffusivity , graph Beltrami flow ( 7 ) is the gradient flow of the discrete Polyakov functional .

Summary: The paper proposes a new graph neural network based on a continuous time flow. Each node is given a pre-processed position from an off-line node embedding method, then both the position and a vector of node features (together called Z) are changed during the flow. The flow changes Z by the average difference of Z of adjacent nodes, weighted by a scaled dot product attention. Additionally, a node-wise initial encoder and final decoder is used. The flow is computed via explicit adaptive or fixed integration and the backward pass via the adjoint method. The authors propose various extensions, including having dynamic adjacency via kNN and using hyperbolic positional space. The authors claim their method is motivated as a discrete analogue of a Beltrami flow.