3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction

Jiaqi Guan; Wesley Wei Qian; Xingang Peng; Yufeng Su; Jian Peng; Jianzhu Ma

Source-linked AI summary

3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction

Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, Jianzhu Ma

arXiv:2303.03543v1Mar 6, 2023q-bio.BMcs.LG

TL;DR

Target-aware molecule generation needs models that represent protein–molecule interactions in 3D without relying on voxelized densities or autoregressive sampling. The paper proposes TargetDiff, a non-autoregressive SE(3)-equivariant diffusion model that jointly generates atom coordinates and types, and reports more realistic structures, better target affinities, and improved affinity ranking and prediction.

Problem
Existing target-aware 3D models rely on voxelized atom densities or autoregressive sampling, while binding-affinity estimation from generative models remains limited.
Method
TargetDiff jointly diffuses continuous atom coordinates and discrete atom types using a protein-conditioned SE(3)-equivariant graph neural network in a non-autoregressive process.
Results
Empirical studies report more realistic 3D structures, better affinities toward protein targets, and improved binding-affinity ranking and prediction without retraining.
Takeaways & Limitations
TargetDiff provides a target-aware generation and evaluation framework that also uses model-derived features or scores for affinity screening.
Takeaways & Limitations
During diffusion, the joint molecule distribution is decomposed into independent coordinate and atom-type distributions, with dependencies handled by the generative model.

Abstract

from arXiv · show

Rich data and powerful machine learning models allow us to design drugs for a specific protein target \textit{in silico}. Recently, the inclusion of 3D structures during targeted drug design shows superior performance to other target-free models as the atomic interaction in the 3D space is explicitly modeled. However, current 3D target-aware models either rely on the voxelized atom densities or the autoregressive sampling process, which are not equivariant to rotation or easily violate geometric constraints resulting in unrealistic structures. In this work, we develop a 3D equivariant diffusion model to solve the above challenges. To achieve target-aware molecule design, our method learns a joint generative process of both continuous atom coordinates and categorical atom types with a SE(3)-equivariant network. Moreover, we show that our model can serve as an unsupervised feature extractor to estimate the binding affinity under proper parameterization, which provides an effective way for drug screening. To evaluate our model, we propose a comprehensive framework to evaluate the quality of sampled molecules from different dimensions. Empirical studies show our model could generate molecules with more realistic 3D structures and better affinities towards the protein targets, and improve binding affinity ranking and prediction without retraining.

1 INTRODUCTION

Target-aware drug generation must model protein–molecule interactions in 3D while avoiding voxelization and autoregressive sampling limitations. TargetDiff addresses these challenges with a non-autoregressive, SE(3)-equivariant diffusion framework and evaluates generated molecules across multiple dimensions.

Target-aware generation is difficult because the synthetically feasible molecular space and binding-pose degrees of freedom are large.
Existing approaches either omit explicit 3D protein–molecule interactions, voxelize 3D space, or use autoregressive atom-by-atom sampling.Voxelization scales cubically with pocket size, while autoregressive methods introduce exposure bias from training–sampling mismatch.
TargetDiff generates target-aware molecules in continuous 3D space using a non-autoregressive diffusion model conditioned on protein atoms.
The model jointly diffuses continuous atom coordinates and discrete atom types with a SE(3)-equivariant graph neural network.Its design targets likelihood invariance to global translation and rotation of the binding complex.
The framework introduces evaluation metrics for assessing generated molecules across multiple dimensions and reports superiority over two representative baselines.
The model can act as a scoring function or unsupervised feature extractor for binding-affinity ranking and prediction.

2 RELATED WORK

Molecular generative models use string, graph, or 3D-structure representations, with recent target-aware methods conditioning generation on protein information. These approaches differ in how directly they represent spatial interactions and how they generate molecules.

