Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization

Authors: siyi  gu
minkai  xu
alexander  powers
weili  nie
tomas  geffner
karsten  kreis
jure  leskovec
arash  vahdat
stefano  ermon
stanford  univeristy
  nvidia
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 00:59:08

Abstract: Generating ligand molecules for specific protein targets, known as structure-based drug design, is a fundamental problem in therapeutics development and biological discovery. Recently, target-aware generative models, especially diffusion models, have shown great promise in modeling protein-ligand interactions and generating candidate drugs. However, existing models primarily focus on learning the chemical distribution of all drug candidates, which lacks effective steerability on the chemical quality of model generations. In this paper, we propose a novel and general alignment framework to align pretrained target diffusion models with preferred functional properties, named AliDiff. AliDiff shifts the target-conditioned chemical distribution towards regions with higher binding affinity and structural rationality, specified by user-defined reward functions, via the preference optimization approach. To avoid the overfitting problem in common preference optimization objectives, we further develop an improved Exact Energy Preference Optimization method to yield an exact and efficient alignment of the diffusion models, and provide the closed-form expression for the converged distribution. Empirical studies on the CrossDocked2020 benchmark show that AliDiff can generate molecules with state-of-the-art binding energies with up to -7.07 Avg. Vina Score, while maintaining strong molecular properties. Code is available at https://github.com/MinkaiXu/AliDiff.

Main paper: 
 introduction: 
 Generating ligand molecules with desirable properties and high affinity to specific protein targets, known as structure-based drug design (SBDD), is a fundamental problem in therapeutic design and biological discovery. It necessitates methods that can produce realistic and diverse drug-like molecules with stable 3D structures and high binding affinities. In the past few years, numerous deep generative models have been proposed to generate molecules in SMILES string representation [Kusner et al., 2017 , Segler et al., 2018] or graph representations [Jin et al., 2018 , Shi et al., 2020 , Guan et al., 2023] . Although these models have shown promise in generating plausible drug-like molecules, they lack sufficient modeling of the 3D protein-ligand interaction with proteins and therefore can hardly be adopted in target-aware molecule generation. As a result, generating ligands conditioned on protein targets remains an open research problem.Recently, with rapid progress in structural biology and the increasing scale of structural data [Francoeur et al., 2020 , Jumper et al., 2021] , numerous target-aware generative models have been proposed to directly generate molecules within the protein targets in 3D. Initial work proposed to sequentially place atoms within the target via autoregressive models [Luo et al., 2021 , Liu et al., 2022 , Peng et al., 2022] , while later work learns diffusion models to jointly design the whole ligand with state-ofthe-art results [Guan et al., 2023 , Lin et al., 2022 , Schneuing et al., 2023 , Huang et al., 2023 , Guan et al., 2024] . Following the biological principle to model the protein-ligand complex interactions, these methods have shown great promise in generating realistic drugs that can bind toward given targets. However, all existing models solely focus on learning the chemical distribution of candidate molecules and treat all training samples equally, while in practice, only the ligand molecules with strong binding affinity and high synthesizability are preferred for real-world therapeutic development. As a result, existing learned models generally lack sufficient steerability regarding the relative quality of model generations and cannot generate faithful samples with the desirable properties. To bridge the gap between existing SBDD models and the necessity for designing ligands with favorable properties, in this paper, we introduce a novel and comprehensive alignment framework to align pretrained target-aware diffusion models with preferred functional properties, named ALIDIFF. ALIDIFF adjusts the target-conditioned chemical distribution toward regions characterized by lower binding energy and structural rationality, as specified by a userdefined reward function, using a preference optimization approach. To this end, we derive a unified variational lower bound to align the likelihoods of both discrete chemical type and continuous 3D coordinate features. We further analyze the winning data overfitting problem commonly associated with preference optimization objectives, and introduce an improved Exact Energy Preference Optimization (E 2 PO) method. E 2 PO analytically ensures a precise and efficient alignment of diffusion models, and we provide a closedform expression for the converged distribution. Our key contributions can be summarized as follows:• We address the challenge of designing favorable target-aware molecules from the perspective of aligning molecule generative models with desirable properties. We introduce the energy preference optimization framework and derive variational lower bounds to align diffusion models for generating molecules with high binding affinity to binding targets. • We analyze the overfitting issue in the preference optimization objective, and propose an improved exact energy optimization method to yield an exact alignment towards target distribution shifted by reward functions. • We conduct comprehensive comparisons and ablation studies on the CrossDocked2020 [Francoeur et al., 2020] benchmark to justify the effectiveness of ALIDIFF. Empirical results demonstrate that ALIDIFF can generate molecules with state-of-the-art binding energies with up to -7.07 Avg. Vina Score, while maintaining strong molecular properties. 
