PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Message Passing for Quantum Chemistry

Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, George E. Dahl

arXiv:1704.01212v2cs.LG

TL;DR

Molecular neural models lacked broad empirical validation despite promising graph-based approaches and limited feature engineering success. This paper unifies these models as Message Passing Neural Networks, explores variants, and achieves state-of-the-art QM9 results, reaching chemical accuracy on 11 of 13 targets.

  • Problem

    Molecular neural architectures had limited empirical evidence of success despite the need to learn features directly from molecular graphs with appropriate invariances.

  • Method

    The paper formulates a general Message Passing Neural Network framework unifying existing graph models and explores novel message, update, output, and scalability variations.

  • Results

    MPNNs achieve state-of-the-art results on all 13 QM9 targets and predict DFT within chemical accuracy on 11 of 13 targets.

  • Takeaways & Limitations

    Appropriately designed MPNNs provide a useful inductive bias for molecular-property prediction, outperforming strong baselines without complicated feature engineering.

  • Takeaways & Limitations

    Generalizing MPNNs to larger molecule sizes remains challenging, particularly when spatial information creates size-dependent distance distributions and fully connected message passing.

Abstract

from arXiv · show

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

Chemical machine learning remains less mature than deep learning in other domains despite rapidly expanding quantum-chemistry and molecular-simulation data. The paper introduces MPNNs as a general graph-learning framework and reports strong QM9 results, including state-of-the-art performance on all 13 targets and chemical-accuracy predictions on 11 out of 13.

  • Motivation: Deep learning has advanced natural-language, audio, image, and video modeling, while machine-learning prediction of molecular and material properties remains in its infancy.The introduction contrasts broad deep-learning success with limited adoption and maturity in chemistry.
  • Motivation: Rapid growth in quantum-chemistry calculations, molecular-dynamics simulations, and high-throughput experiments creates data that classical techniques do not effectively use.The paper argues that more powerful methods require suitable inductive biases informed by atomic-system symmetries.
  • Framework: The paper introduces Message Passing Neural Networks as a general supervised-learning framework that learns graph features directly while remaining invariant to graph isomorphism.MPNNs abstract commonalities among promising neural graph models to clarify their relationships and enable novel variations.
  • Benchmark: QM9 provides the benchmark, comprising 130k molecules, 13 properties, and 13 regression tasks approximated by expensive DFT simulations.Performance is assessed against DFT error and the chemistry community’s chemical-accuracy target.
  • Results: State-of-the-art results are achieved on all 13 targets, while DFT is predicted within chemical accuracy on 11 out of 13 targets.The paper presents this as the performance of its developed MPNN on QM9.
  • Results: Several MPNNs predict DFT within chemical accuracy on 5 out of 13 targets using molecular topology alone, without spatial information.The paper also develops a training method that enlarges node representations without increasing computation time or memory.

2. Message Passing Neural Networks

MPNNs provide a common framework for graph neural models, using learned message passing over node and edge features followed by permutation-invariant graph readout. The framework exposes implementation choices and motivates application-focused efforts to improve modeling and computational efficiency.

  • Framework: MPNNs describe at least eight literature models on undirected graphs with node features xv and edge features evw.The formalism can also be extended to directed multigraphs.
  • Framework: Each MPNN alternates T message-passing time steps with a readout phase that computes a whole-graph feature vector.Message passing uses message functions Mt and vertex update functions Ut.
  • Framework: Learned differentiable functions Mt, Ut, and R define the model, while R must be permutation-invariant over node states to preserve graph-isomorphism invariance.Edge states can also be learned and updated analogously, although only Kearnes et al. (2016) used this idea among existing MPNNs.
  • Existing models: Duvenaud et al. (2015) cannot identify correlations between edge states and node states because its message vector separately sums connected nodes and edges.This limitation follows from the structure of its message-passing scheme.
  • Motivation: The authors pursue MPNNs in a practically important application to identify crucial implementation details and potentially expose the limits of the model family.They also present an MPNN modification intended to improve computational costs, addressing computation time as a limitation of existing approaches.

3. Related Work