Molecule Generation with Different Representations: String-based models commonly use SMILES and adapt language-model architectures such as recurrent neural networks for molecule generation.
Molecule Generation with Different Representations: Graph-based methods operate directly on molecular graphs but are limited in modeling spatial information important for molecular properties and functions.
Molecule Generation with Different Representations: 3D-structure-based methods generate molecules in three-dimensional space to represent molecular geometry more directly.
Target-Aware Molecule Generation: Target-aware methods generate SMILES, molecular graphs, or voxelized 3D molecules conditioned on protein contexts or targets.
Diffusion Models: Diffusion models are latent-variable generative models that learn generation through forward noise injection and reverse denoising processes.

3 METHODS

TargetDiff is a non-autoregressive, SE(3)-equivariant diffusion model that generates target-aware molecules in continuous 3D space while modeling atom coordinates and atom types jointly. Its generative network explicitly incorporates protein–ligand interactions and can also provide representations for affinity ranking and prediction.

3.2 OVERVIEW OF TARGETDIFF: TargetDiff models target-aware molecule generation as a diffusion process with forward noise injection and a reverse process that recovers molecules from noise.Both processes are defined as Markov chains, with the reverse distribution approximated by a neural network.
3.3 MOLECULAR DIFFUSION PROCESS: The model jointly generates continuous atom coordinates and discrete one-hot atom types conditioned on a protein binding site.Gaussian distributions model coordinates, while categorical distributions model atom types and their features.
3.3 MOLECULAR DIFFUSION PROCESS: TargetDiff factorizes the diffusion distribution into independent coordinate and atom-type distributions, while its generative process models their dependencies.This factorization enables concise formulations and efficient sampling of noisy molecular states.
3.4 PARAMETERIZATION OF EQUIVARIANT MOLECULAR GENERATIVE PROCESS: A shifting protein center of mass and an SE(3)-equivariant transition network yield likelihood invariance to global translation and rotation.The proposition applies this construction to the protein–ligand complex during reverse generation.
3.4 PARAMETERIZATION OF EQUIVARIANT MOLECULAR GENERATIVE PROCESS: An SE(3)-Equivariant GNN alternately updates atom embeddings and ligand coordinates using interatomic distances and edge-type features.Protein coordinates remain fixed through a ligand mask, and the final embedding produces denoised atom types through an MLP and softmax.
3.5 TRAINING OBJECTIVE: The model is trained by optimizing a variational bound combining atom-coordinate and atom-type losses.The coordinate KL divergence has a closed form and can be implemented as an MSE between true and predicted clean coordinates.
3.6 AFFINITY RANKING AND PREDICTION AS UNSUPERVISED LEARNER: Under a parameterization that freezes the coordinate-update branch, TargetDiff provides hidden embeddings and denoised atom types for affinity ranking and prediction.The resulting representation is used as an unsupervised signal rather than requiring retraining for the affinity task.

4 EXPERIMENTS

TargetDiff is evaluated on molecular structure realism, target binding affinity, and affinity ranking or prediction using CrossDocked2020 and PDBBind v2020 benchmarks. Across these evaluations, it generally outperforms representative voxelized and autoregressive baselines, while retaining some property and reconstruction limitations.

Evaluation setup: The evaluation framework measures generated molecules across molecular structures, target binding affinity, and molecular properties.The experiments use 100,000 CrossDocked2020 training complexes and 100 novel test complexes, comparing TargetDiff with liGAN, AR, Pocket2Mol, and GraphBP.
Molecular Structures: TargetDiff best matches reference bond-distance distributions, preserving two carbon-carbon distance modes that liGAN, AR, and Pocket2Mol fail to reproduce.Jensen-Shannon divergence is used for distribution comparison, and TargetDiff outperforms other methods across major bond types.
Molecular Structures: TargetDiff generates more consistent rigid fragments after force-field optimization, although its reconstruction algorithm produces a larger proportion of seven-membered rings.liGAN and AR instead generate many three- and four-membered rings, while the seven-member-ring pattern is identified as a limitation and future research direction.
Target Binding Affinity: TargetDiff molecules have the best median Vina energy for 57% of targets, compared with 4% for liGAN, 13% for AR, and 26% for Pocket2Mol.On average, 58.1% of TargetDiff molecules have better binding affinity than the reference molecule, and Vina Score and Vina Min show the same overall advantage.
Target Binding Affinity: TargetDiff can occupy protein pockets more fully than AR in examples where AR covers only part of the pocket, potentially reducing target specificity.The comparison includes binding-pose visualizations and center-of-mass shifts for generated versus reference molecules.
Binding Affinity Ranking and Prediction: TargetDiff-derived features improve affinity ranking and prediction: entropy complements Vina, while hidden embeddings improve supervised correlation and augmented EGNN achieves the best results on all four PDBBind metrics.The ranking analysis uses Spearman correlation on CrossDocked2020, and prediction is evaluated on a time-split PDBBind v2020 benchmark.

