Graph Warp Module: an Auxiliary Module for Boosting the Power of Graph Neural Networks in Molecular Graph Analysis

1 INTRODUCTION . Recently , Graph Neural Network ( GNN ) is a popular choice of model in the analysis of molecular datasets in medicinal and material science . Many molecular datasets consist of molecular graphs with feature vectors associated to each atom , and numerous methods based on GNN has been proposed to date just for learning the features of chemical molecules ( Wu et al. , 2018 ; Duvenaud et al. , 2015 ; Kearnes et al. , 2016 ; Li et al. , 2017 ; Gilmer et al. , 2017 ; Shang et al. , 2018 ) , such as those pertaining to electrical conductivity and toxicity . One problem in the application of GNN to molecular datasets is the difficulty in reducing the training loss . Unlike in the applications of Deep Neural Networks ( DNNs ) to image datasets , the training loss of GNN on molecular dataset does not decrease consistently with the number of layers nor number of nodes per layers ( cf . Fig . 4 in the appendix , thin dashed lines ) , and this seems to happen to numerous GNN architectures that are used in applications today ( Duvenaud et al. , 2015 ; Li et al. , 2016 ; Kipf & Welling , 2017 ; Xu et al. , 2019 ; Busbridge et al. , 2018 ) . Unfortunately , many strong techniques developed for deep CNNs such as ResNet ( He et al. , 2016 ) can not be applied naively to GNN , because the tasks for GNNs are oftentimes fundamentally different in nature from that of standard DNN . For example , each graph data to be passed to the network can differ in size , and it is also often desired that GNN is equivariant ( invariant under the reordering of vertices ) in general . To the best of the authors ’ knowledge , there have not been many studies done to date that directly addressed the problem of training loss . In this study , we propose graph warp module ( GWM ) , a supernode ( Li et al. , 2017 ; Gilmer et al. , 2017 ; Battaglia et al. , 2018 ) based auxiliary module that can be attached to generic GNNs of various types to improve its representation power . The I/O of the auxiliary module is defined independently from the GNN to which it is attached , and the users can install the GWM just by adding a small segment of code . Our GWM consists of three major components . The first component is virtual supernode ( Li et al. , 2017 ; Gilmer et al. , 2017 ) , which communicates with all nodes in the graph and promotes the remote message passing . The second and third components are attention unit ( Vaswani et al. , 2017 ; Veličković et al. , 2018 ) and gating units ( Cho et al. , 2014 ) . These adaptive weighting functions in the module help to adjust the flow of messages and deliver a message of appropriate strength to each node in the graph . Our GWM can consistently improve the performance of various types of GNN on various types of dataset . In Fig . 1 we show the effect of GWM on the performance of four types of GNNs with the same embedding dimension and the same number of layers on four molecular graph datasets . As we Reduction Ratio of Train Loss Re d uc tio n Ra tio o f T es t L os s 0.2 0.4 0.80.6-0.2 0.0 1.0 0.2 0.4 0.6 0.0 0.8 1.0 NFP GGNN RSGCN GIN Tox21 HIV QM9 Lipo DatasetGNN Figure 1 : Train loss reduction and test loss reduction achieved by GWM on various modeldataset pair . The shape of each point presents the dataset used , and the color of each point presents the GNN model used . The horizontal axis denotes the ratio of reduced training loss , and the vertical axis denotes that of the test loss . That the x-coordinate of a point is positive implies that the attachment of GWM improved the train performance for the corresponding model-dataset pair . That the ycoordinate of a point is positive implies that GWM improves the test-performance . can see in the figure , the attachment of GWM reduces both train loss and test loss for all but three model-dataset pairs . The GWM provides not only more representation powers ( less train loss ) , but also better generalization performances ( less test loss ) for various GNNs . As we will show in section 4.4 , we can further improve the positive effect of GWM using hyperparameter optimization softwares such as Optuna ( Akiba et al. , 2019 ) . Our contributions are as follows : 1 . We introduce GWM , an auxiliary module that can help improve the representation power of the GNNs that are designed for the analysis of small graphs . 2 . We show that the attachment of GWM can improve both the representation power and the generalization ability of various GNN models on many popular molecular graph datasets . 2 RELATED WORK . 2.1 VIRTUAL SUPERNODE . A common challenge in the application of GNNs to a graphical dataset is the difficulty in propagating the information across remote parts of graphs . Previously proposed solutions include sub-sampling ( Hamilton et al. , 2017 ) and pooling of neighbor nodes ( Ying et al. , 2018 ) . However , these clustering approaches are not too effective on the analysis of small graphs , in which every node can have a strong influence on the graph label . In this study , we use supernode ( Gilmer et al. , 2017 ; Li et al. , 2017 ; Pham et al. , 2017 ) to promote the global propagation of the information in molecular graphs . By adding a supernode to a graph , we can allow any pair of nodes in the graph to communicate through the supernode in one hop . Battaglia et al . ( Battaglia et al. , 2018 ) discusses a framework of GNNs that generalizes the supernodeaugmented GNNs . One advantage of the supernode-based approach is that we can modify the network architecture while keeping the original GNN model intact . However , naive addition of a supernode to a graph can potentially lead to inadvertent over-smoothing of information propagation ( c.f . ( Li et al. , 2018 ) ) . In our study , we therefore make the supernode a module by combining it with a gated message passing mechanism . This auxiliary module enables us to regulate the amount and type of information that is propagated through the feature space of the supernode . 2.2 MESSAGE PASSING AND ATTENTION/GATE MECHANISM IN GNN . The supernode in our GWM transmits information using the mechanism of message passing neural network ( MPNN ) ( Gilmer et al. , 2017 ) , which is defined recursively as follows by composing multiple layers of the form : h ` , i = F ` ( { h ` −1 , j ; j ∈ N ( i ) ∪ { i } } ) ( 1 ) where i , j are indices of nodes in a graph . h ` , i is the feature vector of the node i at the ` th layer , N ( i ) is the neighborhood of the node i , and F ` is an appropriate choice of function that updates the feature vectors of the previous layer . That is , MPNNs work by passing the information of each node to its neighbors in a recursive manner . Various methods are proposed for the choice of F and for the method of pooling the information of the neighbors of each node ( Schlichtkrull et al. , 2017 ; Kipf & Welling , 2017 ; Li et al. , 2016 ; Bruna & Szlam , 2014 ; Duvenaud et al. , 2015 ) . Attention mechanism is a mechanism that helps the network regulate the importance of each node/edge in message passing ( cf . ( Wang et al. , 2018 ) ) . A Relational GCN ( Schlichtkrull et al. , 2017 ) assumes that the aggregation weights of h ` −1 , j is fixed a priori in all F ` . With such architectures , however , one can not regulate the higher order correlation amongst the outputs from each node . Graph Attention Networks ( GATs ) ( Veličković et al. , 2018 ) introduce a self-attention mechanism ( Vaswani et al. , 2017 ) , which is equipped with a trainable set of weights that controls the importance of edges for each node . The relational graph attention network ( RGAT ) ( Busbridge et al. , 2018 ) also builds upon GAT and constructed multiple types of attentions derived from relation-type-wise intermediate node representations . Finally , a GRU ( Cho et al. , 2014 ) is a gating mechanism originally introduced for recurrent neural networks . Gated Graph sequence Neural Networks ( GGNN ) ( Li et al. , 2016 ) are the first to apply GRUs to the GNNs , and their method aims to introduce a recurrence relation between successive layers . Although not in the form of super-node , a very recent paper ( anonymous , 2019 ) uses a mechanism that combines GGNN and self-attention to capture both local and global relationships of a graph structure of program source codes . Our GWM is equipped with both multirelational attention mechanisms and GRUs to grant the module greater flexibility for the transmission of message between supernode and the bulk nodes . 3 GRAPH WARP MODULE . Our Graph Warp Module ( GWM ) is made of three building blocks : ( 1 ) a supernode , ( 2 ) a transmitter unit , and ( 3 ) a warp gate unit ( Fig . 2 ) . In a GWM-attached GNN , information is propagated across the graph through communication between the supernode and the original ( main ) GNN at each layer . Messages from the supernode and the main GNN are transmitted to the warp gate through the transmitter unit , and the results of the communication are passed back to the module/main network through the warp gate units . In this section , we describe the Graph Warp Module in detail , and present the motivation of our design . 3.1 PREMISE : VANILLA GNN AND ITS I/O . Before describing our GWM , we need to present the I/O notation for the family of GNNs we consider , and explain how they will be used when a GWM is attached to a GNN . We denote an arbitrary graph with the edge set E and the node set V as G = ( V , E ) . We will label the nodes in V as i = 1 , 2 , ... |V | , and represent each edge as a pair of nodes in V . The adjacency matrix A ∈ R|V |×|V | is a matrix whose ( i , j ) th entry is the weight assigned to the edge between the node i and the node j . Each instance of data passed to the GNN is a set of input feature vectors . We denote an input feature vector associated with node i as xi . The type of GNN that we consider computes the output recursively by applying a composition of smooth functions F ` to xis . With the understanding that xj = h0 , j , let h ` , i = F ` −1 , i ( h ` −1 , j ; j ∈ V ) ∈ Rd be the vector of features to be assigned to the ith node by the ` th layer of the GNN . When the GNN is operating on its own without the attachment of a GWM , the GNN updates a feature vector using h ` , i = F ` −1 , i ( h ` −1 , j ; j ∈ V ) ∈ Rd . Finally , the GNN reports some form of the aggregation of { hL , i ; i ∈ V } as the final Readout output . When a GWM is attached to the GNN , the main ( bulk ) GNN is requested to report F ` −1 , i ( h ` −1 , j ; j ∈ V ) ∈ Rd as the message from the ` − 1th main layer to the module , where it is treated as an element in the intermodule hyperspace and is mixed with the transmission from the supernode . The GWM will return the mixed message h ` back to the ` th layer of the main GNN . At the same time , the GWM requests a transmission message from the main GNN to the ` th supernode . The module will mix the transmission and the message from the ` − 1th supernode and return the mixed message g ` to the ` th supernode . The final output is produced by aggregating { hL , i ; i ∈ V } and gL .

Summary: Graph Neural Networks is a popular architecture for the analysis of chemical molecules. The authors propose an auxiliary module that can be attached to a GNN that can boost the representation power of GNNs. The auxiliary module has three building blocks: 1. a supernode, 2. a transmitter unit and 3. a warp gate unit. The authors show through carefully designed experiments that these additions can be attached to any type of GNN and that they are successful in reducing both the training error and test error. A variety of graph regression and graph classification tasks are chosen to show the efficacy of the method.