DiffBP: Generative Diffusion of 3D Molecules for Target Protein Binding

Generating molecules that bind to specific proteins is an important but challenging task in drug discovery. Previous works usually generate atoms in an auto-regressive way, where element types and 3D coordinates of atoms are generated one by one. However, in real-world molecular systems, the interactions among atoms in an entire molecule are global, leading to the energy function pair-coupled among atoms. With such energy-based consideration, the modeling of probability should be based on joint distributions, rather than sequentially conditional ones. Thus, the unnatural sequentially auto-regressive modeling of molecule generation is likely to violate the physical rules, thus resulting in poor properties of the generated molecules. In this work, a generative diffusion model for molecular 3D structures based on target proteins as contextual constraints is established, at a full-atom level in a non-autoregressive way. Given a designated 3D protein binding site, our model learns the generative process that denoises both element types and 3D coordinates of an entire molecule, with an equivariant network. Experimentally, the proposed method shows competitive performance compared with prevailing works in terms of high affinity with proteins and appropriate molecule sizes as well as other drug properties such as drug-likeness of the generated molecules.

Success in reverse pharmacology proved that structurebased drug design (SBDD) is a promising method for faster and more cost-efficient lead compound discovery [10,27]. It directs the discovery of a lead compound which could show potent target inhibitory activity with knowledge about the disease at the molecular level. However, it is underexplored to use machine learning approaches to design molecules that can bind to a specific protein as the target. On one hand, massive data is required to develop machine learning approaches. Such datasets are recently available [13]. On the other hand, this task is considerably challenging in machine learning for the following three reasons. (i) The protein binding site as the conditional context is complicated. Not only do the 3D geometric structures of the target proteins influence the structure and property of the binding molecules, but also other informative contexts should be considered for generating molecules with high affinities, such as the amino acid types. (ii) The desired distribution over molecular chemistry and coordinates has enormous support sets. Unlike the conformation generation task, the chemical formulas as 2D graph constraints are not known, so the sophisticated coupling of element types, continuous 3D coordinates, and other chemical properties or geometries should be captured by a well-designed model. (iii) Geometric symmetries of molecules should be considered for generalization. In the physical 3D space, these symmetries include translations and rotations from the Euclidean group. In particular, if we perform these symmetry operations on a binding site, the generated molecules are expected to be rotated or translated in the same way.
by the former generated ones and the given protein context. However, several problems exist in such an auto-regressive schema. (i) The sequential modeling of a molecule violates real-world physical scenarios, in which the interactions among atoms are global. In other words, each atom's position and element type are affected by all the other atoms in molecular systems. (ii) The auto-regressive sampling for a molecule usually suffers from 'early stopping' problem. To be specific, the model tends to generate molecules with a small number of atoms (also called molecule size), thereby failing to capture the true distribution of atom numbers in drug-like molecules.
To tackle these significant limitations, we propose the target-aware molecular Diffusion model for Protein Binding (DiffBP). By harnessing diffusion denoising generative models' ability to generate samples of high-quality [6,18,52], and great expressivity of equivariant graph neural networks [22,47], our model can generate molecules of good drug properties, with high affinity when binding to the target protein. We summarize our main contributions as: • We first analyze the problems emerging in autoregressive models from a physics and probabilistic perspective. Motivated by these observations, we propose DiffBP, which directly models the full atoms in targetaware molecule generation.
• Delicately-designed procedures are established, considering molecules' center of mass prediction, atom number distribution modeling and other aspects.
• Experiments demonstrate that our DiffBP achieves promising performance compared with previous methods, in terms of appropriate molecule size, higher affinity with target protein and other drug properties.

