SAF: Smart Aggregation Framework for Revealing Atoms Importance Rank and Improving Prediction Rates in Drug Discovery

TL;DR

This paper introduces SAF, a novel weighted aggregation method for graph neural networks in drug discovery, which enhances prediction accuracy and atom importance interpretability by using a Boltzmann distribution-based weighting scheme.

Contribution

SAF presents a new non-linear weighted aggregation approach that improves prediction rates and atom importance ranking in graph neural networks for drug discovery.

Findings

Improved antibiotic activity prediction accuracy.

Recapitulates functional groups in beta-Lactam antibiotics.

Provides a regulated attention mechanism for interpretability.

Abstract

Machine learning, and representation learning in particular, has the potential to facilitate drug discovery by screening a large chemical space in silico. A successful approach for representing molecules is to treat them as a graph and utilize graph neural networks. One of the key limitations of such methods is the necessity to represent compounds with different numbers of atoms, which requires aggregating the atom's information. Common aggregation operators, such as averaging, result in loss of information at the atom level. In this work, we propose a novel aggregating approach where each atom is weighted non-linearly using the Boltzmann distribution with a hyperparameter analogous to temperature. We show that using this weighted aggregation improves the ability of the gold standard message-passing neural network to predict antibiotic activity. Moreover, by changing the temperature hyperparameter, our approach can reveal the atoms that are important for activity prediction in a smooth and consistent way, thus providing a novel, regulated attention mechanism for graph neural networks. We further validate our method by showing that it recapitulates the functional group in beta-Lactam antibiotics. The ability of our approach to rank the atoms' importance for a desired function can be used within any graph neural network to provide interpretability of the results and predictions at the node level.