Title: Utilizing Edge Features in Graph Neural Networks via Variational Information Maximization

Abstract: Graph Neural Networks (GNNs) broadly follow the scheme that the representation vector of each node is updated recursively using the message from neighbor nodes, where the message of a neighbor is usually pre-processed with a parameterized transform matrix. To make better use of edge features, we propose the Edge Information maximized Graph Neural Network (EIGNN) that maximizes the Mutual Information (MI) between edge features and message passing channels. The MI is reformulated as a differentiable objective via a variational approach. We theoretically show that the newly introduced objective enables the model to preserve edge information, and empirically corroborate the enhanced performance of MI-maximized models across a broad range of learning tasks including regression on molecular graphs and relation prediction in knowledge graphs.

1. Introduction:

Many real-world datasets naturally come in the form of graphs, such as citation networks (Kipf & Welling, 2017) , social networks Hamilton et al. (2017) , knowledge graphs (Schlichtkrull et al., 2018) , molecular graphs (Scarselli et al., 2009; Duvenaud et al., 2015) etc., all of which consist of a number of nodes and edges equipped with their inherent features. Recently, impressive performance has been achieved in graph learning tasks with various forms of Graph Neural Networks (GNNs) (Zhou et al., 2018) . Compared to prior works, such as node2vec (Grover & Leskovec, 2016) , GNNs learn the state of a node by recursively aggregating messages from its neighbors: combining the graph structure with node features. Intuitively, edge features should play an important role in graph learning tasks. For example, chemical bonds in a molecule have a high impact on chemical properties of molecules, and edge features in knowledge graphs encode important relations between concepts, data, and entities. Our proposed method focuses on improving the usage of edge features in GNNs.

The expressive power of GNNs largely depends on how the message is passed between nodes. A widely adopted scheme is multiplying neighbor node states with a parameterized transform matrix before aggregation (Gilmer et al., 2017; Xu et al., 2019) . Despite tremendous success of GNNs, existing models do not exhaustively exploit the full potentials of edge features on graphs. For example, many GNNs such as GCN (Kipf & Welling, 2017) , ChebyNet (Defferrard et al., 2016) and GAT (Veličković et al., 2018) do not even consider categorized edge types. To utilize edge features in multi-relational graphs, RGCN (Schlichtkrull et al., 2018) proposes to learn a different transform matrix for each edge type, respectively. However, it does not generalize to edge features in continuous space. MPNN (Gilmer et al., 2017) introduces an edge network that takes edge feature vectors as input and outputs transform matrices, which are used to transform states of neighbor nodes. In principle, the MPNN framework can handle complex edge features. Yet, the lack of maximization of MI between edge and message channels implies that the MPNN may give an edge-independent transform matrix.

In this work, we aim to more efficiently exploit the full potentials of edge features from the perspective of training. We propose the Edge Information maximized Graph Neural Network (EIGNN) that maximizes the Mutual Information (MI) between edge features and the message passing channel which is parameterized as the transform matrix in the widely-accepted message passing framework (transformation and aggregation) (Gilmer et al., 2017; Xu et al., 2019) . Considering the challenge of computing the MI, we adopt a variational approach to reformulate it as an differentiable objective, which can be easily applied as a regularization term. We theoretically show that EIGNN can reduce information loss of edge features. Apart from demonstrating the impressive performance of EIGNN on extensive benchmarks of molecular graphs and knowledge graphs, we also analyze and attribute the enhanced effectiveness of EIGNN to the exploitation of edge features instead of the regularization effects. Notably, attribution analysis on molecular graphs show that EIGNN can capture domain knowledge without human interference.

Preliminaries Let G = (V, E) be a graph with node feature vectors x v ∈ R d for node v ∈ V and edge feature vectors e vw ∈ E for the edge connecting node v and w. In GNNs, the state of each node is updated recursively using neighbor nodes. Let N v be the set of neighbor nodes of v and h (l) v ∈ R d l be the hidden state of v at l-th layer, where d l is the dimension of the hidden layer. For simplicity of notation, we use a single d to denote the dimension such that h

(l) v ∈ R d . We also have h (0) v = x v at the input layer.

2 Related works

2.1. Relational Modeling in Graph Neural Networks:

Single-relational modeling. Many variants such as GCN (Kipf & Welling, 2017) , GAT (Veličković et al., 2018) , ChebyNet (Defferrard et al., 2016) , GraphSAGE (Hamilton et al., 2017) focus on learning node states. These models can assign weight to neighbors, but they can not handle various edge features. A typical neighborhood aggregation scheme is

EQUATION

where σ denotes an activation function, α vw can be a normalization constant or a learned attention coefficient (Veličković et al., 2018) . States of all neighbors are multiplied by the same trainable transform matrix W (l)

1 . Sometimes the self-connection is also treated in the same way, s.t., W

(l) 0 = W (l) 1 .

Multi-relational modeling. A simple strategy to handle multi-relational graphs is assigning each edge type with a separate transform matrix as presented in RGCN (Schlichtkrull et al., 2018 ) and adopted by GGNN (Li et al., 2016) and LNet (Liao et al., 2019) . RGCN updates node states according to the following scheme

EQUATION

where N r v is the collection of neighboring nodes of v with relation r ∈ R and α vw,r is a normalization constant similar as α vw in Eq. (1). Such a scheme faces challenge in handling edge features of continuous space. GGNN and LNet do not focus on the improvement of edge expressibility. GGNN introduces Gated Recurrent Unit (GRU) (Cho et al., 2014) and LNet focus on handling multi-scale connections.

