Title: GRAPH WARP MODULE: AN AUXILIARY MODULE FOR BOOSTING THE POWER OF GRAPH NEURAL NET-WORKS IN MOLECULAR GRAPH ANALYSIS Anonymous authors Paper under double-blind review

Abstract: Graph Neural Network (GNN) is a popular architecture for the analysis of chemical molecules, and it has numerous applications in material and medicinal science. Current lines of GNNs developed for molecular analysis, however, do not fit well on the training set, and their performance does not scale well with the complexity of the network. In this paper, we propose an auxiliary module to be attached to a GNN that can boost the representation power of the model without hindering the original GNN architecture. Our auxiliary module can improve the representation power and the generalization ability of a wide variety of GNNs, including those that are used commonly in biochemical applications.

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) cannot 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; Veličković 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 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:

EQUATION

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 

Graph Warp Module:

Node feature vectors @Layer l Supernode feature vector @Layer l Supernode feature vector @Layer l-1

Node feature vectors @Layer l-1 (molecular graph)

Feature vectors + virtual "supernode"

Node feature vectors @Layer l Node feature vectors @Layer l-1 Before Attaching the Graph Warp Module A GWM consists of a supernode, a transmitter unit, and a warp gate unit. A GWM can be added to the original GNN as an auxiliary module. At each layer, the supernode and the main network communicate through the transmitter and the warp gate. 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 cannot regulate the higher order correlation amongst the outputs from each node. Graph Attention Networks (GATs) (Veličković 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 x i . The type of GNN that we consider computes the output recursively by applying a composition of smooth functions F to x i s. With the understanding that x j = h 0,j , let h ,i = F -1,i (h -1,j ; j ∈ V ) ∈ R d 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 ) ∈ R d . Finally, the GNN reports some form of the aggregation of {h L,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 ) ∈ R d 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 {h L,i ; i ∈ V } and g L .

3.2. SUPERNODE:

A supernode is a special node that is connected to nodes in the original graph to promote global information propagation across the network (Fig. 2 ). A supernode is to be prepared for each th layer of the main GNN, and we associate a feature vector g to the supernode at the th layer. At each layer, the transmitter requests the following from the supernode: (1) a message G (g ) for the + 1th layer and (2) a transmission to the main network, where G is an appropriate choice of a smooth function.

Because a supernode is a superficial variable, we must initialize g 0 manually. For instance, we can use some form of aggregation of the global graph features (e.g. a number of nodes or edges, graph diameter, girth, cycle number, min, max, histogram, or an average of input node features, . . . ). A detailed example of the aggregated feature is presented in the appendix.

3.3. TRANSMITTER UNIT:

The transmitter unit handles the communications between the main GNN module and the GWM (Fig. 3 ). The transmitter module is responsible for translating the messages from the recipient into a form that can be mixed in the intermodule hyperspace. We will use multiple types of messages and thus use a separate attention mechanism for each type of message. Before transmitting messages from the main GNN to the supernode, the transmitter uses the K-head attention mechanism to aggregate messages of each type. We enumerate the components included therein:

• m main→super ,k

: aggregated message of head k from the main GNN to the supernode at layer .

• h main→super : transmission from the main GNN to the supernode, derived from m main→super ,k .

• g super→main : transmission from the supernode to the main at layer .

The transmissions are to be constructed from the following set of equations. For a vector v, we use v m:n ∈ R (m-n)d to denote the concatenation of the vectors v m , v m+1 , ... ∈ R d . Throughout, we use capital letters to denote the trainable coefficients.

EQUATION

EQUATION

EQUATION

) where α ,i,k denotes an attention weight of the ith node at the kth head (type) and the lth layer.

The transmission from the supernode to the main is simply given by:

g super→main = tanh (F g -1 ) ∈ R D .

(5) There is no analogue of m for the supernode because we are not considering a set of messages of different types to be transmitted from the supernode.

3.4. WARP GATE:

The warp gate is responsible for merging the transmitted messages and for passing the results to the supernode and the main network through self recurrent units. The gate uses warp gate coefficients to control the mixing-rate of the messages. The components of the warp Gate are:

• h 0 : inputs to the GRU unit at th layer that transmits the message to the main GNN.

• g 0 : inputs to the GRU unit at th layer that transmits the message to the supernode.

• ĥ ,i : the message

F -1,i (h -1,k ; k ∈ V ) ∈ R D from

the -1th layer of the main network, expressed in the intermodule hyperspace.

