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 extends from graph diffusion convolution (GDC) that used hand-tuned neighborhood size for feature aggregation to adaptively learning the size. ADC can automatically learn the neighborhood radius for each GNN layer and each feature dimension and acts as a plugin to any GNN model. Experimental results and ablation studies show that learning a unique neighborhood radius for each feature channel in each GNN layer consistently and improves the performance of downstream graph mining tasks.