Prior quantum-chemistry learning methods sought to accelerate or approximate DFT and quantum-mechanical calculations, but commonly relied on costly computations or hand-engineered representations with limited generalization and symmetry handling. Other graph-learning approaches instead construct canonical graph representations before classification.

  • Quantum-chemistry approximations: DFT offers speed–accuracy tradeoffs but remains too slow for large systems and has systematic and random errors relative to exact quantum mechanics.DFT scales as O(N_e^3), and a single 9-heavy-atom QM9 molecule takes around an hour on one Xeon E5-2660 core.
  • Quantum-chemistry approximations: Empirical potentials can be fast and accurate, but each new atomic composition requires a first-principles construction.
  • Prior machine-learning approaches: Neural networks and KRR have approximated DFT terms, molecular-dynamics energies and forces, or molecular atomization energies, without improving DFT efficiency in the exchange-correlation approach.Hu et al. used neural networks for the exchange-correlation potential; Behler and Parrinello modeled Silicon-melt energies and forces, while Rupp et al. inferred atomization energies.
  • Prior machine-learning approaches: Hand-engineered molecular features have intrinsic limitations, including difficulty with more than three atom species, failure on novel compositions, or a lack of graph-isomorphism invariance.Behler and Parrinello’s representation was invariant to graph isomorphism but struggled beyond three species and failed to generalize to novel compositions; Rupp et al.’s invariance had to be learned downstream.
  • Graphical-data methods: Some graphical-data methods preprocess inputs into canonical graph representations before passing them to a standard classifier.Examples include Niepert et al. (2016) and Rupp et al. (2012).

4. QM9 Dataset

The QM9 dataset contains about 134k small organic molecules with DFT-derived structures and diverse computed properties, providing a fundamental benchmark for molecular property prediction. These properties span atomization, vibrational, electronic, and electron-distribution categories.

  • Dataset: QM9 contains about 134k drug-like organic molecules composed of H, C, O, N, and F, with up to 9 heavy atoms.Each molecule has a reasonable low-energy DFT structure with atom positions available.
  • Properties: The dataset computes properties grouped into atomization energies, molecular vibrations, electronic states, and electron spatial distributions.These categories cover fundamental chemical quantities used to assess molecular property prediction.
  • Properties: Atomization properties include U0, U, H, and G, measuring energy required to break molecules under different temperatures and pressures.The reported quantities are atomization energy at 0K, atomization energy at room temperature, enthalpy of atomization, and free energy of atomization.
  • Properties: Vibrational properties include the highest fundamental frequency ω1 and zero point vibrational energy ZPVE.The corresponding units are cm−1 for ω1 and eV for ZPVE.
  • Properties: Electronic properties include HOMO, LUMO, and the electron energy gap ∆ε, defined as εHOMO −εLUMO.These describe electron-state energies and their difference.
  • Properties: Electron-distribution properties include electronic spatial extent ⟨R2⟩, dipole moment norm µ, and static polarizability norm α.Their units are Bohr2, Debye, and Bohr3, respectively.

5. MPNN Variants

The paper explores MPNN variants by changing message functions, propagation schemes, readouts, representations, and hyperparameters around GG-NN. It also introduces an efficient multi-copy propagation method to improve scalability while preserving permutation invariance.

  • Experiments began from GG-NN as a strong baseline, varying message and output functions, input representations, and hyperparameters.
  • The model represents molecules with atom features and bond or spatial-distance edge vectors, while initializing padded node states from atom features and using weight tying with GRU updates.
  • Message functions: The Edge Network extends GG-NN by mapping vector-valued edge features to d × d matrices, while Pair Message allows messages to depend on both source and destination states.
  • Propagation: Propagation variants add virtual edges between unconnected node pairs or a latent master node that provides global read-write communication at every message-passing step.
  • Readout functions: The set2set readout operates on projected final node-state and input-feature tuples and produces an order-invariant graph embedding before neural-network output.
  • Scalability: A multi-copy propagation scheme preserves permutation invariance while enabling larger hidden states with fewer parameters; for k = 8, n = 9 and d = 200 it achieved a factor of 2 speedup.Dense-graph message passing requires O(n^2d^2) floating point multiplications, motivating the factorized design.

6. Input Representation

The input representation combines atom-level features with graph encodings of bond types and optional interatomic distances. Distance information is represented either through discrete bins or vector-valued edge features, while explicit hydrogens increase graph size and slow training.

  • Atom and graph features: Atom features capture electronic properties and the bonds each atom participates in, with the complete feature list provided in table 1.The representation includes features for each atom in a molecule.
  • Hydrogen representation: Making hydrogens explicit nodes produces graphs with up to 29 nodes and significantly slows training because the graphs are larger.The alternative is to include hydrogen count as a node feature.
  • Atom and graph features: Without distance information, adjacency entries encode four discrete bond types: single, double, triple, or aromatic.These bond categories define the chemical graph representation.
  • Distance representations: Distance-aware matrix-multiply messages use 10 hand-chosen bins spanning [0, 2], eight uniform bins over [2, 6], and [6, ∞].The bins were selected from a histogram of all distances, yielding an adjacency alphabet of size 14.
  • Distance representations: Vector-valued edge messages represent each adjacency entry with 5 dimensions, including Euclidean distance and a four-dimensional one-hot encoding.The passage specifies the distance dimension and the one-hot encoding dimensions, but its final feature description is truncated.

7. Training

