# neural_message_passing_for_quantum_chemistry__ca3f5c05.pdf

Neural Message Passing for Quantum Chemistry

Justin Gilmer 1 Samuel S. Schoenholz 1 Patrick F. Riley 2 Oriol Vinyals 3 George E. Dahl 1

Abstract Supervised learning on molecules has incredible potential to be useful in chemistry, drug discovery, and materials science. Luckily, several promising and closely related neural network models invariant to molecular symmetries have already been described in the literature. These models learn a message passing algorithm and aggregation procedure to compute a function of their entire input graph. At this point, the next step is to find a particularly effective variant of this general approach and apply it to chemical prediction benchmarks until we either solve them or reach the limits of the approach. In this paper, we reformulate existing models into a single common framework we call Message Passing Neural Networks (MPNNs) and explore additional novel variations within this framework. Using MPNNs we demonstrate state of the art results on an important molecular property prediction benchmark; these results are strong enough that we believe future work should focus on datasets with larger molecules or more accurate ground truth labels.

1. Introduction

The past decade has seen remarkable success in the use of deep neural networks to understand and translate natural language (Wu et al., 2016), generate and decode complex audio signals (Hinton et al., 2012), and infer features from real-world images and videos (Krizhevsky et al., 2012). Although chemists have applied machine learning to many problems over the years, predicting the properties of molecules and materials using machine learning (and especially deep learning) is still in its infancy. To date, most research applying machine learning to chemistry tasks (Hansen et al., 2015; Huang & von Lilienfeld, 2016;

1Google Brain 2Google 3Google Deep Mind. Correspondence to: Justin Gilmer <gilmer@google.com>, George E. Dahl <gdahl@google.com>.

Proceedings of the 34 th International Conference on Machine Learning, Sydney, Australia, PMLR 70, 2017. Copyright 2017 by the author(s).

DFT 103 seconds

Message Passing Neural Net

10 2 seconds

Figure 1. A Message Passing Neural Network predicts quantum properties of an organic molecule by modeling a computationally expensive DFT calculation.

Rupp et al., 2012; Rogers & Hahn, 2010; Montavon et al., 2012; Behler & Parrinello, 2007; Schoenholz et al., 2016) has revolved around feature engineering. While neural networks have been applied in a variety of situations (Merkwirth & Lengauer, 2005; Micheli, 2009; Lusci et al., 2013; Duvenaud et al., 2015), they have yet to become widely adopted. This situation is reminiscent of the state of image models before the broad adoption of convolutional neural networks and is due, in part, to a dearth of empirical evidence that neural architectures with the appropriate inductive bias can be successful in this domain.

Recently, large scale quantum chemistry calculation and molecular dynamics simulations coupled with advances in high throughput experiments have begun to generate data at an unprecedented rate. Most classical techniques do not make effective use of the larger amounts of data that are now available. The time is ripe to apply more powerful and flexible machine learning methods to these problems, assuming we can find models with suitable inductive biases. The symmetries of atomic systems suggest neural networks that operate on graph structured data and are invariant to graph isomorphism might also be appropriate for molecules. Sufficiently successful models could someday help automate challenging chemical search problems in drug discovery or materials science.

In this paper, our goal is to demonstrate effective machine learning models for chemical prediction problems

Neural Message Passing for Quantum Chemistry

that are capable of learning their own features from molecular graphs directly and are invariant to graph isomorphism. To that end, we describe a general framework for supervised learning on graphs called Message Passing Neural Networks (MPNNs) that simply abstracts the commonalities between several of the most promising existing neural models for graph structured data, in order to make it easier to understand the relationships between them and come up with novel variations. Given how many researchers have published models that fit into the MPNN framework, we believe that the community should push this general approach as far as possible on practically important graph problems and only suggest new variations that are well motivated by applications, such as the application we consider here: predicting the quantum mechanical properties of small organic molecules (see task schematic in figure 1).

In general, the search for practically effective machine learning (ML) models in a given domain proceeds through a sequence of increasingly realistic and interesting benchmarks. Here we focus on the QM9 dataset as such a benchmark (Ramakrishnan et al., 2014). QM9 consists of 130k molecules with 13 properties for each molecule which are approximated by an expensive1 quantum mechanical simulation method (DFT), to yield 13 corresponding regression tasks. These tasks are plausibly representative of many important chemical prediction problems and are (currently) difficult for many existing methods. Additionally, QM9 also includes complete spatial information for the single low energy conformation of the atoms in the molecule that was used in calculating the chemical properties. QM9 therefore lets us consider both the setting where the complete molecular geometry is known (atomic distances, bond angles, etc.) and the setting where we need to compute properties that might still be defined in terms of the spatial positions of atoms, but where only the atom and bond information (i.e. graph) is available as input. In the latter case, the model must implicitly fit something about the computation used to determine a low energy 3D conformation and hopefully would still work on problems where it is not clear how to compute a reasonable 3D conformation.

When measuring the performance of our models on QM9, there are two important benchmark error levels. The first is the estimated average error of the DFT approximation to nature, which we refer to as DFT error. The second, known as chemical accuracy, is a target error that has been established by the chemistry community. Estimates of DFT error and chemical accuracy are provided for each of the 13 targets in Faber et al. (2017). One important goal of this line of research is to produce a model which can achieve chemical accuracy with respect to the