Problem Statement
For a protein-molecule (also called protein-ligand) binding system as C, which contains N + M atoms, we represent the index set of molecules as I mol and proteins as I pro , where |I mol | = N and |I pro | = M . To be specific, let a i be the K-dimensional one-hot vector indicating the atom element type of the i-th atom in the binding system and x i be its 3D Cartesian coordinate, and then can be split into two sets as C = M ∩ P, where M = {(a i , x i ) : i ∈ I mol } and P = {(a j , x j ) : j ∈ I pro }. For protein-aware molecule generation, our goal is to establish a probabilistic model to learn the conditional distribution of molecules conditioned on the target proteins, i.e. p(M|P).
Problems in auto-regressive models. Recently proposed auto-regressive models sequentially generate (a i , x i ) by modeling the conditional probability p(a i , j=1 ∪ P is the intermediate binding system at the step i. By this means, the desired probability is modeled as a sequence of conditional distributions, as p(M|P) = N i=1 p(a i , x i |C i−1 ), where C 0 = P. By contrast, in real-world protein-molecule systems, there are force interactions between any pair (or even higher order) of atoms such that the energy function can be decomposed as The stable system reaches an energy-minimal state. From the perspective of energy-based generative models, the corresponding Boltzmann distribution is written as where κ is Boltzmann constant and τ is the temperature, indicating that the modeling of probability is based on joint distributions at a full atom level, instead of sequential conditional distributions. Thus, autoregressive models are likely to violate physical rules. To address it, we employ the following diffusion models for molecule generation at a full-atom level.

Diffusion Models
Diffusion on continuous variables. Diffusion models [50,61] for continuous variables learn the data distribution by manually constructing the forward diffusion process and using a denoising model to gradually remove the noise added in the diffusion process. The latter process is called the reverse denoising process. Denote the input data point by z = z 0 , and the diffusion process adds multivariate Gaussian noise to z t for t = 0, . . . , T , so that where α t ∈ R + is usually monotonically decreasing from 1 to 0, and σ 2 is increasing, which means the retained input signals are gradually corrupted by the gaussian noise along t-axis, leading to q(z T ) = N (0, I). For variancepreserving diffusion process [18], α t = 1 − σ 2 t ; For variance-exploding process [51,52], α t = 1. Following recently proposed Variational Diffusion Models [19,25], where signal-to-noise ratio is defined as The Markov representation of the diffusion process can be equivalently written with transition distribution as where s = t − 1, α t|s = α t /α s and σ 2 t|s = σ 2 t − α 2 t|s σ 2 s . In the true denoising process, when the input signals are given, the transition distribution is also normal, and given by where µ(z 0 , z t ) = In the generative denoising process, because the input signals are not given, it firstly uses a neural network to approximate z 0 byẑ 0 , and then the learned transition distribution which is similar to the Eq. (4) is written as p(z s |z t ) = N (µ(ẑ 0 , z t ), σ 2 t|s I).
Besides, by rewriting the variational lower bound on the likelihood of z 0 , the loss function can be simplified to Instead of directly predictingẑ 0 using neural networks, diffusion models try to predict the added noise in a scorematching way. Specifically, according to Eq. (1), z t = α t z 0 + σ t t , where t ∼ N (0, I), and then the neural network φ predictsˆ = φ(z t , t), so thatẑ 0 = 1 αt z t − σt αtˆ t . Further, the simplified loss function can is written as Diffusion on discrete variables. Differing from gaussian diffusion processes that operate in continuous spaces, diffusion models for discrete variables [2,5] firstly define the diffusion process of random variables z t ∈ {1, . . . , K} with K categories as where z t is the one-hot vector of z t ,Q t = Q 1 Q 2 . . . Q t with [Q t ] ij = q(z t = j|z s = i) denoting diffusion transition probabilities, and [z 0Qt ] i is the probability of z t = i. The true denoising process is given as Here we employ the absorbing discrete diffusion model, which parameterizes the transition matrix Q t as where K + 1 is an absorbing state, usually denoted as [MASK] token in text generation. β t monotonically increases from 0 to 1, means that when t = T , all the discrete variables are absorbed into the K + 1 category. Due to the effectiveness of BERT-style training, a neural network is used to directly predict p(z 0 |z t ) rather than p(z s |z t ), leading to a BERT-like training objective as where λ(t) is the weights at different time step. For example, λ(t) = 1/T leads to equal weights.

Equivariance
For molecular systems, the atoms' positions are represented as 3D Cartesian coordinates, so both equivariance of the transition distributions and invariance of the data distribution of 3D structure are required for generalization w.r.t. the SE(3) group. Formally, f : R 3 → R 3 is an equivariant function w.r.t. SE(3) group, if for any rotation and translation transformation in the group, which is represented by R as orthogonal matrices and t ∈ R 3 respec- In the setting of generative models, the learned distribution p(x) should be invariant to SE(3) group [19]. Kohler et al. [28] showed that an invariant distribution composed with an equivariant invertible function results in an invariant distribution. In the setting of diffusion models, the transition distribution is defined to be SE (3) However, because the translation equivariance cannot be preserved in the diffusion process (Appendix. A.1), we fix the center of mass (CoM) to zero to avoid all the translation transformation (details in Sec. 3.4), and diffuse and denoise the atoms' coordinates in the linear subspace with i∈I mol x i = 0. Further, in the reverse process, because the initial distribution of p(x T ) is set as standard Gaussian, the distribution naturally satisfies invariance to rotation transformation. The zero CoM trick in the corresponding denoising process also circumvents that p(x T + t) = p(x T ), which makes it impossible for p to be a distribution.

Target-aware Diffusion Process
By using the notations of diffusion models, we write the input binding system as C 0 , and the intermediate noisy system as C t , where t = 1, . . . , T . Note that both the diffusion and denoising processes are all performed on M, while the protein system P keeps unchanged in the processes. Therefore, for notation simplicity, we omit the subscript i ∈ I mol in the following and assume that the forward diffusion process as diffusion process is written as Diffusion on discrete types. For the discrete atom element type, we use the absorbing diffusion model, where the noise schedule is chosen as uniform. In detail, assume K + 1 is the absorbing state, and the atom element type a i,t ∈ {1, . . . , K + 1}, corresponding to its one-hot vector

Equivariant Graph Denoiser
To learn the transition distribution p(x i,t−1 , a i,t−1 |C t ), we use the EGNN [47] satisfying rotational-equivariance w.r.t. x i and invariance w.r.t. a i . Specifically, Then a commonly-used SE(3)-EGNN reads where the l-th equivariant convolutional layer is defined as in which i ∈ I mol , and j ∈ I mol ∪ I pro . The input e i is the atom type embedding of a i , e t is the time embedding of t, and V(i) is the neighborhood of i established with KNN according to C t . h (l) i ∈ R D h is the i-th atom's hidden state in the l-th layer which is rotational-invariant. Another alternative architecture of the graph denoiser is Geometric Vector Perceptron [1,22]. In general, we find these two architectures have close performance empirically. In the final layer, a softmax function following a multi-layer percep-

Optimization Objective
Denoising continuous positions. As we set the final layer's output asˆ i,t = x (L) i,t , and E[ i,t ] = 0, we firstly translateˆ i,t , so that iˆ i,t = 0, and then employ Eq. (7) as our loss on atom's continuous position, which reads Denoising discrete types. The rotational-invariant probability vectorp i,0 = softmax(f mlp (h i )) gives the distribution of the i-th atom's element type. We use Eq. (11) as the training loss, to recover the atoms of the absorbing type K +1 back to its true type with the uniform weights of time.
where CE(·, ·) is the cross-entropy loss, which is only calculated on the absorbing types.
Avoiding intersection for binding. In the binding site, chances are that the noisy atom positions x t go across the surface of a protein, leading to the intersection of proteins and molecules. In the reverse denoising process, we hope to include the inductive bias of non-intersection. To achieve it, we add an intersection loss as a regularization term in protein docking [14]. We turn to previous works on the surface of proteins and point cloud reconstruction [54,57], where the surface of a protein point cloud In this way, {x ∈ R 3 : S(x) < γ} is the interior of the protein, and the atoms of the binding molecule should be forced to lay in {x ∈ R 3 : S(x) > γ}. As a result, the inductive bias as a regularization loss function reads wherex i,0 is the approximated positions at t = 0, asx i,0 = 1 αt x i,t − σt αtˆ i,t . γ and ρ are predefined, where we give the choice of them in Sec. 5.3.
Reconstruction on other attributes. In the denoising step of t = 1, our equivariant graph denoiser can also be used to recover other attributes of the binding systems, which is based on x i,0 ≈ x i,1 as the SNR(1) → ∞. For example, in order to predict the binary atom attribute z aro of 'is aromatic', another prediction head g mlp : i,1 )). And the binary cross-entropy loss can be used to train the denoiser, leading to where y aro i is the i-th atom's ground-truth label on the attribute of 'is aromatic'.

