Structure-based Drug Design with Equivariant Diffusion Models

Introduction

The rational design of molecular drug-like compounds remains an outstanding challenge in biopharmaceutical research. Structure-based drug design (SBDD) aims to generate small-molecule ligands that bind to a specific 3D protein structure with high affinity and specificity (Anderson, 2003). Traditionally, SBDD campaingns are initiated either by high-throughput experimental or virtual screening (Lyne, 2002; Shoichet, 2004) of large chemical databases. Not only is this expensive and time consuming but it restricts the exploration of chemical space to the historical knowledge of previously studied molecules, with a further emphasis usually placed on commercial availability (Irwin & Shoichet, 2005). Moreover, the optimization of initial lead molecules is often a biased process, with significant reliance on human intuition (Ferreira et al., 2015).

Recent advances in geometric deep learning, especially in modelling geometric structures of biomolecules (Bronstein et al., 2021; Atz et al., 2021), provide a promising direction for structure-based drug design (Gaudelet et al., 2021). Despite remarkable progress in the use of deep learning as surrogate docking models (Lu et al., 2022; Stärk et al., 2022; Corso et al., 2022), deep learning-based design of ligands that bind to target proteins remains an open problem. Early attempts have been made to represent molecules as atomic density maps, and variational auto-encoders were utilized to generate new atomic density maps corresponding to novel molecules (Ragoza et al., 2022). However, it is nontrivial to map atomic density maps back to molecules, necessitating a subsequent atom-fitting stage. Follow-up work addressed this limitation by representing molecules as 3D graphs with atomic coordinates and types which circumvents the post-processing steps. Li et al. (2021) proposed an autoregressive generative model to sample ligands given the protein pocket as a conditioning constraint. Peng et al. (2022) improved this method by using an $E(3)$-equivariant graph neural network which respects rotation and translation symmetries in 3D space. Similarly, Drotár et al. (2021); Liu et al. (2022) used autoregressive models to generate atoms sequentially and incorporate angles during the generation process. Li et al. (2021) formulated the generation process as a reinforcement learning problem and connected the generator with Monte Carlo Tree Search for protein pocket-conditioned ligand generation. However, the main premise of sequential generation methods may not hold in real scenarios, since it imposes an artificial ordering scheme in the generation process and, as a result, the global context of the generated ligands may be lost. In addition, previous methods often focus on de novo ligand generation and are not readily applicable to other drug design tasks and downstream molecular optimization.

In this work, we propose an $SE(3)$-equivariant 3D-conditional diffusion model for SBDD (DiffSBDD) that respects translation, rotation, and permutation symmetries. We evaluate our model in the following ways:

We benchmark our model on the de novo molecule generation task, for which we introduce two strategies, protein-conditioned generation, which considers the protein as a fixed context, and learning the joint distribution of protein and ligand atoms. In the latter case, conditional samples can be produced with a modified sampling scheme.

We further demonstrate how pre-trained DiffSBDD models can be used out-of-the-box to optimize arbitary molecular properties via a noise/denoise scheme in combination with an evolutionary algorithm.

We present an inpainting-inspired approach to accomplish challenging molecular (re)design problems, such as scaffold hopping/elaboration and fragment growing/merging, all training a new model.

Finally, we curate an experimentally determined binding dataset derived from Binding MOAD (Hu et al., 2005), which supplements the commonly used synthetic CrossDocked (Francoeur et al., 2020) dataset to validate our model performance under realistic binding scenarios. The experimental results demonstrate that DiffSBDD is capable of generating novel, diverse and drug-like ligands with high predicted binding affinities to given protein pockets, and case studies underline the method’s flexibility as a tool for candidate molecule refinement.

Background

Denoising diffusion probabilistic models (DDPMs) (Sohl-Dickstein et al., 2015; Ho et al., 2020) are a class of generative models inspired by non-equilibrium thermodynamics. Briefly, they define a Markovian chain of random diffusion steps by slowly adding noise to sample data and then learning the reverse of this process (typically via a neural network) to reconstruct data samples from noise.

adds noise to the data $\bm{z}_{\text{data}}$ and produces a latent noised representation $\bm{z}_{t}$ for $t=0,\dotsc,T$. $\alpha_{t}$ controls the signal-to-noise ratio $\text{SNR}(t)=\alpha_{t}^{2}/\sigma_{t}^{2}$ and follows either a learned or pre-defined schedule from $\alpha_{0}\approx 1$ to $\alpha_{T}\approx 0$ (Kingma et al., 2021). We also choose a variance-preserving noising process (Song et al., 2020) with $\alpha_{t}=\sqrt{1-\sigma_{t}^{2}}$.

Since the noising process is Markovian, we can write the denoising transition from time step $t$ to $s<t$ in closed form as

with $\alpha_{t|s}=\frac{\alpha_{t}}{\alpha_{s}}$ and $\sigma_{t|s}^{2}=\sigma_{t}^{2}-\alpha_{t|s}^{2}\sigma_{s}^{2}$ following the notation of Hoogeboom et al. (2022). This true denoising process depends on the data sample $\bm{z}_{\text{data}}$, which is not available when using the model for generating new samples. Instead, a neural network $\phi_{\theta}$ is used to approximate the sample $\hat{\bm{z}}_{\text{data}}$. More specifically, we can reparameterize Equation (1) as $\bm{z}_{t}=\alpha_{t}\bm{z}_{\text{data}}+\sigma_{t}\bm{\epsilon}$ with $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ and directly predict the Gaussian noise $\hat{\bm{\epsilon}}_{\theta}=\phi_{\theta}(\bm{z}_{t},t)$. Thus, $\hat{\bm{z}}_{\text{data}}$ is simply given as $\hat{\bm{z}}_{\text{data}}=\frac{1}{\alpha_{t}}\bm{z}_{t}-\frac{\sigma_{t}}{\alpha_{t}}\hat{\bm{\epsilon}}_{\theta}$.

