Guided multi-objective generative AI for drug design

This post was contributed by Kevin Ryczko, Amit Kadan, Takeshi Yamazaki, Arman Zaribafiyan, Ross Pivovar, and Ilan Gleiser

In this post, we’ll discuss how SandboxAQ is transforming computer-aided drug design (CADD) by deploying their workloads in the cloud on AWS. CADD has the potential to accelerate the drug discovery process – something that otherwise generally takes more than a decade.

Identifying compounds with a high binding affinity (ligands) to a target protein is one of the most crucial steps in CADD for designing drugs that bind their targets selectively and specifically. But this is an extremely challenging task because of the large chemical space we need to explore.

The traditional approach has been to virtually screen billions of compounds from a large pre-enumerated dataset with various methodologies that evaluate and filter the compounds, leaving only tens to hundreds for experimental testing and validation. This has a high computational cost, and success isn’t certain if an optimal ligand exists outside of the data being examined.

Enter artificial intelligence

Based on its remarkable success in text and image generation, generative artificial intelligence has recently emerged in the scientific community to perform de novo drug discovery. At a high level, a generative machine learning (ML) model learns the underlying distribution of a data set and encodes it into a high-dimensional latent space. The latent space allows the model to represent chemical space numerically, encoding molecules as vectors. By encoding discrete data points from a molecular dataset into a continuous numerical space, the model must extrapolate – thereby exploring uncharted chemical space not captured in the data.

A generative model proposes new compounds by sampling the latent space and transforming a latent vector into a molecular representation via the decoding layers of a neural network. Conventionally, it samples the vector from the latent space randomly. But, there’s no guarantee that a generative model produces an optimal ligand when performing random sampling. A hurdles remains: reducing the computational cost to produce a molecule with the desired properties.

Figure 1: Our generative AI, IDOLPro, utilizes gradient information from differentiable scores to efficiently traverse the latent space of a generative model. Optimized points in the latent space yield molecules with desired properties.

Our approach to de novo drug design

We developed IDOLpro (Inverse Design of Optimal Ligands for Protein pockets), a guided multi-objective generative AI to ensure that optimal ligands are produced. IDOLpro incorporates feedback (gradients) from a scoring function to update the latent vectors (Figure 1). Standard machine learning libraries like Pytorch and Tensorflow automatically generate a computational graph so that the model weights can be efficiently updated to optimize a particular loss function. We take advantage of this computational graph during inference to find a latent space vector which optimizes an objective function to quantify a set of desired attributes. Doing this allows us to use robust optimization algorithms to traverse the latent space, quickly and efficiently finding novel ligands with optimal expected properties.

In the case of in silico drug discovery, one of the most important objectives is to find ligands with high binding affinities to a target protein. However, optimizing binding affinity alone produces ligands that are useless in practice due to overlooking other important criteria like synthesizability, toxicity, and solubility. To address this issue, we built guided multi-objective generative AI for structure-based drug discovery, optimizing both binding affinity and synthesizability simultaneously, producing synthesizable compounds that bind well to the target protein.

To validate our approach, we used the Vina score to approximate the binding affinity and synthetic accessibility (SA) score to approximate the synthesizability of the generated molecules, while guiding the latent variables of a state-of-the-art generative model for protein-conditioned ligand generation – DiffSBDD.

Building guided multi-objective generative AI for drug design with cloud computing

AWS was essential for us to develop our optimization pipeline. Development of the pipeline, model training and testing, and experimentation (inference) were all performed on Amazon Elastic Compute Cloud (Amazon EC2) instances that AWS ParallelCluster spun up for us (Figure 2 shows the architecture). The open-source Slurm scheduling software built into AWS ParallelCluster allowed us to request hundreds of GPU-based Amazon EC2 instances in parallel, enabling rapid experimentation of our optimization pipeline.

Our optimization process requires defining a computational graph that begins with the latent vectors of the generative model and ends with the scores of interest. These computational graphs can become large and we need to store them in GPU memory for maximum efficiency. For this reason, we chose NVIDIA A10G Tensor Core GPUs with 24 GB of VRAM (g5.2xlarge).

To have a differentiable synthesizability score, we trained an ensemble of equivariant neural networks (each on a different random split of data) to predict the synthetic accessibility given a molecular structure. To train, validate, and test our ensemble, we used NVIDIA Tesla V100 GPUs (p3.2xlarge). We tracked all our training runs using MLFlow, which we set up on a t2.large EC2 instance. We also configured the MLFlow backend store to use a PostgreSQL relational database, which ran on a db.t3.medium instance in Amazon Relational Database Service (Amazon RDS). And we stored the datasets and training metrics in an Amazon Simple Storage Service (Amazon S3) bucket, all managed by MLFlow.