related work: 
 Structure-Based Drug Design. With increasing amount of structural data becoming accessible, generative models have attracted growing attention for structure-based molecule generation. Early research [Skalic et al., 2019] proposes to generate SMILES representations from protein contexts by sequence generative models. Inspired by the progress in 3D and geometric modeling, many works proposed to solve the problem directly in 3D space. For instance, Ragoza et al. [2022] voxelizes molecules within atomic density grids and generates them through a Variational Autoencoder framework. Luo et al. [2021] , Peng et al. [2022] , Liu et al. [2022] , Powers et al. [2023] developed autoregressive models to generate molecules by sequentially placing atoms or chemical groups within the target. Following the autogressive backbone, FLAG [Zhang et al., 2023] and DrugGPS [Zhang and Liu, 2023] take advantage of chemical priors of molecular fragments to generate ligand molecules piece by piece, leading to more realistic substructures. More recently, diffusion models achieved exceptional results in synthesizing high-quality images and texts, which have also been successfully used for ligand molecule generation [Guan et al., 2023 , Lin et al., 2022 , Schneuing et al., 2023 , Huang et al., 2023 , Guan et al., 2024] . These models generate molecules by progressively denoising atom types and coordinates while maintaining physical symmetries with SE(3)-equivariant neural networks. While the existing works focus on designing molecules using various deep generative models, they often struggle with generating molecules that exhibit different desirable properties, e.g., strong binding affinity, high synthesizability, and low toxicity. Real-word drug discovery projects almost always seek to optimize or constrain these properties [D Segall, 2012 , Bickerton et al., 2012] .In this work, we aim to address the challenge with a novel and general preference optimization framework.Reinforcement learning from human feedback (RLHF). Recently, significant efforts have been devoted to aligning generative models with human preferences. The use of reinforcement learning to incorporate feedback from humans and AI into finetuning large language models is exemplified by Reinforcement Learning from Human Feedback (RLHF) [Ziegler et al., 2020 , Ouyang et al., 2022] . Research works have incorporated human feedback to improve performance across various domains, such as machine translation [Nguyen et al., 2017] , summarization [Stiennon et al., 2020] , and also diffusion models [Uehara et al., 2024b,a] . Notably, Rafailov et al. [2023] designed a new preference paradigm that enables training language models to satisfy human preferences directly without reinforcement learning. This algorithm was later applied to diffusion models for text-to-image generation tasks [Wallace et al., 2023] . Concurrent work [Zhou et al., 2024] attempts to apply DPO for designing antibodies with rationality and functionality. To the best of our knowledge, we are the first alignment approach for target-aware ligand design, where the conditional distribution is shifted toward desirable properties. 