Algorithm 1 Training DiffBP (Sampling in Appendix. A.2)
Input: Zero-centered molecule {a j , x j } N j=1 , protein {a j , x j } M j=1 , and graph denoiser φ θ Sample t ∼ U(0, . . . , T ), i,t ∼ N (0, I) Subtract center of mass from i,t Compute To sum up, the overall loss function used in the training process reads

Generative Denoising Process
Generating atom positions. In generating the atom positions, we firstly sample where the parameters is calculated by Generating element types. For atoms' element types, {a i,T } are all set as the absorbing type of K + 1 at time T . Then, in each step at t, we randomly select (T − t)/T of atoms to predict their element types. Besides, if the element type of an atom has been recovered in the past steps, it would not change, which indicates that the recovery of the element type of each atom is only performed once.
, which indicates that E[ˆ i,t ] ≈ 0 for any i ∈ I mol and 0 ≤ t ≤ T . Combining with E[x i,T ] = 0, we can easily obtain that E[x i,t ] ≈ 0. Therefore, we choose the mass center of the binding molecules as zero by translating the protein-molecule binding system such that the origin of the global coordinate system coincides with the molecule's CoM. Otherwise, it is challenging for the generated molecules to be located in the binding site.
Another problem raised by this translation is that the CoM of the molecules is unknown in the generation process, since only the information of proteins is the pre-given context as conditions. Besides, as a full-atom generation method, it is necessary to assign the atom numbers of the binding molecules before both diffusion and denoising process. To address these problems, an additional graph neural network ϕ ω (P) ∈ R 3 × Z + with the same intermediate architectures as shown in Eq. (17) is pre-trained as a pregeneration model, which aims to generate the atom numbers as well as the molecules' CoM before the denoising process. In detail, where x mol c is finally obtained by an SE(3)-equivariant pooling as mean({x ) is used to predict the distribution of atom numbers with u mlp : R D h → R Ns , where N s is the number of different molecule sizes which can be statistically obtained by the training set, and u mlp transforms the latent features into logits for calculating the probability. A dictionary is used to map the predicted classn s to Z + . For example, Dict = {1 : 18, 2 : 27} means whenn s = 1, the atom number N mol will be assigned as 18. Once the molecule's CoM and atom number are obtained, we translate the protein system by x j := x j − x mol c for j = 1, . . . , M , and set N := N mol as the molecule size for generation of the diffusion model. The overall workflow of DiffBP is presented in Fig. 1.
Structure-based drug design. Increasing scientific interests have been raised in SBDD, which aims to generate both 2D chemical formulas and 3D structures conditioned on target protein structure as contextual constraints. For example, LiGAN [37] and 3DSBDD [34] are grid-based models which predict whether the grid points are occupied by specific atoms. Further, Pocket2Mol [42] and GraphBP [30] generate atoms auto-regressively and directly model the probability of the next atom's type as discrete categorical attribute and position as continuous geometry, achieving success in generating the molecules of high validity and affinity with the target proteins.