The neural network is trained to maximise the likelihood of observed data by optimizing a variational lower bound on the data, which is equivalent to the simplified training objective (Ho et al., 2020; Kingma et al., 2021) $\mathcal{L}_{\text{train}}=\frac{1}{2}||\bm{\epsilon}-\phi_{\theta}(\bm{z}_{t},t)||^{2}$ up to a scale factor (see Appendix A for details).

E​(n)𝐸𝑛E(n)-equivariant Graph Neural Networks

In our case, since the nodes have both geometric atomic coordinates $\boldsymbol{x}$ as well as atomic type features $\boldsymbol{h}$, we can use a simple implementation of EGNN proposed by Satorras et al. (2021), in which the updates for features $\bm{h}$ and coordinates $\bm{x}$ of node $i$ at layer $l$ are computed as follows:

where $\phi_{e}$, $\phi_{\text{att}}$, $\phi_{h}$ and $\phi_{h}$ are learnable Multi-layer Perceptrons (MLPs) and $d_{ij}$ and $a_{ij}$ are the relative distances and edge features between nodes $i$ and $j$ respectively. In Section 3.1, we propose a modified coordinate update that breaks symmetry to reflections.

Equivariant Diffusion Models for SBDD

We use an equivariant DDPM to generate molecules and binding conformations jointly with respect to a specific protein target. We represent protein and ligand point clouds as graphs that are further processed by EGNNs (Satorras et al., 2021). We consider two distinct approaches to 3D pocket conditioning: (1) a conditional DDPM that receives a fixed pocket representation as context in each denoising step, and (2) a model that is trained to approximate the joint distribution of ligand-pocket pairs combined with a modified sampling procedure at inference time.

We parameterize the noise predictor $\hat{\bm{\epsilon}}_{\theta}=\phi_{\theta}(\bm{z}^{(L)}_{t},\bm{z}^{(P)}_{\text{data}},t)$ with an EGNN (Satorras et al., 2021; Hoogeboom et al., 2022). To process ligand and pocket nodes with a single GNN, atom types and residue types are first embedded in a joint node embedding space by separate learnable MLPs. We employ the same message-passing scheme outlined in Equations (3)-(5), however, following (Igashov et al., 2022), we do not update the coordinates of nodes that belong to the pocket to ensure the three-dimensional protein context remains fixed throughout the EGNN layers.

Equivariance to the orthogonal group $O(3)$ (comprising rotations and reflections) is achieved because we model both prior and transition probabilities with isotropic Gaussians where the mean vector transforms equivariantly w.r.t. rotations of the context (see Hoogeboom et al. (2022) and Appendix E). Ensuring translation equivariance, however, is not as easy because the transition probabilities $p(\bm{z}_{t-1}|\bm{z}_{t})$ are not inherently translation-equivariant. In order to circumvent this issue, we follow previous works (Köhler et al., 2020; Xu et al., 2022; Hoogeboom et al., 2022) by limiting the whole sampling process to a linear subspace where the center of mass (CoM) of the system is zero. In practice, this is achieved by subtracting the center of mass of the system before performing likelihood computations or denoising steps. Since equivariance of the transition probabilities depends on the parameterization of the noise predictor $\hat{\bm{\epsilon}}_{\theta}$, we can make the model sensitive to reflections with a simple additive term in the EGNN’s coordinate update:

using the cross product which changes sign under reflection. Here, $\bar{\bm{x}}^{l}$ denotes the center of mass of all nodes at layer $l$. $\phi_{x}^{\times}$ is an additional MLP. This modification is discussed in more detail in Appendix F.

2 Joint generation of pocket and small molecule

In addition to the pocket-conditional model presented above, we also train a DDPM for generating molecular systems consisting of pocket atoms and ligand atoms. This unconditional model thus approximates the joint distribution $p(\bm{z}^{(L)}_{\text{data}},\bm{z}^{(P)}_{\text{data}})$. The training procedure is identical to the equivariant molecule generation procedure developed by Hoogeboom et al. (2022) aside from the fully-connected neural networks that embed protein and ligand node features in a common space as described in Section 3.1, and the breaking of reflection symmetry using Equation (6).

When combined with inpainting, as described in the following section, this model can be used for target-specific molecule generation. Here, we do not learn the conditional distribution during training but delegate conditioning entirely to the sampling algorithm.

3 Inpainting

Generating compounds or parts thereof conditioned on the molecular context is reminiscent of inpainting, a technique originally introduced for completing missing parts of images (Song et al., 2020; Lugmayr et al., 2022) but also adopted in other domains, including biomolecular structures (Wang et al., 2022). An unconditional diffusion model can generate approximate conditional samples if the given context is injected into the sampling process by modifying the probabilistic transition steps. This idea has been successfully applied to image imputation (Song et al., 2020; Lugmayr et al., 2022) but can be readily extended to three-dimensional point cloud data.