method: 
 In this section we present ALIDIFF, a general framework for aligning target-aware diffusion models with various molecular functionalities. We first provide an overview of the target-aware ligand diffusion model and our Reinforcement Learning from Feedback formulation (section 3.1). Next, we introduce the energy optimization approach for aligning the diffusion model and analyze the potential limitations of the framework (section 3.2). We then further introduce an exact energy optimization method from a distribution matching perspective to align the generative model efficiently and exactly (section 3.3). A visualization of the framework is shown in Figure 2 .Notation. We focus on aligning molecule generative models for structure-based drug design, which can be abstracted as generating molecules that can bind to a given protein target. Following the convention in the related literature [Luo et al., 2021 , Guan et al., 2023] , the molecule and target protein are represented as M = {(xEQUATIONP )} N P i=1 , respectively, where N M and N P denote the number of atoms of the molecule M and the protein P. x ∈ R 3 and v ∈ R K denote the atomic 3D position and chemical type, respectively, with K being the dimension of atom types. For brevity, we denote the molecule as a matrix m = [x M , v M ] where x M ∈ R N M ×3 and v M ∈ R N M ×K , and denote the protein as a matrix p = [x P , v P ] where x P ∈ R N P ×3 and v P ∈ R N P ×K . The task can then be formulated as modeling the conditional distribution p(m|p).Preliminaries. Diffusion Models have been previously used to model the joint distribution of atomic types and positions [Guan et al., 2023 , Schneuing et al., 2023 , Lin et al., 2022] . This approach consists of a forward diffusion process and a reverse generative (denoising) process. Both processes are only defined on the ligand molecules m, with fixed proteins p. In the forward process, small Gaussian and categorical noises are gradually injected on atomic coordinates x and types v as follows:EQUATIONwhere N and C stand for the Gaussian and categorical distribution respectively, and β t corresponds to a (fixed or learnable) variance schedule. Note that, in certain recent work q process can be learnable with dependence on the conditioning p [Huang et al., 2023] . We omit the subscript M for the ligand molecule without ambiguity here and denote the atom positions and types at time step t as x t and v t . Using Bayes theorem, the posterior conditioned on m 0 can be computed in closed form:EQUATIONwhere μ(x t , x 0 ) = Ho et al., 2020 , Austin et al., 2021] . At timestep T , q converges to the prior with Gaussians on coordinates and uniforms on atom types. The reverse process, also known as the generative process, learns a neural network parameterized by θ to recover data by iterative denoising. The denoising step can be approximated with predicted Gaussians µ θ and categorical distributions c θ as follows:√ ᾱt-1βt 1-ᾱt x 0 + √ α t (1-ᾱt-1) 1-ᾱt x t , βt = 1-ᾱt-1 1-ᾱt β t , α t = 1 -β t , ᾱt = t s=1 α s , c(v t , v 0 ) = c * K k=1 c * k , and c * (v t , v 0 ) = [α t v t + (1 -α t )/K] ⊙ [ᾱ t-1 v 0 + (1 -ᾱt-1 )/K] [EQUATIONwhere [x 0 , v0 ] = ϵ θ ([x t , v t ], t, p) are predictions from a denoising network ϵ θ . Importantly, the denoising network here is specifically parameterized by equivariant neural networks, resulting in an SE(3)-invariant likelihood p θ (m|p) on the protein-ligand complex [Xu et al., 2022] .Overview. As ligand molecules with desirable properties, e.g., high binding affinity and synthesizability, are required for real-world therapeutic development, we aim to align the ligand diffusion model with these preferences. Such preferences can be defined as a reward model r(•) : M × P → R calculated from various cheminformatics software, e.g., binding affinity, drug-likeness, synthesizability, or their combinations. We fine-tune and align the pre-trained diffusion model with the reinforcement learning framework. Specifically, given a dataset D containing given protein targets, inspired by RLHF [Ouyang et al., 2022] , this fine-tuning is achieved by maximizing the reward:EQUATIONwhere p θ and p ref are the distributions induced by the fine-tuned and pre-trained models, respectively. In this work, p θ and p ref are the fine-tuned and pre-trained molecule diffusion models, as introduced above. β is a hyperparameter controlling the KL divergence regularization. Note that, here the reward is a known black-box function, unlike typical RLHF where it is unknown and has to be estimated from preferences. In the following section, we elaborate on how the alignment objective is rewritten with diffusion forward and reverse processes defined on atomic types and coordinates.Though the reward function is known, evaluating reward values such as binding affinity is computationally expensive and we instead resort to aligning with a labeled offline dataset. We start with a dataset D = {(p, m w , m l )} where p denotes the protein condition and m w ≻ m l is a pair of winning and losing ligands with respect to certain specified energy, e.g., binding energy.The optimal solution to the RLHF objective from Equation (4) can be written in closed-form p * θ (m|p) ∝ p ref (m|p) exp( 1 β r(m, p)) [Peters and Schaal, 2007] . Following the preference optimization algorithm [Rafailov et al., 2023] , we use the Bradley Terry (BT, [Bradley and Terry, 1952] ) model p(m 0 1 ≻ m 0 2 |p) = σ(r(m 0 1 , p) -r(m 0 2 , p)) to reformulate the RLHF objective as:EQUATION)Due to the intractability of p θ (m|p) for diffusion models, we instead follow recent work on diffusionbased preference optimization [Wallace et al., 2023] to align the whole reverse process and utilize Jensen's inequality to optimize its negative evidence lower bound optimization (ELBO):EQUATIONwhere we omit the conditioning on the protein target p for compactness. We further approximate the reverse process p θ (m 1:T |m 0 ) with the forward process q(m 1:T |m 0 ) for efficient sampling of m 1:T , and obtain the following expression after some derivations [Wallace et al., 2023] :LDPO-Diffusion (θ) = -E (p,m w 0 ,m l 0 )∼D,t∼[0,T ],m w t ∼q,m l t ∼q log σ -βT D KL (q(m w t-1 |m w 0,t )∥p θ (m w t-1 |m w t )) -D KL (q(m w t-1 |m w 0,t )∥p ref (m w t-1 |m w t )) -D KL (q(m l t-1 |m l 0,t )∥p θ (m l t-1 |m l t )) + D KL (q(m l t-1 |m l 0,t )∥p ref (m l t-1 |m l t )) (7) Let [x 0 ,v0 ] be the predicted atom position and type, which are fed into Equation ( 3) to obtain the posterior distributions. With the joint diffusion processes Equations ( 1) to (3) on both continuous x and discrete v features, the above KL divergences can be decomposed and calculated as: θ , v0,θ and x0,ref , v0 ,ref be the predictions from the fine-tuned and from the original pretrained model, respectively. Then, we can further obtain the preference optimization loss on x and v, respectively, as follows:D KL (q(m t-1 |m 0,t )∥p(m t-1 |m t )) = D x,t-1 KL (q(x t-1 |x 0,t )∥p(x t-1 |x t )) + D v,t-1 KL (q(c t-1 |c 0,t )∥p(c t-1 |c t )), D x,t-1 KL (q(x t-1 |x 0,t )∥p(x t-1 |x t )) = 1 βt ∥ μ(x t , x 0 ) -μ(x t , x0 )∥ 2 + C = γ t ∥x 0 -x0 ∥ 2 + C, D v,t-1 KL (q(c t-1 |c 0,t )∥p(c t-1 |c t )) = k c(v t , v 0 ) k log c(v t , v 0 ) k c(v t , v0 ) k , (8) where γ t = ᾱt-1β 2 t 2σ 2 t (1-ᾱt) 2 and C is a constant. Let x0,EQUATIONWith Jensen's inequality and the convexity of -log σ, we can derive the final objective as a (weighted) sum of atom coordinate and type preference lossesL x t-1 + L v t-1, which turns the sum of the KL terms outside -log σ and serves as an upper bound of Equation ( 7):EQUATIONwhere the preference is assigned separately to atom types v and coordinates x. The loss decomposition imposes a fine-grained preference assignment on chemical elements and geometric strcutures and enables us to choose weights to balance the training of the two variables [Guan et al., 2023 [Guan et al., , 2024]] .The overall training and sampling algorithms of ALIDIFF are summarized in Appendix B.Although DPO enjoys the advantage of efficient fine-tuning without fitting a reward function, recent theoretical investigations reveal that it is highly vulnerable to overfitting by pushing all the probability mass on the winning sample [Azar et al., 2024] . Specifically, the non-linear transformation log σ of Equation ( 5) pushes the log p θ (m w |p) -log p θ (m l |p) towards infinity, completely removing the likelihood for the losing sample regardless of any regularization in the original RLHF setup Equation ( 4) [Azar et al., 2024 , Tang et al., 2024] . Let us analyze the problem with an example consisting of two ligand molecules m w and m l with their rewards measured as r w and r l (e.g., calculated from binding energy). The DPO objective in Equation ( 10) tends to just greedily maximize towards p(m w ≻ m l |p) → 1. However, the optimal preference probability can be calculated by the BT model [Bradley and Terry, 1952] as p(m w ≻ m l |p) = σ(r wr l ), and our alignment goal is to shift the distribution to align with this p instead of greedy maximization. To address the over-optimization issue, we introduce an improved objective with regularization on preference maximization, named Exact Energy Preference Optimization (E 2 PO). Let Lx t (θ) and Lv t (θ) denote terms for reverse preference optimization:EQUATIONOur E 2 PO objective function takes a cross-entropy form to align the distributions p θ (m w ≻ m l |p) towards p(m w ≻ m l |p). Formally, it is given by:EQUATIONwhere the second term Lx t-1 + Lv t-1 weighted by 1 -σ(r wr l ) helps to alleviate the overfitting on the winning data sample. Notably, for r w >> r l , we have σ(r wr l ) ≈ 1, indicating that the regularized objective in Equation ( 12) will still change back to the original objective in Equation ( 10), where overfitting on the extremely better data is expected. In principle, with the regularization objective, we have: .The theorem provides an analytical guarantee for the optimal shifted distribution after alignment that avoids over-optimization. Assuming we achieve convergence on the KL divergence, we have that p ref  