Setup
Datasets. For evaluation of our model, we follow the previous works [30,37] and use the CrossDocked2020 [13] dataset to generate ligand molecules that bind to target proteins' pockets based on the pocket structures, with the same split of training and test set. For atom-level contexts of protein pockets, we employ element types and the amino acids the pocket atoms belong to as their SE (3) Methods. Three state-of-the-art methods including 3DS-BDD [34], Pocket2Mol [42] and GraphBP [30] are employed as benchmarks. These three methods are all autoregressive generative models. Pocket2Mol and GraphBP directly generate the continuous coordinates of atoms. In comparison, 3DSBDD generates the atoms' position on regular grids. We generate 100 molecules for each protein pocket to calculate metrics for further evaluation.
Metrics. Following GraphBP and LiGAN, for the calculation of Affinity Score, we adopt Gnina [38,44] which is an ensemble of CNN scoring functions that were trained on the CrossDocked2020 dataset. Empirical studies show such CNN predicted affinity is more accurate than Autodock Vina empirical scoring function. Quantitatively, two metrics on Affinity Score are reported: (i) ∆Binding as the percentage of generated molecules that have higher predicted binding affinity than the corresponding reference molecules. It measures the maximum performance that the model can achieve, and the larger the better.  smaller MPBG represents the better binding affinity of the generated molecules on average. Besides, other scores are used for further comparison including (i) QED as quantitative estimation of drug-likeness; (ii) SA as synthetic accessibility score; (iii) Sim as the average Tanimoto similarities of the multiple generated molecules in one pocket, reviling the diversity of generation; (iv) LPSK as the ratio of the generated drug molecules satisfying the Lipinski's rule of five. These chemical properties are calculated by RD-Kit package [29]. After the molecules are generated by the model, Openbabel [41] is used to construct chemical bonds between atoms and Universal Force Field (UFF) minimization [45] is used for refinement before the metric calculation. Note that in Pocket2Mol, refinement with UFF is not performed, so we report the metric of Pocket2Mol-without-UFF in Appendix. B. And the reported metrics in this part are all based on Pocket2Mol-with-UFF.

