[Paper Review] Neural Message Passing with Edge Updates for Predicting Properties of Molecules and Materials

Peter Bjørn Jørgensen, Karsten W. Jacobsen+1 more|arXiv (Cornell University)|Jun 8, 2018

Machine Learning in Materials Science21 references56 citations

TL;DR

The paper extends SchNet with an edge update network that conditions edge messages on the receiving atom, improving formation energy predictions across QM9, Materials Project, and OQMD, with k-nearest neighbor graphs outperforming fixed cutoffs.

ABSTRACT

Neural message passing on molecular graphs is one of the most promising methods for predicting formation energy and other properties of molecules and materials. In this work we extend the neural message passing model with an edge update network which allows the information exchanged between atoms to depend on the hidden state of the receiving atom. We benchmark the proposed model on three publicly available datasets (QM9, The Materials Project and OQMD) and show that the proposed model yields superior prediction of formation energies and other properties on all three datasets in comparison with the best published results. Furthermore we investigate different methods for constructing the graph used to represent crystalline structures and we find that using a graph based on K-nearest neighbors achieves better prediction accuracy than using maximum distance cutoff or the Voronoi tessellation graph.

Motivation & Objective

• Motivate fast, accurate prediction of molecular and material properties as an alternative to expensive DFT calculations.
• Extend neural message passing to include edge updates that depend on the receiving atom.
• Evaluate on QM9, The Materials Project, and Open Quantum Materials Database (OQMD).
• Investigate graph construction choices (cutoff vs. K-nearest vs. Voronoi tessellation) for crystalline structures.

Proposed method

• Embed atomic numbers to initial node states and expand interatomic distances with radial basis functions.
• Introduce an edge update network E_t that makes edge features depend on the sending and receiving atom states (e_vw^t) and the previous edge state (e_vw^{t}).
• Define a message function M_t that uses the updated edge features and neighbor states to aggregate messages.
• Update node states with a residual graph update S_t(h_v^t, m_v^{t+1}).
• Use a readout R that sums (or averages) atom-wise contributions to predict properties.
• Benchmark against SchNet and enn-s2s on QM9, Materials Project, and OQMD; compare graph construction approaches (constant cutoff, K-nearest, Voronoi).

Experimental results

Research questions

• RQ1 Does incorporating an edge update network that conditions edge messages on both sender and receiver atom states improve predictive accuracy for molecular and material properties?
• RQ2 How does the choice of graph construction (cutoff distance, K-nearest neighbors, Voronoi tessellation) affect predictive performance for crystalline materials?
• RQ3 Can the edge-updated model achieve state-of-the-art results on QM9, Materials Project, and OQMD datasets?
• RQ4 What properties (e.g., formation energy) benefit most from edge updates and how do results compare to existing models?

Key findings

• The edge update network yields significant improvements in prediction accuracy over baseline SchNet on all three datasets.
• QM9, Materials Project, and OQMD show lower mean absolute errors with the proposed model compared to V-RF and SchNet ( Table 2 ).
• On QM9, the proposed model achieves lower MAEs for 9 out of 12 properties and matches the remaining 3 (Table 3).
• For crystalline materials, a K-nearest neighbor graph (e.g., K=24) with 13.7 meV/atom MAE achieved the best performance, outperforming fixed cutoff and Voronoi approaches (Figure 4).
• The edge updates are most influential in conditioning the interaction on the receiving atom, as seen in the learned filter visualizations (first-layer dependence on receiving vs. sending atom).