experiment setup: 
 Dataset. We train and evaluate ALIDIFF using the CrossDocked2020 dataset [Francoeur et al., 2020] . Following the common setup in this field [Luo et al., 2021 , Guan et al., 2023] , we refined the initial 22.5 million docked protein binding complexes by selecting docking poses with RMSD lower than 1Å with the ground truth and diversifying proteins with a sequence identity below 30%. To apply ALIDIFF, we further preprocess our data and construct a dataset of the form D = {(p, m w , m l )}, where p denotes the protein, m w denotes the preferred molecules, and m l denotes rejected molecules based on the user-defined reward. In our setting, we choose two ligand molecules per pocket site and label the preference by a certain reward, e.g. binding energy for our main benchmark. We provide ablations with more reward functions in Section 4.3. Details of preference pair selection are presented in Appendix E. The final dataset uses a train and test split of 65K and 100.Baselines. We compare our model with the following baselines: liGAN [Ragoza et al., 2022] is a conditional VAE model that utilizes a 3D CNN architecture to both encode and generate voxelized representations of atomic densities; AR [Luo et al., 2021] , Pocket2Mol [Peng et al., 2022] and GraphBP [Liu et al., 2022] are autoregressive models that learn graph neural networks to generate 3D molecules atom by atom sequentially; TargetDiff [Guan et al., 2023] and DecompDiff [Guan et al., 2024] are diffusion-based approaches for generating atomic coordinates and types via a joint denoising process; IPDiff [Huang et al., 2023] is the most recent state-of-the-art diffusion-based approach that further integrates the interactions between the target protein and the molecular ligand into the generation process.Evaluation metrics. We evaluate the generated molecules by comparing binding affinity with the target and critical molecular properties. We analyze the generated molecules across 100 test proteins, reporting the mean and median for affinity-based metrics (Vina Score, Vina Min, Vina Dock, and High Affinity) and molecular property metrics (drug-likeness QED [Bickerton et al., 2012] , synthesizability SA [Ertl and Schuffenhauer, 2009] , and diversity). We use AutoDock Vina [Eberhardt et al., 2021] to estimate binding affinity scores, using the common setup described by Luo et al. [2021] , Ragoza et al. [2022] . Specifically, Vina Score estimates binding affinity from the generated 3D structures, Vina Min refines the structure through local minimization before estimation, Vina Dock uses a re-docking procedure to reflect the optimal binding affinity, and High Affinity gauges the percentage of generated molecules that bind better than reference molecules per protein.Binding Affinity and Molecular Properties. We compare the performance of our proposed method ALIDIFF against the above baseline methods. Our model is fine-tuned from IPDiff, the ligand generative model. We report the results in Table 1 , and leave more implementation details in Appendix D. As shown in the results, ALIDIFF significantly outperforms all non-diffusion-based models in binding-related metrics, and also surpasses our base model IPDiff in all binding affinity related metrics by a notable margin. In particular, ALIDIFF increases the binding-related metrics Avg. Vina Score, Vina Min, and Vina Dock by 10.1%, 8.56%, and 3.9% compared with IPDiff. Our superior performance in binding-related metrics demonstrates the effectiveness of energy preference We also notice a trade-off between binding affinity and property-related metrics. While we achieve state-of-the-art performance on all binding affinity metrics, the performance on QED and SA metrics slightly decreases. This phenomenon has been commonly observed in previous studies where achieving high binding affinity can often sacrifice other molecular metrics [Guan et al., 2023 , Huang et al., 2023] . This is because the highest affinity can potentially only be achieved by rather specific and unique molecules, which are harder to synthesize than simple molecules, and hence these trade-offs are expected. Besides, in real-world drug discovery, binding affinity is typically a more critical metric as molecules with more stable interaction with the pocket site are important, whereas QED and SA work mainly as rough filters [Guan et al., 2023] . For these reasons, we believe the deterioration in molecular properties is well compensated by the improvement in binding affinity, especially with such little deterioration in property metrics. In addition, in the following section (Table 2 ), we further discuss incorporating molecular properties into the reward, which shows slightly lower performance gain on affinity but archives improvements also on molecular properties.Effect of reward objectives. To further explore the potential of ALIDIFF, we evaluate the effect of combining optimization objectives (r = r affinity + r SA ; r = r affinity + r QED ) and investigate whether such a combined reward function can lead to better molecular properties to counter the trade-off we discussed before. As shown in Table 2 , the results indicate that finetuning solely with binding affinity apparently achieves better performance in terms of binding affinity metrics. However, ALIDIFF-Affinity+SA generates compounds with better drug-likeness (QED) and synthetic accessibility (SA). Both models exhibit similar performance in terms of structural diversity. This suggests that while ALIDIFF-Affinity is superior for binding affinity, incorporating synthetic accessibility considerations (Affinity + SA) results in compounds that are more drug-like and easier to synthesize, enabling more efficient multi-objective drug development. Moreover, ALIDIFF-Affinity+QED achieves better binding affinity compared with ALIDIFF-Affinity, while the improvement in QED is relatively minimal. Thus, balancing these objectives highlights the potential for overcoming trade-offs in molecular optimization. Comparison with Supervised Fine-Tuning. Supervised Fine-Tuning (SFT) serves as an alternative method for generating molecules with user-defined optimization objectives. We select the top 50% protein-ligand samples with higher quality in user-defined reward from the training dataset and fine-tune the baseline model with the same training and sampling setting. The results in Table 3 show that SFT did not show improvement over the baseline, and ALIDIFF demonstrates significantly superior results compared to SFT.Effect of preference optimization methods. As discussed in Section 3.3, the original DPO objective is vulnerable to overfitting and we propose to avoid it with regularization by weighting preference losses with the user-defined rewards. We compare the direct use of energy preference optimization by ranking molecule pairs (ALIDIFF-DPO) and exact energy optimization with user-defined reward function (ALIDIFF-E 2 PO) in Table 3 . The results show that ALIDIFF-E 2 PO achieves superior performance over ALIDIFF-DPO in all binding affinity metrics (Vina Score, Vina Min, and Vina Dock) while maintaining competitive scores in QED, SA and diversity. In terms of drug-likeness and structural diversity, ALIDIFF-E 2 PO performs competitively, indicating that while it prioritizes binding affinity, it still maintains favorable drug-like properties and diversity. This further supports our previous hypothesis regarding the trade-off between binding affinity and molecular properties.An additional ablation study on the effect of exact energy optimization is presented in Appendix E. General applicability to ligand diffusion models. We further justify the general applicability of the proposed approach by finetuning another diffusion-based SBDD model, TargetDiff [Guan et al., 2023] , with exact energy optimization (ALIDIFF-T), As shown in Table 4 , ALIDIFF-T surpasses TargetDiff on all binding affinity and molecular properties, with a 6.2%, 16.6%, 6.9%, 2.8% increase in Avg. Vina Score, QED, SA, and diversity, respectively. The results further justify that our approach is generally applicable to diffusion-based SBDD models. Notably, ALIDIFF-T archives even better QED and SA compared with ALIDIFF, which allows users to choose the model based on the specific purpose for molecular properties. Also, we notice the percentage of improvement of binding affinity from Target-Diff to ALIDIFF-T is slightly lower than that from IPDiff to ALIDIFF. This can be explained as preference optimization is more effective when the model distribution is more similar to the preference data distribution, and IPDiff is shown to fit CrossDocked data better than TargetDiff [Huang et al., 2023] .Strength of β. We further evaluate ligand molecules generated by ALIDIFF trained with varying β values in Figure 5 . Recall that β influences the scale of energy preference optimization and regularization with respect to the reference model. The results indicate a clear trade-off between binding affinity and molecular properties with varying β. Lower β values (e.g., 0.01) significantly enhance binding affinity metrics (Vina Score, Vina Min, Vina Dock), but at the cost of lower drug-likeness (QED) and diversity. Conversely, higher β values improve QED, suggesting that these configurations generate more drug-like compounds while maintaining consistent synthetic accessibility and diversity. We believe β reaches an equilibrium around β = 1, where binding affinity is maximized without sacrificing too much loss in molecular properties. This ablation study demonstrates that the parameter β can offer a useful tool to train ALIDIFF models with different desired trade-offs between binding affinity and useful molecular properties, which can vary for different drug development use cases. 