1By comparison, the inference time of the neural networks discussed in this work is 300k times faster.

true targets as measured by an extremely precise experiment. The dataset containing the true targets on all 134k molecules does not currently exist. However, the ability to fit the DFT approximation to within chemical accuracy would be an encouraging step in this direction. For all 13 targets, achieving chemical accuracy is at least as hard as achieving DFT error. In the rest of this paper when we talk about chemical accuracy we generally mean with respect to our available ground truth labels.

In this paper, by exploring novel variations of models in the MPNN family, we are able to both achieve a new state of the art on the QM9 dataset and to predict the DFT calculation to within chemical accuracy on all but two targets. In particular, we provide the following key contributions:

We develop an MPNN which achieves state of the art results on all 13 targets and predicts DFT to within chemical accuracy on 11 out of 13 targets.

We develop several different MPNNs which predict DFT to within chemical accuracy on 5 out of 13 targets while operating on the topology of the molecule alone (with no spatial information as input).

We develop a general method to train MPNNs with larger node representations without a corresponding increase in computation time or memory, yielding a substantial savings over previous MPNNs for high dimensional node representations.

We believe our work is an important step towards making well-designed MPNNs the default for supervised learning on modestly sized molecules. In order for this to happen, researchers need to perform careful empirical studies to find the proper way to use these types of models and to make any necessary improvements to them, it is not sufficient for these models to have been described in the literature if there is only limited accompanying empirical work in the chemical domain. Indeed convolutional neural networks existed for decades before careful empirical work applying them to image classification (Krizhevsky et al., 2012) helped them displace SVMs on top of handengineered features for a host of computer vision problems.

2. Message Passing Neural Networks

There are at least eight notable examples of models from the literature that we can describe using our Message Passing Neural Networks (MPNN) framework. For simplicity we describe MPNNs which operate on undirected graphs G with node features xv and edge features evw. It is trivial to extend the formalism to directed multigraphs. The forward pass has two phases, a message passing phase and a readout phase. The message passing phase runs for T

Neural Message Passing for Quantum Chemistry

time steps and is defined in terms of message functions Mt and vertex update functions Ut. During the message passing phase, hidden states ht v at each node in the graph are updated based on messages mt+1 v according to

w N(v) Mt(ht v, ht w, evw) (1)

ht+1 v = Ut(ht v, mt+1 v ) (2)

where in the sum, N(v) denotes the neighbors of v in graph G. The readout phase computes a feature vector for the whole graph using some readout function R according to

ˆy = R({h T v | v G}). (3)

The message functions Mt, vertex update functions Ut, and readout function R are all learned differentiable functions. R operates on the set of node states and must be invariant to permutations of the node states in order for the MPNN to be invariant to graph isomorphism. In what follows, we define previous models in the literature by specifying the message function Mt, vertex update function Ut, and readout function R used. Note one could also learn edge features in an MPNN by introducing hidden states for all edges in the graph ht evw and updating them analogously to equations 1 and 2. Of the existing MPNNs, only Kearnes et al. (2016) has used this idea.

Convolutional Networks for Learning Molecular Fingerprints, Duvenaud et al. (2015)

The message function used is M(hv, hw, evw) = (hw, evw) where (., .) denotes concatenation. The vertex update function used is Ut(ht v, mt+1 v ) = σ(Hdeg(v) t mt+1 v ), where σ is the sigmoid function, deg(v) is the degree of vertex v and HN t is a learned matrix for each time step t and vertex degree N. R has skip connections to all previous

hidden states ht v and is equal to f

v,t softmax(Wtht v)

where f is a neural network and Wt are learned readout matrices, one for each time step t. This message passing scheme may be problematic since the resulting message vector is mt+1 v = (P ht w, P evw) , which separately sums over connected nodes and connected edges. It follows that the message passing implemented in Duvenaud et al. (2015) is unable to identify correlations between edge states and node states.

Gated Graph Neural Networks (GG-NN), Li et al. (2016)

The message function used is Mt(ht v, ht w, evw) = Aevwht w, where Aevw is a learned matrix, one for each edge label e (the model assumes discrete edge types). The update function is Ut = GRU(ht v, mt+1 v ), where GRU is the Gated

Recurrent Unit introduced in Cho et al. (2014). This work used weight tying, so the same update function is used at each time step t. Finally,

v V σ i(h(T ) v , h0 v) j(h(T ) v ) (4)

where i and j are neural networks, and denotes elementwise multiplication.

Interaction Networks, Battaglia et al. (2016)

This work considered both the case where there is a target at each node in the graph, and where there is a graph level target. It also considered the case where there are node level effects applied at each time step, in such a case the update function takes as input the concatenation (hv, xv, mv) where xv is an external vector representing some outside influence on the vertex v. The message function M(hv, hw, evw) is a neural network which takes the concatenation (hv, hw, evw). The vertex update function U(hv, xv, mv) is a neural network which takes as input the concatenation (hv, xv, mv). Finally, in the case where there is a graph level output, R = f( P

v G h T v ) where f is

a neural network which takes the sum of the final hidden states h T v . Note the original work only defined the model for T = 1.

Molecular Graph Convolutions, Kearnes et al. (2016)