Complex-relational modeling. The relation in a graph can be quite complex, expressed as a general feature vector e. MPNN (Gilmer et al., 2017) introduces an edge network which takes edge feature vectors as input and outputs transform matrices. A single edge network is shared in a MPNN model. The forward propagation is formalized as

EQUATION

where f : e → W denotes the edge network. Recently, some research works treat a multirelational problem as the complex-relational one by introducing a continuous edge embedding vector for each edge type, so as to handle increasing number of relations (Nathani et al., 2019) .

Although the MPNN architecture allows the usage of arbitrary edge features, this advantage is not utilized in practice. MPNN can actually learn an edge-independent transform matrix.

A GNN model that efficiently utilizes edge features is yet to emerge.

2.2. Readout functions:

After several forward propagations, GNNs yield final states of all nodes, which are suitable for node/edge classification or regression. For graph classification or regression tasks, we can apply a readout function (Ying et al., 2018; Vinyals et al., 2015) such that

EQUATION

) where h L v is the state of v at the last layer, R is the readout function that outputs a graph-level representation y, e.g., summing up the final node states, applying hierarchical pooling (Ying et al., 2018) or using the set2set model (Vinyals et al., 2015) .

3 Our method

3.1. The Usage of Mutual Information:

In probability theory and information theory, MI is a measure of mutual dependence between two random variables. Our method proposes to preserve edge information in GNNs, which is important in many real-world graph structures such as molecules -apart from node (atom) features, attributes of edges (bonds) are equally important for predicting properties of molecules. To this end, we maximize I(e; W ) -the MI between the edge feature vector e and the message passing channel, i.e., the transform matrix W which is used to transform neighbor node states in the forward propagation. Our method can be easily generalized to directly maximize the MI between edge features and the message itself in methods that do not explicitly have the transform matrix, e.g., the message from node w to node v can be expressed as f (h v , e vw , h w ) rather than f (e vw )h w , which is shown in Section 4.3.

An general principle of maximum MI is described for unsupervised learning task by (Linsker, 1988) and MI inspired objective functions have long been adopted in unsupervised learning (Bridle et al., 1992; Barber & Agakov, 2006; Veličković et al., 2019; Hjelm et al., 2019) , semi-supervised learning (Krause et al., 2010) and generative adversarial networks (Chen et al., 2016) . Specifically, DGI (Veličković et al., 2019 ) also applies MI to GNNs. DGI proposes to learn node-wise representations in an unsupervised manner by maximizing the MI between node representations and corresponding high-level summaries of graphs, using adversarial learning and negative sampling. The node representations may then be retrieved and used for downstream tasks, such as node classification. DGI can be used to pre-train GNNs, as demonstrated in Hu et al. (2019) . Our EIGNN also proposes information maximization but targets a completely different objective and adopts a quite different approach.

3.2. A Variational Approach to Maximize Mutual Information:

Computing I(e; W ) itself is intractable in practice, needless to say that training a model requires the derivative. Thus, we adopt a variational approach (Agakov, 2004) to reformulate I(e; W ) as a differentiable objective. We show that our objective is an approximated lower bound of I(e; W ) and notably, optimizing our objective does lead to maximizing I(e; W ). Following MPNN (Gilmer et al., 2017) , we use an edge network to parameterize the transform matrix W and relate it to edge features. Therefore, the prior

EQUATION

where δ(•) is the Dirac delta function. The posterior p(e|W ) is intractable, so we define a variational distribution q(e|W ), which can be obtained by defining a neural network g : W → e. Specifically, q(e|W ) substitutes to some distribution (such as Gaussian distribution) with parameter g(W ). In this way, f and g are similar to the probabilistic encoder and decoder in the Variational Auto-Encoder (VAE) (Kingma & Welling, 2013) . Then we can approximate I(e; W ) with a differentiable objective L I (f, g; e) as follows. Theorem 1. Let e be the edge feature vector, W be the transform matrix with conditional distribution p(W |e) specified by the probabilistic encoder f as shown in Eq. ( 5) and q(e|W ) be the variational distribution specified by the probabilistic decoder g, then we have

EQUATION

where L I (f, g; e) = log q(e|f (e)) and H(•) denotes the entropy.

Proof. Let D KL (• •) denote the KL-divergence, which should be nonnegative, then we have

I(e; W ) = H(e) -H(e|W ) = H(e) + E W ∼p(W ) [E e∼p(e|W ) [log p(e|W )]] = H(e) + E W ∼p(W ) [E e∼p(e|W ) [log p(e|W ) -log q(e|W ) + log q(e|W )]] = H(e) + E W ∼p(W ) [D KL (p(e|W ) q(e|W )) + E e∼p(e|W ) [log q(e|W )]] ≥ H(e) + E W ∼p(W ) [E e∼p(e|W ) [log q(e|W )]] = H(e) + E e∼p(e),W ∼p(W |e) [log q(e|W )] (a) = H(e) + E e∼p(e) [log q(e|f (e))]

where the equality (a) follows from Eq. ( 5).

According to Theorem 1, we can maximize the variational lower bound for I(e; W ). The bound becomes tight when the variational distribution q(e|W ) approaches the true posterior p(e|W ). Moreover, H(e) is a constant because the distribution of edge feature vector e is fixed for given graphs, hence we can equivalently maximize L I (f, g; e). We choose the widely accepted Gaussian distribution as the prior distribution for the probabilistic decoder g,

EQUATION

Then we have

L I (f, g; e) = log q(e|f (e)) = log N (e; g(f (e)), σ 2 I) = -λ e -g(f (e)) 2 2 (8)