As an example, Figure 2 outlines the sampling procedure for molecule generation with a fixed scaffold. A subset of all atoms is fixed and serves as the molecular context we want to condition on. All other atoms are generated by the DDPM. To the this end, we diffuse the fixed atoms at each time step and predict a new latent representation $z_{t-1}^{\text{gen}}$ with the neural network. We then replace the generated atoms corresponding to fixed nodes with their forward noised counterparts:

where $\mathcal{M}$ denotes the set of mask indices used to uniquely identify nodes corresponding to fixed atoms. In this manner, we traverse the Markov chain in reverse order from $t=T$ to $t=0$ to generate conditional samples. Because the noise schedule decreases the noising process’s variance to almost zero at $t=0$ (Equation (1)), the final sample is guaranteed to contain an unperturbed representation of the fixed atoms. This approach can be applied to pocket-conditioned ligand-inpainting by fixing all pocket nodes when sampling from the joint distribution model (Section 3.2). However, it is much more general and allows us to mask and replace arbitrary parts of the ligand-pocket system without retraining—an option we explore in Section 4.4.

Resampling

Trippe et al. (2022) show that this simple replacement method inevitably introduces approximation error that can lead to inconsistent inpainted regions. In our experiments, we observe that the inpainting solution sometimes generates dislocated molecules that are not properly positioned in the target pocket (see Figure 7 for an example). Trippe et al. (2022) propose to address this limitation with a particle filtering scheme that upweights more consistent samples in each denoising step. We, however, choose to adopt the conceptually simpler idea of resampling (Lugmayr et al., 2022), where each latent representation is repeatedly diffused back and forth before advancing to the next time step as reviewed in Appendix C.1. This enables the model to harmonize its prediction for the generated part and the noisy sample from the fixed part (Eq. (7)), which does not include any information about the generated part. We choose $r=10$ resamplings per denoising step based on empirical results discussed in Appendix C.1.

Experiments

We use the CrossDocked dataset (Francoeur et al., 2020) and follow the same filtering and splitting strategies as in previous work (Luo et al., 2021; Peng et al., 2022). This results in 100,000 high-quality protein-ligand pairs for the training set and 100 proteins for the test set. The split is done by 30% sequence identity using MMseqs2 (Steinegger & Söding, 2017).

Binding MOAD

We also evaluate our method on experimentally determined protein-ligand complexes found in Binding MOAD (Hu et al., 2005) which are filtered and split based on the proteins’ enzyme commission number as described in Appendix D. This results in 40,344 protein-ligand pairs for training and 130 pairs for testing.

Baselines

We compare with three recent deep learning methods for structure-based drug design. 3D-SBDD (Luo et al., 2021), Pocket2Mol (Peng et al., 2022) and GraphBP (Liu et al., 2022) are auto-regressive schemes relying on graph representations of the protein pocket and previously placed atoms to predict probabilities based on which new atoms are added. 3D-SBDD uses heuristics to infer bonds from generated atomic point clouds while Pocket2Mol directly predicts them during the sequential generation process. GraphBP incorporates relative locations with respect to a local coordinate system when placing each new atom. Additionally, we include a concurrently developed diffusion model for target-aware molecule design, TargetDiff (Guan et al., 2023), in the comparison. TargetDiff is conceptually very similar to our conditionally trained model (DiffSBDD-cond) but employs a different diffusion formalism for the categorical atom types. For 3D-SBDD, Pocket2Mol and TargetDiff, we re-evaluate the generated ligands on the CrossDocked dataset from the original papers. For GraphBP, we train the model with both CrossDocked and Binding MOAD datasets with the default hyperparameters. We also attempted to train Pocket2Mol on Binding MOAD, but did not manage to robustly train the model on this dataset due to instability during training.

Evaluation

For every experiment, we evaluated all combinations of all-atom and $C_{\alpha}$ level graphs with conditional and inpainting-based approaches respectively. Full details of model architecture and hyperparameters are given in Appendix C. We sampled 100 valid moleculesDue to occasional processing issues the actual number of available molecules is slightly lower on average (see Appendix G.1). for each target pocket and removed all atoms that are not bonded to the largest connected fragment. Ligand sizes were sampled from the training distribution as described in Appendix C. This procedure yields significantly smaller molecules than the reference ligands from the test set. We therefore increase the mean ligand size by 5 atoms for CrossDocked and 10 atoms for Binding MOAD, respectively, to sample approximately equally sized molecules. This correction improves the observed docking scores which are highly correlated with the ligand size (see Figure 6).

We employ widely-used metrics to assess the quality of our generated molecules (Peng et al., 2022; Li et al., 2021): (1) Vina Score is an empirical estimation of binding affinity between small molecules and their target pocket; (2) QED is a simple quantitative estimation of drug-likeness combining several desirable molecular properties; (3) SA (synthetic accessibility) is a measure estimating the difficulty of synthesis; (4) Lipinski measures how many rules in the Lipinski rule of five (Lipinski et al., 2012), which is a loose rule of thumb to assess the drug-likeness of molecules, are satisfied; (5) Diversity is computed as the average pairwise dissimilarity (1 - Tanimoto similarity) between all generated molecules for each pocket; (6) Inference Time is the average time to sample 100 molecules for one pocket across all targets. All docking scores and chemical properties are calculated with QuickVina2 (Alhossary et al., 2015) and RDKit (Landrum et al., 2016).

2 De novo molecular design