• ĝ : the message G -1 (g -1 ) ∈ R D from the -1th supernode, where G is an appropriate smooth function with outputs in the intermodule hyperspace. • z ,i : tensor of warp gate coefficients for the transmission from the supernode to the main GNN.

• z (S):

,i : tensor of warp gate coefficients for the transmission from the main GNN to the supernode.

The module then mixes the transmissions and the messages from the previous layer by applying the following gated interpolations:

EQUATION

EQUATION

EQUATION

)

where σ is a nonlinear function whose range lies in [0, 1]. Finally, the warp gate returns the mixed messages to the main GNN and the supernode through gated recurrent unit (GRU) s:

h ,i = GRU h -1,i , h 0 ,i ∈ R D , g = GRU g -1 , g 0 ∈ R D . (9)

As for the structure of GRU, we used the original design introduced by (Cho et al., 2014) .

As we will show with ablation studies (Sec. 4.4), every component of GWM is essential in making the module work. The Attention coefficients (Eqs.(3,4)) are important because the amount and the type of information that must be transmitted to remote nodes may differ for different nodes. The Gating coefficients (Eqs.(6-8)) are important because we want to regulate the transmission from each node in the graph to the supernode and vice versa. We use different recurring units (Eq.9) for the transmission from the module to the supernode and the transmission from the module to the main network for each layer because the amount of the information that must be reinforced may differ for the main network and the supernode at each layer.

3.4.1. COMPUTATIONAL COMPLEXITY:

Let V be the vertex set, K be the number of attention heads, and D be the dimension of the node embedding. Then the additional computational cost incurred by the attachment of GWM is at most O(|V |KD 2 ). As for the actual computation time, GWM attached module consumes approximately double the time of the original unattached version.

4. EXPERIMENTS:

In this section, we present our experimental results on multiple molecular graph datasets, testing the efficacy of the GWM for graph regression tasks and graph classification tasks.

4.1. DATASETS:

We used four datasets collected in MoleculeNet (Wu et al., 2018) . These datasets are described in the SMILES string format, which admits the graph representations we described above. For details, please see (Wu et al., 2018) .

For the graph regression tasks, we used the QM9 dataset and the Lipophilicity (LIPO) dataset. QM9 is a dataset with numerical labels, containing about 133K drug-like molecules with 12 important chemical-energetic, electronic, and thermodynamic properties, such as HOMO, LUMO, and electron gaps. The LIPO dataset is another numeric-valued dataset, containing the solubility values of roughly 4K drug molecules. Each instance of data in these datasets is a pair of a molecular graph and a numerical value(s): the 12 chemical properties in the QM9 dataset, and the solubility in the LIPO dataset. For both datasets, the task is to predict the numerical value(s) from the molecular graph. We evaluated the performance of the models using mean absolute errors (MAEs). We report the averaged MAE over 12 sub-tasks (properties) for QM9.

For the graph classification tasks, we used the Tox21 and the HIV datasets. The Tox21 dataset contains about 8K pairs of molecular graph and 12 dimensional binary vector that represent the experimental outcomes of toxicity measurements on 12 different targets. The HIV dataset contains roughly 42K pairs of molecular graph and binary label that represent the medicinal effect of the molecule. For these datasets, the task is to predict the binary label(s) from the molecular graph.

For these tasks, we use ROC-AUC values as a measure of performance. We report the averaged ROC-AUC over 12 sub-tasks (targets) for Tox21.

Throughout, we used the train/validation/test data splits of the "scaffold" type, which is considered by (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014) as the difficult type for test predictions. Please find the appendix for details.

4.2. CHOICES OF THE MAIN GNN MODELS AND IMPLEMENTATIONS:

We test GWMs on various GNN models. Neural Fingerprints (NFP) (Duvenaud et al., 2015) and Weavenet (Kearnes et al., 2016) are relatively classical baselines. A Gated Graph Neural Network (GGNN) (Li et al., 2016 ) is a strong GRU-based GNN. Renormalized Spectral Graph Convolutional Network (RSGCN1 ) (Kipf & Welling, 2017) , a popular GNN model approximating a CNN for graphs (Defferrard et al., 2016) . The relational graph attention network (RGAT) (Busbridge et al., 2018) uses multiple attention mechanisms for a set of edge types. Graph Isomorphism Network (GIN) (Xu et al., 2019) employs multi layer perceptrons within each layer for richer transformations.

We implement all models in Chainer (Tokui et al., 2015) . In the readout layer, we first aggregate all information from the main nodes in the same way as in the original paper, concatenated the result with the features from the supernode, and passed the concatenated tensor to a fully connected layer.