Property Evaluation
There are two problems we want to figure out in this part: (i) Can DiffBP generate molecules of high binding affinity with target protein at a full-atom level?
(ii) Is the distribution of atom number of molecules generated by DiffBP close to drug-like molecules'?
Guiding by these two questions, we first define the generated molecule sizes at three levels: small, medium, and large, according to the reference molecules' sizes, as whereM is the sampled molecules. We set the threshold δ as 6 in practice. Table. 1 presents the two quantitative binding metrics of the four methods, where 'Ratio' is the proportion of molecules of different sizes to the total number of molecules generated. Conclusions can be reached as: • 3DSBDD and Pocket2Mol both suffer from 'earlystopping' problem of auto-regressive generating process, in which the small ratio is at a high value in comparison with the other two. In contrast, 75.19% of the molecules generated by DiffBP are of medium size, achieving the lowest MPBG and highest ∆Binding.
• It is notable that the calculation of binding scores tends to give larger molecules a higher affinity because larger molecules are more likely to make more interactions with the target proteins, which is a common problem of preference of existing scoring functions. † Therefore, the overall binding score of GraphBP benefits from large proportions of generated molecules with large sizes. However, it also suffers from the distribution shift of atom sizes, i.e. the medium ratio is very small while large molecules account for a substantial part of generated samples. Although the reference molecules are not the gold standard in the task of SBDD, they can reflect some properties of the binding site, e.g. the site size can be reflected by the reference molecule size to some degree. Therefore, using the reference molecule size to define the appropriate atom numbers of generated molecules is reasonable, and comparing the metrics in the groups of different sizes shows the preference for the molecules' sizes of the scoring functions of drugs' properties. Other metrics for drug properties on generated molecules are shown in Table. 2. It shows that all the metrics are competitive for the molecules generated by DiffBP. However, the scoring functions of different metrics also have their preference for molecule sizes, we give the details in Appendix. B. Fig. 2 are two generation examples.

Sub-structure Analysis
We evaluate the ratio of different bond types and molecules containing different rings of the generated molecules and reference, to give a comparison of different methods' generation in terms of sub-structure, shown † https://github.com/gnina/gnina/issues/167 in Table. 3. For atom type analysis, see Appendix. B. As presented, DiffBP tends to generate more bonds generally, as well as 'Pent' rings, which is an immediate consequence of larger atom numbers of the generated molecules. Pocket2Mol has the advantage to generate fewer 'Tri' rings and more 'Hex' rings, in coordination with experimental analysis in [42].