This work deviates slightly from other MPNNs in that it introduces edge representations et vw which are updated during the message passing phase. The message function used for node messages is M(ht v, ht w, et vw) = et vw. The vertex update function is Ut(ht v, mt+1 v ) = α(W1(α(W0ht v), mt+1 v )) where (., .) denotes concatenation, α is the Re LU activation and W1, W0 are learned weight matrices. The edge state update is defined by et+1 vw = U t(et vw, ht v, ht w) = α(W4(α(W2, et vw), α(W3(ht v, ht w)))) where the Wi are also learned weight matrices.

Deep Tensor Neural Networks, Sch utt et al. (2017)

The message from w to v is computed by

Mt = tanh W fc((W cfht w + b1) (W dfevw + b2))

where W fc, W cf, W df are matrices and b1, b2 are bias vectors. The update function used is Ut(ht v, mt+1 v ) = ht v + mt+1 v . The readout function passes each node independently through a single hidden layer neural network and sums the outputs, in particular

v NN(h T v ).

Laplacian Based Methods, Bruna et al. (2013); Defferrard et al. (2016); Kipf & Welling (2016)

Neural Message Passing for Quantum Chemistry

These methods generalize the notion of the convolution operation typically applied to image datasets to an operation that operates on an arbitrary graph G with a real valued adjacency matrix A. The operations defined in Bruna et al. (2013); Defferrard et al. (2016) result in message functions of the form Mt(ht v, ht w) = Ct vwht w, where the matrices Ct vw are parameterized by the eigenvectors of the graph laplacian L, and the learned parameters of the model. The vertex update function used is Ut(ht v, mt+1 v ) = σ(mt+1 v ) where σ is some pointwise non-linearity (such as Re LU).

The Kipf & Welling (2016) model results in a message function Mt(ht v, ht w) = cvwht w where cvw = (deg(v)deg(w)) 1/2 Avw. The vertex update function is U t v(ht v, mt+1 v ) = Re LU(W tmt+1 v ). For the exact expressions for the Ct vw and the derivation of the reformulation of these models as MPNNs, see the supplementary material.

2.1. Moving Forward

Given how many instances of MPNNs have appeared in the literature, we should focus on pushing this general family as far as possible in a specific application of substantial practical importance. This way we can determine the most crucial implementation details and potentially reach the limits of these models to guide us towards future modeling improvements.

One downside of all of these approaches is computation time. Recent work has adapted the GG-NN architecture to larger graphs by passing messages on only subsets of the graph at each time step (Marino et al., 2016). In this work we also present a MPNN modification that can improve the computational costs.

3. Related Work

Although in principle quantum mechanics lets us compute the properties of molecules, the laws of physics lead to equations that are far too difficult to solve exactly. Therefore scientists have developed a hierarchy of approximations to quantum mechanics with varying tradeoffs of speed and accuracy, such as Density Functional Theory (DFT) with a variety of functionals (Becke, 1993; Hohenberg & Kohn, 1964), the GW approximation (Hedin, 1965), and Quantum Monte-Carlo (Ceperley & Alder, 1986). Despite being widely used, DFT is simultaneously still too slow to be applied to large systems (scaling as O(N 3 e ) where Ne is the number of electrons) and exhibits systematic as well as random errors relative to exact solutions to Schr odinger s equation. For example, to run the DFT calculation on a single 9 heavy atom molecule in QM9 takes around an hour on a single core of a Xeon E5-2660 (2.2 GHz) using a version of Gaussian G09 (ES64L-G09Rev D.01) (Bing et al., 2017). For a 17 heavy atom molecule, computation time is

up to 8 hours. Empirical potentials have been developed, such as the Stillinger-Weber potential (Stillinger & Weber, 1985), that are fast and accurate but must be created from scratch, from first principles, for every new composition of atoms.

Hu et al. (2003) used neural networks to approximate a particularly troublesome term in DFT called the exchange correlation potential to improve the accuracy of DFT. However, their method fails to improve upon the efficiency of DFT and relies on a large set of ad hoc atomic descriptors.

Two more recent approaches by Behler & Parrinello (2007) and Rupp et al. (2012) attempt to approximate solutions to quantum mechanics directly without appealing to DFT. In the first case single-hidden-layer neural networks were used to approximate the energy and forces for configurations of a Silicon melt with the goal of speeding up molecular dynamics simulations. The second paper used Kernel Ridge Regression (KRR) to infer atomization energies over a wide range of molecules. In both cases hand engineered features were used (symmetry functions and the Coulomb matrix, respectively) that built physical symmetries into the input representation. Subsequent papers have replaced KRR by a neural network.

Both of these lines of research used hand engineered features that have intrinsic limitations. The work of Behler & Parrinello (2007) used a representation that was manifestly invariant to graph isomorphism, but has difficulty when applied to systems with more than three species of atoms and fails to generalize to novel compositions. The representation used in Rupp et al. (2012) is not invariant to graph isomorphism. Instead, this invariance must be learned by the downstream model through dataset augmentation.

In addition to the eight MPNNs discussed in Section 2 there have been a number of other approaches to machine learning on graphical data which take advantage of the symmetries in a number of ways. One such family of approaches define a preprocessing step which constructs a canonical graph representation which can then be fed into into a standard classifier. Examples in this family include Niepert et al. (2016) and Rupp et al. (2012). Finally Scarselli et al. (2009) define a message passing process on graphs which is run until convergence, instead of for a finite number of time steps as in MPNNs.

4. QM9 Dataset