where λ > 0 is a constant determined by σ and the dimension of e, taken as a tunable parameter. The following Theorem 2 shows that maximizing the objective in Eq. ( 8) does lead to the maximization of I(e; W ), hence enables the model to preserve edge information. Theorem 2. Assume the optimal solution of maximizing L I (f, g; e) is f and g , then f also maximizes I(e; W ).

Proof. Note that H(e) is a constant when the graphs are given. In information theory, we have H(g(f (e))) ≤ H(f (e)) ≤ H(e). Since f and g is the optimal solution of maximizing L I (f, g; e) presented in Eq. ( 8), it is not difficult to see that e = g (f (e)), ∀e ∈ E. In this case, the inequalities in Eq. 9 become equalities, i.e., H(g (f (e))) = H(f (e)) = H(e). Therefore, we have I(e; W ) = H(e), i.e., the maximum is attained in this case.

3.3. Edge Information Maximized Graph Neural Networks:

Our EIGNN is derived by implementing our MI objective in GNNs. As a concrete example, the forward propagation of our model follows the formulation in Eq. ( 3), where the dege network f : e → W is expressed as a multi-layer perceptron (MLP). According to theoretical analysis presented in Sec. 3.2 , we introduce another MLP g : W → e as the decoder.

For graph regression or classification tasks, the model outputs a prediction y for each graph G, which has label ŷ. Without MI maximization, we denote the vanilla loss as L 0 (ŷ, y; G). Common choice of L 0 includes Mean Square Error (MSE), Mean Absolute Error (MAE) and Cross Entropy (CE). For a graph G = (V, E), EIGNN maximizes L I (f, g; e) and minimizes L 0 (ŷ, y; G) using the following loss function

EQUATION

where E e∈E [•] denotes taking the mean over all edges in G = (V, E) and λ is the regularization parameter. When EIGNN is trained using mini-batches, L 0 (ŷ, y; G) is averaged over all graphs in the batch while L I (f, g; e) is averaged over all edges of all graphs in the batch.

Similarly, for relational prediction tasks in knowledge graphs, EIGNN directly yields nodelevel representations h v for each node v ∈ V and edge-level representations e for each relationship. The objective function of EIGNN can be derived from the translational scoring function Bordes et al. (2013) , which learns embedding such that for a given valid triple t vw = (h v , e vw , h w ) from the valid set S, the condition

d tvw = h v + e vw -h w ≈ 0 holds. Let L 0 = E tvw∈S E t vw ∈S max{d t vw -d tvw + γ, 0}

, where S denotes a set of invalid triples and γ is a margin hyper-parameter, EIGNN can be trained by minimizing the following loss,

EQUATION

4 Experiments

In this section, we first conduct experiments on a large quantum chemistry benchmark QM9, which is challenging for most baselines. Then we evaluate EIGNN on several useful molecule benchmarks and use attribution analysis to show that EIGNN increases the impact of edges and captures domain knowledge without human interference. Finally, we adopt our method to large-scale knowledge graphs and evaluate the performance on challenging relation prediction tasks using a wide variety of real-world datasets. All experimental results demonstrate a clear and substantial improvement of EIGNN over the state-of-the-art methods.

4.1. Quantum Chemistry:

QM9 (Ramakrishnan et al., 2014 ) is a large benchmark containing 134k molecules with 12 quantum chemistry regression properties, which have been show to be quite challenging for many GNNs (Gilmer et al., 2017) . Feature engineering of nodes and edges exactly follows (Gilmer et al., 2017) such that molecules are preprocessed as graphs according to atom features and bond features. We compare our EIGNN with nine state-of-the-art baselines which can be categorized into three groups according to the ability of handling edge features: i) GCN, ChebyNet, GAT and GIN (Xu et al., 2019) we use the same set2set (Vinyals et al., 2015) readout, which has been demonstrated to work well in (Gilmer et al., 2017) .

In Table 1 , we list regression results for all methods. We report individual MAE for each target in their original scale, averaged MAE (Avg. MAE) over 12 properties, and averaged normalized MAE (Avg. nMAE; averaged over normalized target properties since different targets have different units and ranges). Our EIGNN achieves the best performance for each metric and each target. Now we are ready to answer the following research questions. i) Are edge features important? Yes. The error has a trend of decreasing with increasing edge features. The comparison between sMPNN (using edge types) and MPNN (using edge types and distance) directly verifies the importance of edge features. It is also consistent with the expert knowledge that distances between pairwise atoms are closely related to quantum properties. For example, the smaller the distance between the two atoms, the stronger the bond is, and consequently a higher bond energy is associated with this atom pair. ii) Does the EIGNN work? Yes. EIGNN achieves the best performance on each target, outperforming the strong baseline MPNN. Moreover, the advantage of EIGNN over MPNN is consistent over 3 runs and the standard deviation on this task is quite small. Detailed results are shown in Table 4 of Appendix A. iii) How does the EIGNN work? Our MI objective is easily implemented on top of vanilla loss function. We have shown that our objective enables preserving of edge information. Fig. 1 demonstrates that regularization such as L 2 weight decay can increase training error while our objective does not. Moreover, the validation performance verifies that regularization itself does not reduce the validation error. Thus, the effectiveness of EIGNN is due to exploiting edge features rather than the regularization effect. We further run an ablation study where we concatenate edge features to node representations (i.e., MPNN+concat in Fig. 1 ) in message passing. Concatenation is unable to identify correlations between edges and nodes (Gilmer et al., 2017) and our results show that it slightly reduces the mean validation error but increases the variance.

4.2. More Molecule Benchmarks with Potential Applications:

We further evaluate EIGNN on three molecule benchmarks: Lipophilicity (Wu et al., 2018) , ESOL (Delaney, 2004) and FreeSolv (Mobley & Guthrie, 2014) . These datasets contain fewer molecules, and have potential usages in applications such as chemistry, drug discovery, and materials science. For example, the property lipophilicity is an important feature of drug molecules that affects both membrane permeability and solubility. The dataset Lipophilicity contains 4200 compounds. ESOL provides water solubility data for 1128 compounds. FreeSolv contains hydration free energy of 642 small molecules in water. We conduct graph regression experiments on these benchmarks. All datasets are split into training, validation and test according to a proportion of 0.8/0.1/0.1. MPNN and our EIGNN share the same architecture with 3 layers of message passing and 3 steps of set2set. We repeat each experiment 3 times with different random seeds. Results of testing root mean square error (RMSE) in Table 2 verify the effectiveness of our method. Our EIGNN outperforms MPNN on each dataset and each run. Detailed results for each run are presented in Appendix B.

Attribution analysis. To understand how our EIGNN reduces the regression error, we conduct attribution analysis, i.e., attributing the prediction of a deep network to its input features, which usually builds up on the standard gradient operator (Sundararajan et al., 2017) . For an output y, i.e., the prediction of a GNN, we define its sensitivity to an edge with features e as (2019) . Compared with MPNN, we can observe an increasing of overall edge attribution under our EIGNN and a decreasing of prediction error in both cases. Interestingly, the attribution under EIGNN is similar to the expert knowledge of chemists: halogen atoms such as {Cl, Br, I} and their bond with the carbon atom greatly effect the lipophilicity of a molecule Wilcken et al. (2013) , while atoms {O, N} also have high impact on the lipophilicity but usually in a negative way (Augustijns & Brewster, 2007) . In Fig. 2 (a), the attribution of the halogen bond C-Cl and the pair {O, N} under our EIGNN is much higher than the one under MPNN, which is consistent with the expert knowledge. In Fig. 2 

4.3. Predicting Relations in Knowledge Graphs:

In this subsection, we adopt EIGNN to tackle the problem of relation prediction in knowledge graphs (KGs), which entails predicting whether a given triple is valid or not. For example, a triple (London, capital of, United Kingdom) should be classified as valid or London should be predicted as the capital of United Kingdom. KGs represent human knowledge as a directed graph, and have been widely used in practical applications, such as semantic search, dialogue generation, question answering etc. Recovering missing relations in KGs have been a major task for practical usages of KGs. We evaluate our methods on three benchmark datasets, WN18RR (Dettmers et al., 2018) , FB15k-237 (Toutanova et al., 2015) and NELL-995 (Xiong et al., 2017) . Without the reversible relation problem (Dettmers et al., 2018) , WN18RR includes 11 relations scraped from WordNet for 40, 943 synsets. FB15k-237 is a subset of Freebase, and contains 14, 541 entities associated with 237 types of edge. NELL-995 is constructed from the 995 th iteration of NELL system, containing 75, 492 entities and 200 types of edge.

Table 3 : Experimental results on WN18RR, FB15K-237 and NELL-995 test sets. Hits@N values are in percentage. The best score is in bold and second best score is underlined.

Dataset WN18RR:

FB15K-237 NELL-995 Hit@N % Hit@N % Hit@N % MRR @1 @3 @10 MRR @1 @3 @10 MRR @1 @3 @10 DistMult A critical issue of applying Eq. ( 3) to KGs is that high-dimensional embedding vectors are required to distinguish massive amount of entities and relations, leading to a rapid growth in number of parameters in EIGNN. To address this issue, we adopt the architecture of (Nathani et al., 2019) and learn graph attention based embeddings that target relation prediction on KGs as follows,

EQUATION

Compared with Eq. ( 3), where m vw = f (e vw )h w , the above equation absorbs the transform matrix f (e vw ) into f (h

EQUATION

w ) and reduces the model parameters. In the following experiments, we implement f using a MLP as in previous experiments, and maximizes the MI between e vw and m vw by introducing another MLP with λ = 0.01. Multi-head attention is further introduced to stabilize the learning process and encapsulate more information about neighbors according to (Veličković et al., 2018) . After training EIGNN, ConvKB (Nguyen et al., 2018 ) is adopted as a regression function for a given triple by analyzing the global embedding properties across each dimension.

In the relation prediction task, the aim is to predict a triple (v, e vw , w) with v or w missing. We can generate a set of candidate triples for each missing entity v by randomly replacing it with an arbitrary one. Scores can be calculated by ConvKB for all triples, and we find the rank of a correct triple by sorting all scores in ascending order. Thus, the performance of relation prediction task can be evaluated by mean reciprocal rank (MRR) and the proportion of correct entities in the top N ranks (Hits@N) for N = 1, 3, and 10 ( Bordes et al., 2013) . We compare our EIGNN with seven state-of-the-art baselines focusing on this task: DistMult (Yang et al., 2014) , ComplEx (Trouillon et al., 2016) , ConvE (Dettmers et al., 2018) , TransE (Bordes et al., 2013 ), ConvKB (Nguyen et al., 2018) , RGCN (Schlichtkrull et al., 2018) and KBGAT (Nathani et al., 2019) . As shown in Table 3 , our EIGNN achieves the best performance for each metric on FB15K-237 and NELL-995, and achieves the best performance on WN18RR with Hit@3 and 10 metrics. The results of KBGAT are reproduced following the official implementation1 , and the results of other methods can be found in the previous peer-reviewed publications, i.e. (Nathani et al., 2019) .

5. Conclusions:

In this work, to make better use of edge features in GNNs, we proposed the edge information maximized graph neural network (EIGNN) that maximizes the mutual information between edge feature vectors and message passing channels. We reformulated the mutual information as a differentiable objective by adopting a variational approach. We have theoretically proved that our proposed objective enables EIGNN to preserve edge information and empirically evaluated EIGNN's performance on a variety of benchmarks incorporating an array of challenging molecular datasets and knowledge graphs. These results clearly manifested a substantial improvement of EIGNN over the prior state-of-the-art methods. Apart from demonstrating the impressive performance of EIGNN, we also showed that its effectiveness is due to exploitation of edge features instead of the regularization effect. Notably, attribution analysis on molecular graphs show that EIGNN can capture domain knowledge in an end-to-end fashion.

A Full Results on QM9 We present full results of quantum property regressions on QM9 with average performance and standard deviation in Table 4 . The (10 -3 ) in the parentheses indicates that the values in the corresponding raw of the table should multiply by 10 -3 . This is simply for clear presentation of the values. We also present a detailed descriptions on the target properties in Table 5 for your reference.

In our results, we directly report MAE instead of Error Ratio [(MAE)/(Chemical Accuracy) (Gilmer et al., 2017) ], because it is common to report MAE in terms of chemical unit. This practice has been widely adopted not only in computational chemistry but also in the machine learning community working on molecular graphs (Wu et al., 2018; Morris et al., 2019) . Still, we add Table 6 which contains the Error Ratio for MPNN and our EIGNN. In this comparison, we follow (Gilmer et al., 2017) and train models separately to predict each target. Our EIGNN consistently outperforms MPNN. 

C More Examples of Attribution:

In this section, we present more examples of attribution analysis on Lipophilicity. As a supplement for Fig. 2 in the main text, the observation here is similar. Compared with MPNN, we can observe an increasing of overall edge attribution under our EIGNN and a decreasing of prediction error in both cases. Notably, our EIGNN is able to capture the expert knowledge. 

:

https://github.com/deepakn97/relationPrediction

Figures and Tables:

(9) I(e; W ) is upper bounded by H(e), I(e; W ) = I(e; f (e)) = H(f (e)) -H(f (e)|e) = H(f (e)) ≤ H(e).

Figure 1: Ablation study. Training and validation error on QM9. The shadow area indicates mean ± std over 3 runs.

Figure 2: Attribution analysis. The color indicates the impact of an edge/atom on the output, i.e., the regression result. EIGNN i) increases the edge attribution, ii) reduces the prediction error and iii) can learn domain knowledge without human interference.

S e = | ∂y ∂e |, where |•| denotes the L 1 norm. Similarly, the sensitivity to an atom with features x is S x = | ∂y ∂x |. Then S e and S x are used as the metrics of attribution in our experiments. As an example, we show in Fig. 2 the attribution for two molecules in Lipophilicity: (a) Nc1nonc1C(=NO)Nc2ccc(F)c(Cl)c2 and (b) Oc1c2ncc3ccccc3c2nn1c4ccccc4 . Molecule in (a) has the potential to be used to treat, prevent and/or diagnose cancer Prinz et al.

(b), similarly, the attribution of atoms {O, N} and the bond C-O under EIGNN is much higher. More examples are presented in Appendix C.

(a) The molecule is CN [C@@H](C)C(= O)N [C@@H](C1CCCCC1)C(= O)N [C@H]2CCCN (CCc3ccc(F )cc3)C2.The attribution of {O, N} and the halogen atom F is higher under EIGNN. (b) The molecule is CC(C)N 1CCN [C@H](C1)C(= O)N 2CCN (CC2)C(= O)N c3ccc(Cl)c(Cl)c3. Our EIGNN successfully captures the importance of two critical halogen atoms Cl and several atoms N. (c) The molecule is COc1ccc(cc1)C(= O)N 2CCCC2 = O. The attribution of two atoms O at the top is

For a fair comparison, we repeat all experiments 3 times with different random seeds while during each run, all methods share the same random seed. We randomly choose 10k molecules for validation, 10k molecules for testing, and keep the rest for training. Each target property is normalized to zero mean and unit variance for training. Each model is trained to predict the 12 target properties simultaneously. λ is naively set to 1 for EIGNN. We use mean square error (MSE) loss to train the models for at most 300 epochs till convergence, and the performance is measured by mean absolute error (MAE). For LNet and GGNN, implementation of the readout function follows the original paper. While for all other models, Quantum property regressions for 12 targets and overall performance (top two raws) on QM9. We repeat all experiments 3 times with different random seeds and report the average performance. Full results with standard deviation are presented in Appendix A, e.g., for MPNN and EIGNN, we have Avg. nMAE 0.0398 ± 0.0002 and 0.0357 ± 0.0005.

