Adaptive Diffusion in Graph Neural Networks

1 Introduction . Graph neural networks ( GNNs ) are a type of neural networks that can be directly coupled with graph-structured data [ 30 , 41 ] . Specifically , graph convolution networks [ 12 , 19 ] ( GCNs ) generalize the convolution operation to local graph structures , offering attractive performance for various graph mining tasks [ 15 , 32 , 37 ] . The graph convolution operation is designed to aggregate information from immediate neighboring nodes into the central node , which is also referred to as message passing [ 14 ] . To propagate information between nodes that are further away , multiple neural layers can be stacked * Now at Meta AI and work done when at Microsoft . †Jie Tang is the corresponding author . * Code is available at https : //github.com/abcbdf/ADC 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to go beyond the immediate hop of neighbors . To directly collect high-order information , spectral based GNNs leverage graph spectral properties to collect signals from global neighbors [ 6 , 12 , 17 ] . Though generating promising results , both strategies are limited to a pre-determined and fixed neighborhood for passing and receiving messages . Essentially , these methods have an implicit assumption that all graph datasets share the same size of receptive field during the message passing process . To break this , graph diffusion convolution ( GDC ) [ 21 ] was recently proposed to extend the discrete message passing process in GCN to a diffusion process , enabling it to aggregate information from a larger neighborhood . For each input graph , GDC hand-tunes the best neighborhood size for feature aggregation by grid-searching the parameters on the validation set , making its practical application limited and sensitive . To eliminate the manual search process of the optimal propagation neighborhood in GDC , we propose the adaptive diffusion convolution ( ADC ) strategy that supports learning the optimal neighborhood from the data automatically . ADC achieves this by formalizing the task as a bilevel optimization problem [ 11 ] , enabling the customized learning of one optimal propagation neighborhood size for each dataset . In other words , all GNN layers and feature channels ( dimensions ) share the same neighborhood size during message passing on each graph . To further this direction , we also enable ADC to automatically learn a customized neighborhood size for each GNN layer and each feature channel from data . By learning a unique propagation neighborhood for each layer , ADC can empower GNNs to capture neighbors ’ information from diverse graph structures , which is fully dependent on the data and downstream learning objective . Similarly , by learning distinct neighborhood size for each feature channel , GNNs are then capable of selectively modeling each neighbor ’ s multiple feature signals . Altogether , ADC makes GNNs fully coupled with the graph structures and all feature channels . By design , ADC is a general plugin that can be directly applied to existing GNN models . By plugging it on several GNNs , we show that the upgraded GNNs can offer significant performance advances over their vanilla versions across a wide range of datasets . Furthermore , experimental results also show that by learning the propagation neighborhood size automatically , ADC can consistently outperform GDC , which customizes this for each dataset by grid search . Finally , we demonstrate that GNNs ’ model capacity can benefit from the better coupling between the its architecture , graph structures , and feature channels , that is , by learning dedicated neighborhood size for each GNN layer and feature channel . 2 Neighborhood Radius in GNNs . We focus on the problem of semi-supervised node classification . The input includes an undirected network G = ( V , E ) , where the node set V contains of n nodes { v1 , ... , vn } and E is the edge set , and A ∈ Rn×n is the symmetric adjacency matrix of graph G. Given the input feature matrix X and a subset of node label Y , the task is to predict the labels of remaining nodes . 2.1 Neighborhood Radius in Message Passing Networks . The convolution operation on graphs can be described as the process of neighborhood feature aggregation or message passing [ 14 ] . The message passing graph convolutional networks can be simply defined as : H ( l ) = γ ( l ) ( ϕ ( l ) ( H ( l−1 ) ) , G ) ( 1 ) where H ( l ) denotes the hidden feature of layer l with H ( 0 ) = X and X as the input feature , ϕ ( · ) denotes feature transformation and γ ( · ) denotes feature propagation . Take GCN [ 19 ] for example . The feature transformation and feature propagation functions correspond to ϕ ( H ) = HW , γ ( Ĥ , G ) = D̃− 1 2 ÃD̃− 1 2 Ĥ , respectively , in which D is the diagonal degree matrix with D̃ii = ∑ j Ãij , and Ĥ denotes hidden feature after transformation . Note that GCN uses the adjacency matrix A with self loop , so it actually uses Ã = I+A . To simplify the notations , we use T to denote D̃− 1 2 ÃD̃− 1 2 . Straightforwardly , the feature transformation function ϕ ( · ) describes how features transform inside each node and the feature propagation function γ ( · ) describes how features propagate between nodes . Essentially , how good a GNN model can utilize graph structures heavily depends on the design of the feature propagation function . Neighborhood radius r. Most graph-based models can be represented as γ ( l ) ( Ĥ , G ) = f ( T ) Ĥ , where f ( T ) is a matrix that can be generated by T. So f ( T ) can be represented as f ( T ) =∑∞ k=0 θkT k. To quantify how far each node could aggregate features from , we define the neighbor- hood radius of a node as r : r = ∑∞ k=0 θkk∑∞ k=0 θk ( 2 ) Here , θk denotes the influence from k-step-away nodes . For a large r , this means the model puts more emphasis on long distance nodes , i.e. , global information . For a small r , this means the model amplifies local information . Neighborhood radius r in GCN . For GCN , the neighborhood radius r = 1 , which is just range of nodes directly connected to it . To collect information beyond direct connections , it is required to stack multiple GCN layers to reach high-order neighborhoods . Neighborhood radius r in multi-hop models . There are attempts to improve GCN ’ s feature propagation function from first-hop neighborhood to multi-hop neighborhood , such as MixHop [ 2 ] , JKNet [ 38 ] , and SGC [ 35 ] . For example , SGC [ 35 ] uses feature propagation function γ ( Ĥ , G ) = TKĤ , where T = D̃− 1 2 ÃD̃− 1 2 . In other words , the neighborhood radius r = K for SGC , which is the range of neighborhoods to collect information from each GNN layer . However , for all multi-hop models , the discrete nature of hop numbers makes r non-differentiable . 2.2 Neighborhood Radius in Graph Diffusion Convolution . Recently , a line of work has been focused on generalizing feature propagation from discrete hops to continuous graph diffusion [ 21 , 36 , 40 ] . Notably , graph diffusion convolution ( GDC ) addresses this by the following propagation setup [ 21 ] : γ ( l ) ( Ĥ , G ) = ∞∑ k=0 θkT kĤ ( 3 ) where k is summed from 0 to infinity , making each node aggregate information from the whole graph . In Eq.3 , the weight coefficients should satisfy ∑∞ k=0 θk = 1 such that the signal strength is not amplified nor reduced through the propagation . The two commonly-used sets of weight coefficients [ 21 , 36 , 40 ] are generated from personalized PageRank ( θk = α ( 1− α ) k ) [ 27 ] and the heat kernel ( θk = e−t t k k ! ) [ 22 ] , respectively . In this work , we focus on heat kernel . Heat kernel . Heat kernel incorporates prior knowledge into the GNN model , which means the feature propagation between nodes follows Newton ’ s law of cooling [ 34 ] , i.e. , the feature propagation speed between two nodes is proportional to the difference between their features . Formally , this prior knowledge can be described as : dxi ( t ) dt = − ∑ j∈N ( i ) Ãij ( xi ( t ) − xj ( t ) ) ( 4 ) where N ( i ) denotes the neighborhood of node i , xi ( t ) represents the feature of node i after diffusion time t. This differential equation can be solved as : X ( t ) = HtX ( 0 ) ( 5 ) where X ( t ) means the feature matrix after diffusion time t and Ht = e− ( I−T ) t is the heat kernel . Neighborhood radius rh in diffusion models . According to the definition of neighborhood radius Eq . 2 , the heat kernel version of the GDC has neighborhood radius rh as rh = ∑∞ k=0 θkk∑∞ k=0 θk = ∑∞ k=0 e −t tk k ! k∑∞ k=0 e −t tk k ! = e−t ∑∞ k=0 tk k ! k e−t ∑∞ k=0 tk k ! = e−t ( ett ) e−tet = t ( 6 ) This suggests that t is the neighborhood radius for the heat kernel based GDC , that is , t becomes a perfect continuous substitute for the hop number in multi-hop models . 3 Adaptive Diffusion Convolution . Recall that the heat kernel version of graph diffusion convolution ( GDC ) has the following feature propagation function as γ ( l ) ( Ĥ , G ) = e−LtĤ = ∞∑ k=0 e−t tk k ! TkĤ ( 7 ) where the Laplacian matrix L = I−T . For each graph dataset , it requires the manual grid search step to determine the neighborhood radius related parameter t. Moreover , t is fixed for all feature channels and propagation layers in each dataset . In this work , we explore how to adaptively learn the neighborhood radius from data for each graph and further examine the potential to generalize it for different feature channels and GNN layers . 3.1 Training Neighborhood Radius . Diffusion convolution enables us to replace GNNs ’ discrete feature propagation function with the continuous heat kernel . Instead of hand-tuning t , we can calculate the gradient of t and update t to converge to an optimal neighborhood , which is the same to learning other weight and bias parameters in the model . Figure 1 shows the training process of learning t. With more epochs , both t and training loss decrease when learning on the training set ( blue ) . Meanwhile , the validation and test accuracies drop dramatically as t tends to zero ( more epochs ) —representing each node could only use its own features to predict the label . That is , learning t directly on the training set causes overfitting . This phenomenon has also been observed in GDC [ 21 ] . The authors found that strong regularization on the difference of θk+1 − θk could help overcome the overfitting issue . However , that would require hand-tuning the regularization factor for every dataset , which is similar to hand-tuning t itself , e.g. , by grid search , further limiting the generalization of the model . To address this issue , we propose a method of training t by using the gradient of the model on the validation set . The goal for the model is to find t∗ that minimizes the validation loss Lval ( t , w∗ ) , where w denotes all the other trainable parameters in the feature transformation function and w∗ denotes the set of parameters that minimize the training loss Ltrain ( t , w ) . This strategy can be formalized as a bilevel optimization problem [ 3 , 11 ] : t∗ = argmin t Lval ( t , w∗ ( t ) ) ( 8 ) w∗ ( t ) = argmin w Ltrain ( t , w ) ( 9 ) By doing so , every time we update t , we need to make w converge to the optimal value , which is too expensive to train . An approximation method is to update t every time we update w , that is , w ( e+1 ) = w ( e ) − α1 5w Ltrain ( t ( e ) , w ( e ) ) ( 10 ) t ( e+1 ) = t ( e ) − α2 5t Lval ( t ( e ) , w ( e+1 ) ) ( 11 ) where e denotes the number of training epochs , α1 and α2 denote the learning rate on the training and validation sets , respectively . The similar idea has been proposed in the gradient-based hyperparameter tuning [ 24 ] and neural architecture search [ 23 ] . Figure 1 shows that using this method ( green lines ) helps avoid overfitting and thus offers better generalization , as there is no sign of test accuracy drop . Meanwhile , t does not diminish to zero , indicating the meaningful learning of this parameter— neighborhood radius . 3.2 Training Neighborhood Radius for Each Layer and Channel Conventional GNNs use the predetermined neighborhood radius for feature propagation . GDC proposes to use different neighborhood radius t for different datasets by hand-tuning the values . The above method furthers this direction by automatically learning the radius t from the given graph . This implies that one t for one dataset , that is , the same t for all GNN layers and all feature channels ( dimensions ) . Adaptive diffusion convolution ( ADC ) . The natural question arises here is whether we can have a unique t for each layer and channel , making them adaptive for the final learning objective . The obstacle that prevents previous models from achieving this lies in the infeasible challenge of hand-tuning or grid-searching the propagation function separately for each feature channel and GNN layer , given that as the number of parameters increases , the time complexity increases exponentially . However , the aforementioned strategy for updating t during the training of the model empowers us to adaptively learn specific t for all layers and all feature channels . Straightforwardly , we have the adaptive diffusion convolution ( ADC ) by extending the feature propagation function in Eq . 7 for each layer and channel , that is , from t to t ( l ) i , γ ( l ) ( Ĥ , G ) i = ∞∑ k=0 e−t ( l ) i ( t ( l ) i ) k k ! TkĤi ( 12 ) where t ( l ) i denotes the neighborhood radius t for the l-th layer and i-th channel , Ĥi represents the i-th column of the hidden feature Ĥ , i.e. , the feature on channel i , and γ ( l ) i denotes the feature propagation function on the l-th layer and i-th channel . This feature propagation function enables the GNN to train a separate t for each feature channel and layer , which is illustrative in Figure 2 . In addition , Figure 1 ( black lines ) also shows that there is no overfitting caused by the increase of the number of hyperparameters ( t ) . Generalized adaptive diffusion convolution ( GADC ) . By now , we introduce the adaptive diffusion convolution based on heat kernel . Without loss of generality , we can have ADC extended to a generalized ADC ( GADC ) , that is , not limiting the weight coefficients θk as heat kernel . Therefore , we have the feature propagation of GADC as : γ ( l ) ( Ĥ , G ) i = ∞∑ k=0 θ ( l ) ki T kĤi ( 13 ) where θ ( l ) ki denotes the weight coefficient for k-hop neighbors on l-th layer and i-th channel . We restrict ∑∞ k=0 θ ( l ) ki = 1 during training . Implementation Details . As we operate differently on each channel , whether we propagate before or after the feature transformation function actually matters . Empirically , we find that propagating on the input channels generates better results than propagating on the output channels ( Cf . Figure 7 for details ) . Therefore , we swap the propagation and transformation steps in the original message passing networks from Eq . 1 to : H ( l ) = ϕ ( l ) ( γ ( l ) ( H ( l−1 ) , G ) ) ( 14 ) Additionally , calculating e−Lt directly is infeasible for large graphs . Practically , we need to use the top-K truncation to approximate the heat kernel , making ADC ( Eq . 12 ) and GADC ( Eq.13 ) respectively updated as : γ ( l ) ( Ĥ , G ) i = K∑ k=0 e−t ( l ) i ( t ( l ) i ) k k ! TkĤi , γ ( l ) ( Ĥ , G ) i = K∑ k=0 θ ( l ) ki T kĤi ( 15 ) Similar to GDC , ADC and GADC are flexible components that can be directly plugged into existing GNN models , enabling them to adaptively learn the neighborhood radius for feature aggregation .

Summary: This paper proposes an adaptive diffusion convolution strategy that can automatically learn the optimal radius of message passing in graph neural networks and addresses the issue of selecting the optimal radius by hand-tuning. Additionally, this mechanism can serve as a plugin for graph neural networks and improve the model capacity and performance. Eventually, the paper validates this claim on the task of semi-supervised node classification.