To investigate the success of MPNNs on predicting chemical properties, we use the publicly available QM9 dataset (Ramakrishnan et al., 2014). Molecules in the dataset consist of Hydrogen (H), Carbon (C), Oxygen (O), Nitrogen (N), and Flourine (F) atoms and contain up to 9 heavy (non Hydrogen) atoms. In all, this results in about 134k drug-

Neural Message Passing for Quantum Chemistry

like organic molecules that span a wide range of chemistry. For each molecule DFT is used to find a reasonable low energy structure and hence atom positions are available. Additionally a wide range of interesting and fundamental chemical properties are computed. Given how fundamental some of the QM9 properties are, it is hard to believe success on more challenging chemical tasks will be possible if we can t make accurate statistical predictions for the properties computed in QM9.

We can group the different properties we try to predict into four broad categories. First, we have four properties related to how tightly bound together the atoms in a molecule are. These measure the energy required to break up the molecule at different temperatures and pressures. These include the atomization energy at 0K, U0 (e V), atomization energy at room temperature, U (e V), enthalpy of atomization at room temperature, H (e V), and free energy of atomization, G (e V).

Next there are properties related to fundamental vibrations of the molecule, including the highest fundamental vibrational frequency ω1 (cm 1) and the zero point vibrational energy (ZPVE) (e V).

Additionally, there are a number of properties that concern the states of the electrons in the molecule. They include the energy of the electron in the highest occupied molecular orbital (HOMO) εHOMO (e V), the energy of the lowest unoccupied molecular orbital (LUMO) εLUMO (e V), and the electron energy gap ( ε (e V)). The electron energy gap is simply the difference εHOMO εLUMO.

Finally, there are several measures of the spatial distribution of electrons in the molecule. These include the electronic spatial extent R2 (Bohr2), the norm of the dipole moment µ (Debye), and the norm of static polarizability α (Bohr3). For a more detailed description of these properties, see the supplementary material.

5. MPNN Variants

We began our exploration of MPNNs around the GG-NN model which we believe to be a strong baseline. We focused on trying different message functions, output functions, finding the appropriate input representation, and properly tuning hyperparameters.

For the rest of the paper we use d to denote the dimension of the internal hidden representation of each node in the graph, and n to denote the number of nodes in the graph. Our implementation of MPNNs in general operates on directed graphs with a separate message channel for incoming and outgoing edges, in which case the incoming message mv is the concatenation of min v and mout v , this was also used in Li et al. (2016). When we apply this to undirected

chemical graphs we treat the graph as directed, where each original edge becomes both an incoming and outgoing edge with the same label. Note there is nothing special about the direction of the edge, it is only relevant for parameter tying. Treating undirected graphs as directed means that the size of the message channel is 2d instead of d.

The input to our MPNN model is a set of feature vectors for the nodes of the graph, xv, and an adjacency matrix A with vector valued entries to indicate different bonds in the molecule as well as pairwise spatial distance between two atoms. We experimented as well with the message function used in the GG-NN family, which assumes discrete edge labels, in which case the matrix A has entries in a discrete alphabet of size k. The initial hidden states h0 v are set to be the atom input feature vectors xv and are padded up to some larger dimension d. All of our experiments used weight tying at each time step t, and a GRU (Cho et al., 2014) for the update function as in the GG-NN family.

5.1. Message Functions

Matrix Multiplication: We started with the message function used in GG-NN which is defined by the equation M(hv, hw, evw) = Aevwhw.

Edge Network: To allow vector valued edge features we propose the message function M(hv, hw, evw) = A(evw)hw where A(evw) is a neural network which maps the edge vector evw to a d d matrix.

Pair Message: One property that the matrix multiplication rule has is that the message from node w to node v is a function only of the hidden state hw and the edge evw. In particular, it does not depend on the hidden state ht v. In theory, a network may be able to use the message channel more efficiently if the node messages are allowed to depend on both the source and destination node. Thus we also tried using a variant on the message function as described in (Battaglia et al., 2016). Here the message from w to v along edge e is mwv = f (ht w, ht v, evw) where f is a neural network.

When we apply the above message functions to directed graphs, there are two separate functions used, M in and an M out. Which function is applied to a particular edge evw depends on the direction of that edge.

5.2. Virtual Graph Elements

We explored two different ways to change how the messages are passed throughout the model. The simplest modification involves adding a separate virtual edge type for pairs of nodes that are not connected. This can be implemented as a data preprocessing step and allows information to travel long distances during the propagation phase.

Neural Message Passing for Quantum Chemistry

We also experimented with using a latent master node, which is connected to every input node in the graph with a special edge type. The master node serves as a global scratch space that each node both reads from and writes to in every step of message passing. We allow the master node to have a separate node dimension dmaster, as well as separate weights for the internal update function (in our case a GRU). This allows information to travel long distances during the propagation phase. It also, in theory, allows additional model capacity (e.g. large values of dmaster) without a substantial hit in performance, as the complexity of the master node model is O(|E|d2 + nd2 master).

5.3. Readout Functions

We experimented with two readout functions. First is the readout function used in GG-NN, which is defined by equation 4. Second is a set2set model from Vinyals et al. (2015). The set2set model is specifically designed to operate on sets and should have more expressive power than simply summing the final node states. This model first applies a linear projection to each tuple (h T v , xv) and then takes as input the set of projected tuples T = {(h T v , xv)}. Then, after M steps of computation, the set2set model produces a graph level embedding q t which is invariant to the order of the of the tuples T. We feed this embedding q t through a neural network to produce the output.