Model Name (emulating) Attention Gatings GRUs Simple supernode (Li et al., 2017; Pham et al., 2017) ---NoGate GWM (Gilmer et al., 2017) -(supernode only) Proposed GWM Table 1 : Supernode-based models validated in the Experiment 4.4. Note that the "simple supernode" ablation model and the model used in (Pham et al., 2017) allow bi-directional message passings between original nodes and the supernode while the one used in (Li et al., 2017) only allows directional messages from the original nodes to the supernode.

For evaluation, we used a softmax cross entropy for classification tasks and a mean squared error for regression tasks. We use a fixed set of hyperparameters throughout the study.

All models were trained with Adam (Kingma & Ba, 2015) . We report the results of the model snapshots of the epoch for which the best validation score was achieved. For implementation details including readouts and hyperparameters, please read the appendix.

4.3. TRAINING AND TEST LOSS REDUCTION:

We explain the details of the experiment that produced the result presented in Fig. 1 . For this experiment, we reported the average r of the loss reduction ratio r = L -L (+) L for both training loss and test loss over 10 runs. L denotes the loss of the vanilla model, and L (+) denote the loss of the GWM-installed model. rtrain denotes a reduction ratio of the training loss, and rtest denotes a reduction ratio of the test loss.

Fig. 1 is the scatter plot of the (r train , rtest ), with the dotted slope representing rtrain = rtest . For all GNNs and datasets, we set L = 3 and D = 50. As we can see in the plot, rtrain s were negative for the two GNNs in the HIV dataset (blue and brown circles). rtrain for GGNN on QM9 (blue cross) was a very small negative value). For the other 13 (model-dataset) pairs, the attachment of GWM consistently reduces the training loss( i.e. improves the fit to the training graph datasets. ) Remarkably, 15 out of 16 pairs had positive rtest s: the GWM improved generalization performances in most cases. It is worthy of note that rtrain and rtest are positively correlated in this scatter plot. We were able to obtain similar results for all other choices of hyperparameters we tested. For the results with different hyperparameter values, see the appendix. This result implies that our GWM has the general effect of improving the generalization performance by augmenting the representation power of the main GNN.

4.4. EFFECT OF THE GWM ON THE REPRESENTATION POWER OF MODEL SPACE:

In the second experiment, we studied the effect of GWM on the representation power of GNN models. To compare the GWM-augmented GNNs with their vanilla GNN counterparts on fair grounds, we used Bayesian optimization to optimize the number of layers and the dimension of feature vectors for each model-dataset pair we tested. Hyperparameters are optimized via the Optuna (Akiba et al., 2019) library.

We conducted a set of ablation studies to investigate the effect of (1) attention mechanism, (2) Gating mechanism, and (3) the Recurrent unit. We used two ablation models. Simple supernode model is a supernode without attention, gatings, and GRU functions. This ablation model can be considered a variant of (Pham et al., 2017 )'s supernode model in which the supernode and the bulk networks communicate with each other in bidirectional manner. This model can be also considered as an extended version of supernode proposed in (Li et al., 2017) .

NoGate GWM is a GWM without gatings, and it lacks GRU for the main GNN.

Table 1 summarizes the details of our ablation studies. For the detailed formulations of the two ablation models, please see the appendix. for each (dataset, GNN) pair. Full tables with standard deviations are presented in the appendix. As we can see in the tables, the proposed GWM improves the generalization performances for 23 out of 24 (model-dataset) pairs. These results suggest that the proposed (full) GWM can improve GNNs' performances irrespective of the choice of GNN models and the dataset.

A few words of caution are in order here. Two ablation models did not improve the generalization of GNNs for the QM9 and the HIV datasets (see the column "# Improved"). This suggests an appropriate combination of attentions, gatings, and GRUs is essential in making the supernode effective for the analysis of molecular graph datasets.

We shall emphasize that the goal of this study is not to find the specific network architecture that achieves the states of the art performance for selected datasets2 . Instead, the goal of our work is propose an attachable module that improves the representation power and a generalization performance irrespective of the choice of GNN architecture. As we can see in the presented result, the attachment of GWM improves the result in most cases; one of the results of our GWM attached model for Tox21 is actually SOTA (0.787, achieved by GWM attached RGAT).

5. CONCLUSION:

For a generic DNN, numerous effective installable modules have been proposed for the improvement of the model (e.g. (Srivastava et al., 2014; Ioffe & Szegedy, 2015; Miyato et al., 2018; He et al., 2016) ). The proposed GWM is the first of its kind to be installed to a generic GNN as an auxiliary module. Experimental results show that the GWM can generally improve the representation power as well as the generalization performance of a GNN, irrespective of the choice of GNN architecture and the molecular graph datasets. We would like to emphasize that the choice of the internal structure of GWM is not limited to the ones we described in this study, and that there are possibly numerous ways to construct a GWM-like module. For example, there is no provable justification for the use of a linear transformation in the transmissions or a bilinear form in the attention coefficients α. 

A OUR FORMULATION OF RGAT:

Apart from the original RGAT (Busbridge et al., 2018) , we have developed a similar GNN in a slightly different formulation. Followings are our RGAT formulation:

EQUATION

EQUATION

EQUATION

W, F, G, A are the coefficient matrix to be tuned. is the index of the layer up to L, k is the index of the attention head up to i, j are the index of the nodes up to N , e i,j = r is the index of the edge type up to R.

The main point is the edge type information in Eq.13. The edge type e i,j = r switches the weight matrix of the attention similarity function, a. This means that the associations between nodes should be computed dependent on the edge type. This is a natural assumption for chemical molecular graphs. Typically we have multiple bond types between nodes = atoms: single-bond, double-bound, triple-bond, and the aromatic ring. It is natural to assume that interactions between atoms are affected by the bond types among the atoms.

The main differences from the original RGAT lie in the Eq.11. The original RGAT assumes that the weight matrix G is also dependent on the edge type (G ,k,ei,j ) while we omit this dependency. Also, the original RGAT does not provide a self-link weight matrix F while we do. We made these changes based on our preliminary experiments. We found our formulation is better than the original RGAT formulation in the MoleculeNet dataset, in terms of the training stability and the generalization performances.

Another difference is the choice of the attention function. In our formulation, the attention similarity measure a(•) is defined by the general attention in (Luong et al., 2015) while the original GAT (Veličković et al., 2018) and the RGAT (Busbridge et al., 2018) employed a simpler concat attention.

B OUR IMPLEMENTATION OF GIN:

We implement the simplest GIN: 2-layer MLP with ReLU activation for each layer and a bias parameter fixed at 0. We regularize GIN with dropout (Srivastava et al., 2014) , instead of batchnormalization (Ioffe & Szegedy, 2015) . This is because the batch-normalization of the Chainer-Chemistry library did not correctly treat the padded node elements in the minibatches when we conducted the experiments.

C EXPERIMENTS DETAILS: GENERAL ISSUES C.1 GRAPH DATA REPRESENTATION:

All datasets used in our experiments are taken from the MoleculeNet (Wu et al., 2018) . Four used datasets are provided in the SMILES format. A SMILES format is a line notation for describing the structure of chemical compounds. We decode a SMILES molecular data into a graph representation of the molecule. A node in the graph corresponds to an atom. Each atom node is associated with the symbolic label of the atom name ("H", "C", ...). An edge in the graph corresponds to a bond between atoms. Each bond edge is associated with the bond type information (single-, double-, ....).

Given the graph, we extract input feature vectors for node x i and that of supernode x . x i , the input feature vector for the node i is a D-dimensional continuous vector, which is an embedded vector of the one-hot atom label vector with a trainable linear transformation. X , The input feature vector for the supernode is a D -dimensional continuous vector, which again is an embedded vector of some graph-global features with a trainable linear transformation. Choices for the graph-global features are detailed in the following section.

The edge information is converted in an adjacency matrix, A .

C.2 EXPLICIT FEATURES FOR SUPERNODE:

Since the supernode does not exist in the original graph G, we have no observable cues for the supernode. For simplicity, we propose to use an aggregation of node features, such as:

• Histograms of discrete labels attached to original nodes

• Averages, maximums, minimums, or medians of numerical attributes attached to original nodes • Histograms of edge types if the graph is multiple relational graph.

• Number of nodes, graph diameters, modularity, and other simple statistics for graph structure.

We can augment the super feature vector x if some additional information about the graph is provided. Essentially, these simple aggregations of the feature vectors do not bring new information into the network. However we found that the graph-wise super feature input boosts the performance of the learned network model.

C.3 DATA SPLITS:

