AI Agents for Molecular Property Prediction
Using natural language to train ML models for ADMET screening
[Update: we received helpful feedback from Pat Walters from OpenADMET pointing out that the datasets and benchmarks we’re using have several flaws, as he’s pointed out before on his Practical Cheminformatics blog. We’ve added some notes to the text below for clarity around some of the points he made about benchmarking and the datasets that we use. Importantly, our goal with this blog was to illustrate how easy it is to train models in Axon without writing any code, not on developing world-beating ADMET models or rigorously benchmarking them. We encourage readers and users to read his blog post and ensure that they’re developing the best ADMET models that they can!]
ADMET screening is an important step in early stages of computational analysis of drug candidates, where prospective drug molecules are assessed for five properties:
Absorption - how well a molecule can be taken up by a person
Distribution - how the molecule becomes distributed throughout the body
Metabolism - how the molecule is broken down after absorption
Excretion - how the molecule (and its metabolites) leave the body
Toxicity - whether or not the molecule causes damage to natural bodily functions
This information is critical to have before trying to take a drug to trial—even if you think a molecule interacts with a target protein exactly how you want it to interact, that molecule might be functionally useless if it can’t get to the protein in the first place or is toxic to the person taking it!
As with most drug discovery tasks, ADMET screening has changed significantly since it first began. Early instances of ADMET predictions relied on linear regression to costly in vivo experiments and later evolved to incorporate in vitro experimental data (see this article). Now, ADMET screening is done computationally on a much larger scale based on troves of these historical data, newer high-throughput experimental results, and public datasets like the Therapeutic Data Commons (TDC). Even in silico methods have changed significantly for ADMET predictions: early linear regression methods were replaced by quantitative structure-activity relationships (QSAR), and followed by more complex machine learning (ML) methods like random forests (RF) and support vector machines (SVM) (see this article). Most recently, these approaches were superseded by deep learning using tools like ChemProp, a directed message passing neural network (D-MPNN) model for property prediction from SMILES strings. Now, there are easily deployed, open-source foundation models like ADMET-AI that are built on ChemProp and trained on data from TDC to make common ADMET predictions for drug molecules. These models can potentially supplement models trained on internal, private ADMET data that pharmaceutical companies might have gathered throughout the years.
At Mirror Physics, we built Axon to tackle multiple steps in the computer-aided drug discovery (CADD) pipeline, including ADMET screening. Axon can help you easily train and develop bespoke models using private data or simply deploy pre-trained models like ADMET-AI. You can request that Axon either use ADMET-AI to evaluate a library of compounds or train a new ChemProp model (or other “classical” ML models like RF and SVM) using natural language. Axon has access to the TDC datasets and ADMET datasets from MoleculeNet, which can be used for (pre-)training before using transfer learning on internal datasets or as a way to test generalizability of a model trained on internal data. Below, we show some examples of how Axon accomplishes these tasks.
Training models from scratch
Plenty of companies that produce small molecule drugs have large private ADMET datasets from historical in-house experiments that might serve as useful training data, in conjunction with publicly available ADMET datasets from TDC. These results can be used to train robust and accurate models to make ADMET predictions, but doing so often requires significant time and effort from chemists with years of training in cheminformatics. Axon can be easily deployed to prepare datasets and train models for ADMET screening, including state-of-the-art D-MPNN models like the most recent version of ChemProp and classical ML models like RF and SVM. Axon has access to classical ML models available in scikit-learn and to tools like rdkit and mordred that can be used to construct molecular fingerprints and descriptors as inputs to these models.
We asked Axon to train several types of models on the blood-brain barrier (BBB) data from TDC, including ChemProp, RF, and SVM models. We let Axon build two versions each of the RF and SVM models that used Morgan fingerprints and Mordred descriptors as inputs, respectively. We also asked Axon to analyze the results and compare them to the results on the public TDC leaderboard for this property and to the results in a whitepaper from Schrödinger on their DeepAutoQSAR model. These binary classification models are evaluated based on the area under the receiver-operator characteristic curve (AUROC), which can be anywhere from 0 to 1 and measures how well each model classifies the molecule as one that does or does not cross the blood-brain barrier. The closer AUROC is to 1, the more accurately the model classifies molecules in the test datasets. Each metric on the TDC leaderboard provides five different splits of the dataset on which individual models are trained and evaluated. We trained a new model from scratch for each split. Axon generated the plot below, showing the average AUROC values and the standard deviation across five different splits of the dataset for the models on the leaderboard, the published value of DeepAutoQSAR, and the models that it just trained. Many of the models did a relatively good job of classifying these molecules as permeable or impermeable, with AUROC values competitive with other models on the public leaderboard. [Update: For full transparency, we omitted the molecules with invalid SMILES strings in this dataset. It has several other errors, as noted elsewhere, and the models we trained here will probably not be particularly useful for real-world deployment and the performance shown in the plots below might not be particularly meaningful. However, we are focused in this post on the process of training these models—we always encourage using the best data available and rigorously testing models before relying on their outputs!]
The results from the models we trained here are not groundbreaking—we want to focus on how easy Axon made the process of gathering the data, training a model on those data, and testing and comparing the model to other results. Axon helps to make the process of training and evaluating the models that cheminformatics experts use even easier and completes its tasks quickly from only two prompts. Bench chemists who aren’t familiar with CADD tools but who want to create models quickly based on data that they’ve gathered can do so reliably without needing to scour documentation. Meanwhile, cheminformatics experts benefit from the speed with which Axon can clean data, write short training scripts, read documentation, and use its compute infrastructure to train and apply these models. If you want to provide more details for how to organize, train, and use these models, you can do so with plain, conversational language and provide as much—or as little—technical detail as you desire. Axon also shows you what it’s doing as it completes tasks, promoting traceability so expert users can check its work. Axon will work with you to produce something tailored to your property prediction needs.
Transfer learning with different data
Scientific data collection never stops after an ML model is trained; instead, those models are used to guide subsequent data collection and investigation. Gathering more data means that you’ll have more information that you can use to improve models in the future. The spread of automated high-throughput experimentation and methods that help automatically guide data collection means that it’s possible for these prediction models to be re-trained regularly with enriching new results. While some classical ML models are relatively cheap to train and deploy, NN models can be more expensive to train—you might not want a large model that you trained a few weeks ago to go to waste every time you want to update it with new data. Alternatively, you might have a NN model that was trained on data for one task and you want to see if it’s possible to tweak the model so that it works for another, similar task. For example, we trained a model to classify molecules as being able or unable to permeate the blood-brain barrier—a binary classification task. That model architecture can be used for other binary classification tasks, like whether or not a drug ingested orally is bioavailable. Both of these situations can be addressed using transfer learning, where a pre-trained model is used as a starting point and its parameters updated to fit new data. Transfer learning can be a nice alternative to training a model from scratch because it typically requires much shorter training times to produce a model that performs well. Notably, the transfer learning process for both of these situations—updating a model with new data for the same task or re-training a model for a new task—is the same. Unfortunately, actually doing transfer learning well can be a bit thorny, and requires organizing model weights, checking model compatibility, updating training scripts, and occasionally doing things like moving large files between different systems.
Axon can ease this transfer learning process. We continued from the BBB permeability session and shifted our focus to the bioavailability benchmark on TDC. Axon trained the same models as it did above from scratch (RF and SVM models from fingerprints and descriptors and a ChemProp model), but it also used transfer learning for another ChemProp model. We prompted Axon to take a ChemProp model that it trained on one of the dataset splits for the BBB task and re-train it on the bioavailability data and compare all of its models to the public leaderboard. You can see that none of our models did particularly well compared to the top models on the leaderboard, but for a first pass, they perform reasonably—bioavailability is trickier to predict than BBB permeability! Notably, our transfer-learned ChemProp model (the orange bar in the plot below) outperformed the ChemProp model that we trained from scratch. That improvement is important because the transfer-learned model was trained on the bioavailability data for a fraction of the amount that the model from scratch was (12 training epochs vs. 50 epochs, respectively).
Altogether, this process took less than an hour and required relatively little input from us to produce usable models. Axon was able to corral the data we needed, read documentation for the models we wanted to train, run the training, test the models, and produce these comparison plots from only a handful of prompts. This quick deployment is only a small example of what Axon is capable of doing for these ML models—our tests here neglected things like hyperparameter optimization, for example, which might help us to improve our predictions on the bioavailability task. Axon can also handle these other tasks, walk users through suggested approaches, and produce simple models like these or more complicated, deeper models as a user might need for their specific application.
Deploying pre-trained models
Finally, if you don’t want to train a model from scratch, you can also easily deploy pre-trained models using natural language with Axon. We’ve made the ADMET-AI model available through Axon to make ADMET prediction a little simpler. ADMET-AI is intended as an ADMET foundation model and is trained on random splits for all of the tasks across TDC. If you don’t want to worry about the cost of training a model from scratch, you can simply use ADMET-AI instead. You can also mix and match different approaches as needed—use ADMET-AI for initial screening, and train and apply a purpose-built model like ChemProp on a smaller dataset for a particular chemical class when the time comes. Ultimately, Axon is as flexible as its users need it to be for these molecular property prediction tasks.
Final thoughts
We conceived of Axon as a platform to make chemistry easier. Axon can help you with chemical property prediction by training new models, updating models using transfer learning, or simply plugging in existing foundation models to help screen molecules using natural language prompts. The ADMET modeling that we showed here is one important step in the drug discovery process that Axon can facilitate, but it can do so much more and links steps of the CADD process together seamlessly. You can request a demo for Axon here if you’re interested in trying it for yourself!
We’re excited about Axon and what it can do. There’s much more about it that we’re working to share with you. Feel free to send me a message at alex@mirrorphysics.com if you want to talk more about Axon’s capabilities.