De novo molecular design refers to designing new molecules from scratch. In our context, we sample new ligands that bind to specific protein pockets. We present quantitative benchmarks covering a range of molecular properties and in silico docking scores (Table 1) and provide examples of generated molecules for a qualitative assessment (Figure 3).

Overall, the experimental results in Table 1 suggest that DiffSBDD can generate diverse small-molecule compounds with predicted high binding affinity, outperforming state-of-the-art methods. DiffSBDD’s advantage is even more apparent when only the best molecules from each pocket are compared (see Top-10% column and Table 8). Not surprisingly, however, molecules generated by TargetDiff have extremely similar properties to DiffSBDD molecules. Among our models, we observe that the joint probability approach with inpainting, while being the more complex case, achieves higher Vina scores than analogous conditional models and often yields better ligands. DiffSBDD aims to approximate the probability distribution of ligands interacting with their binding pockets and does not explicitly optimize any molecular property. It should therefore generate molecules with scores similar to the distribution observed for reference ligands from the test set. We see that this is indeed the case for all metrics except the synthesizability score (SA). While it is unclear why the model fails to better approximate the distribution of synthetic accessibility scores, simple techniques can be used for downstream optimization of this property once promising candidates are found (Section 4.3).

Generally, presenting the full atomic context to the model improves Vina scores as well as the agreement of predicted conformations with docked poses (Appendix G.5) compared to the $C_{\alpha}$-only models. This can be seen as evidence of better protein-ligand interaction modelling. Interestingly, it does not seem to happen at the expense of reduced diversity despite stricter constraints. It is also worth noting that DiffSBDD cannot reconstruct the full atomic detail of protein pockets well in the challenging case of joint distribution learning in full atom mode. Nevertheless, with inpainting we enforce correct pocket structures externally, which suffices to achieve competitive results empirically. A representative selection of molecules for two targets (4tos and 2v3r) are presented in Figure 3. This set is curated to be representative of our high scoring molecules, with both realistic and non-realistic motifs shown. Random selections of generated molecules are presented in Figure 10.

Binding MOAD

Results for the Binding MOAD dataset with experimentally determined binding complex data are reported in Table 1 as well. DiffSBDD generates highly diverse molecules but, unlike the CrossDocked case, docking scores are lower on average than corresponding reference ligands from this dataset. We believe the reason to be twofold: the Binding MOAD training set is much smaller and also contains more challenging ground-truth ligands (native binders) whereas CrossDocked complexes can have unrealistic protein-ligand interactions. This hypothesis is supported by the lower Vina scores of reference molecules from the synthetic dataset.

Generated molecules for two representative targets (PDB: 6c0b and PDB: 3gt9) are shown in Figure 3. The first target (PDB: 6c0b) is a human receptor which is involved in microbial infection (Chen et al., 2018) and possibly tumor suppression (Ding et al., 2016). The reference molecule, a long fatty acid (see Figure 3, bottom left panel) that aids receptor binding (Chen et al., 2018), has too high a number of rotatable bonds and low a number of hydrogen bond donors/acceptors to be considered a suitable drug (QED of 0.36). Our model however, generates drug-like (QED between 0.29-0.62) and suitably sized molecules by adding aromatic rings connected by a small number of rotatable bonds, which allows the molecules to adopt a complementary binding geometry and is entropically favourable (by reducing the degrees of freedom), a classic technique in medicinal chemistry (Ritchie & Macdonald, 2009).

3 Molecule Optimization

We use our model to optimize exciting candidate molecules, a common task in drug discovery called lead optimization. This is when we take a compound found to have high binding affinity and optimize it for better ‘drug-like’ properties. We first noise the atom features and coordinates for $t$ steps (where $t$ is small) using the forward diffusion process. From this partially noised sample, we can then denoise the appropriate number of steps with the reverse process until $t=0$. This allows us to sample new candidates of various properties whilst staying in the same region of active chemical space, assuming $t$ is small (Appendix Figure 12). This approach is inspired by (Luo et al., 2022) but note this does not allow for direct optimization of specific properties. Instead, it can be regarded as an exploration around the local chemical space whilst maintaining high shape and chemical complementary via the conditional DDPM.

We extend this idea by combining the partial noising/denoising procedure with a simple evolutionary algorithm that optimizes for specific molecular properties. We find that our model performs well at this task out-of-the-box without requiring additional fine-tuning. As a showcase, we optimize a molecule in the test set targeting PDB 5ndu, a cancer therapeutic (Barone et al., 2020), which has low SA and QED scores, 0.31 and 0.35 respectively, but high binding affinity. Over a number of rounds of optimization, we can observe significant increases in QED (from 0.35 to mean of 0.43) whilst still maintaining high similarity to the original molecule (Figure 4a). We can also rescue the low synthetic accessibility score of the seed molecule by producing a battery of highly accessible molecules when selecting for SA. Finally, we observe that we can perform significant optimization of binding affinity after only a few rounds of optimization. Figures 4b-d show three optimization runs alongside representative molecules with substantially optimized scores (QED, SA or Vina) whilst maintaining comparable binding affinity and globally similar structures. Full details are provided in Appendix G.8.

4 Flexible molecule design with inpainting

In drug discovery it is very common to design molecules around previously identified potent substructures. For example, we may wish to design a scaffold around a set of functional groups (scaffold hopping) or extend an existing fragment to make a whole molecule (fragment growing). We can realise a number of these sub-tasks through molecular inpainting (see Section 3.3), whereby we can inpaint in and around fixed regions of the molecule to design whole molecules. Unlike previous methods, using DiffSBDD in this a way does not require retraining a new model on any specialized, usually synthetic, datasets. Rather, the simple definition of an arbitrary binary mask allows the method to generalize to any inpainting task whilst using a conditional DDPM trained on all protein-drug data. Four case studies are presented in Figure 5 as a proof-of-concept, with full methodological descriptions in Appendix C.2.