5.4. Multiple Towers

One issue with MPNNs is scalability. In particular, a single step of the message passing phase for a dense graph requires O(n2d2) floating point multiplications. As n or d get large this can be computationally expensive. To address this issue we break the d dimensional node embeddings ht v into k different d/k dimensional embeddings ht,k v and run a propagation step on each of the k copies separately to get temporary embeddings { ht+1,k v , v G}, using separate message and update functions for each copy. The k temporary embeddings of each node are then mixed together according to the equation

ht,1 v , ht,2 v , . . . , ht,k v = g ht,1 v , ht,2 v , . . . , ht,k v (5)

where g denotes a neural network and (x, y, . . .) denotes concatenation, with g shared across all nodes in the graph. This mixing preserves the invariance to permutations of the nodes, while allowing the different copies of the graph to communicate with each other during the propagation phase. This can be advantageous in that it allows larger hidden states for the same number of parameters, which yields a computational speedup in practice. For example, when the message function is matrix multiplication (as in GG-NN) a propagation step of a single copy takes O n2(d/k)2 time, and there are k copies, therefore the

Table 1. Atom Features

Feature Description

Atom type H, C, N, O, F (one-hot) Atomic number Number of protons (integer) Acceptor Accepts electrons (binary) Donor Donates electrons (binary) Aromatic In an aromatic system (binary) Hybridization sp, sp2, sp3 (one-hot or null) Number of Hydrogens (integer)

overall time complexity is O n2d2/k , with some additional overhead due to the mixing network. For k = 8, n = 9 and d = 200 we see a factor of 2 speedup in inference time over a k = 1, n = 9, and d = 200 architecture. This variation would be most useful for larger molecules, for instance molecules from GDB-17 (Ruddigkeit et al., 2012).

6. Input Representation

There are a number of features available for each atom in a molecule which capture both properties of the electrons in the atom as well as the bonds that the atom participates in. For a list of all of the features see table 1. We experimented with making the hydrogen atoms explicit nodes in the graph (as opposed to simply including the count as a node feature), in which case graphs have up to 29 nodes. Note that having larger graphs significantly slows training time, in this case by a factor of roughly 10. For the adjacency matrix there are three edge representations used depending on the model.

Chemical Graph: In the abscence of distance information, adjacency matrix entries are discrete bond types: single, double, triple, or aromatic.

Distance bins: The matrix multiply message function assumes discrete edge types, so to include distance information we bin bond distances into 10 bins, the bins are obtained by uniformly partitioning the interval [2, 6] into 8 bins, followed by adding a bin [0, 2] and [6, ]. These bins were hand chosen by looking at a histogram of all distances. The adjacency matrix then has entries in an alphabet of size 14, indicating bond type for bonded atoms and distance bin for atoms that are not bonded. We found the distance for bonded atoms to be almost completely determined by bond type.

Raw distance feature: When using a message function which operates on vector valued edges, the entries of the adjacency matrix are then 5 dimensional, where the first dimension indicates the euclidean distance between the pair of atoms, and the remaining four are a one-hot encoding of

Neural Message Passing for Quantum Chemistry

the bond type.

7. Training

Each model and target combination was trained using a uniform random hyper parameter search with 50 trials. T was constrained to be in the range 3 T 8 (in practice, any T 3 works). The number of set2set computations M was chosen from the range 1 M 12. All models were trained using SGD with the ADAM optimizer (Kingma & Ba (2014)), with batch size 20 for 3 million steps ( 540 epochs). The initial learning rate was chosen uniformly between 1e 5 and 5e 4. We used a linear learning rate decay that began between 10% and 90% of the way through training and the initial learning rate l decayed to a final learning rate l F, using a decay factor F in the range [.01, 1].

The QM-9 dataset has 130462 molecules in it. We randomly chose 10000 samples for validation, 10000 samples for testing, and used the rest for training. We use the validation set to do early stopping and model selection and we report scores on the test set. All targets were normalized to have mean 0 and variance 1. We minimize the mean squared error between the model output and the target, although we evaluate mean absolute error.

In all of our tables we report the ratio of the mean absolute error (MAE) of our models with the provided estimate of chemical accuracy for that target. Thus any model with error ratio less than 1 has achieved chemical accuracy for that target. In the supplementary material we list the chemical accuracy estimates for each target, these are the same estimates that were given in Faber et al. (2017). In this way, the MAE of our models can be calculated as (Error Ratio) (Chemical Accuracy). Note, unless otherwise indicated, all tables display result of models trained individually on each target (as opposed to training one model to predict all 13).

We performed numerous experiments in order to find the best possible MPNN on this dataset as well as the proper input representation. In our experiments, we found that including the complete edge feature vector (bond type, spatial distance) and treating hydrogen atoms as explicit nodes in the graph to be very important for a number of targets. We also found that training one model per target consistently outperformed jointly training on all 13 targets. In some cases the improvement was up to 40%. Our best MPNN variant used the edge network message function, set2set output, and operated on graphs with explicit hydrogens. We were able to further improve performance on the test set by ensembling the predictions of the five models with lowest validation error.