5 CONCLUSION

The conclusion presents TargetDiff as a 3D equivariant diffusion model for target-aware molecule generation and binding-affinity prediction, while identifying bond generation and fragment-based generation as future directions.

TargetDiff is a 3D equivariant diffusion model for target-aware molecule generation and enhancing binding affinity prediction.
Future work could incorporate bond generation into diffusion to skip the bond inference algorithm.
Fragment-based molecule-generation techniques are another proposed future direction.
The implementation, experimental data, and model checkpoints are available online.

A PROOF OF SE(3)-EQUIVARIANCE OF GENERATIVE MARKOV TRANSITION

The proof establishes SE(3)-equivariance of the generative Markov transition by combining invariant input features, distance preservation, equivariant network outputs, and center-of-mass preprocessing.

The generative transition must preserve equivariance under global rotation and translation of ligand and protein coordinates.
Pairwise distances remain unchanged under SE(3) transformations because rotations preserve norms and translations cancel.
Atom and edge features are initially SE(3)-invariant, supporting invariant hidden representations through the equivariant layers.
The network predicts an SE(3)-equivariant denoised coordinate output, which is used to obtain the posterior mean for the previous diffusion step.
Moving the protein center of mass to zero removes translation dependence, leaving rotation equivariance for the Markov transition.

B ANALYSIS OF INVARIANT INITIAL DENSITY

The invariant initial density is constructed by centering protein coordinates at zero, making the large-time diffusion distribution standard Gaussian and invariant to rigid transformations.

At a sufficiently large diffusion timestep, the noisy protein-conditioned coordinates are assumed to follow a Gaussian centered at the protein center with unit standard deviation.
Centering protein coordinates at zero produces an SE(3)-invariant initial density and simplifies diffusion and generation to standard Gaussian sampling.
Complex coordinates are translated into the protein zero-center-of-mass subspace before density evaluation.
Because the projection preserves norms under rotations and the base distribution is isotropic Gaussian, the resulting density is invariant under SE(3).

C PROOF OF INVARIANT LIKELIHOOD

The proof of invariant likelihood relies on an invariant prior and an equivariant transition, with categorical diffusion posteriors derived from symmetric transition probabilities.

An invariant initial density composed with an equivariant transition function yields an invariant distribution.
The likelihood proof is organized around the invariant-prior condition for transformed protein-ligand coordinates.
The categorical forward and cumulative transitions define noisy atom-type distributions used to derive the posterior over atom types.
Symmetry of the categorical transition function follows because matching and nonmatching atom-type probabilities are treated consistently in either direction.

E OVERALL TRAINING AND SAMPLING PROCEDURES

TargetDiff is trained and sampled through a joint diffusion process over ligand coordinates and atom types, conditioned on protein binding-site structure. Its equivariant network operates on dynamically constructed protein–ligand graphs throughout the process.