Conclusion

In this work, we propose DiffSBDD, an $SE(3)$-equivariant 3D-conditional diffusion model for structure-based drug design. We demonstrate the effectiveness of DiffSBDD in generating novel and diverse ligands with predicted high-affinity for given protein pockets on two datasets. We also show that a joint model in combination with an inpainting-based sampling approach can achieve competitive results to direct conditioning on a wide range of molecular metrics. Finally, we provide case studies that showcase the model’s versatility and applicability to downstream drug design tasks. Future work will attempt to study the performance in these domains further. We hope that this flexible approach to using DiffSBDD allows medicinal chemists to augment their domain expertise with a powerful data-driven diffusion model.

Acknowledgements

We thank Xingang Peng and Shitong Luo for providing us generated molecules of the Pocket2Mol and 3D-SBDD methods. We thank Hannes Stärk and Joshua Southern for valuable feedback and insightful discussions. This work was supported by the European Research Council (starting grant no. 716058), the Swiss National Science Foundation (grant no. 310030_188744), and Microsoft Research AI4Science. Charles Harris is supported by the Cambridge Centre for AI in Medicine Studentship which is in turn funded by AstraZeneca and GSK. Michael Bronstein is supported in part by ERC Consolidator grant no. 724228 (LEMAN).

References

Appendix A Variational Lower Bound

To maximise the likelihood of our training data, we aim at optimising the variational lower bound (VLB) (Kingma et al., 2021; Hoogeboom et al., 2022)

during training. The prior loss should always be close to zero and can be computed exactly in closed form while the reconstruction loss must be estimated as described in Hoogeboom et al. (2022). In practice, however, we simply minimise the mean squared error $\mathcal{L}_{\text{train}}=\frac{1}{2}||\bm{\epsilon}-\hat{\bm{\epsilon}}||^{2}$ while randomly sampling time steps $t\sim\mathcal{U}(0,\dots,T)$, which is equivalent up to a multiplicative factor.

Appendix B Note on Equivariance of the Conditional Model

The 3D-conditional model can achieve equivariance without the usual “subspace-trick”. The coordinates of pocket nodes provide a reference frame for all samples that can be used to translate them to a unique location (e.g. such that the pocket is centered at the origin: $\sum_{i}\bm{x}^{(P)}_{i}=\bm{0}$). By doing this for all training data, translation equivariance becomes irrelevant and the CoM-free subspace approach obsolete. To evaluate the likelihood of translated samples at inference time, we can first subtract the pocket’s center of mass from the whole system and compute the likelihood after this mapping. Similarly, for sampling molecules we can first generate a ligand in a CoM-free version of the pocket and move the whole system back to the original location of the pocket nodes to restore translation equivariance. As long as the mean of our Gaussian noise distribution $p(\bm{z}_{t}|\bm{z}^{(P)}_{\text{data}})=\mathcal{N}(\bm{\mu}(\bm{z}^{(P)}_{\text{data}}),\sigma^{2}\bm{I})$ depends equivariantly on the pocket node coordinates $\bm{x}^{(P)}$, $O(3)$-equivariance is satisfied as well (Appendix E). Since this change did not seem to affect the performance of the conditional model in our experiments, we decided to keep sampling in the linear subspace to ensure that the implementation is as similar as possible to the joint model, for which the subspace approach is necessary.

Appendix C Implementation Details

As part of a sample’s overall likelihood, we compute the empirical joint distribution of ligand and pocket nodes $p(N_{L},N_{P})$ observed in the training set and smooth it with a Gaussian filter ($\sigma=1$). In the conditional generation scenario, we derive the distribution $p(N_{L}|N_{P})$ and use it for likelihood computations.

For sampling, we can either fix molecule sizes manually or sample the number of ligand nodes from the same distribution given the number of nodes in the target pocket:

For the experiments discussed in the main text, we increase the mean size of sampled molecules by 5 (CrossDocked) and 10 (Binding MOAD) atoms, respectively, to approximately match the sizes of molecules found in the test set. This modification makes the reported QuickVina scores more comparable as the in silico docking score is highly correlated with the molecular size, which is demonstrated in Figure 6. Average molecule sizes after applying the correction are shown in Table 2 together with corresponding values for generated molecules from other methods.

Preprocessing

All molecules are expressed as graphs. For the $C_{\alpha}$ only model the node features for the protein are set as the one hot encoding of the amino acid type. The full atom model uses the same one hot encoding of atom types for ligand and protein nodes. We refrain from adding a categorical feature for distinguishing between protein and ligand atoms in this case and continue using two separate MLPs for embedding the node features instead.

Noise schedule

We use the pre-defined polynomial noise schedule introduced in (Hoogeboom et al., 2022):

A tiny offset $\epsilon=10^{-5}$ is used to avoid numerical problems at $t=0$ defining the final noise schedule:

Feature scaling

We scale the node type features $\bm{h}$ by a factor of 0.25 relative to the coordinates $\bm{x}$ which was empirically found to improve model performance in previous work (Hoogeboom et al., 2022). To train joint probability models in the all-atom scenario, it was necessary to scale down the coordinates (and corresponding distance cutoffs) by a factor of 0.2 instead in order to avoid introducing too many edges in the graph near the end of the diffusion process at $t=T$.