In table 2 we compare the performance of our best MPNN variant (denoted with enn-s2s) and the corresponding ensemble (denoted with enn-s2s-ens5) with the previous state of the art on this dataset as reported in Faber et al. (2017). For clarity the error ratios of the best non-ensemble models are shown in bold. This previous work performed a comparison study of several existing ML models for QM9 and we have taken care to use the same train, validation, and test split. These baselines include 5 different hand engineered molecular representations, which then get fed through a standard, off-the-shelf classifier. These input representations include the Coulomb Matrix (CM, Rupp et al. (2012)), Bag of Bonds (Bo B, Hansen et al. (2015)), Bonds Angles, Machine Learning (BAML, Huang & von Lilienfeld (2016)), Extended Connectivity Fingerprints (ECPF4, Rogers & Hahn (2010)), and Projected Histograms (HDAD, Faber et al. (2017)) representations. In addition to these hand engineered features we include two existing baseline MPNNs, the Molecular Graph Convolutions model (GC) from Kearnes et al. (2016), and the original GG-NN model Li et al. (2016) trained with distance bins. Overall, our new MPNN achieves chemical accuracy on 11 out of 13 targets and state of the art on all 13 targets.

Training Without Spatial Information: We also experimented in the setting where spatial information is not included in the input. In general, we find that augmenting the MPNN with some means of capturing long range interactions between nodes in the graph greatly improves performance in this setting. To demonstrate this we performed 4 experiments, one where we train the GG-NN model on the sparse graph, one where we add virtual edges, one where we add a master node, and one where we change the graph level output to a set2set output. The error ratios averaged across the 13 targets are shown in table 3. Overall, these three modifications help on all 13 targets, and the Set2Set output achieves chemical accuracy on 5 out of 13 targets. For more details, consult the supplementary material. The experiments shown tables 3 and 4 were run with a partial charge feature as a node input. This feature is an output of the DFT calculation and thus could not be used in an applied setting. The state of art numbers we report in table 2 do not use this feature.

Towers: Our original intent in developing the towers variant was to improve training time, as well as to allow the model to be trained on larger graphs. However, we also found some evidence that the multi-tower structure improves generalization performance. In table 4 we compare GG-NN + towers + set2set output vs a baseline GGNN + set2set output when distance bins are used. We do this comparison in both the joint training regime and when training one model per target. The towers model outperforms the baseline model on 12 out of 13 targets in both

Neural Message Passing for Quantum Chemistry

Table 2. Comparison of Previous Approaches (left) with MPNN baselines (middle) and our methods (right) Target BAML BOB CM ECFP4 HDAD GC GG-NN DTNN enn-s2s enn-s2s-ens5

mu 4.34 4.23 4.49 4.82 3.34 0.70 1.22 - 0.30 0.20 alpha 3.01 2.98 4.33 34.54 1.75 2.27 1.55 - 0.92 0.68 HOMO 2.20 2.20 3.09 2.89 1.54 1.18 1.17 - 0.99 0.74 LUMO 2.76 2.74 4.26 3.10 1.96 1.10 1.08 - 0.87 0.65 gap 3.28 3.41 5.32 3.86 2.49 1.78 1.70 - 1.60 1.23 R2 3.25 0.80 2.83 90.68 1.35 4.73 3.99 - 0.15 0.14 ZPVE 3.31 3.40 4.80 241.58 1.91 9.75 2.52 - 1.27 1.10 U0 1.21 1.43 2.98 85.01 0.58 3.02 0.83 - 0.45 0.33 U 1.22 1.44 2.99 85.59 0.59 3.16 0.86 - 0.45 0.34 H 1.22 1.44 2.99 86.21 0.59 3.19 0.81 - 0.39 0.30 G 1.20 1.42 2.97 78.36 0.59 2.95 0.78 .842 0.44 0.34 Cv 1.64 1.83 2.36 30.29 0.88 1.45 1.19 - 0.80 0.62 Omega 0.27 0.35 1.32 1.47 0.34 0.32 0.53 - 0.19 0.15 Average 2.17 2.08 3.37 53.97 1.35 2.59 1.36 - 0.68 0.52

Table 3. Models Trained Without Spatial Information Model Average Error Ratio GG-NN 3.47 GG-NN + Virtual Edge 2.90 GG-NN + Master Node 2.62 GG-NN + set2set 2.57

Table 4. Towers vs Vanilla GG-NN (no explicit hydrogen) Model Average Error Ratio GG-NN + joint training 1.92 towers8 + joint training 1.75 GG-NN + individual training 1.53 towers8 + individual training 1.37

individual and joint target training. We believe the benefit of towers is that it resembles training an ensemble of models. Unfortunately, our attempts so far at combining the towers and edge network message function have failed to further improve performance, possibly because the combination makes training more difficult. Further training details, and error ratios on all targets can be found in the supplementary material.

Additional Experiments: In preliminary experiments, we tried disabling weight tying across different time steps. However, we found that the most effective way to increase performance was to tie the weights and use a larger hidden dimension d. We also early on found the pair message function to perform worse than the edge network function. This included a toy pathfinding problem which was originally

2As reported in Sch utt et al. (2017). The model was trained on a different train/test split with 100k training samples vs 110k used in our experiments.

