Graph neural networks for drug discovery

The problem

Screening drug candidates in the lab is slow and expensive. If a model can predict a molecule’s pharmacokinetic properties up front, chemists can spend lab time on the candidates most likely to succeed.

Approach

At CSIRO I developed and evaluated an AI application that predicts these properties, and entered the approach in the OpenADMET and ExpansionRx blind challenge.

• Graph representations. Molecules are featurised as graphs and modelled with Chemprop message-passing networks. I also experimented with adding fingerprint features (Morgan, Avalon, ErG and MACCS).
• Multi-task learning. The model learns all endpoints jointly, supplemented with public datasets (Galapagos’ Polaris set and a merged Novartis set of roughly 300K molecules) after chemical curation and SMILES canonicalisation.
• Target transforms. Skewed endpoints are log-transformed, with a logit transform for percentage endpoints.
• Ensembling. Models are trained on five random splits and the submission averages all five checkpoints.
• Tuning at scale. Training and hyperparameter search ran on HPC with SLURM. Hand-picked hyperparameters ended up beating automated Tree-structured Parzen Estimator search.

Results

• Top 13 of 370+ global participants.
• 0.588 macro-averaged relative absolute error.

Working across disciplines

The work was shared with AI, chemistry and bioengineering teams to judge practical applicability, with findings written up as technical reports for stakeholders.