Table content: <table><tr><td>Method</td><td>GCN</td><td>ChebyNet</td><td>GAT</td><td>GIN</td><td>RGCN</td><td>GGNN</td><td>LNet</td><td>sMPNN</td><td>MPNN</td><td>EIGNN</td></tr><tr><td>Avg. nMAE</td><td>0.135</td><td>0.121</td><td>0.137</td><td>0.100</td><td>0.102</td><td>0.099</td><td>0.099</td><td>0.089</td><td>0.040</td><td>0.036</td></tr><tr><td>Avg. MAE</td><td>5.306</td><td>4.303</td><td>5.470</td><td>3.480</td><td>3.817</td><td>3.661</td><td>3.653</td><td>3.161</td><td>0.693</td><td>0.633</td></tr><tr><td>mu</td><td>0.568</td><td>0.518</td><td>0.567</td><td>0.478</td><td>0.506</td><td>0.518</td><td>0.472</td><td>0.472</td><td>0.110</td><td>0.097</td></tr><tr><td>alpha</td><td>0.881</td><td>0.793</td><td>0.891</td><td>0.621</td><td>0.632</td><td>0.608</td><td>0.623</td><td>0.528</td><td>0.332</td><td>0.294</td></tr><tr><td>HOMO(10 -3 )</td><td>5.451</td><td>4.775</td><td>5.429</td><td>4.183</td><td>4.453</td><td>4.483</td><td>3.889</td><td>3.854</td><td>2.481</td><td>2.230</td></tr><tr><td>LUMO(10 -3 )</td><td>6.400</td><td>5.674</td><td>6.331</td><td>4.796</td><td>5.138</td><td>5.153</td><td>4.194</td><td>4.549</td><td>2.862</td><td>2.593</td></tr><tr><td>gap(10 -3 )</td><td>8.201</td><td>7.097</td><td>8.193</td><td>6.096</td><td>6.500</td><td>6.602</td><td>5.813</td><td>5.634</td><td>3.620</td><td>3.275</td></tr><tr><td>R2</td><td>53.56</td><td>41.95</td><td>54.52</td><td>34.65</td><td>40.10</td><td>39.68</td><td>35.27</td><td>33.49</td><td>6.064</td><td>5.646</td></tr><tr><td>ZPVE(10 -3 )</td><td>2.533</td><td>2.527</td><td>2.271</td><td>1.744</td><td>1.477</td><td>1.292</td><td>1.438</td><td>1.345</td><td>0.679</td><td>0.612</td></tr><tr><td>U0</td><td>2.042</td><td>1.984</td><td>2.290</td><td>1.422</td><td>1.059</td><td>0.697</td><td>1.806</td><td>0.791</td><td>0.416</td><td>0.357</td></tr><tr><td>U</td><td>2.042</td><td>1.984</td><td>2.290</td><td>1.422</td><td>1.059</td><td>0.697</td><td>1.755</td><td>0.791</td><td>0.416</td><td>0.357</td></tr><tr><td>H</td><td>2.042</td><td>1.984</td><td>2.290</td><td>1.422</td><td>1.059</td><td>0.697</td><td>1.796</td><td>0.791</td><td>0.416</td><td>0.357</td></tr><tr><td>G</td><td>2.042</td><td>1.984</td><td>2.290</td><td>1.422</td><td>1.059</td><td>0.696</td><td>1.778</td><td>0.791</td><td>0.416</td><td>0.357</td></tr><tr><td>Cv</td><td>0.473</td><td>0.420</td><td>0.479</td><td>0.309</td><td>0.317</td><td>0.315</td><td>0.312</td><td>0.262</td><td>0.134</td><td>0.121</td></tr></table>

Testing RMSE on Lipophilicity, ESOL and FreeSolv.

Table content: <table><tr><td>Dataset</td><td colspan="2">Lipophilicity</td><td colspan="2">ESOL</td><td colspan="2">FreeSolv</td></tr><tr><td>Method</td><td>MPNN</td><td>EIGNN</td><td>MPNN</td><td>EIGNN</td><td>MPNN</td><td>EIGNN</td></tr><tr><td>mean±std</td><td>0.678±0.042</td><td>0.653±0.025</td><td>0.805±0.064</td><td>0.776±0.071</td><td>1.398±0.081</td><td>1.273±0.137</td></tr></table>

Full results of quantum property regressions for 12 targets and overall performance (nMAE and MAE in top two raws) on QM9. We repeat all experiments 3 times with different random seeds and report the average performance and standard deviation. This is a supplement for Table1. in the main text.