With this combination, when running IDOLpro on AWS, we efficiently generated 100 novel ligands per protein pocket in an hour using just a single GPU.

Fig. 2: Outline of AWS technologies used in our generative AI pipeline. We used AWS ParallelCluster to launch training and inference tasks, with artifacts stored in AWS S3. We used an MLFlow server set up on an AWS EC2 instance to track experiments, with data stored in an AWS RDS instance.

Surpassing the binding affinity of experimentally verified molecules

We investigated the performance of our pipeline on two commonly used benchmark datasets, CrossDocked and Binding MOAD. The CrossDocked dataset contains the computationally-docked protein-ligand pair structures. The Binding MOAD dataset contains experimentally determined pocket-ligand pairs. We compared the results of our tool to other existing tools in the literature that also perform de novo ligand generation.

Both datasets contain the structural information for a large number of protein-ligand interactions. The ligands in these interactions are useful to identify a binding pocket within the protein target and hence serve as the input to a generative model that generates ligands directly into the protein pocket.

Fig. 3: IDOLpro produces molecules with better binding affinity than other generative tools on two benchmark sets – CrossDocked (left) and the Binding MOAD (right). IDOLpro also yields average Vina scores that are better than those of the reference molecules found in each dataset. For CrossDocked these reference molecules were from virtual screening, while for the Binding MOAD, these reference molecules are from experiments.

Our generative AI, IDOLpro, outperformed all other generative models in these benchmark sets as we show in Figure 3. For CrossDocked, we found a 10% (0.73 kcal/mol) improvement compared to state-of-the-art (DiffSBDD) and a 17% improvement compared to the ligands in the test set (labeled Test-set).

To put these energy differences into perspective, we can use them to estimate the IC50 value, which directly informs the efficacy of an inhibitor drug. When considering the performance of our tool compared to the next-best tool, this difference in Vina scores corresponds to a 3.4x improvement in the IC50 value.

Amongst generative ML tools, IDOLpro is producing – to the best of our knowledge – the best binding affinity ligands to date. With our pipeline, for the first time, we can generate ligands with better binding affinities than the experimentally verified ligands. Relative to those same experimentally verified ligands, when optimizing both binding affinity and synthesizability, we were able to produce more easily synthesizable ligands with better binding affinity 95% of the time.

Fine-tuning reference ligands with generative AI

Fig. 4: Example of using IDOLpro for performing lead optimization. IDOLpro produces molecules that have slight changes in their scaffold and functional groups relative to the reference molecules. IDOLpro is able to produce many molecules with better binding affinity and synthetic accessibility scores.

In addition to de novo generation, our pipeline can perform lead optimization in which the initial ligand is optimized towards having more optimal properties. In this case, we start with a reference ligand, identify its place in the latent space in the generative model, and then begin the optimization procedure. Lead optimization often fixes a large part of the molecule, i.e., the scaffold while optimizing the rest of the molecule. As a simple demonstration, we performed optimization without holding any of the molecules fixed in the present work. However, as seen in Fig. 4 our optimization also produces ligands with a similar scaffold.

When performing lead optimization on the reference ligands from a testing subset of CrossDocked, we found that 95% of the ligands produced from our pipeline have better binding affinity and synthesizability scores than the references. If we compare this to our de novo generation strategy, only 76% of the ligands have both better binding affinity and synthesizability scores.

We found that de novo generation produces the lowest binding affinity scores, and lead optimization produces more synthesizable ligands. Performing lead optimization with our pipeline strikes a fine balance between binding affinity and synthesizability, producing viable, next-generation drug molecules.

Conclusion

IDOLpro is a guided multi-objective generative AI to accelerate structure-based drug design with AWS. Our method outperformed all other generative models on two commonly used benchmark datasets for protein-ligand binding. In addition, for the first time, we generated ligands with better binding affinities than experimentally verified ligands.

Our future work is to integrate other important metrics into IDOLpro such as toxicity and solubility to ensure the generation of more suitable ligands, along with the consideration of more accurate binding affinity metrics such as free energy perturbation (FEP) based affinity. The modular architecture of our guided multi-objective generative AI is general and allows us to expand it to other industry-relevant applications like catalyst and battery material discovery.

We think that guided multi-objective generative AI will become the foundation for many drug and material discovery activities.  IDOLpro is SandboxAQ’s cloud-based simulation package and it includes several methods to compute different properties of molecules and materials. We’ve made IDOLpro available on QEMIST Cloud.

The content and opinions in this blog are those of the third-party author and AWS is not responsible for the content or accuracy of this blog.