Models were trained with a uniform 50-trial hyperparameter search and a fixed SGD/ADAM optimization schedule. For QM-9, training used the remaining molecules after 10,000-sample validation and test splits, with normalized targets and MSE optimization.

  • Optimization and hyperparameters: Each model–target combination used a uniform random hyperparameter search with 50 trials, with T constrained to 3 ≤T ≤8 and M sampled from 1 ≤M ≤12.In practice, any T ≥3 worked.
  • Optimization and hyperparameters: Models used SGD with ADAM, batch size 20, and 3 million steps (540 epochs), with initial learning rates uniformly sampled between 1e−5 and 5e−4.
  • Dataset and objectives: QM-9 contains 130462 molecules; 10000 samples were randomly assigned to validation, 10000 to testing, and the remainder to training.Validation supported early stopping and model selection, while scores were reported on the test set.
  • Dataset and objectives: Targets were normalized to mean 0 and variance 1, training minimized mean squared error, and evaluation used mean absolute error.

8. Results

The results establish how error ratios relate to chemical accuracy and identify input, training, and architectural choices that improve MPNN performance on QM9. Experiments further favor spatial or long-range information, tower structures, tied weights with larger hidden dimensions, and edge-network messages.

  • Evaluation metric: Error ratios below 1 indicate chemical accuracy, with model MAE computed as Error Ratio × Chemical Accuracy.The reported ratio divides model MAE by the provided chemical-accuracy estimate for each target.
  • Model and training choices: Including complete edge features and explicit hydrogen nodes was important for several targets, while per-target training outperformed joint training by up to 40%.Complete edge features comprise bond type and spatial distance; joint training covered all 13 targets.
  • Training Without Spatial Information: When spatial information was omitted, augmenting MPNNs with long-range interactions substantially improved performance.The experiments compared sparse GG-NN graphs, virtual edges, a master node, and a modified graph construction.
  • Towers: The multi-tower structure improved generalization in comparisons against vanilla GG-NN, although combining towers with the edge-network message function did not further improve performance.The comparison used distance bins under both joint and per-target training; the authors suggest towers resemble an ensemble.
  • Additional Experiments: Tied weights with a larger hidden dimension were more effective than disabling weight tying, and edge-network messages outperformed pair messages on 11 of 13 jointly trained targets.The edge-network function achieved an average error ratio of 1.53 versus 3.98 for pair messages, so pair messages were not pursued further.

9. Conclusions and Future Work

MPNNs with suitable message, update, and output functions provide a useful inductive bias for molecular-property prediction, outperforming strong baselines without complicated feature engineering. Future work should improve scalability and generalization to larger graphs, especially when using spatial information.

  • Conclusions: MPNNs with appropriate message, update, and output functions outperform several strong baselines while eliminating complicated feature engineering.The models provide a useful inductive bias for predicting molecular properties.
  • Conclusions: Long-range interactions between graph nodes are important, either through a master node or the set2set output.The conclusion identifies these mechanisms as important for model performance.
  • Conclusions: The towers variation improves model scalability, although additional improvements are needed.The supplied passage states that towers makes the models more scalable but ends before specifying further improvements.
  • Future Work: Future MPNNs should generalize effectively to larger graphs or use benchmarks that expose generalization issues across graph sizes.Generalizing to larger molecule sizes is particularly challenging when spatial information is used.
  • Future Work: Spatial generalization is difficult because pairwise distance distributions depend heavily on the number of atoms.The passage also notes that the most successful spatial-information methods create a further challenge, but the supplied text is truncated.

10. Appendix

The appendix relates graph-convolution models to MPNNs, explains their weighted-neighbor message-passing interpretation, defines the QM9 target properties, and summarizes architectural findings on data efficiency, long-range interactions, and message functions.

  • Graph-convolution models: Graph-convolution models from Defferrard et al., Bruna et al., and Kipf & Welling can be interpreted as MPNNs and typically use simple message functions on larger graphs.These models generalize convolutions to graphs with N nodes.
  • Graph-convolution models: Kipf & Welling’s layer is an MPNN that adds self-loops and computes a weighted average of neighboring node states at each time step.The weights are scalar entries of the normalized adjacency-based matrix L.
  • Benchmark targets: The benchmark targets cover atomization energies, vibrational properties, electronic states, and spatial electron distributions.Examples include U0, U, H, G, ω1, ZPVE, HOMO, LUMO, ∆ε, ⟨R2⟩, µ, and α.
  • Architecture and scaling: 11k samples matched or exceeded the best baseline’s performance with 110k samples on both R2 and Omega.This result used the edge network + set2set output architecture.
  • Architecture and scaling: Allowing long-range interactions was important, while the edge network consistently outperformed the pair message function across most targets.The comparison included sparse graphs, virtual edges, master nodes, and set2set outputs.
Loading 1704.01212v2…