Table content: <table><tr><td>Method</td><td>GCN</td><td>ChebyNet</td><td>GAT</td><td>GIN</td></tr><tr><td>Avg. nMAE</td><td>0.1350±0.0046</td><td>0.1206±0.0084</td><td>0.1367±0.0050</td><td>0.1001±0.0007</td></tr><tr><td>Avg. MAE</td><td>5.3063±0.1964</td><td>4.3032±0.4814</td><td>5.4698±0.2040</td><td>3.4799±0.0402</td></tr><tr><td>mu</td><td>0.5679±0.0078</td><td>0.5180±0.0131</td><td>0.5670±0.0102</td><td>0.4783±0.0041</td></tr><tr><td>alpha</td><td>0.8811±0.0308</td><td>0.7932±0.0778</td><td>0.8913±0.0314</td><td>0.6209±0.0028</td></tr><tr><td>HOMO(10 -3 )</td><td>5.4510±0.0790</td><td>4.7750±0.2040</td><td>5.4290±0.1660</td><td>4.1830±0.0370</td></tr><tr><td>LUMO(10 -3 )</td><td>6.4000±0.1230</td><td>5.6740±0.2670</td><td>6.3310±0.2390</td><td>4.7960±0.0520</td></tr><tr><td>gap(10 -3 )</td><td>8.2010±0.2150</td><td>7.0970±0.3620</td><td>8.1930±0.2910</td><td>6.0960±0.0420</td></tr><tr><td>R2</td><td>53.563±1.0319</td><td>41.950±4.8289</td><td>54.519±1.5992</td><td>34.647±0.2167</td></tr><tr><td>ZPVE(10 -3 )</td><td>2.5330±0.1070</td><td>2.5270±0.3560</td><td>2.2710±0.1450</td><td>1.7440±0.0100</td></tr><tr><td>U0</td><td>2.0422±0.3281</td><td>1.9842±0.2035</td><td>2.2899±0.2000</td><td>1.4215±0.0857</td></tr><tr><td>U</td><td>2.0422±0.3281</td><td>1.9842±0.2035</td><td>2.2899±0.2000</td><td>1.4215±0.0857</td></tr><tr><td>H</td><td>2.0422±0.3281</td><td>1.9842±0.2035</td><td>2.2899±0.2000</td><td>1.4215±0.0857</td></tr><tr><td>G</td><td>2.0423±0.3281</td><td>1.9842±0.2036</td><td>2.2899±0.2000</td><td>1.4215±0.0857</td></tr><tr><td>Cv</td><td>0.4730±0.0209</td><td>0.4199±0.0499</td><td>0.4787±0.0164</td><td>0.3093±0.0035</td></tr><tr><td>Method</td><td>RGCN</td><td>GGNN</td><td>LNet</td><td>sMPNN</td></tr><tr><td>Avg. nMAE</td><td>0.1021±0.0016</td><td>0.0992±0.0013</td><td>0.0992±0.0061</td><td>0.0888±0.0014</td></tr><tr><td>Avg. MAE</td><td>3.8175±0.0605</td><td>3.6608±0.0723</td><td>3.6527±0.3417</td><td>3.1610±0.0697</td></tr><tr><td>mu</td><td>0.5056±0.0048</td><td>0.5179±0.0076</td><td>0.4717±0.0063</td><td>0.4718±0.0096</td></tr><tr><td>alpha</td><td>0.6321±0.0145</td><td>0.6077±0.0092</td><td>0.6225±0.0508</td><td>0.5278±0.0106</td></tr><tr><td>HOMO(10 -3 )</td><td>4.4530±0.1290</td><td>4.4830±0.0650</td><td>3.8889±0.1617</td><td>3.8540±0.0410</td></tr><tr><td>LUMO(10 -3 )</td><td>5.1380±0.1210</td><td>5.1530±0.0890</td><td>4.1935±0.2205</td><td>4.5490±0.0810</td></tr><tr><td>gap(10 -3 )</td><td>6.5000±0.1390</td><td>6.6020±0.1400</td><td>5.8132±0.6456</td><td>5.6340±0.0570</td></tr><tr><td>R2</td><td>40.102±0.7428</td><td>39.685±0.8212</td><td>35.275±3.0531</td><td>33.489±0.6562</td></tr><tr><td>ZPVE(10 -3 )</td><td>1.4770±0.0090</td><td>1.2920±0.0340</td><td>1.4376±0.0769</td><td>1.3450±0.0260</td></tr><tr><td>U0</td><td>1.0589±0.0231</td><td>0.6969±0.0413</td><td>1.8058±0.2533</td><td>0.7914±0.0446</td></tr><tr><td>U</td><td>1.0589±0.0231</td><td>0.6966±0.0418</td><td>1.7555±0.2196</td><td>0.7914±0.0446</td></tr><tr><td>H</td><td>1.0589±0.0231</td><td>0.6974±0.0408</td><td>1.7964±0.2428</td><td>0.7914±0.0446</td></tr><tr><td>G</td><td>1.0589±0.0231</td><td>0.6961±0.0421</td><td>1.7780±0.2458</td><td>0.7914±0.0446</td></tr><tr><td>Cv</td><td>0.3170±0.0152</td><td>0.3146±0.0125</td><td>0.3124±0.0303</td><td>0.2625±0.0040</td></tr><tr><td>Method</td><td>MPNN</td><td>EIGNN</td><td/><td/></tr><tr><td>Avg. nMAE</td><td>0.0398±0.0002</td><td>0.0357±0.0005</td><td/><td/></tr><tr><td>Avg. MAE</td><td>0.6929±0.0212</td><td>0.6331±0.0298</td><td/><td/></tr><tr><td>mu</td><td>0.1095±0.0014</td><td>0.0974±0.0026</td><td/><td/></tr><tr><td>alpha</td><td>0.3318±0.0026</td><td>0.2939±0.0054</td><td/><td/></tr><tr><td>HOMO(10 -3 )</td><td>2.4810±0.0200</td><td>2.2300±0.0310</td><td/><td/></tr><tr><td>LUMO(10 -3 )</td><td>2.8620±0.0370</td><td>2.5930±0.0440</td><td/><td/></tr><tr><td>gap(10 -3 )</td><td>3.6200±0.0180</td><td>3.2750±0.0520</td><td/><td/></tr><tr><td>R2</td><td>6.0637±0.2511</td><td>5.6464±0.3098</td><td/><td/></tr><tr><td>ZPVE(10 -3 )</td><td>0.6790±0.0140</td><td>0.6120±0.0170</td><td/><td/></tr><tr><td>U0</td><td>0.4164±0.0225</td><td>0.3574±0.0100</td><td/><td/></tr><tr><td>U</td><td>0.4164±0.0225</td><td>0.3575±0.0100</td><td/><td/></tr><tr><td>H</td><td>0.4164±0.0225</td><td>0.3574±0.0100</td><td/><td/></tr><tr><td>G</td><td>0.4164±0.0225</td><td>0.3575±0.0101</td><td/><td/></tr><tr><td>Cv</td><td>0.1339±0.0013</td><td>0.1208±0.0027</td><td/><td/></tr></table>

Regression targets on QM9.

