Equivariant Graph Neural Networks for Toxicity Prediction

Abstract

Predictive modeling for toxicity is a crucial step in the drug discovery pipeline. It can help filter out molecules with a high probability of failing in the early stages of de novo drug design. Thus, several machine learning (ML) models have been developed to predict the toxicity of molecules by combining classical ML techniques or deep neural networks with well-known molecular representations such as fingerprints, InChi keys, SMILES strings, or 2D-graphs. Since molecules live in the physical 3D space, it is expected that their more natural and accurate representation would be defined in this geometric space as in ab initio methods for computing physicochemical properties. In such space, these molecules are also subject to geometric symmetries that only adequate 3D molecular representations can capture. Hence, the implications of using 3D structures in toxicity predictive modeling is an essential yet open question that needs further investigation. Recent studies showed great success in using equivariant graph neural networks (EGNNs) for representation learning based on 3D structures to predict potential energy surfaces and physicochemical properties of organic molecules and materials. Inspired by this, we investigate the performance of EG-NNs to construct reliable ML models for toxicity prediction. We use the equivariant transformer (ET) model in TorchMD-NET, an SE(3)-equivariant attention-weighted message-passing neural network. Eleven different toxicity datasets taken from Molecu-leNet, Therapeutic Data Commons (TDC), and ToxBenchmark have been considered to evaluate the capability of ET for toxicity prediction. Our results show that ET adequately learns 3D representations of molecules that can be successfully correlated with toxicity activity, achieving good accuracies on most datasets comparable to state-of-the-art 2D graph-based models and significantly better than a SMILES-based trans-former model. However, exhaustive (single- and multi-) conformational analysis has revealed a lack of correlation between the conformational samples and the predictive performance, which we find is only slightly affected by the conformer selection. Moreover, combining toxicity prediction in a multi-task setting with a neural network potential (energy prediction), or adding the total energy of the respective molecule as an additional node feature (besides atom types), does not benefit or even decrease the model's performance throughout all datasets. This suggests that a physicochemical property, such as the total energy, can not be related to toxicity activity with the selected ML method. We also provide an attention weight analysis for helping to explain the black-box toxicity prediction of our model in 3D space and, thus, increase reliability in ML models. In summary, our findings highlight the relevance of considering 3D geometry information in EGNNs and provide a straightforward way to integrate molecular conformers into ML-based pipelines for toxicity prediction.