designed to benefit from using pair messages. Also, when trained jointly on the 13 targets the edge network function outperforms pair message on 11 out of 13 targets, and has an average error ratio of 1.53 compared to 3.98 for pair message. Given the difficulties with training this function we did not pursue it further. For performance on smaller sized training sets, consult the supplementary material.

9. Conclusions and Future Work

Our results show that MPNNs with the appropriate message, update, and output functions have a useful inductive bias for predicting molecular properties, outperforming several strong baselines and eliminating the need for complicated feature engineering. Moreover, our results also reveal the importance of allowing long range interactions between nodes in the graph with either the master node or the set2set output. The towers variation makes these models more scalable, but additional improvements will be needed to scale to much larger graphs.

An important future direction is to design MPNNs that can generalize effectively to larger graphs than those appearing in the training set or at least work with benchmarks designed to expose issues with generalization across graph sizes. Generalizing to larger molecule sizes seems particularly challenging when using spatial information. First of all, the pairwise distance distribution depends heavily on the number of atoms. Second, our most successful ways of using spatial information create a fully connected graph where the number of incoming messages also depends on the number of nodes. To address the second issue, we believe that adding an attention mechanism over the incoming message vectors could be an interesting direction to explore.

Neural Message Passing for Quantum Chemistry

Acknowledgements

We would like to thank Lukasz Kaiser, Geoffrey Irving, Alex Graves, and Yujia Li for helpful discussions. Thank you to Adrian Roitberg for pointing out an issue with the use of partial charges in an earlier version of this paper.

Battaglia, Peter, Pascanu, Razvan, Lai, Matthew, Rezende, Danilo Jimenez, and Kavukcuoglu, Koray. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems, pp. 4502 4510, 2016.

Becke, Axel D. Density-functional thermochemistry. iii. the role of exact exchange. The Journal of Chemical Physics, 98(7):5648 5652, 1993. doi: 10.1063/1. 464913. URL http://dx.doi.org/10.1063/1. 464913.

Behler, J org and Parrinello, Michele. Generalized neural-network representation of high-dimensional potential-energy surfaces. Phys. Rev. Lett., 98: 146401, Apr 2007. doi: 10.1103/Phys Rev Lett.98. 146401. URL http://link.aps.org/doi/10. 1103/Phys Rev Lett.98.146401.

Bing, Huang, von Lillenfeld, O. Anatole, and Bakowies, Dirk. personal communication, 2017.

Bruna, Joan, Zaremba, Wojciech, Szlam, Arthur, and Le Cun, Yann. Spectral networks and locally connected networks on graphs. ar Xiv preprint ar Xiv:1312.6203, 2013.

Ceperley, David and Alder, B. Quantum monte carlo. Science, 231, 1986.

Cho, Kyunghyun, Van Merri enboer, Bart, Bahdanau, Dzmitry, and Bengio, Yoshua. On the properties of neural machine translation: Encoder-decoder approaches. ar Xiv preprint ar Xiv:1409.1259, 2014.

Defferrard, Micha el, Bresson, Xavier, and Vandergheynst, Pierre. Convolutional neural networks on graphs with fast localized spectral filtering. In Advances in Neural Information Processing Systems, pp. 3837 3845, 2016.

Duvenaud, David K, Maclaurin, Dougal, Iparraguirre, Jorge, Bombarell, Rafael, Hirzel, Timothy, Aspuru Guzik, Al an, and Adams, Ryan P. Convolutional networks on graphs for learning molecular fingerprints. In Advances in neural information processing systems, pp. 2224 2232, 2015.

Faber, Felix, Hutchison, Luke, Huang, Bing, Gilmer, Justin, Schoenholz, Samuel S., Dahl, George E., Vinyals, Oriol, Kearnes, Steven, Riley, Patrick F., and von

Lilienfeld, O. Anatole. Fast machine learning models of electronic and energetic properties consistently reach approximation errors better than dft accuracy. https://arxiv.org/abs/1702.05532, 2017.

Hansen, Katja, Biegler, Franziska, Ramakrishnan, Raghunathan, Pronobis, Wiktor, von Lilienfeld, O. Anatole, Mller, Klaus-Robert, and Tkatchenko, Alexandre. Machine learning predictions of molecular properties: Accurate many-body potentials and nonlocality in chemical space. The journal of physical chemistry letters, 6(12):2326 2331, 2015. doi: 10.1021/acs.jpclett. 5b00831. URL http://dx.doi.org/10.1021/ acs.jpclett.5b00831.

Hedin, Lars. New method for calculating the one-particle green s function with application to the electron-gas problem. Phys. Rev., 139:A796 A823, Aug 1965. doi: 10.1103/Phys Rev.139.A796. URL http://link. aps.org/doi/10.1103/Phys Rev.139.A796.

Hinton, Geoffrey, Deng, Li, Yu, Dong, Dahl, George E., Mohamed, Abdel-rahman, Jaitly, Navdeep, Senior, Andrew, Vanhoucke, Vincent, Nguyen, Patrick, Sainath, Tara N, et al. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IEEE Signal Processing Magazine, 29 (6):82 97, 2012.

Hohenberg, P. and Kohn, W. Inhomogeneous electron gas. Phys. Rev., 136:B864 B871, Nov 1964. doi: 10.1103/ Phys Rev.136.B864. URL http://link.aps.org/ doi/10.1103/Phys Rev.136.B864.