Hyperparameters

Postprocessing

For postprocessing of generated molecules, we use a similar procedure as in (Luo et al., 2021). Given a list of atom types and coordinates, bonds are first added using OpenBabel (O’Boyle et al., 2011). We then use RDKit to sanitise molecules and filter for the largest molecular fragment.

C.1 Resampling

Here, we briefly recapitulate the resampling algorithm introduced in Ref. (Lugmayr et al., 2022). The key intuition is that inpainting with the replacement method combines a generated part with an independently sampled latent representation of the known part. Even though the neural network tries to reconcile these two components in every step of the denoising trajectory, it cannot succeed because the same issue reoccurs in the following step. Lugmayr et al. (2022) thus propose to apply the neural network several times before proceeding to the next noise level, allowing the DDPM to preserve more conditional information and move the sample closer to the data distribution again. The procedure is summarized in Algorithm 1 including our center of mass adjustment discussed in Section 3.3.

To empirically study the effect of the number of resampling iterations applied, we generated ligands for all test pockets with $r=1$, $r=5$, and $r=10$ resampling steps, respectively. Because the resampling strategy slows down sampling approximately by a factor of $r$, we used the striding technique proposed by Nichol & Dhariwal (2021) and reduced the number of denoising steps proportionally to $r$. Nichol & Dhariwal (2021) showed that this approach reduces the number of sampling steps significantly without sacrificing sample quality. In our case, it allows us to retain sampling speed while increasing the number of resampling steps.

To gauge the effect of resampling for molecule generation we show the distribution of RMSD values between the center of mass of reference molecules and generated molecules in Figure 7. The unmodified replacement method ($r=1$) produces molecules that are clearly farther away from the presumed pocket center than the conditional model. Increasing $r$ moves the mean distance closer to the average displacement of molecules from the conditional method. This effect seems to saturate at $r=10$ which is in line with the results obtained for images (Lugmayr et al., 2022).

Table 4 shows that neither the additional resampling steps nor the shortened denoising trajectory degrade the performance on the reported molecular metrics. The average docking scores even improve slightly which might reflect better positioning of generated ligands in the pockets prior to docking. The same model trained with $T=500$ diffusion steps was used in all three cases.

C.2 Flexible molecule design with inpainting

Molecular scaffolds, also called Bemis-Murcko scaffolds, are often a defined as the molecular core of a compound to which R (functional) groups are attached via attachment points (Bemis & Murcko, 1996). Defining the scaffold allows us to study the key functional groups mediating receptor binding and uncover structure-activity relationships (SAR). Once understood, we can fix the molecular decorations and build new scaffolds to support those functional groups in place, thus allowing us to find structurally distinct compounds that have the same/similar activity. Such a task is known as scaffold hopping and is a central process in medicinal chemistry which allows us to discover structurally novel compounds (Böhm et al., 2004).

In our case study (Figure 5a) we perform scaffold hopping around the functional groups of PDB:2gm1. Such a task could be simplified by partially noising then denoising the existing high performing scaffold and may even be desirable if there are medicinal or synthetic chemistry constraints in place. However, to demonstrate the power of our approach, we instead choose to sample a brand new scaffold de novo by sampling $n$ (where $n$ is the number of atoms present in the original scaffold) new atoms from $\mathcal{N}(\bm{0},\bm{I})$ whilst fixing the functional groups during inpainting.

Scaffold elaboration

Scaffold elaboration can be seen as the inverse problem of scaffold hopping. Here, we have a known scaffold of ideal structure and we wish to enumerate/optimize the functional groups at or around the attachment points to improve activity (Abel & Bhat, 2017).

For our case study (Figure 5b), we perform optimisation of all the R-groups attached to the molecule in PDB:5ndu. $n$ new atoms are sampled de novo as above while the scaffold is kept fixed. Note, the attachment points can be moved if deemed suitable by our learnt model.

Fragment merging

Fragment merging is the task of combining fragments with an overlapping binding site (Li, 2020).

For the fragment merging example, instead of masking existing molecules, we instead take fragments from an experimental fragment screen against the non structural protein 3 (NSP3) from SARS-CoV-2 (Schuller et al., 2021). Using two fragments as input (PDB:5rue and PDB:5rsw) we successfully replicate the fragment merge undertaking in Gahbauer et al. (2023) which was accomplished using the chemoinformatics-based approach Fragmensteinhttps://github.com/matteoferla/Fragmenstein.

Fragment growing

Fragment growing is the task of adding chemical groups/whole other fragments to a seed fragment to improve properties (Li, 2020).

For our case study (Figure 5d), we perform fragment growing in two directions around the central fragment in PDB:5ndu. The molecule is broken into three fragments using the Breaking Retrosynthetically Interesting Chemical Substructures (BRICS) (Degen et al., 2008) protocol. Out of these, the middle fragment is fixed whilst the two outer fragment, plus accompanying linkers, are redesigned. Here, we elect to showcase the second regime of partially noising/denoising the inpainted region, allowing us to partially redesign the two peripheral fragments whilst fixing the central one, which could be useful for a number of retrosynthetic and activity reasons.

Appendix D Binding MOAD Dataset

Appendix E Proofs

Here we used $||\bm{R}\bm{x}||_{2}=||\bm{x}||_{2}$ for orthogonal $\bm{R}$.