In chemical datasets, a totally random shuffling of samples into train/val/test subsets is not always a valid way of data splitting. Therefore MoleculeNet provides several ways of data splitting. The "random" split is the random sample shuffling that are most familiar to the machine learning community. The "scaffold" split separate samples based on the molecular two-dimensional structure. Since the scaffold split separates structurally different molecules into different subsets, "it offers a greater challenge for learning algorithms than the random split" (Wu et al., 2018) . Throughout the paper, we adopt the scaffold split to assess the full potential of the GWM-attaching GNNs.

The actual construction of the scaffold split train/validation/test subsets has a freedom of algorithm choices. We basically adopted the algorithm provided by the deepchem3 library, which is the standard split algorithm for many papers. However, for the experiment of train/test loss comparison, we adopted the algorithm provided by the Chainer Chemistry library.

C.4 READOUT LAYER:

In many applications of GNNs users may expect a single fixed-length vector representing the characteristics of the graph G. So we add the 'readout' layer to aggregate the original node hidden states {H } and the global node hidden states {g }.

A main issue in the readout unit is how to aggregate the original nodes, whose number varies for each graph. A simple way is to take an arithmetic average (sum) of the hs at the L-th layer, but we can also use a DNN to compute (non-linear) "average" of hs (Li et al., 2016; Gilmer et al., 2017) . After the aggregation of the node hidden states, we simply concatenate it with gs and apply some transformations to achieve the readout vector, r:

EQUATION

In the above equation, DNN r1 is a multi-layer perceptron or a fully connected layer to mix the concatenated hidden vectors. We adopted a simple fully-connected layer for DNN r1 in this paper. DNN r2 is a specific readout unit accompanied with a original GNN to aggregate variable-length H L .

C.5 OPTIMIZER All models were trained with Adam (Kingma & Ba, 2015) , α = 0.001, β 1 = 0.9, and β 2 = 0.999.

C.6 ABLATION MODELS FORMULATION:

Here we detail the formulation of the ablation models used in the main comparison experiments (Sec. 4.4).

As written in the main manuscript, we formulate the two ablation models (Table 1 ) as follows. A simple supernode model is an ablation model without attentions, gates, nor the GRUs, and it can be considered as a variant of (Li et al., 2017; Pham et al., 2017) with bidirectional communication between the supernode and the bulk, omitting all attentions and gates. First, there is no attention for the Transmitter. So the message from the main nodes to the supernode is just a sum of hidden vectors, h -1,: :

EQUATION

)

Originally there is no messages from the supernode to the main GNN in (Li et al., 2017) , but we allow such a simple message in this ablation model:

EQUATION

Messages are merged by simple linear combinations, instead of gates and GRUs, following (Li et al., 2017) :

EQUATION

EQUATION

)

We find it difficult to fully recover the supernode of (Gilmer et al., 2017) since their description on the supernode is quite limited. Thus, a NoGate GWM model, which surrogates (Gilmer et al., 2017) , only capture the essence of their supernode: no gatings for merger, and GRU is not installed for the nodes of the main GNN. In this model, we use the same attention-based Transmitter unit as in Eqs.

(2-6). We reduce the adaptive gatings in the Warp unit by a simple averaging, and omit the GRU for h ,i .

EQUATION

EQUATION

EQUATION

EQUATION

C.7 HYPERPARAMETER

We fix a part of hyperparameters throughout the experiments, which does not influence performances so much: the number of heads in all multi-head attention mechanisms to K = 8, and used R = 4 edge types for the multi-relational mechanism in all models. Also, at every layer, we set the dimension of the supernode feature to be the same as that of the features of the nodes in the main GNN.

In the next section, we list the other hyperparameters (the number of layers L, the dimension of feature vectors D(= D )) used in several experiments/figures, as well other experimental/implementation details.

C.8 COMPUTATIONAL ENVIRONMENT:

We use a single GPU (mainly nvidia Tesla V100) for an experimental run. A run roughly takes 1 hour to 1 day, depending on the hyperparameters and the GNN models. 

D EXPERIMENTS DETAILS: FOR EACH EXPERIMENT:

In this section, we report details for each experiment, including the chosen hyperparameters and additional results.

D.1 TRAINING LOSS BEHAVIORS OVER NUMBER OF LAYERS AND EMBEDDING DIMENSIONS:

We used the Tox21 dataset to confirm behaviors of training losses, mentioned in the first paragraph of the introduction (Fig. 4 ).

To study the effect of the number of layers L (the left panel), we fixed the dimension D = 32. To study the effect of the feature vector dimensions D (the right panel), we fixed the number of layers L = 4.

All models are trained for 30 epochs.

