Chemical Properties from Graph Neural Network-Predicted Electron Densities
Ethan M. Sunshine, Muhammed Shuaibi, Zachary W. Ulissi, John R. Kitchin, “Chemical Properties from Graph Neural Network-Predicted Electron Densities”, The Journal of Physical Chemistry C 127(48), 23459-23466 (2023). https://doi.org/10.1021/acs.jpcc.3c06157
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Note on the source text: the extracted text is the arXiv preprint (arXiv:2309.04811v1), so wording may differ from the version of record. Figures are not available in the text, so results shown only graphically are described qualitatively here.
Question
Machine learning models trained on adsorption energies in the Open Catalyst 2020 (OC20) data set had not reached errors below about 0.3 eV, which the authors note corresponds to more than five orders of magnitude uncertainty in a rate constant at 300 K. Since density functional theory holds that every ground-state property follows from the electron density, the paper asks whether a graph neural network (GNN) can be trained to predict the electron density itself for diverse catalyst surfaces, and whether established physics-based tools (Bader partitioning, dipole calculations, a DFT code) can then extract useful properties from that predicted density without the model ever being trained on those properties (Sunshine et al. 2023).
Approach
The authors use a pointwise (probe) approach following the DeepDFT model of Jorgensen and Bhowmik: grid points where the density is wanted are added to the atomic graph as nodes that receive but do not send messages, which suits planewave codes that already store the density on a grid. Two changes make the predictions physically usable. Probes beyond the cutoff radius of every atom are forced to zero density, removing a spurious uniform negative charge in vacuum regions, and the model takes the total valence electron count as input and scales each random batch of probes to carry its proportional share, so the full predicted density integrates exactly to the right number of electrons. The model is a SchNet variant built inside the Open Catalyst Project (OCP) code, which also allows reuse of weights from machine-learned potentials. A modified neighbor search that skips probe-probe edges cuts graph construction cost from O((A+P)2) to O(A(A+P)) for A atoms and P probes. The data set takes 32,794 OC20 training structures, 897 validation and 100 test structures from the out-of-distribution split (no shared adsorbates or surfaces), with RPBE pseudo-valence densities from VASP CHGCAR files, about 5 million grid points per structure and well over 500 GB in total. Training minimizes a normalized absolute density error with at least 104 probes per step for about 20 epochs. Predicted densities are written back to CHGCAR files, with PAW augmentation occupancies taken from a zero-step VASP run, and compared against the superposition of atomic charge densities (SACD) that VASP uses by default.
Key findings
On the test set (4.94 million probes per structure) the model reaches a normalized density MAE of 2.54%, versus 8.30% for SACD. It is clearly better in the high-density regime but cannot distinguish densities below about 10-4 electrons per cubic angstrom, which barely affect the training metric. A learning curve over smaller training subsets shows accuracy is still limited by data rather than architecture, with extrapolation suggesting an MAE near or below 2% with one to two orders of magnitude more data. Bader charges computed from predicted densities have roughly an order of magnitude lower error than SACD charges, and dipole moments about half the SACD error, both improving with training set size, even though the model never saw charges or dipoles. Using predicted densities to start VASP self-consistency calls, the number of Hamiltonian evaluations falls as training data grows, but no model consistently beats the SACD starting guess, partly because wavefunctions and augmentation occupancies must also converge. Predicting a full density grid takes 1-2 GPU-minutes, much slower than machine-learned potentials, and the densities are not accurate enough to give forces directly.
Significance
The work proposes an alternative route for atomistic machine learning in catalysis: learn the density once, then compute charges, dipoles and potentially other properties with existing, interpretable methods that enforce charge conservation. It also lays out the practical obstacles, namely petabyte-scale data if extended to all of OC20, the need for architectures whose cost is linear in the number of probes (ruling out many-body models such as GemNet), and better initial guesses for PAW occupancies and wavefunctions before density prediction can speed up DFT (Sunshine et al. 2023).