Further Analysis on Modules
Several modules are not necessary to fill the task in our workflow. Here we attempt to figure out whether they can bring improvements to the performance. The first is the intersection loss as regularization as proposed in Sec. 3.3. The regularization term affects the success rate (a.k.a. validity) greatly. Here we give the validity of the molecules obtained by the trained model v.s. different loss parameters, as shown in Figure. 3. It shows that the regularization loss can bring improvements to generation validity slightly. However, when the two coefficients are too large, the validity decreases. The rationale for it is that the loss term can be viewed as a repulsive force from protein atoms to molecule atoms, and when the repulsion is too large, the positions of atoms are likely to collapse into a small region, causing the number of each atom's neighbors too large, i.e. the valences of atoms are larger than allowed.   Secondly, the CoMs and atom numbers obtained by pregeneration models in our workflows can be replaced by the ones given by reference molecules, leading to the variant of DiffBP(Pre-Ref). If the generated CoMs stray too far off the binding sites, or the molecule sizes are too large, the contextual information of protein pockets will hardly affect the generative denoising process effectively. As shown in Table. 4, our pipeline of DiffBP(Pre-Gen) consisting of a pregeneration model and a diffusion generative model works well since the molecules generated by DiffBP(Pre-Gen) achieve comparable affinity scores to DiffBP(Pre-Ref).

Conclusion
A 3D molecule generative model based on diffusion denoising models called DiffBP is proposed, to generate druglike molecules of high binding affinities with the target proteins. The generation process is in line with the laws of physics, i.e. generate molecules at a full atom level. Limitations still exist, including the relatively low validity and deviation from the real drug molecule in the substructures.

A.1. Analysis on invariance in diffusion models
In the forward diffusion process for positions, which reads If a translation t is performed on the input positions and the added noise, then By contrast, for a rotation transformation, As i,t ∼ N (0, 1), by this way and thus For the true position distribution, q(Rx i,0 ) = q(x i,0 ) In this way, the intermediate states diffusion process satisfies the rotation invariance. And to learn the reverse process, we only require p(Rx i,0 |Rx i,t ) to be rotationalequivariant. By Eq. (23), the probability of each generated intermediate state will also be rotation invariance.

A.2. Details in diffusion models
Diffusion Schedule. As α t and σ t are two parameters changing gradually along t-axis. Lots of noise schedules proposed, such as Linear [18] and Cosine [39], here we choose a simple polynomial schedule and variance-preserving process as Sampling Process. We give the detailed pseudo-codes of sampling algorithm of DiffBP in Algorithm 2.

. Binding Affinity
As in Pocket2Mol, the UFF refinement is not mentioned, so here we give the Binding Affinity of Pocket2Molwithout-UFF in Table. B1.

B.2. Detailed drug properties
Since there is a high correlation between Binding Scores and molecule sizes, here we list the other score functions of drug properties to figure out their preferences in Table. B2.
It shows that score functions of QED, SA and LPSK have their preference for molecule sizes. QED and LPSK tend to give smaller molecules higher scores, while SA prefers larger ones.

B.3. Atom type analysis
Expect the sub-structure, we give the ratio of different element types of atoms, for further comparison, as shown in Table. B3. Note that we only count the percentage of heavy atoms, in line with the generative models. It shows that for the most common elements like 'C','O','N', DiffBP and Pocket2Mol generate the average ratio that most closely resembles the true distribution. However, for the uncommon elements, such as 'S' and 'Cl', DiffBP generates more than reference.

B.4. Implementation details
Hyperparameters. In practice, we found that GVP is more numerical stable than EGNN, so we choose 12-layers GVP for the pre-generation model and 9-layers GVP for Graph denoiser. The hidden dimensions of vertices are fixed as 256 and 64 for invariant node attributes (atom types) and equivariant geometry positions, respectively. In training, the batch size is set as 8, and the initial learning rate is set as 10 −4 . AdamW optimizer is used, together with the Exponential Moving Average training trick.
Platforms. We use a single NVIDIA A100(81920MiB) GPU for a trial. The codes are implemented in Python 3.9 mainly with Pytorch 1.12, and run on Ubuntu Linux.