Thinner dashed lines in Fig. 4 plot the training losses of original networks (w/o GWM). Thick solid lines in Fig. 4 plot the training losses of networks augmented with GWM. In general, GWM has an effect of decreasing the training loss for most choices of numbers of layers and dimensions of the node feature vectors, for all GNNs. In general, similar results are obtained from the other three datasets so we omit these figures.

D.2 SECTION 4.3: TRAIN AND TEST LOSS REDUCTION:

In the experiment in the Section 4.3 (Fig. 1 ), we align the hyperparameters including L, D among a vanilla GNN and its GWM-installed counterpart, to compare the loss function values. We used four datasets. For each dataset, we fixed L and D for all GNN models to compare the loss reduction performances.

In Fig. 1 , L = 3 and D = 50 for all GNNs and datasets. In preliminary experiments, we manually changed Ls (∈ {2, 3, 4}) and Ds (∈ 32, 50, 100, 150) in some extent, but found the overall tendency of the scatter plots does not dramatically change. This is partially understood from the Figure 1 in the main manuscript: the loss curves of the vanilla GNNs and their GWM-augmented counterparts evolve in roughly parallel. This implies the ratios of loss reductions are not so much dependent on the hyperparameter choices.

Here we show the results of other hyperparameter settings. All cases we observe the similar plot patterns. Chosen Ls and Ds are presented in the Table 4 .

Note that GWM-attached GNNs tend to perform better with deeper architecture (i.e larger L), when the vanilla GNNs without GWM do not change its overall performance with the number of layers (Table 4 ). The performance of the GWM-attached models often improve with the number of layers until the model depth reaches quite a large number ( 8) (e.g. [QM9, GGNN+Proposed GWM], [HIV, GIN+Proposed GWM]). Out of 24 pairs (=4 datasets × 6 GNN models), 11 pairs perform better with deeper architecture when deployed with GWM.

For the QM9 dataset, we trained the models for 50 epochs. For the HIV and the Tox21 dataset, we trained the models for 100 epochs. For the LIPO dataset, we trained the models for 200 epochs. Finally we present the detailed reports on the sub-tasks of QM9 and Tox21. These datasets consists of 12 sub-tasks and so far we reported the sub-task-averaged scores. Below we report the scores changes brought by GWM-attached GNNs in Table 7 (QM9) and Table 8 (Tox21). For QM9, we report the relative reductions of MAE in percentage(%). For Tox21, we report the improvements of binary classification accuracy in percentage (%). Numbers of sub-tasks improved by GWM (positive value slots) are in good accordance with the score gains in Table 2 and Table 3 

:

Referred as RSGCN in Chainer-Chemistry package, but often simply referred as "GCN" in several papers.

For example, Yang et al. (2019) andLi et al. (2017) achieve the SotA AUCs (0.776) of GNNs on the HIV dataset.

https://deepchem.io/

Figures and Tables:

Figure1: 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.

Figure2: The overview of the proposed Graph Warp Module (GWM). A GWM consists of a supernode, a transmitter unit, and a warp gate unit. A GWM can be added to the original GNN as an auxiliary module. At each layer, the supernode and the main network communicate through the transmitter and the warp gate.

Figure 3: Details of the GWM computations.

Figure 4: Training losses of various GNN models on a molecule graph dataset (Tox21). The horizontal axis denotes the number of GNN layers (the left panel) or the dimension of the node feature vectors (the right panel). Color denotes the GNN model. Thinner dashed lines are the losses of the vanilla GNNs, while thicker solid lines show the losses of the GNNs attached with the proposed Graph Warp Module ("GWM"). Scores are partially unavailable due to memory shortages.

Fig.5is the scatter plot of the (r train , rloss ), L = 3 and D = 32. In this case, two HIV dataset plus one case for QM9 reported the increase of the training loss. For other 13 pairs, the GWM successfully reduce the training losses. It is remarkable that all 16 plots have positive test loss reduction rates: namely, the generalization performances are improved by the GWM in this choice of the hyperparameters.

Fig.6is the scatter plot of the (r train , rloss ), L = 4 and D = 100. In this case, only two cases of the HIV datasets record the negative rtrain values. For other 14 pairs, the GWM successfully reduce the training losses and 13 out of these 14 pairs have positive rtest s.All models were trained for 30 epochs.

Table 3 respectively present the MAEs on the regression tasks and the ROC-AUCs for the classification tasks, averaged over 10 random runs. In the tables, bold faces indicate the improvements from the vanilla GNN and asterisks indicate the best model among supernode models