Training Procedure: Training samples a diffusion time, centers protein coordinates, perturbs ligand coordinates and atom types, and predicts denoised coordinates and atom types.The loss combines coordinate MSE with a posterior atom-type KL loss.
Equivariant Network: At each equivariant layer, the model dynamically constructs a k-nearest-neighbor graph from known protein coordinates and current ligand coordinates.The graph uses protein and ligand atom features together with distance and bond-related edge features.
Sampling Procedure: Sampling first draws ligand size from a pocket-size-conditioned prior, initializes noisy coordinates and atom types, and iteratively samples reverse posteriors.The output is a generated ligand molecule intended to bind the protein pocket.
Equivariant Network: The implementation uses 9 equivariant Transformer layers with hidden dimension 128 and 16 attention heads, alongside separate coordinate and atom-type noise schedules.The coordinate schedule is sigmoid-based, while the atom-type schedule is cosine-based, with 1000 diffusion steps.

F.3 TRAINING DETAILS

Training uses pocket-size-aware ligand-atom priors and practical optimization settings, with pocket size estimated robustly from protein-atom distances. Generation samples ligand sizes from these learned priors rather than matching reference molecules.

Training Details: Training uses Adam with learning rate 0.001, batch size 4, gradient clipping at 8, and atom-type loss weight α = 100.Protein coordinates receive Gaussian noise with standard deviation 0.1 during training.
Training Details: The model converges within 24 hours and 200k steps on one NVIDIA GeForce GTX 3090 GPU.
Pocket Size Estimation: Pocket size is estimated as the median of the top 10 farthest pairwise protein-atom distances.The median is selected for robustness.
Prior Distribution: Ligand-atom priors are estimated from histograms within 10 quantile bins of training pocket sizes, reflecting a positive correlation between pocket and ligand size.
Prior Distribution: During generation, ligand atom counts are randomly sampled from pocket-size-conditioned priors rather than required to match reference molecules.Reference atom counts are provided during training but are unnecessary during generation.

G ADDITIONAL EVALUATION RESULTS

Additional evaluations reproduce TargetDiff’s favorable binding-affinity trend with QVina and reveal that generated-molecule SA-score differences are confounded by molecular size.

Binding Affinity: TargetDiff-generated molecules achieve state-of-the-art binding affinity under QVina, showing a trend similar to the AutoDock Vina evaluation.QVina is described as faster but less accurate than AutoDock Vina.
Synthetic Accessibility: Pearson R=−0.56, p ≤10−80 indicates a strong negative correlation between SA score and molecular size.The passage states that SA-score differences between generated molecules could therefore be artifacts of size differences.
Evaluation Summary: Table S1 summarizes properties of reference molecules and molecules generated by TargetDiff and other baselines.

H SAMPLING TIME ANALYSIS

TargetDiff is designed to scale more favorably with molecule size than autoregressive models by processing additional atoms in parallel. Its average wall-clock time is intermediate among the compared models, while supplementary examples cover both favorable and unfavorable target-specific outcomes.

Sampling Time Analysis: TargetDiff processes additional atoms in parallel, whereas autoregressive models require additional generation steps for larger molecules.This is presented as a major scaling advantage over autoregressive models such as AR.
Sampling Time Analysis: The time study compares inference-time ratios relative to generating a 10-atom molecule across five randomly selected binding pockets.The comparison includes AR, Pocket2Mol, GraphBP, and TargetDiff and generates 100 molecules per pocket.
Sampling Time Analysis: For 100 valid molecules, average wall-clock times are 7785s for AR, 2544s for Pocket2Mol, 105s for GraphBP, and 3428s for TargetDiff.TargetDiff is described as having moderate sampling efficiency relative to AR and Pocket2Mol; GraphBP is fastest but produces lower-quality molecules.
Binding-Pose Examples: Supplementary visualizations show best, median, and worst targets with generated poses and calculated Vina energies.The selected targets are 1H36 A, 1DXO A, and 2GNS A.
Binding-Pose Examples: For targets 4KCQ A and 1R1H A, supplementary figures compare three AR-generated molecules with three TargetDiff-generated molecules where AR has better Vina-estimated affinity.