[Paper Review] TAGMol: Target-Aware Gradient-guided Molecule Generation

Vineeth Dorna, D. Subhalingam+6 more|arXiv (Cornell University)|Jun 3, 2024

Advanced biosensing and bioanalysis techniquesBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

TAGMol proposes a target-aware, gradient-guided molecular generation framework that decouples molecular generation from property prediction, using an SE(3)-equivariant diffusion model and a multi-objective guide network to steer generation toward desired properties. It achieves a 22% improvement in average Vina Score while optimizing for drug-likeness (QED) and synthesizability (SA), outperforming state-of-the-art baselines in structure-based drug design.

ABSTRACT

3D generative models have shown significant promise in structure-based drug design (SBDD), particularly in discovering ligands tailored to specific target binding sites. Existing algorithms often focus primarily on ligand-target binding, characterized by binding affinity. Moreover, models trained solely on target-ligand distribution may fall short in addressing the broader objectives of drug discovery, such as the development of novel ligands with desired properties like drug-likeness, and synthesizability, underscoring the multifaceted nature of the drug design process. To overcome these challenges, we decouple the problem into molecular generation and property prediction. The latter synergistically guides the diffusion sampling process, facilitating guided diffusion and resulting in the creation of meaningful molecules with the desired properties. We call this guided molecular generation process as TAGMol. Through experiments on benchmark datasets, TAGMol demonstrates superior performance compared to state-of-the-art baselines, achieving a 22% improvement in average Vina Score and yielding favorable outcomes in essential auxiliary properties. This establishes TAGMol as a comprehensive framework for drug generation.

Motivation & Objective

• To address the limitations of existing generative models in structure-based drug design that overly focus on binding affinity and neglect broader drug-like properties.
• To overcome the challenge of training on low-quality or incomplete datasets by decoupling generation from property prediction and injecting property signals during diffusion sampling.
• To develop a method that simultaneously optimizes for target binding and key pharmacological properties such as QED and SA, even when these signals are sparse in training data.
• To enable multi-property guidance in 3D molecular generation by training separate guide models for each property and combining their gradients during sampling.
• To reduce reliance on post-optimization techniques by embedding property constraints directly into the generative process.

Proposed method

• Decompose the drug generation task into two components: a base SE(3)-equivariant graph diffusion model for generating realistic 3D molecular geometries.
• Train a time-dependent guide model using SE(3)-invariant GNNs to predict molecular properties (e.g., QED, SA, binding affinity) from noisy molecular inputs.
• Apply classifier guidance principles by using the gradient of the guide model to steer the latent space during diffusion sampling, guiding the generation toward desired properties.
• Enable multi-objective guidance by training individual guide models per property and combining their gradients during sampling to explore high-performing regions of the latent space.
• Use docking scores (Vina) as a proxy for binding affinity, leveraging the geometric inductive bias of the guide to improve coordinate-level generation.
• Integrate property guidance into the denoising process without modifying the base generative model, preserving diversity and reconstruction quality.

Figure 1: Overview of TAGMol . (a) Training a property-oriented guide using existing data. (b) Utilizing the trained guide and diffusion model to steer the diffusion sampling process towards the optimal regions of the property of interest.

Experimental results

Research questions

• RQ1 Can a decoupled generation and guidance framework improve molecular generation in structure-based drug design beyond binding affinity alone?
• RQ2 How effectively can gradient-based guidance from property predictors steer diffusion sampling toward molecules with improved drug-like properties?
• RQ3 To what extent can a single generative model achieve simultaneous optimization of binding affinity, QED, and SA scores through multi-property guidance?
• RQ4 Does the use of noisy inputs during guide training improve robustness when training data lacks sufficient high-quality property signals?
• RQ5 Can the framework maintain performance across multiple properties when guiding for all simultaneously, without significant degradation in any individual metric?

Key findings

• TAGMol achieves a 22% relative improvement in average Vina Score compared to state-of-the-art baselines, demonstrating superior binding affinity prediction.
• The model maintains or improves upon key pharmacological properties, including QED and SA scores, even when not explicitly optimized for them.
• Multi-property guidance results in the highest overall hit rate and substantial improvements across all metrics, with only a minor, acceptable drop in SA scores.
• Single-property guidance for binding affinity yields the largest performance gain, attributed to the guide’s geometric inductive bias aligning with docking-based affinity scoring.
• The method shows robustness: guiding for one property does not significantly degrade performance in others, indicating effective collaboration between the denoising and guide models.
• Visualizations confirm that generated molecules are structurally valid and adopt reasonable binding poses, validating the quality of generated outputs.

Figure 2: Comparison of distance distributions between all-atom distances of reference molecules in the test set (Reference) and distances in model-generated molecules. The Jensen-Shannon divergence (JSD) between these two distributions is calculated and reported.