Table content: <table><tr><td>Dataset</td><td>GNN model</td><td>NFP</td><td colspan="4">Weave RGAT GGNN RSGCN</td><td>GIN</td><td># Improved</td></tr><tr><td>LIPO</td><td>vanilla GNN +Simple Supernode</td><td>.677 .693</td><td>1.19 1.01</td><td>.753 .740</td><td>.582 .604</td><td>.801 .775</td><td>.844 .819</td><td>-4/6</td></tr><tr><td/><td>+NoGate GWM</td><td>.675</td><td>.721</td><td>.688</td><td>.576</td><td>.787</td><td>.847</td><td>5/6</td></tr><tr><td/><td>+Proposed GWM</td><td>.672*</td><td>.688*</td><td>.659*</td><td>.569*</td><td>.752*</td><td>.784*</td><td>6/6</td></tr><tr><td>QM9</td><td>vanilla GNN +Simple Supernode</td><td>6.16 7.68</td><td>6.38 5.51</td><td>8.96 9.00</td><td>4.92 5.41</td><td>15.2 14.6</td><td>14.0 11.5*</td><td>-3/6</td></tr><tr><td/><td>+NoGate GWM</td><td>6.84</td><td>5.40*</td><td>9.21</td><td>5.52</td><td>12.5</td><td>12.9</td><td>3/6</td></tr><tr><td/><td>+Proposed GWM</td><td>6.64*</td><td>5.90</td><td>8.39*</td><td>4.88*</td><td>11.9*</td><td>11.8</td><td>5/6</td></tr></table>

MAEs on the LIPO dataset and QM9 dataset. Smaller values are better. Scores on QM9 are the average MAEs over 12 sub-tasks. The score of select models are presented in the appendix.

Table content: <table><tr><td>Dataset</td><td>GNN model</td><td>NFP</td><td colspan="4">Weave RGAT GGNN RSGCN</td><td>GIN</td><td># Improved</td></tr><tr><td>HIV</td><td>vanilla GNN +Simple supernode</td><td>.724 .707</td><td>.670 .676</td><td>.707 .704</td><td>.746 .764*</td><td>.746 .728</td><td>.729 .729</td><td>-2/6</td></tr><tr><td/><td>+NoGate GWM</td><td>.714</td><td>.680</td><td>.726</td><td>.744</td><td>.742</td><td>.739</td><td>3/6</td></tr><tr><td/><td>+Proposed GWM</td><td>.731*</td><td>.681*</td><td>.748*</td><td>.762</td><td>.758*</td><td>.755*</td><td>6/6</td></tr><tr><td>Tox21</td><td>vanilla GNN +Simple supernode</td><td>.763 .770</td><td>.710 .750</td><td>.764 .787*</td><td>.757 .790</td><td>.760 .770*</td><td>.740 .763</td><td>-6/6</td></tr><tr><td/><td>+NoGate GWM</td><td>.775*</td><td>.764</td><td>.786</td><td>.792*</td><td>.759</td><td>.766</td><td>5/6</td></tr><tr><td/><td>+Proposed GWM</td><td>.769</td><td>.767*</td><td>.787*</td><td>.785</td><td>.766</td><td>.768*</td><td>6/6</td></tr></table>

ROC-AUCs on the HIV dataset and Tox21 dataset. Larger values are better. Scores on Tox21 are the average MAEs over 12 sub-tasks. Scores of select models are presented in the appendix.

Table content: <table/>

Effective choices of supernode features are also open to further research. Our study can possibly open an entirely new avenue for the architectural study of GNNs. Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, Andrew Palmer, Volker Settels, Tommi Jaakkola, Klavs Jensen, and Regina Barzilay. Are Learned Molecular Representations Ready For Prime Time? arXiv, pp. 1904.01561, 2019.

Table content: <table><tr><td>Rex Ying, Jiaxuan You, Christopher Morris, Xiang Ren, William L. Hamilton, and Jure Leskovec.</td></tr><tr><td>Hierarchical Graph Representation Learning with Differentiable Pooling. In Advances in Neural</td></tr><tr><td>Information Processing Systems 31 (Proceedings of NeurIPS), 2018.</td></tr></table>

MAEs on the LIPO dataset and QM9 dataset. The number of layers and the dimension of the feature vectors are defined via Bayesian Optimization for each method and each dataset. Averages (standard deviations) over 10 random runs. Smaller values are better.