E.2 O​(3)𝑂3O(3)-equivariance of the transition probabilities

The denoising transition probabilities from time step $t$ to $s<t$ are defined as isotropic normal distributions:

Therefore, $p_{\theta}(\bm{x}^{(L)}_{t-1}|\bm{x}^{(L)}_{t},\hat{\bm{x}}^{(L)},\bm{x}^{(P)})$ is $O(3)$-equivariant by a similar argument to Section E.1 if $\bm{\mu}_{t\rightarrow s}$ is computed equivariantly from the three-dimensional context.

Recalling the definition of $\bm{\mu}_{t\rightarrow s}=\frac{\alpha_{t|s}\sigma_{s}^{2}}{\sigma_{t}^{2}}\bm{x}^{(L)}_{t}+\frac{\alpha_{s}\sigma_{t|s}^{2}}{\sigma_{t}^{2}}\hat{\bm{x}}^{(L)}$, we can prove its equivariance as follows:

where $\hat{\bm{x}}^{(L)}$ defined as $\hat{\bm{x}}^{(L)}=\frac{1}{\alpha_{t}}\bm{x}_{t}^{(L)}-\frac{\sigma_{t}}{\alpha_{t}}\hat{\bm{\epsilon}}$ is equivariant because:

E.3 O​(3)𝑂3O(3)-equivariance of the learned likelihood

and the sample’s likelihood is $O(3)$-equivariant:

Appendix F S​E​(3)𝑆𝐸3SE(3)-equivariant Graph Neural Network

Chiral molecules cannot be superimposed by any combination of rotations and translations. Instead they are mirrored along a stereocenter, axis, or plane. As chirality can significantly alter a molecule’s chemical properties, we use a variant of the $E(3)$-equivariant graph neural networks (Satorras et al., 2021) presented in Equations (3)-(5) that is sensitive to reflections and hence $SE(3)$-equivariant. We change the coordinate update equation, Equ. (5), of standard EGNNs in the following way

where $\bar{\bm{x}}^{l}$ denotes the center of mass of all nodes at layer $l$. This modification makes the EGNN layer sensitive to reflections while staying close to the original formalism. Since the resulting graph neural networks are only equivariant to the $SE(3)$ group, we will hereafter call them SE(3)GNNs for short.

This result shows that the output coordinates are only equivariantly transformed if $\bm{R}$ is orientation preserving, i.e. $\det(\bm{R})=1$. If $\bm{R}$ is a reflection ($\det(\bm{R})=-1$), coordinates will be updated with an additional summand that breaks the symmetry.

The learnable coefficients $\phi_{x}^{d}(\bm{h}^{l}_{i},\bm{h}^{l}_{j},d_{ij}^{2},a_{ij})$ and $\phi_{x}^{\times}(\bm{h}^{l}_{i},\bm{h}^{l}_{j},d_{ij}^{2},a_{ij})$ only depend on relative distances and are therefore $E(3)$-invariant. Their arguments are represented with the “$\cdot$” symbol for brevity. Likewise, the normalization factor $||(\bm{x}_{i}^{l}-\bar{\bm{x}}^{l})\times(\bm{x}_{j}^{l}-\bar{\bm{x}}^{l})||$ is abbreviated as $Z^{\times}_{ij}$. Already in the first line we used the fact that the mean transforms equivariantly. Furthermore, we use $\bm{R}\bm{a}\times\bm{R}\bm{b}=\det(\bm{R})\bm{R}(\bm{a}\times\bm{b})$ in the second step, which can be derived as follows:

F.2 Empirical Results

To show the effectiveness of this architecture on a simple toy example, we repeat the classification experiment by Adams et al. (2021) who train neural networks to classify tetrahedral chiral centers as right-handed (rectus, ‘R’) or left-handed (sinister, ‘S’). We closely follow their data split and experimental set-up and only replace the classifier with EGNN and SE(3)GNNs, respectively. The results in Table 5 clearly demonstrate that the $SE(3)$-equivariant EGNN is capable of solving this task (without any hyperparameter optimization) whereas the $E(3)$-equivariant version does not do better than random guessing.

Appendix G Extended results

The numbers of available molecules differ slightly between different methods due to computational issues or missing molecules in the available baseline sets. More precisely, on average $100$, $99.7$, $100$ and $100$ molecules have been evaluated per pocket for DiffSBDD-cond, Diff-SBDD-inpaint, DiffSBDD-cond ($C_{\alpha}$), and DiffSBDD-inpaint ($C_{\alpha}$), respectively. For Pocket2Mol, $98.4$ molecules are available per pocket. The set of 3D-SBDD molecules does not contain generated ligands for two test pockets. For the remaining $98$ pockets, $89.9$ molecules are available on average. For GraphBP, $86.1$ molecules are evaluated per pocket. For Binding MOAD, $100$ molecules were available per pocket for evaluation of all DiffSBDD models except for the conditional model with full atomic pocket representation ($99.9$ on average). We could further evaluate $88$ molecules from GraphBP per pocket.

G.2 Additional Molecular Metrics

In addition to the molecular properties discussed in Section 4 we assess the models’ ability to produce novel and valid molecules using four simple metrics: validity, connectivity, uniqueness, and novelty. Validity measures the proportion of generated molecules that pass basic tests by RDKit–mostly ensuring correct valencies. Connectivity is the proportion of valid molecules that do not contain any disconnected fragments. We convert every valid and connected molecule from a graph into a canonical SMILES string representation, count the number unique occurrences in the set of generated molecules and compare those to the training set SMILES to compute uniqueness and novelty respectively.