conclusion: 
 In this paper, we present ALIDIFF, a novel framework to align pretrained target-aware molecule diffusion models with desired functional properties via preference optimization. Our key innovation is the Exact Energy Preference Optimization method, which enables efficient and exact alignment of the diffusion model towards regions of lower binding energy and structural rationality specified by user-defined reward functions. Extensive experiments on the CrossDocked2020 benchmark demonstrate the strong performance of ALIDIFF. By incorporating user-defined reward functions and an improved Exact Energy Preference Optimization method, ALIDIFF successfully achieves state-of-the-art performance in binding affinity while maintaining competitive molecular properties.In the future, we plan to explore more expressive molecular reward function classes within our framework and extend ALIDIFF to real-world prospective drug design settings by integrating it into online drug discovery pipelines. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: General machine learning method without specific concern in our mind.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: General machine learning method without specific concern in our mind.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: General machine learning method without specific concern in our mind.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All assets get credited.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: While ALIDIFF exhibits promising performance, there are still potential limitations to our current approach. For example, ALIDIFF takes binding affinity as our reward function which is computed by AutoDock Vina [Eberhardt et al., 2021] in this work. However, computing binding energy via software is an approximation and sometimes can be very inaccurate. In the future, we plan to explore experiment-measured energy or ensemble different binding affinity calculation software, e.g., GlideScore [Friesner et al., 2004] In addition, in this work, we focus on an offline learning setting where the preference pairs are off-the-shelf. This is because computing binding affinity is computationally expensive. An important future direction to extend the work toward real-world drug discovery scenarios could be incorporating the online setting but with a limited number of query.The pseudo-code for ALIDIFF and ALIDIFF-T are provided below. Sampling procedures are the same as Guan et al. [2023] and Huang et al. [2023] .Algorithm 1 Training Procedure ALIDIFF 1: Input: Protein-ligand binding dataset {P, M w , M l } N 1 , pre-trained neural network ϕ θ , reference network ϕ ref , learnable neural network ψ θ2 and pretrained interaction prior network ψ IP . 2: while ϕ θ and ψ θ2 not converge do 3:[[p 0 , m w 0 , m l 0 ]] ∼ {P, M w , M l } N i=1 where m w 0 = {x w 0 , v w 0 }, m l 0 = {x l 0 , v l 0 } 4:Obtain r w , r l for m w , m l , respectively.t ∼ U (0, . . . , T )6:Move the complex to make CoM of protein atoms zero 7:Obtain shifts[[s Mw 0 , s M l 0 ]] and interactions [[f Mw 0 , f M l 0 , f P 0 ]]from ψ IP and ψ θ2 according to [Huang et al., 2023] . ϵ ∼ N (0, I) 10:x w t = √ ᾱt x w 0 + s Mw t + √ 1 -ᾱt ϵ, x l t = √ ᾱt x l 0 + s M l t + √ 1 -ᾱt ϵ 11: Perturb v w 0 , v l 0 to obtain v w t , v l t 12:g ∼ Gumbel(0, 1)13:log c w = log( ᾱt v w 0 + (1 -ᾱt /K), log c l = log( ᾱt v l 0 + (1 -ᾱt /K) 14: v w t = onehot(arg max i (g i + log c w i )), v l t = onehot(arg max i (g i + log c l i ))15:Embed v w t , v l t into hw,0 t , hM l ,0 t , and embed v P 0 into hP,0 t 16: Obtain features [[h w,0 t , h M l ,0 t , h P,0 t ]] through prior-conditioning 17: Predict (x w 0|t , vw 0|t ) from ϕ θ ([[h Mw,0 t , h P,0 t ]], [[f Mw 0 , f P 0 ]]) 18: Predict (x l 0|t , vl 0|t ) from ϕ θ ([[h M l ,0 t , h P,0 t ]], [[f M l 0 , f P 0 ]]) 19: Predict (x w 0|t,ref , vw 0|t,ref ) from ϕ ref ([[h Mw,0 t , h P,0 t ]], [[f Mw 0 , f P 0 ]]) 20: Predict (x l 0|t,ref , vl 0|t,ref from ϕ ref ([[h M l ,0 t , h P,0 t ]], [[f M l 0 , f P 0 ]]) 21: Compute loss L with (x w 0|t , vw 0|t ), (x l 0 , v l 0 ), (x w 0|t,ref , vw 0|t,ref ), (x l 0|t,[[p, m w 0 , m l 0 ]] ∼ {P, M w , M l } N i=1 where m w 0 = {x w 0 , v w 0 }, m l 0 = {x l 0 , v l 0 } 4:Sample diffusion time t ∼ U (0, . . . , T )5:Move the complex to make CoM of protein atoms zero 6:Perturb x w 0 , x l 0 to obtain x w t , x l t :x w t = α t x w 0 + (1 -α t )ϵ, x l t = α t x l 0 + (1 -α t )ϵ, where ϵ ∼ N (0, I) 7: Perturb v w 0 , v l 0 to obtain v w t , v l t : 8: logc w = log(α t v w 0 + (1 -α t )/K) 9: logc l = log(α t v l 0 + (1 -α t )/K) 10: v w t = one_hot(arg max[g i + logc w i ])11:v l t = one_hot(arg max[g i + logc l i ]), where g ∼ Gumbel(0, 1)EQUATIONCompute L(θ) = L x (θ) + αL v (θ) according to Equation ( 10) We prove the theorem with Lemmas C.1 and C.2. Lemma C.1 justifies the least square objective is the variational upper bound for preference optimization, and Lemma C.2 shows that regularized preference optimization corresponds to exact KL divergences between the optimal and parameterized distributions. A version of similar proof can be found in Wallace et al. [2023] and Ji et al. [2024] , Chen et al. [2024] respectively, and to be self-contained we incorporate these proofs here. Compared with Wallace et al. [2023] , we introduce an additional term into the diffusion optimization. And compared with Ji et al. [2024] , Chen et al. [2024] , we explicitly drop the assumption for drawing infinite samples m for each pocket p.Lemma C.1. The objective function Equation (12) L ALIDIFF-E 2 PO (θ) = -E (p,m w 0 ,m l 0 )∼D,t∼[0,T ],m w t ∼q,m l t ∼q (σ(r w -r l ))(L x t-1 + L v t-1 ) + (1 -σ(r w -r l ))( Lx t-1 + Lv t-1) is a variational upper bound of:EQUATIONWe refer readers to Appendix S2 of Diffusion-DPO [Wallace et al., 2023] for the full proof. The bound is derived from Jensen's inequality and the convexity of the function -log σ. Proof. First of all, we can rewrite the objective Equation ( 13) in the following form, expanding the sigmoid function: Then we can substitute this equation and rewrite L E 2 PO (θ): Then we can plug the above results into Equation ( 15) and further simplify L E 2 PO : ,m) with p * and further simplify the above equation: L E 2 PO (θ) = E p∼D E pref(m1:2|p) - 2 i=1 e r(p,mi)L E 2 PO (θ) = E p∼D E pref(m1:2|p) - 2 i=1 e r(p,mi) 2 j=1 e r(p,mj ) log p θ (mi|p) pref(mi|p) β 2 j=1 p θ (mj |p) pref(mj |p) β = E p∼D E pref(m1:2|p) - 2 i=1 e r(p,mi)L E 2 PO (θ) = E p∼D E pref(m1:2|p) - 2 i=1 e r(p,mi) 2Z(p) log pθ (m i |p) 2p ref (m i |p) = E p∼D E pref(m1:2|p) -L E 2 PO (θ) = E p∼D - 1 2 2 i=1 e r(pData. Following [Guan et al., 2023] , proteins and ligands are expressed with atom coordinates and a one-hot vector containing the atom types. For proteins, each atom type is represented by a one-hot vector covering 20 distinct amino acids. Ligand atoms are encoded using a one-hot vector that discriminates among several elements, specifically H, C, N, O, F, P, S, Cl. Additionally, a one-dimensional binary flag is incorporated to differentiate whether atoms are part of the protein or the ligand. We further apply two separate single-layer Multi-Layer Perceptrons (MLPs) to transform the input data into 128-dimensional latent spaces, providing a compact and informative representation for subsequent computational stages.Preference Pair Generation. For each synthetic molecule, we first locate its corresponding protein binding site and compute reward according to user-defined reward function for all synthetic molecules of the corresponding the binding site. We select a losing sample with lower reward and construct the preference. The selection process is detailed in Appendix E.Architecture. We follow the same architecture as IPDiff [Huang et al., 2023] , which includes a learnable diffusion denoising model ϕ θ1 , learnable neural network ϕ θ2 and pretrained interaction prior network IPNET. The architecture of all models used in our method is the same as IPDiff.Pretraining Details. Following existing work, we adopted the Adam optimizer with a learning rate of 0.001 and parameters β values of (0.95, 0.999). The training was conducted with a batch size of 4 and a gradient norm clipping value of 8. To balance the losses for atom type and atom position, we applied a scaling factor λ of 100 to the atom type loss. Additionally, we introduced Gaussian noise with a standard deviation of 0.1 to the protein atom coordinates as a form of data augmentation. Our parameterized diffusion denoising model, IPDiff was trained on a single NVIDIA A6000 GPU and achieved convergence within 200k steps.Training Details. For finetuning, the pre-trained diffusion model is further fine-tuned via the gradient descent method Adam with init learning rate=5e-6, betas=(0.95,0.999). We keep other setting the same as pretraining. We use β = 5 in Equation ( 5). We trained our model with one NVIDIA GeForce GTX A100 GPU, and it could converge within 30k steps.Effect of diffusion steps. In fig. 6 , we present a comprehensive ablation study examining the impact of diffusion steps on the optimization of molecular properties using our novel ALIDIFF framework. The visualizations at the top of the figure showcase the progressive refinement of molecular structures across increasing diffusion steps (t = 200 to t = 1000). These images clearly illustrate how our model gradually enhances the molecular fitting within the target binding site, which is critical for improving drug efficacy. The plotted data below provides a quantitative analysis of QED, SA, Vina Lipinski. We further compared Lipinski's Rule of Five [Lipinski et al., 2012] across all comparison methods. Lipinski's Rule of Five is another measurement for assessing drug-likeness besides QED, and we would like to incorporate this metric to validate our performance in generating drug-like molecules. The results of Lipinski's scores are reported in Table 5 . The results are consistent with our evaluation using QED score, as all diffusion-based models are not achieving high drug-likeness. We maintain similar drug-likeness as our backbone models targetDiff and IPDiff. Choice of m l . Our generated dataset is obtained by directly transforming a standard labeled dataset into a pairwise preference dataset. Yet the binding affinity labels are continuous values where sometimes the difference between preferred and dispreferred is minimal. Therefore, the effect of energy preference optimization is highly sensitive to the overall data quality. Justification: We have provided computer resources details in Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 


Decision: This paper contributes a novel diffusion-based model, focusing on generating ligand molecules with high affinity to their protein targets. To this end, Authors proposed a improved Exact Energy Preference Optimization method based on reinforcement raining. 

The reviewers have pointed out that the proposed approach has limited computational novelty, but still is of high practical utility. The rebuttal also discussed that the optimization of the affinity is associated with lower drug likeliness scores, which however seems a property of the diffusion models, and was traded for higher drug likeliness by fragment-based methods.

The paper seems of interest to the NeurIPS community.

Reviews: review: paper_summary: paper_summary: This paper proposes a novel alignment framework, known as ALIDIFF, to align pretrained target diffusion model with preferred functional properties for structure-based drug design. ALIDIFF shifts the target-conditioned chemical distribution towards regions with higher binding affinity and structural rationality, specified by user-defined reward functions, via the preference optimization approach. The experiments show that this method can generate ligands with better binding energy and maintain competitive molecular properties.
strength_weakness: strength_weakness: - Lipinski metric should also be reported in Table 1. This is a commonly used metric for measure drug-likeness in previous work. 
- Fragment-based methods, such as FLAG and DrugGPS, is missing in Table 1. For comprehensive comparison, I recommend add those two baselines. 
- Generated molecules have inferior drug-likeness properties. Isn't it possible to also align the pre-trained diffusion model with high drug-likeness region?
questions: questions: Please refer to the weakness part.
limitations: limitations: Limitation is discussed.
paper_summary: paper_summary: This paper proposes a novel and general alignment framework to align pretrained target diffusion models with preferred functional properties, named AliDiff. AliDiff adjusts the target-conditioned chemical distribution toward regions characterized by lower binding energy and structural rationality. AliDiff can generate molecules with lower binding energies.
strength_weakness: strength_weakness: - The experiments are only conducted on the CrossDocked2020 dataset. Are there any alternative datasets?
- In addition to well-established benchmarks, is it possible to show several practical case studies to verify the effectiveness of the proposed algorithms. For example, generating ligands for the target of SARS-CoV-2.
questions: questions: See weaknesses
limitations: limitations: None
paper_summary: paper_summary: This paper presents ALIDIFF, a framework that aligns pre-trained target-aware molecule diffusion models with desired functional properties using preference optimization. The key contribution of ALIDIFF is the Exact Energy Preference Optimization method, which precisely aligns diffusion models to regions with lower binding energy and better structural rationality based on user-defined reward functions. Extensive experiments on the CrossDocked2020 benchmark showcase ALIDIFF's strong performance, generating molecules with state-of-the-art binding energies and maintaining competitive molecular properties. By incorporating user-defined reward functions and improving the Exact Energy Preference Optimization method, ALIDIFF achieves significant advancements in binding affinity.
strength_weakness: strength_weakness: * The method leverages IPDiff and further optimizes it using RL, which makes performance improvements easier. Consequently, the methodological novelty appears marginal.
* The paper lacks specific information on the time required for pre-training, fine-tuning, and binding energy computations via AutoDock Vina. This absence makes it difficult to fully assess the advantages and limitations of the proposed method.
  - AutoDock Vina is known to be time-consuming, and RL also requires significant computational time. Including these time details would help illustrate both the strengths and constraints of the proposed approach.
questions: questions: * Train/test split is 65K/100. It seems that the test set is too small compared to the train set. Is there any reason for that?
* How much time does it take to pre-train, fine-tune and binding energy computations via AutoDock Vina?
limitations: limitations: Yes
paper_summary: paper_summary: This paper focuses on generating ligands with desired properties, such as high binding affinity, for the protein-conditioned ligand generation task. The authors propose a Preference Optimization (PO)-based fine-tuning method for pre-trained generative models. They extend DPO (Data-driven Preference Optimization) to propose exact energy optimization that directly utilizes reward values in the loss function. Experiments demonstrate that the proposed method can generate ligand molecules with the desired properties.
strength_weakness: strength_weakness: * The practical impact of the achieved improvement needs to be clearly quantified.
* There is no consideration or results for comparison or combination with methods other than DPO, such as IPO or SLiC.
questions: questions: * The proposed objective function seems to be heavily influenced by the scale of the reward value r. Was any adjustment necessary? In particular, was simple addition sufficient for the Affinity + SA experiment that combined multiple rewards? Additionally, is it easy to incorporate other rewards, such as diversity?
* The paper states that if $r^w \gg r^l$, then $\sigma(r^w - r^l) \approx 1$. However, wouldn't it be more accurate to say that if $r^w - r^l \ge 1$, then $\sigma(r^w - r^l) \approx 1$? What is the actual distribution of $r^w - r^l$ in the training data?
limitations: limitations: The proposed objective function cannot be used when reward values are not explicitly available, making it unsuitable for cases involving manually binary-labeled results.

Comments: title: Looking forward to feedback during the reviewer-author discussion period
Looking forward to feedback during the reviewer-author discussion period
Thanks
Thanks for the clarification
comment: Thanks again for your time spent reviewing our work and the constructive comments and suggestions. We provided evaluations of two related works and our results showcase that the proposed method consistently outpeforms all baselines on binding affinity metrics, which have also been acknowledged by other reviewers. We also added lipsinki as another measure of drug-likeness and conducted new experiments on aligning molecules with drug-likeness propoerties, as suggested.

As the discussion period draws to a close, we would be grateful for your further feedback on our clarifications and recognition of our research's contributions. Thank you again for your valuable suggestions!

Best,

Authors
Thank you very much for your timely reply and recognition of our efforts! We appreciate your valuable suggestions and will incorporate the discussions and results to the final version. Please let me know if you have any other questions!

Sincerely, Authors
Dear reviewer 5bxb,


We would like to first express our sincere gratitude for your time and effort in reviewing our paper. We appreciate your valuable suggestions and will incorporate the discussions and results to the final version.


Considering the author-reviewer discussion is ending very soon, could you kindly check our responses and let us know if you have further concerns? We are more than willing to address any other concerns or questions.


We would greatly appreciate it if the reviewer would consider adjusting the score, on the basis of our response and other review comments.
 

Thanks again for your constructive reviews!
 

Sincerely, 
Authors
Thank you for the detailed response. It has addressed my concerns.
I still believe that this work is worth presenting.

> Mathematically, for the sigmoid function, ...

Sorry, it is my misunderstanding. You are correct.
Dear reviewer bNcB,

Thank you very much for your timely reply and recognition of our efforts! We appreciate your valuable suggestions and will incorporate the discussions and results to the final version. Please let me know if you have any other questions!

Sincerely,
Authors
Dear reviewer fWJ4,

Thank you very much for your timely reply and recognition of our efforts! We appreciate your valuable suggestions and will incorporate the discussions and results to the final version. Please let me know if you have any other questions!

Sincerely,
Authors
I'm glad to see that the author's response has addressed my concerns and I maintain the score, which tends to acceptance.
Thank you for the detailed rebuttal. I have read it thoroughly. Based on the clarified points, I believe that while the methodological novelty may be limited, the practical value is significant. Therefore, I will raise my score to 5. Thank you for the excellent research.