Hu, Li Hong, Wang, Xiu Jun, Wong, Lai Ho, and Chen, Guan Hua. Combined first-principles calculation and neural-network correction approach for heat of formation. The Journal of Chemical Physics, 119(22):11501 11507, 2003.

Huang, Bing and von Lilienfeld, O. Anatole. Communication: Understanding molecular representations in machine learning: The role of uniqueness and target similarity. The Journal of Chemical Physics, 145(16): 161102, 2016. doi: 10.1063/1.4964627. URL http: //dx.doi.org/10.1063/1.4964627.

Kearnes, Steven, Mc Closkey, Kevin, Berndl, Marc, Pande, Vijay, and Riley, Patrick. Molecular graph convolutions: Moving beyond fingerprints. Journal of Computer-Aided Molecular Design, 30(8):595 608, 2016.

Kingma, Diederik and Ba, Jimmy. Adam: A method for stochastic optimization. ar Xiv preprint ar Xiv:1412.6980, 2014.

Neural Message Passing for Quantum Chemistry

Kipf, T. N. and Welling, M. Semi-Supervised Classification with Graph Convolutional Networks. Ar Xiv eprints, September 2016.

Krizhevsky, Alex, Sutskever, Ilya, and Hinton, Geoffrey E. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, pp. 1097 1105, 2012.

Li, Yujia, Tarlow, Daniel, Brockschmidt, Marc, and Zemel, Richard. Gated graph sequence neural networks. ICLR, 2016.

Lusci, Alessandro, Pollastri, Gianluca, and Baldi, Pierre. Deep architectures and deep learning in chemoinformatics: the prediction of aqueous solubility for drug-like molecules. Journal of chemical information and modeling, 53(7):1563 1575, 2013.

Marino, Kenneth, Salakhutdinov, Ruslan, and Gupta, Abhinav. The more you know: Using knowledge graphs for image classification. ar Xiv preprint ar Xiv:1612.04844, 2016.

Merkwirth, Christian and Lengauer, Thomas. Automatic generation of complementary descriptors with molecular graph networks. Journal of chemical information and modeling, 45(5):1159 1168, 2005.

Micheli, Alessio. Neural network for graphs: A contextual constructive approach. IEEE Transactions on Neural Networks, 20(3):498 511, 2009.

Montavon, Gr egoire, Hansen, Katja, Fazli, Siamac, Rupp, Matthias, Biegler, Franziska, Ziehe, Andreas, Tkatchenko, Alexandre, von Lilienfeld, O. Anatole, and M uller, Klaus-Robert. Learning invariant representations of molecules for atomization energy prediction. In Advances in Neural Information Processing Systems, pp. 440 448, 2012.

Niepert, Mathias, Ahmed, Mohamed, and Kutzkov, Konstantin. Learning convolutional neural networks for graphs. In Proceedings of the 33rd annual international conference on machine learning. ACM, 2016.

Ramakrishnan, Raghunathan, Dral, Pavlo O, Rupp, Matthias, and Von Lilienfeld, O Anatole. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1, 2014.

Rogers, David and Hahn, Mathew. Extended-connectivity fingerprints. Journal of chemical information and modeling, 50(5):742 754, 2010.

Ruddigkeit, Lars, Van Deursen, Ruud, Blum, Lorenz C, and Reymond, Jean-Louis. Enumeration of 166 billion organic small molecules in the chemical universe

database gdb-17. Journal of chemical information and modeling, 52(11):2864 2875, 2012.

Rupp, Matthias, Tkatchenko, Alexandre haand M uller, Klaus-Robert, and von Lilienfeld, O. Anatole. Fast and accurate modeling of molecular atomization energies with machine learning. Physical review letters, 108 (5):058301, Jan 2012. URL http://dx.doi.org/ 10.1103/Phys Rev Lett.108.058301.

Scarselli, Franco, Gori, Marco, Tsoi, Ah Chung, Hagenbuchner, Markus, and Monfardini, Gabriele. The graph neural network model. IEEE Transactions on Neural Networks, 20(1):61 80, 2009.

Schoenholz, Samuel S., Cubuk, Ekin D., Sussman, Daniel M, Kaxiras, Efthimios, and Liu, Andrea J. A structural approach to relaxation in glassy liquids. Nature Physics, 2016.

Sch utt, Kristof T, Arbabzadah, Farhad, Chmiela, Stefan, M uller, Klaus R, and Tkatchenko, Alexandre. Quantumchemical insights from deep tensor neural networks. Nature Communications, 8, 2017.

Stillinger, Frank H. and Weber, Thomas A. Computer simulation of local order in condensed phases of silicon. Phys. Rev. B, 31:5262 5271, Apr 1985. doi: 10.1103/ Phys Rev B.31.5262. URL http://link.aps.org/ doi/10.1103/Phys Rev B.31.5262.

Vinyals, Oriol, Bengio, Samy, and Kudlur, Manjunath. Order matters: Sequence to sequence for sets. ar Xiv preprint ar Xiv:1511.06391, 2015.

Wu, Yonghui, Schuster, Mike, Chen, Zhifeng, Le, Quoc V., Norouzi, Mohammad, Macherey, Wolfgang, Krikun, Maxim, Cao, Yuan, Gao, Qin, Macherey, Klaus, et al. Google s neural machine translation system: Bridging the gap between human and machine translation. ar Xiv preprint ar Xiv:1609.08144, 2016.