Table 6 shows that only a small fraction of all generated molecules is invalid and must be discarded for downstream processing. A much larger percentage of molecules is fragmented but, since we can simply select and process the largest fragments in these cases, low connectivity does not necessarily affect the efficiency of the generative process. Moreover, all models produce diverse sets of molecules unseen in the training set.

G.3 Octanol-water partition coefficient

The octanol-water partition coefficient ($\log P$) is a measure of lipophilicity and is commonly reported for potential drug candidates (Wildman & Crippen, 1999). We summarize this property for our generated molecules in Table 7.

G.4 Vina scores for subsets of generated molecules

Table 8 summarizes Vina docking scores for different subsets of molecules with the highest scores among all molecules generated for a given pocket. While there is a significant number of poorly scoring molecules among all generated candidates, these results show that DiffSBDD also produces a large population with high predicted binding affinity. Furthermore, the performance gap between DiffSBDD and competing methods increases as the selection becomes stricter.

G.5 Agreement of generated and docked conformations

G.6 Random generated molecules

Randomly selected molecules generated with our method and 2 baseline methods (3D-SBDD and Pocket2Mol) when trained with CrossDocked are presented in Figure 10.

G.7 Distribution of docking scores by target

We present extensive evaluation of the docking scores for our generated molecules in Figure 11. We evaluate all models trained with a given dataset first against all targets (Figure 11A+C) and 8 randomly chosen targets (Figure 11B+D). We note that the all-atom model trained using CrossDocked data outperforms all other methods. Unsurprisingly, model performance is highly target dependent, likely varying with properties like pocket geometry, size, charge, and hydrophbicity, which would affect the propensity of generating high affinity molecules.

G.8 Optimization

We demonstrate the effect the number of noising/denoising steps ($t$) has on various molecular properties in Figure 12. We test all values of $t$ at intervals of 10 steps and 200 molecules are sampled at every timestep. Note this does not allow for explicit optimization of any particular property unless combined with the evolutionary algorithm.

During the evolutionary algorithm, at the end of every generation the top 10 docking molecules are used to seed the next population. Every seed molecule is elaborated into 20 new candidates with a randomly chosen $t$ between 10 and 150. To make the first population, we start with the single reference molecule and sample 200 new molecules with $t$ chosen as above.

Appendix H Related Work

Inspired by non-equilibrium thermodynamics, diffusion models have been proposed to learn data distributions by modeling a denoising (reverse diffusion) process and have achieved remarkable success in a variety of tasks such as image, audio synthesis and point cloud generation (Kingma et al., 2021; Kong et al., 2021; Luo & Hu, 2021). Recently, efforts have been made to utilize diffusion models for molecule design (Du et al., 2022b). Specifically, Hoogeboom et al. (2022) propose a diffusion model with an equivariant network that operates both on continuous atomic coordinates and categorical atom types to generate new molecules in 3D space. Torsional Diffusion (Jing et al., 2022) focuses on a conditional setting where molecular conformations (atomic coordinates) are generated from molecular graphs (atom types and bonds). Similarly, 3D diffusion models have been applied to generative design of larger biomolecular structures, such as antibodies (Luo et al., 2022) and other proteins (Anand & Achim, 2022; Trippe et al., 2022).

Structure-based Drug Design

Structure-based Drug Design (SBDD) (Ferreira et al., 2015; Anderson, 2003) relies on the knowledge of the 3D structure of the biological target obtained either through experimental methods or high-confidence predictions using homology modelling (Kelley et al., 2015). Candidate molecules are then designed to bind with high affinity and specificity to the target using interactive software (Kalyaanamoorthy & Chen, 2011) and often human-based intuition (Ferreira et al., 2015). Recent advances in deep generative models have brought a new wave of research that model the conditional distribution of ligands given biological targets and thus enable de novo structure-based drug design. Most of recent work consider this task as a sequential generation problem and design a variety of generative methods including autoregressive models, reinforcement learning, etc., to generate ligands inside protein pockets atom by atom (Drotár et al., 2021; Luo et al., 2021; Li et al., 2021; Peng et al., 2022).

Geometric Deep Learning for Drug Discovery

Geometric deep learning refers to incorporating geometric priors in neural architecture design that respects symmetry and invariance, thus reduces sample complexity and eliminates the need for data augmentation (Bronstein et al., 2021). It has been prevailing in a variety of drug discovery tasks from virtual screening to de novo drug design as symmetry widely exists in the representation of drugs. One line of work introduces graph and geometry priors and designs message passing neural networks and equivariant neural networks that are permutation- and translation-, rotation-, reflection-equivariant, respectively (Duvenaud et al., 2015; Gilmer et al., 2017; Satorras et al., 2021; Lapchevskyi et al., 2020; Du et al., 2022a). These architectures have been widely used in representing biomolecules from small molecules to proteins (Atz et al., 2021) and solving downstream tasks such as molecular property prediction (Schütt et al., 2018; Klicpera et al., 2020), binding pose prediction (Stärk et al., 2022) or molecular dynamics (Batzner et al., 2022; Holdijk et al., 2022). Another line of work focuses on generative design of new molecules (Du et al., 2022b, c). More specifically, molecule design is formulated as a graph or geometry generation problem, following either a one-shot generation strategy that generates graphs (atom and bond features) in one step or attempting to generate them sequentially.