Table content: <table><tr><td>RSGCN GIN</td><td>.746 (.011) .729 (.020)</td><td>.728 (.009) .729 (.025)</td><td>.742 (.019) .739 (.012)</td><td>.758* (.022) .755* (.009)</td><td>.760 (.005) .740 (.007)</td><td>.770* (.006) .763 (.006)</td><td>.759 (.007) .766 (.008)</td><td>.769 (.006) .768* (.007)</td></tr><tr><td>GGNN</td><td>.746 (.018)</td><td>.764* (.014)</td><td>.744 (.006)</td><td>.762 (.016)</td><td>.764 (.009)</td><td>.790 (.006)</td><td>.792* (.008)</td><td>.785 (.005)</td></tr><tr><td>NFP WeaveNet RGAT</td><td>.724 (.017) .670 (.020) .707(.039)</td><td>.707 (.023) .676 (.031) .704 (.019)</td><td>.714 (.018) .680 (.043) .726 (.024)</td><td>.731* (.020) .681* (.010) .748* (.019)</td><td>.763 (.004) .710 (.029) .744 (.008)</td><td>.770 (.008) .750 (.020) .787* (.005)</td><td>.775* (.007) .764 (.011) .786 (.009)</td><td>.769 (.007) .767* (.021) .787* (.009)</td></tr><tr><td>GNN model</td><td>vanilla GNN</td><td>+Simple supernode</td><td>+NoGate GWM</td><td>+Proposed GWM</td><td>vanilla GNN</td><td>+Simple supernode</td><td>+NoGate GWM</td><td>+Proposed GWM</td></tr><tr><td>Dataset</td><td colspan="2">HIV</td><td/><td/><td colspan="2">Tox21</td><td/><td/></tr></table>

ROC-AUCs on the HIV dataset and Tox21 dataset. The number of layers and the dimension of feature vectors are defined via Bayesian Optimization for each method and each dataset. Averages (standard deviations) over 10 random runs. Larger values are better.

Table content: <table/>

. Relative reduction of MAE (%) GWM, 12 tasks in QM9. Larger values are better.

Table content: <table><tr><td>GNN/Tasks</td><td colspan="5">mu alpha HOMO LUMO gap</td><td>r2</td><td>zpve</td><td>cv</td><td>u0</td><td colspan="3">u298 h298 g298</td></tr><tr><td>GGNN</td><td>-0.6</td><td>-0.9</td><td>1.0</td><td>1.6</td><td colspan="2">4.6 2.9</td><td>19</td><td colspan="2">0.3 -3.9</td><td>-4.8</td><td>-8.3</td><td>-4.2</td></tr><tr><td>WeaveNet</td><td>-2.2</td><td>-0.3</td><td>27</td><td>-35</td><td>52</td><td>6.7</td><td>107</td><td>-11</td><td>84</td><td>91</td><td>86</td><td>87</td></tr><tr><td/><td/><td>GNN/Tasks</td><td colspan="7">Task1 Task2 Task3 Task4 Task5 Task6</td><td/><td/><td/></tr><tr><td/><td/><td>RGAT</td><td>0.9</td><td>0.2</td><td>0.7</td><td colspan="2">0.5</td><td>1.2</td><td>0.2</td><td/><td/><td/></tr><tr><td/><td/><td>GGNN</td><td>0.0</td><td>3.7</td><td>2.1</td><td colspan="2">1.1</td><td>1.5</td><td>16</td><td/><td/><td/></tr><tr><td/><td colspan="2">GNN/Tasks</td><td colspan="8">Task7 Task8 Task9 Task10 Task11 Task12</td><td/><td/></tr><tr><td/><td/><td>RGAT</td><td>-0.5</td><td>1.2</td><td>-0.0</td><td>0.2</td><td/><td>1.9</td><td>0.4</td><td/><td/><td/></tr><tr><td/><td/><td>GGNN</td><td>5.9</td><td>2.8</td><td>0.7</td><td>1.0</td><td/><td>2.4</td><td>1.5</td><td/><td/><td/></tr></table>

Improvements of binary classification accuracy (%) by GWM, 12 tasks in Tox21. Larger values are better.

Table content: <table/>

Peer review (iclr2020):
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. The paper is well written and easy to follow. The modification suggested by the authors is novel and useful. Experiments are well designed. An aspect that is not clear from the paper is the ability of these models to overfit the data as you increase the representation power of the network. While the authors claim that to be one of the shortcomings of existing GNNs, it is not clear whether the proposed method solves that problem. For example, in figure 4. the training loss hardly decreases as the number of layers are increased. It would be good if the authors can share any insights on this point. Overall, I think this is a good paper that the community will benefit from.