Partial Charge Prediction and Pattern Extraction from a AttentiveFP Graph Neural Network
Marc Lehner, Paul Katzberger, Greg Landrum and Sereina Riniker
Lab. für Physikalische Chemie, Eidgenössische Technische Hochschule Zürich (ETHZ), Switzerland
Molecular dynamics (MD) simulations enable the time-resolved study of bio-molecular processes. The quality of MD simulations is, however, highly dependent on the set of interaction parameters used, so-called force fields. The accurate partial-charge assignment of all simulated atoms is hence a crucial part of every MD simulation. Due to the slowly decaying nature of the Coulomb interactions, the effects of different partial-charge assignments can be observed over long distances and can have drastic effects on the stability of a MD simulation. Therefore, many schemes have been developed over the last decades to improve partial-charge assignment: Classical tabulated values, ab initio calculations, or the prediction with machine learning models. However, all these approaches have some shortcomings in either accuracy, speed, or interpretability.
Here, we present an option to combine the accuracy of ab initio calculations, the speed of machine learning models, and the interpretability of tabulated assignments. An attention-based graph neural network (GNN) is trained on a diverse dataset to predict high-quality atom-in-molecule (AIM) partial charges, calculated on a def2-tzvp/TPSSh Quantum Mechanical level of theory with an implicit solvent (polarizable continuum model). Then a model-agnostic approach was used to extract the most important sub-graph on an atomistic level to provide the user with the same level of interpretability as for tabulated values.
For the graph neural network the AttentiveFP architecture is used with a modification to include the physics of partial charges. The final network layer rescales the sum of individual partial charges to the formal charge. It can be shown, that this physics informed network then increased the accuracy of partial charge predictions, compared to the unmodified AttentiveFP model. The network is then used together with the GNNExplainer from Pytorch-Geometric to extract the relevant attention of nodes in a sub-graph. While these sub-graphs are already in good agreement with general “chemical knowledge”, satisfying the human interpretability of the model, they can also be used to generate a simpler, rule based model for charge assignment. The attention values can be used to prioritize decisions in the re-construction of a sub-graph, therefore linearizing the otherwise complex scaling of partial charge assignment in complex molecular graphs. Making this method well suited, even for large bio-molecular systems or screening of large databases.
The method is then compared to state of the art models like the semi-empirical am1-bcc, other machine learning based approaches or tabulated empirical values from common force fields in accuracy, assignment speed, and human interpretability.