Table content: <table><tr><td>Target property</td><td>Description</td><td>Unit</td></tr><tr><td>mu</td><td>Dipole moment</td><td>D</td></tr><tr><td>alpha</td><td>Isotropic polarizability</td><td>a 3 0</td></tr><tr><td>HOMO</td><td>Highest occupied molecular orbital energy</td><td>E h</td></tr><tr><td>LUMO</td><td>Lowest unoccupied molecular orbital energy</td><td>E h</td></tr><tr><td>gap</td><td>Gap between HOMO and LUMO</td><td>E h</td></tr><tr><td>R2</td><td>Electronic spatial extent</td><td>a 2 0</td></tr><tr><td>ZPVE</td><td>Zero point vibrational energy</td><td>E h</td></tr><tr><td>U0</td><td>Internal energy at 0K</td><td>E h</td></tr><tr><td>U</td><td>Internal energy at 298.15K</td><td>E h</td></tr><tr><td>H</td><td>Enthalpy at 298.15K</td><td>E h</td></tr><tr><td>G</td><td>Free energy at 298.15K</td><td>E h</td></tr><tr><td>Cv</td><td>Heat capavity at 298.15K</td><td>cal molK</td></tr></table>

Error Ratio [(MAE)/(Chemical Accuracy)(Gilmer et al., 2017)] on QM9. Note that the energy values of {U0, U, H, G} are per molecule rather than per atom. Following(Gilmer et al., 2017), models are separately trained on each target.In this section, we present experimental results on Lipophilicity, ESOL and FreeSolv with detailed results for each run. The results in Table7verify that our EIGNN consistently outperforms MPNN.

Table content: <table><tr><td>Target</td><td>MPNN</td><td>EIGNN</td><td>Improvement (%)</td></tr><tr><td>mu</td><td>0.87</td><td>0.80</td><td>8.24</td></tr><tr><td>alpha</td><td>2.64</td><td>2.44</td><td>7.57</td></tr><tr><td>HOMO</td><td>1.54</td><td>1.39</td><td>10.2</td></tr><tr><td>LUMO</td><td>1.31</td><td>1.28</td><td>2.13</td></tr><tr><td>gap</td><td>2.20</td><td>1.97</td><td>10.4</td></tr><tr><td>R2</td><td>1.00</td><td>0.72</td><td>28.4</td></tr><tr><td>ZPVE</td><td>5.31</td><td>4.43</td><td>16.5</td></tr><tr><td>U0</td><td>64.3</td><td>35.8</td><td>44.4</td></tr><tr><td>U</td><td>56.6</td><td>23.6</td><td>58.3</td></tr><tr><td>H</td><td>77.4</td><td>35.2</td><td>54.5</td></tr><tr><td>G</td><td>46.6</td><td>34.4</td><td>26.1</td></tr><tr><td>Cv</td><td>1.69</td><td>1.79</td><td>-6.22</td></tr><tr><td colspan="4">B Full Results on Lipophilicity, ESOL and FreeSolv</td></tr></table>

Testing RMSE on Lipophilicity, ESOL and FreeSolv. This is a supplement for Table2in the main text.

Table content: <table><tr><td>Dataset</td><td colspan="2">Lipophilicity</td><td colspan="2">ESOL</td><td colspan="2">FreeSolv</td></tr><tr><td>Seed</td><td>MPNN</td><td>EIGNN</td><td>MPNN</td><td>EIGNN</td><td>MPNN</td><td>EIGNN</td></tr><tr><td>0</td><td>0.718</td><td>0.676</td><td>0.770</td><td>0.718</td><td>1.396</td><td>1.109</td></tr><tr><td>1</td><td>0.696</td><td>0.664</td><td>0.750</td><td>0.733</td><td>1.299</td><td>1.265</td></tr><tr><td>2</td><td>0.620</td><td>0.619</td><td>0.894</td><td>0.876</td><td>1.499</td><td>1.443</td></tr><tr><td>mean±std</td><td>0.678±0.042</td><td>0.653±0.025</td><td>0.805±0.064</td><td>0.776±0.071</td><td>1.398±0.081</td><td>1.273±0.137</td></tr></table>

Peer review (iclr2020):
This paper proposed an auxiliary loss based on mutual information for graph neural network. Such loss is to maximize the mutual information between edge representation and corresponding edge feature in GNN ‘message passing’ function. By assuming a gaussian distribution of edge feature given edge representation, the training can be done efficiently with tractable density. Experiments on molecule regression and knowledge graph completion show better performance than MPNN. Overall the paper is written in a clear way which is easy to follow. The idea of using mutual information as some kind of regularization is also interesting. However, there are some concerns I have with the paper: Regarding formulation 1. The derivation up to Eq(8) looks fine to me, where the assumptions are reasonable. However from Eq(8) one can see this is reduced to an ‘auto-encoder’ type of regularization, where one can have a trivial solution for reconstruction -- the identity network, when the hidden dimension is larger than input dimension. And in this paper, dimension of W should always be larger than the dimension of e (for example, in molecules e should be low dimension vector with the bond type, distance, etc., while W should have dimension that matches the node embeddings). I think the original loss (i.e., the supervised MSE, cross entropy etc) would help a bit with such degenerated case, but it is possible that the learned f(e) contains both identity mapping (or equivalent) and the representation that contributes to original loss. 2. Actually I’m also not sure if I get the motivation here. If one needs to do this regularization for edges, why don’t we consider this auxiliary loss for node embeddings as well? As in molecules, atoms have more interesting features than bonds, which should account more if the mutual information loss is needed. Regarding experiment 1. In Figure 1, the training loss of EIGNN is better than MPNN. This is a bit counterintuitive to me, as I think the auxiliary is a kind of regularization -- which might help with generalization but not necessarily the training loss. 2. The original paper of MPNN reports the relative MAE. Is it possible to report the results using the same metric as previous paper? It would make the comparison more consistent -- though showing the improvement in current way is not too bad. 3. I think one simple ablation study would to concat the edge feature directly inside the ‘message passing’ procedure, or have some ‘residual’ type of connection for edge features.