Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
Authors: Seungbeom Lee
Munsun Jo
Jungseul Ok
Dongwoo Kim
Venue: ICML 2025 poster (2025)
Keywords: Structure-Based Drug Design
Deep Learning
Direct Property Optimization
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We propose a multi-reward Direct Preference Optimization (DPO) method to fine-tune models toward preferred properties, such as binding affinity, validity, and drug-likeness.

Abstract: Deep learning-based Structure-based drug design aims to generate ligand molecules with desirable properties for protein targets. While existing models have demonstrated competitive performance in generating ligand molecules, they primarily focus on learning the chemical distribution of training datasets, often lacking effective steerability to ensure the desired chemical quality of generated molecules. To address this issue, we propose a multi-reward optimization framework that fine-tunes generative models for attributes, such as binding affinity, validity, and drug-likeness, together. Specifically, we derive direct preference optimization for a Bayesian flow network, used as a backbone for molecule generation, and integrate a reward normalization scheme to adopt multiple objectives. Experimental results show that our method generates more realistic ligands than baseline models while achieving higher binding affinity, expanding the Pareto front empirically observed in previous studies.

Enhancing Ligand Validity and Affinity in Optimization
Structure-Based Drug Design with
Seungbeom Lee" "1
Munsun Jo 2
Jungseul Ok"
Dongwoo Kim 13
Abstract
0.75
MegletMol
Poc
Deep learning- based structure-baser drug design aims to generate ligand molecules with desirable properties for protein targets. While existing mod- els have demonstrated competitive performance in generating ligand molecules, they primarily focus on learning the chemical distribution of training datasets, often lacking effective steerability to en- sure the desired chemical quality of generated molecules. To address this issue, we propose a multi-reward optimization framework that fine- tunes generative models for attributes, such binding affinity, validity, and drug-likeness, to- gether. Specifically, we derive direct preference optimization for a Bayesian flow network, used as backbone for molecule generation, and integrate reward normalization scheme to adopt multi- ple objectives. Experimental results show that method generates more realistic ligands than baseline models while achieving higher binding affinity, expanding the Pareto front empirically observed previous studies.
Raference
Churs
PLAG AR
Referepe
FLAG
McLCHAFT
10
MelCCrlft
DecompDiff
0.65
DecompDii
IPDiT
TargetDiff
0.6
TargeDDfff
Decompppp
Opcompppt
IPI
AADiI
0.55
30
Alroir
(a) Vina Dock vs. SA
(b) Vina Dock vs. SE
Figure I. Comparison of median Vina docking score (lower better) against Synthetic accessibility (SA, higher is better) and Strain Energy (SE, lower better). The dashed blue line represents Pareto front obtained from previous works.
els (Huang et al., 2024; Qu et al., 2024; Guan et al., 2023), they often neglect other crucial factors like ease of conforma- tion stability such as strain energy or synthesis. In practice, drug discovery demands a careful balance of many proper- ties, making single-objective or narrowly focused solutions insufficient.
In current deep generative modeling for SBDD, standard strategies generally concentrate on minimizing the discrep- ancy hetween a generated ligand and a training set of protein- ligand complexes. Although this can lead to strong bind- ing affinities, improving multiple objectives beyond bind- ing, e-g., synthetic accessibility, usually requires careful tuning of loss terms elaborate constraints. As result, purely likelihood-based approaches can struggle to expand the Pareto front, which captures trade-offs among distinct objectives.
1. Introduction
Designing ligand molecules that simultaneously optimize multiple objectives - such as binding affinity, synthetic ac- cessibility, and strain energy structure-based drug design (SBDD). Conventional deep- learning approaches typically address only one two ob- jectives at time, most commonly focusing on high binding affinity and basic chemical feasibility, e.g., through simple validity checks via the valency rule. While these methods have advanced the generative capabilities of molccular mod-
central challenge in
Figure illustrates the trade-offs among Vina docking score measuring binding affinity to target protein, synthetic acces- sibility (SA), and strain energy (SE) using results from nine different generative models, along with reference dataset included in the test set. Although recently proposed meth- ods such as MolCRAFT (Qu et al., 2024), IPDiff (Huang et al., 2024), and AliDiff (Gu et al., 2024), outperform the reference data in terms of Vina docking score, they under- perform on SA and SE. Ideally, design process would push
'Equal contribution 'Graduate School of Artificial Intelli- gence, POSTECH, South Korea ²KT Corporation, South Korea DDppartment of Computer Science & Engineering, POSTECHH South Korea. Correspondence to: Dongwoo Kim <dong- woo.kim @ kc..s
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
the Pareto frontier toward both lower docking scores and higher SA (or lower SE). In contrast, these findings show
generative modeling task, leveraging protein-ligand com- plex data to learn a distribution over valid 3D conforma- tions. Autoregressive approaches place atoms or fragments within protein pockets in sequential manner (Luo et al., 2021; Peng et al., 2022), while diffusion models gener- ate candidate molecules by denoising from random noise, learning continuous atom positions and discrete atom types jointly (Guan al., 2023; Huang et al., 2024; Guan et al., 2024; Zhou et al., 2024). Most recently, Qu et al. (2024) introduced MOICRAFT, which uses Bayesian Flow Net- works (BFNs) (Graves et al., 2023) to sample in parameter space rather than directly in data space, leading to improved binding affinity and more stable 3D molecular structures. To further improve the quality of binding affinity, Gu et al. (2024) introduced AliDiff, which fine-tunes a diffusion- based generative model via DPO.
that existing models remain overly focused on improving affinity, thus calling into question the practical utility of their generated molecules.
Our Approach Building recent developments in preference-based reinforcement learning, we introduce multi-reward Direct Preference Optimization (DPO) frame- work adapted for Bayesian Flow Networks (BFNs), which recently show promising performance on ligand genera- tion tasks (Qu et al., 2024). DPO refines a pretrained gen- erative model using feedback from external software or metrics, enabling fine- aineggraind control over multiple proper- ties (Rafailov et al., 2024; Kim et al., 2024). In our setting, we treat each property, such as binding affinity, synthetic accessibility, and strain energy, as a separate reward. We then apply softmax- -based normalization to ensure commen- surability of metrics with different ranges and adopt an ensemble strategy to penalize high variance in multi-reward predictions. This multi-reward scheme expands the trade-off frontier between objectives, thereby improving molecular candidates across multiple dimensions.
On the other hand, Harris et al. (2023) raises questions about the evaluation practices with the deep generative models for SBDD. They find that many generated molecules from re- cent generative models violate physical constraints and call for an expanded evaluation of SBDD. A similar observation is made for deep learning based binding affinity prediction in Buttenschoen al. (2024). Our approach builds upon this foundation by adopting multi-reward perspective that jointly optimizes for binding affinity and other molecu- lar properties (e.g., synthetic accessibility. strain energy), thereby expanding the Pareto frontier in multi-objective SBDD.
encertainty-regulariae
2. Related Work
Preference Alignment and Direct Preference Optimiza- tion (DPO) Preference alignment has emerged as power- ful paradigm for adjusting pretrained models to better reflect desired behaviors. Originally adopted in large language models (LLMs) to align with human feedback (Ouyang et al., 2022). DPO (Rafailov et al., 2024) provides a more direct means of incorporating user or expert preferences compared to traditional reinforcement learning. Subse- quently, DPO-based methods have been extended to text- to-image generative models (Wallace et al., 2023), show- ing that preference alignment is not limited to language domains. In multi-objective settings, Kim et al. (2024) pro- pose a confidence-basee or variance-penalizee approach to combine heterogeneous reward signals, further enhancing stability and interpretability.
3. Preliminaries
We first introduce the problem formulation of structure- based drug design with notations used throughout this manuscript. We then provide an overview of Bayesian Flow Networks (BFNs) et al., 2023) and MolCRAFT(QU et al., 2024), where the latter leverages BFNs to generate target molecules.
3.1. Problem Formulation
In structure-based drug design (SBDD), the goal is to gen- erate ligand molecules conditioned on given protein tar- get. We denote a protein target by P = ((x (d) ) where x R is the three-dimensional coordinate of (t)
Structure-Based Drug Design Structure-based drug de- sign (SBDD) focuses on generating or modifying small- molecule ligands to exhibit strong binding affinity for a target protein while maintaining favorable pharmacological properties. Classical computational approaches rely de novo design or docking-based virtual screening (Schneider & Fechner, 2005) to identify candidate molecules, followed by iterative refinement and validation through experimental assays.
the i-th atom, and vip E Dr is a one-hot vector en- coding the atom type. A molecule is analogously repre- sented as M ((x [( (()) where x(() R3 and VM
(6) E PM denote the diree-dimensional coordinate and the atom-type vector, respectively. For convenience, we denote each molecule by m = [x, v] where [s is the con- catenation of x E NuMx and v E R N Pixs given target protein as p.
Recent deep learning methods cast SBDD as a conditional
2

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
SBDD aims to discover a molecule m given protein and its binding site p. During training, set of protein-reference molecule pairs is provided. As primary performance met- ric, external software such as AutoDock Vina (Eberhardt et al., 2021) is employed to estimate the binding affinity between the protein and generated molecules, Along with affinity scores, additional molecular properties (e.g., valid- ity, synthesizability, and
Receiver distribution At time the receiver uses pa- rameters Q--1 from the previous step and the noisy sam- ple i from the sender to reconstruct a denoised m. Let + (0.-1.P. be an equivaaaui neural network that outputs rn from parameters Q,-1, the protein pocket p. and time t. Given that the distribution and noise schedule of the sender distribution is known, the receiver distribution defined
can be evaluated to assess the quality of the generated compounds further.
PR PR m 0.-1, = ps m 4 (0.-I,,,0),, (2)
3.2. Bayesian Flow Networks
Hence the receiver distribution becomes close to the sender distribution as d accurately reconstructs the orig- inal molecule m.
Bayesian flow network (BFN) (Graves al.. 2023) is class of generative models that integrate core ideas from both diffusion models (Ho et al., 2020) and Bayesian inference. Unlike traditional diffusion models that operate directly in the data space, BFNs perform their forward and reverse "diffusion"" in the parameter space of the data distribution. This process ensures the generative process is continuous and differentiable regardless of whether data discrete or continuous, making BFNs particularly suitable for molecule generation, where atom coordinates continuous, and atom types are discrete. MolCRAFT (Qu et al., 2024) adopts the BFN as a backbone for conditional molecule generative model given proteins. Below, we outline the details of BFNs in the context of MoICRAFT.
Bayesian update To ensure that distribution parameters 0, can be updated in tractable way, BFNs employ Bayesian inference. Given the previous parameter 0,- and the noisy sample m,, the updated parameter 0g can be obtained via a Bayesian update function h(0,-1 me, 0+)) Or. The update function infers the posterior after observing noisy sample m, from the sender distribution ps mm, ax) and an input distribution L1 (ml0,-1), often selected conju- gate prior of the sender distribution.
This induces the Bayesian update distribution pu. which is mush-forward distribution of ps under transformation h. i.e., Pu (f0) == ps([m | (mm) = 0')), where we omit the conditions for brevity. By marginalizing out 01 one obtains the Bayesian flow distribution:
At a high level, BENs can be conceptualized as an iterative message exchange between two components: a sender and receiver. In each round, the sender transmits a noisy sample to the receiver, akin to the forward diffusion process in a standard diffusion model. The receiver then attempts to infer the sender's s parameters from the noisy input, effec- tively performing a reverse "denoising" step analogous to the reverse diffusion phase. This iterative cycle continues until reaching the final estimate of the data distribution in parameter space.
Pr (0 m. p.t. = P Pu (0, 0o m. p. ag+),
(3)
where On is the initial parameter used to define a simple distribution such as the standard normal for continuous or uniform for discrete. A key advantage of BFNs is that, with suitable noise injections and the input distribution, O, can be calculated recursively without explicit simulation.
Sender distribution Formally, given a molecule m [x, v], the sender injects continuous noise for both coordi- nates and types independently according to a predefined noise schedule. We denote the resulting noise-injected molecule by m [x, v). The sender distribution is then factorized as:
Training objective Training proceeds by minimizing the Kullback-Leibler divergence between the sender distribu- tion ps and the receiver distribution pR:
DKL (ps |RR]] ((m. p) = (E,, mi 0.- 0i
(4)
For the continuous-time extension of BFNs and additional technical details, we refer readers to Graves et al. (2023).
Nu
NM
ps(m m),) II p(x x
II
p 1( af),
(1)
3.3. Direct Preference Optimization
Reinforcement learning with human feedback (RLHF) has been proposed to align the large language models with human preference (Ouyang et al., 2022). Let D = (c,x",,')) be preference dataset, where x" and x' are winning and losing responses given prompt c. In standard
where a₁ is an accuracy parameter such that the sender samples are entirely uninformative about x when a = 0 and become more informative as a increases. One can use different schedules for the coordinate ax and atom type o..
3

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
RLHF pipelines, one first trains a reward model capture these preferences and then uses policy optimization meth- ods (e.g.. PPO) to fine-tune the model such that it produces higher-scoring outputs.
where m, is noisy sample drawn at time step (the "re- ceiver" output), governed by the BFN's sender/receiver dis- tributions, and O, are parameters sampled from the Bayesian flow distribution pr, which performs Bayesian updates in parameter space.
Direct Preference Optimization (DPO) offers an alternative that bypasses the need for an explicit reward-model approx- imation in the training loop (Rafailov al., 2024). Let (e(x) be model to be fine-tuned. Given a reference model Pret, DPO is derived from the optimal solution to the RLHF objective and can be expressed as
Unlike Diffusion-DPOL the reconstruction at time t depends on the parameter 0:-1. Expanding the ratio into sum over time steps with Jensen's 's inequality, one obtains an upper bound involving differences of KL terms at each time step. By approximating the receiver distribution in expectation with the sender distribution, we can obtain the following loss:
LDPO
Po (xx
Pe (a') Poet (x²) ].
-Ee s!! logo Blog
Poet (xx)) - B log
LEEN (OONNDPO(@) = - EE L
Dy
(5)
which is separated into two parts: atom loss 4 and type loss CY. A complete derivation is provided in Appendix A.
where we omit the condition e for brevity. Minimizing this loss encourages the fine-tuned model po to assign higher relative likelihood the winning samples than the losing samples, while remaining close Pret- Recently, Diffusion- DPO extends DPO to diffusion-based generative models by expanding the reward model over the entire diffusion trajectory, enabling preference-driven fine-tuning in multi- step generative processes (Wallace et al., 2023).
Let x and x' be the clean coordinates for the winning and losing molecules at time t = 0. Then the atom loss has the form of
Cx = E log (- (BT o7 2
1 - -b1]].
(8)
with
4. Multi-Reward Optimization for BFNS
w
| 10
$5 (0(-))1)) - ||x=0 D2 (0--1)1
In this section, we present our Direct Preference Optimiza- tion for Bayesian Flow Networks (BFN-DPO), an end-to- end framework for controllable molecule generation. Our method draws on recent advances in DPO and adapts these ideas to BFNs. We then extend DPO-BFN to handle multiple reward signals via a softmax based normalization strategy and an imcertainty-regulariiedd ensemble, allowing fair com- parisons among rewards with different numerical scales.
and analogously for 4 Here, (-) and 4 generate coordinate predictions under the fine-tuned and reference BFNs, and at is a concentration parameter for the coordi- nates at step t.
For the discrete atom types, we have similar construction based on KL divergences. Each "winning" and "losing" noisy type vi" and v is compared to the true one-hot type vector v or The derivation yields:
4.1. Direct Preference Optimization on BFNs
We consider a preference dataset D {(p, m", m')), where each triplet consists of protein target p. a "winning" molecule m and "losing" molecule m'. These prefer- ences are determined by pre-computed reward values (c.g., binding affinity). We denote by (ra (.) and (.).)) the reference and fine-tuned BFN models, respectively. Draw- ing on Diffusion-DPe (Wallace et al., 2023), we define the BFN-DPO objective for a single preference pair (m", mm')
CT = - E 12[1og((-30) (-8T (A3) - (4%)..)] (9)
where
(A
In Imps(ve v ,,x) - Inpa(V; () ax) - -Inps(VV .a)) - -npp( (0 of)],
and similarly for (4))-- In practice, these log- logelikeli terms could be implemented as cross-entropy or KL diver- gences for categorical predictions.
Po (in 1
CBEN-DFO(D) = -E log a ( log Pest (mrr.
-1)
pe (mirr 0 0
0 00:T-1 )}
-B log
By minimizing (f)) the model Pa is driven to gen- crate noisy samples ri that better align with automated preferences than m. all while remaining close to Pref-
Pret (mirr
4

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
4.2. Multi-Reward Normalization
Fw f becomes too large (close 1), the second term is nullified; thus, over-optimizing already "high-scoring" win- ners is discouraged. In effect, EºPO enables more balanced training and helps prevent mode collapse onto narrow set of winning samples.
While one can fine-tune Bayesian Flow Network (BFN) solely using single reward signal (e.g., validity), doing may degrade other desirable properties like binding affin- ity. To address this limitation, we adopt a multi-reward BFN-DPO scheme inspired by Kim et al. (2024), where multiple reward signals (c.g., binding affinity, validity) are simultaneously optimized.
5. Experiment
In this section, we provide the result of the multi-object optimization approach with empirical experiments. Due to the space limit, missing implementation details are provided in Appendix B.
Suppose we have K reward functions that evaluate proteule-protein pair (p, m). Since these rewards may have different numerical ranges and may be preferred in dif- ferent directions (e.g., negative for binding affinity, positive for strain energy), direct averaging can be misleading. To handle scale discrepancies, we normalize each reward via a softmax over a mini-batch size B:
5.1. Experimental Setting
Datasets We fine-tune pretrained model on the Cross- Docked dataset (Francoeur et al., 2020). The CrossDocked dataset comprises 100K training pairs and 100 target pro- teins for testing. For each protein, win-lose pairs are selected based on multi-reward scores described in Section 4.2 from molecules generated by the pretrained model. These pairs serve as input for fine-tuning via proposed BFN-DPO procedure. We generate 100 new molecules per target pro- tein in the test set for evaluation.
(()) exp(Is(1())) Ee (xp())(r))))
(10)
where (())) is the j-th reward for the i-th molecule, and (-()) is simple transformation based on the j-th reward, so that higher values consistently indicate better rewards across different metrics. This maps each reward to the (0, 1) interval, ensuring that distinct reward signals become comparable despite inherently different numerical ranges.
Baselines We compare our method to the following base- lines: AR (Luo et al., 2021) and Poket2Mol (Peng et al., 2022) generate molecules autoregressive, and FLAG (Zhang & Liu, 2023) generates at the fragment level. TargetD- iff (Guan al., 2023). DecompDiff (Guan et al., 2024), De- compOpt (Zhou et al., 2024), and IPDiff (Huang et al., 2024) are diffusion-based generative models. IPDiff incorporates protein-ligand interactions as a priori in the forward process. MolCRAFT (Qu et al., 2024) employs BFNs, where noise and denoise processes are performed in a fully continuous parameter space. AliDiff (Gu et al., 2024) applies a single reward optimization to improve the binding affinity based on IPDiff.
After normalization, one still must combine the m rewards into a single score. Following Kim et al. (2024), we reduce the influence of high-variance reward profiles by subtracting variance-basee penalty. Specifically, let μ4s, be the mean of the normalized rewards for molecule i. The final multi- reward score for that molecule is
MA. 7
> (1 m
(11)
where y is hyperparameter controlling how strongly to penalize high reward variance. This strategy promotes molecules that achieve balanced performance across all re- ward dimensions, rather than excelling at some while failing at others. More details on the hyperparameters can be found in Appendix B.
Evaluation metrics We evaluate the generated molecules from three different perspectives: binding affinity, conforma- tion validity, and molecular properties. For binding affinity, we use Vina Score, Vina Min, and Vina Dock measured by AutoDock Vina (Eberhardt et al., 2021) following the setup used in Ragoza et al. (2022). The Vina Score measures the binding affinity of the generated molecule as is. Vina Min measures the affinity after optimizing the molecular structure without changing the docking position. Vina Dock measures the binding affinity after re-docking the molecule; hence, reconsider the molecule's position and orientation. Low scores indicate stronger protein-ligand binding affinity. We report average and median scores for all Vina metrics. For conformation validity, we report strain energy (SE) mea- suring molecule conformation energy and protein-ligand
We further stabilize training by adopting the E*PO, a vari- ant of DPO (Gu et al., 2024), to avoid excessive focus winning samples when their advantage is already large. Con- cretely, we replace the standard BEN-DPO loss DEEN-DOO with
LRINE'O = - F (F - -Fi) LIENDRO
+ (1 o( - -ri) (2 LBFN.DRO)
(12)
where For and Fj denote the final multi newlid scores for the winning and losing molecules, respectively. If the difference
5

---

Enhancing Ligand Validity and Affinity in Structure-Bascl Drug Design with Multi-Reward Optimization
Table Performance comparison different reward methods. We mark the best, and the second-best performances in bold and underline, respectively. (L)(1) indicate whether a smaller/larger number is better, respectively.
Binding Affinity
Conformation Stability
Molecular Properties
Method
Vina Score (L)
Vina Min (L)
Vina Dock (U)
SE (L)
Clash (L)
SA (t)
QED (f)
Avg. size
Avg.
Med.
Avg.
Med.
Avg.
Med.
25%
50%
75%
Avg.
Avg.
Avg.
Reference
-6.36
-6.46
-6.11
-6.49
-7.45
-7.26
1.98
3.98
6.23
5.51
0.73
0.48
22.8
AR
-5.75
-5.64
-6.18
-5.88
-6.75
-6.62
2.79
7.39
16.34
4.46
0.64
0.51
17.7
Pocket2Mol
-5.14
-4.70
-6.42
-5.82
-7.15
-6.79
0.22
1.50
4.82
6.24
0.74
0.57
17.7
FLAG
16.48
4.53
1.21
-4.04
-5.63
-6.61
1.32
4.96
16.84
40.83
0.70
0.49
21,5
TifgetDiff'
-5.47
-6.30
-6.64
-6.83
-7.80
-7.91
7.48
16.85
33.70
10.84
0.58
0.48
24.2
DecompDiff
-5.19
-5.27
-6.03
-6.00
-7.03
-7.16
4.17
10.81
22.77
8.16
0.66
0.51
21.2
DecompOpt
-5.67
-6.04
-7.04
-7.09
-8.39
-8.43
12.34
21.37
37.69
14.56
0.61
0.45
29.4
IPDiff
-6.70
-7.53
-7.63
-7.71
-8.41
-8.39
10.30
27.05
65.49
8.79
0.61
0,52
24.3
AliDiff
-6.97
-7.75
-7.87
-7.94
-8.65
-8.64
10.09
27.60
8.76
0.56
0.50
24.4
MolCRAFT
-6.59
-7.04
-7.27
-7.26
-7.92
-8.01
3.10
7.92
14.87
7.09
0.69
0.50
22.7
Ours
-7.18
-7.38
-7.89
-7.77
-8.62
-8.64
2.05
5.56
11.62
6.69
0.74
0.55
22.8
clash measuring the number of potential overlaps between ligand atoms and protein (Harris et al., 2023). For molecular properties, we use synthetic accessibility (SA) (Ertl & Schuf- fenhauer, 2009) and quantitative estimation of likgelikeness (QED) (Bickerton al., 2012).
Rocket2Moll
0.76
Pucker
Ours
0.55
Reference
TPDuT
0.7
Decompthin
MolCRAFT
FLAG
AR
MolCran
ARDiff
0.65
DecompDift
IPTHI
TingeiDin'
0.6
DecomgOOt
Multi-reward optimization We choose MolCRAFT (Qu al., 2024) as a backbone for the multi-reward optimization. In the following analysis, we report the results of the fine- tuned model with Vina Score, SE, and QED. Although it is possible to optimize all seven evaluation metrics together with BFN-DPO, we find that this generally degrades the model performance. A more detailed analysis of reward selection is provided in Section 5.4.
0.45
Decomp
ARTHIT
TurgetDoff
0.50
30
20
10
(a) Vina Dock vs. QED
(b) SE vs.
SA
Figure 2. Comparison between the median Vina docking score and QED, and SE and SA.
Figure 1 illustrates how different models balance their aver- age Vina Dock scores against SA and SE. In both subplots, previous approaches occupy curve where improving one metric typically comes at the expense of the other. In con- trast, our multi-reward method moves beyond this frontier by achieving stronger binding affinity while simultaneously preserving higher SA and lower SE. This outcome high- lights the effectiveness of our framework in expanding the Pareto front for multi-objective structure-based drug design.
5.2. Results
Table 1 shows the overall performance of ligand molecules generated by our method and baselines. Our approach achieves four best and four second-best results in eight metrics, indicating robust and well-rounded performance in SBDD. While autoregressive models ensure stable confor- mations and molecular properties but struggle with binding affinity, and biffusion-based models achieve high binding affinity at the cost of conformation stability and synthetic ac- cessibility, MoICRAFT with BFNs balances both effectively. Qu et al. (2024) observes that binding affinity scores can artificially boosted by generating larger molecules, which present more surface area for binding. The average size of generated molecules remains comparable to those produced by MOICRAFT, indicating that no affinity hacking occurs. Note that smaller molecules often excel in SE and SA, which helps explain why AR and Pocket2Mol, both of which generate relatively smaller ligands, report particularly high SE and SA scores.
We further illustrate the comparison between Vina Dock and QED in Figure 2a and SE and SA in Figure 2b. A similar trend can be observed from the Vina Dock and QED, but no clear trade-off between SE and SA is observable. Instead, we observe a strong correlation between SE and SA.
5.3. Analysis
Validity of generated molecules Although the confor- mation stability and molecular property metrics show the validness of generated molecules by our method, we con-
6

---

Enhancing Ligand Validity and Affinity in Structure- Drug Design with Multi-Reward Optimization
100
of passing valldity test.
84%885%
Model
71%
AIDIff
IPDIf!
60
59%
Difg
Decompopt
485,50%
Rocket2Mal
25X225%
MolCRAFT
Ours
20
moi preed jnaded
mai sand toaded
panitzationn
nen, convertiiee
all attne, connected
bund jengths
Internai, preni clasn
bond angies
aromati,, ring farness
douilng Fand, Hamess
protein marimum, dofance minimom, disance, to protoin
minimum, distanca to organic, cofartors miniuum dicturce to horganic, colactors
interaal energy
minuuum distance, waters
volumo overtay, with organic calartors volume, overiag with marguric, corazzors
volume overian, with presein
volume overiae weators
AllDiff
Decompopt
IPDiff
TargetDiff
AR DecompDiff Pocket MolCRAFT Ours
(a) The percentage valid molecules.
(b) Cumulative chart showing the remaining percent valid molecules after each validity test.
Figure 3. Validity check of the generated molecules from different models through PoseBuster (Buttenschoen et al., 2024).
120
Table 2. Comparison between different opitmization approaches. GS is short for gradient surgery. A model fine-tuned with multi- rewards shows better binding affinity scores than the model with gradient surgery.
Strain Energy
80
60
Binding Aflinit
Conformation Stabiiit
Molecular Properties
SE
SA Avp. (*)
QED
20
H
Med. (i)
Med. (1)
Avg. (1)
Avg. (1)
Reference
-7.66
3.96
5.51
0.73
0.48
ANDIFF
IPDif
Decompppt
TargetDiff
DecompDiff
MolERAFT
Ours
FLAG
Refarennc
Packet2Mal
Pretrained
7.62
GS
-6.00
9.69
7.18
0.66
0.50
6 GS
-6.87
8.62
7.11
0.69
0.52
Multi-Stage
Ours
-7.38
5.56
6.69
0.74
0.55
Figure Comparison of strain energy (SB) distributions for molecules generated by different methods.
this issue, Figure presents the SE distributions for different methods. Lower SE values correspond more physically plausible, stable 3D conformations. Our method clearly generates molecules with lower SE distributions, indicating structurally sound ligands. These improvements do not come at the expense of binding affinity.
duct more comprehensive validity check through the 20 validity tests proposed by PoseBuster (Buttenschoen et al., 2024). The test considers both intra-molecular and inter- molecular validity with the protein. Figure shows the percentage of molecules that pass validity tests on the generated 1OK molecules for the test proteins. Our multi- reward method generates significantly more valid ligands compared to other generative models. AliDiff, fine-tuned IPDiff (Huang et al., 2024) with Vina Dock as a single re- ward, is known to achieve the best binding performance so far. However, the model fails to generate valid molecules in many cases. Figure shows which test each model fails satisfy the validity over 20 validity tests. Notably, many models, including IPDiff and AliDiff, fail to satisfy proper bond angles between atoms of generated molecules. Some models fail to keep proper distance from the protein.
Synthetic Accessibility (SA) Figure presents the me- dian SA scores of molecules generated by AliDiff, IPDiff, and MolCRAFT with reference molecules across target pro- teins in the test The compared models recent state- of of-tth-art models for binding affinity. The results show that our model consistently generates molecules with high SA across all targets. We also report the proportion of proteins for which each model generated molecules with the highest median SA score. Our model obtains the best SA scores for 44% of proteins.
Strain Energy (SE) Recent studies (Buttenschoen et al., 2024; Harris et al., 2023) caution that deep learning-base based SBDD models may artificially inflate binding scores by producing ligands with unrealistically high SE. To examine
5.4. Ablation Studies
Different optimization approaches Multi-reward opti- mization can framed multi-task learning problem

---

Enhancing Ligand Validity and Affinity in Structure-Bascl Drug Design with Multi-Reward Optimization
protein
protein 21
protein
protein 61
protein
0.9
0.8
SA (Med.)
0.7
0.5
0.4
0.3-
The proportion of the highest SA (Med.) Reference (54.0%)
IPDiff (0.0%)
ALIDIFF (0.0%) MoiCRAFT (2.0%)
Ours (44.0%)
Figure 5. The median Synthetic Accessibility (SA) scores of the molecules generated for each of the 100 target proteins in the test set. We compare our model to the top three models from the highest median binding affinity (AliDiff, IPDiff, MolCRAFT), as well as the reference. The values in parentheses represent the proportion of the best SA score for each protein
Table Comparison based on the number of rewards. A model fine-tuned with three rewards shows better binding affinity scores than the model with all seven rewards.
Binding Affinity
Conformation Stability
Drug-like Properties
Rewards
Vina Score
Vina Min
Vina Dock
SE Med. (4)
Clash
SA Avg. (t)
QED
Avg. (4)
Avg. ())
Avg. (l)
Avg. (J)
Avg. (t)
No
Pretrained
-6.59
-7.27
-7.92
7.62
7.09
0.69
0.50
Score
-6.63
-7.36
-8,01
6.85
7.25
0.71
0.53
Min
-6.11
-7.43
-8.02
6.88
7.36
0.71
0.53
Dock
-6.63
-7.40
-8.1
6.99
7.34
0.71
0.53
Single
SE
-6.51
-7.25
-7.79
6.79
7.40
0.70
0.50
Clash
-6.44
-7.09
-7.58
9.87
7.16
0.66
0.48
QED
-6.45
-7.31
-7.80
6.13
7.38
0.72
0.54
SA
-6.50
-7.29
-7.80
7.33
7.47
0.71
0.53
Multi (Ours)
Dock+SE+QED
-7.18
-7.89
-8.62
5,56
6.69
0.74
0.55
7 rewards
-6.50
-7.28
-7.86
7.22
7.28
0.53
0.70
when each reward is treated as a separate task. To explore this perspective, we adopt the widely used multi-task op- timization technique known as gradient surgery (Yu al., apply the gradient surgery to DPO losses of Vina score, SE, and QED (3GS). Since the DPO loss consists of atom and type losses as in Equation (7), we further apply the gradient surgery to these losses separately, resulting in six losses in total (6GS). addition, we exam- ine multi-stee stage optimization procedure, wherein the model iteratively fine-tuned using newly generated molecules from each preceding iteration. We report the result of the two-stage approach.
approach outperforms our uncertainty-base reward opti- mization. The gradient surgery methods perform worse than the reference model in all metrics except QED.
2020). Specifically,
Multi-reward selection Table 3 show the performances with different choices metrics as rewards. Trying to op- timize all seven rewards at once underperforms in several categories, implying the complex loss landscape with all rewards. In general, optimizing Vina metrics shows a strong correlation, i.e., optimizing one Vina metric improve the performance of the other two Vina metrics. The optimiza- tion over QED improves the performance of SA, leading to our final choice of rewards. Note that Clash cannot be improved even by the corresponding single reward. We leave the improvement of Clash for future work.
Table 2 compares the performance of each approach to our uncertainty- reward optimization. We only report the Vina Score among all Vina metrics, but they share similar trends. Neither the gradient surgery nor the multi-stage
8

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
6. Conclusion
sibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of chem- informatics, 1:1-11, 2009.
In this work, we propose a multi-reward optimization frame- work to enhance deep learning-base -ased structure-based drug design. Our method addresses the limitations of previous approaches by fine-tuning generative models to optimize multiple attributes, including binding affinity, validity, and drug-likonn -ikeniieeees.. By utilizing direct preference optimization within Bayesian flow network and integrating a reward normalization scheme, we enable the generation of more realistic ligands with improved binding affinity. Experimen- tal results demonstrate that our approach outperforms the baseline models, offering more robust solution for drug discovery tasks.
Francoeur, G., Masuda, T., Sunseri, J., Jia, A., Iovanisci, R. B., Snyder, I., and Koes, D. R. Three-dimensional convolutional neural networks and cross-docked data set for structure- baasd drug design. Journal of chemical information and modeling, 6009):4200-4215, 2020.
Graves, A., Srivastava, R. K., Atkinson, T., and Gomez, F. Bayesian flow networks. arXiv preprint arXiv:2308.0000 07037, 2023.
Gu, S., Xu, M., Powers, A., Nie, W., Geffner, T., Kreis, K., Leskovec, J., Vahdat, A., and Ermon, S. Align- ing target-aware molecule diffusion models with ex- act energy optimization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024, URL https:: /tttt::///penrevieww forum?? id=EMcvxXtzNu.
Acknowledgements
This work was partly supported by Institute of Information & communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (RS-2019-
11191906: Artificial Intelligence Graduate School Program (POSTECH), RS-2024-00457882: Artificial Intelligence Research Hub Project) and Korea Evaluation Institute of Industrial Technology (KEIT) grant funded by the Korea government (MOTIE, Korea) and the National Research Foundation of Korea (NRF) grant funded by the Korea gov- ermment (MSIT) (RS-2024-00337955, RS-2023-00217286)
Guan, J., Qian, W. W., qiang liu, Ma, W.-Y., Ma, J., and Peng, J. Energy-inspired molecular conformation op- timization. In International Conference on Learning Representations, 2022. URL httpt:://////.. /ttp::///ppeeeeeeeeeeee net met/forummidd=7FIIM
Guan, J., Qian, W. W., Peng, X., Su, Y., Peng, J., and Ma, J. 3d equivariant diffusion for target-aware molecule gener- ation and affinity prediction. In International Conference Learning Representations, 2023.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Guan, J., Zhou, X., Yang, Y., Bao, Y., Peng, J., Ma. J.. Liu, Q., Wang. L.,, and Gu, Q. Decompdiff: diffusion models with decomposed priors for structure-basee drug design. arXiv preprint 240.::::000009022 2024.
References
Harris, C., Didi, K., Jamasb, A. R., Joshi, C. K., Mathis, S. V., Lio, P., and Blundell, T. Benchmarking generated poses: How rational is cructure-based drug design with generative models? arXiv preprint arXiv:2308.074111 2023.
Ba, L.,, Kiros, J. R., and Hinton, G. Layer normaliza- tion, 2016. URL wwwphh:ttsttiiiieeeeeeee http//aiii....gaa/a// wrgg:/s//6/// /167 06450.
Bickerton, G. R., Paolini, G. V., Besnard, J., Muresan, S., and Hopkins, A. L. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):00 9800009982222 2012.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6851, 2020.
Buttenschoen, M., Morris, G. M., and Deane, C. M. Pose- busters: Ai-based docking methods fail to generate physi- cally valid poses or generalise to novel sequences. Chem- ical Science, 15(9):313-3311 3139, 2024.
Huang, Z., Yang. L., Zhou, X., Zhang, Z., Zhang, W., Zheng, X., Chen, J., Wang, Y., Bin, C., and Yang, W. Protein- ligand interaction prior for binding-aware 3d molecule diffusion models. In The Twelfth International Confer- ence on Learning Representations, 2024.
Eberhardt, J., Martins, D., Tillack, A. F., and Forli, S. Autodock vina 1.2. O: New docking methods, expanded force field, and python bindings. Journal of chemical information and modeling, 61(8):38913389, 2021.
Kim, K., Jeong, J., An, M., Ghavamzadeh, M., Dvijotham, K., Shin, J., and Lee, K. Confidence-anee aware reward op- timization for fine-tunig text-to-image models. arXiv preprint orXiv:2404.0/ 2024.
Ertl, P. and Schuffenhauer, A. Estimation of synthetic acces-

---

Enhancing Ligand Validity and Affinity in Structure-Basel Drug Design with Multi-Reward Optimization
Luo, S., Guan, J., Ma, J., and Peng, J. A 3d generative model for structure-based drug design. Advances in Neural Information Processing Systems, 34:6229-653 6239, 2021.
Zhou, X., Cheng, X., Yang, Y., Bao, Y., Wang, L., and Gu, Q. Decompopt: Controllable and decomposed diffusion models for structure-based molecular optimization. arXiv preprint arXiv:2403. 2024.
Ouyang, I.., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27785 2022.
Peng, X., Luo, S., Guan, J., Xie, Q., Peng, J., and Ma, Pocket2Mol: Efficient molecular sampling based 3D protein pockets. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Pro- ceedings of the 39th International Conference on Ma- chine Learning, volume 162 Proceedings of Machine Learning Research, pp. 17644-17655, PMLR, 17-23 Jul 2022. URL https:::: http///poceeng.....ee mlr press/ v1 v162/peng225 html.
Qu, Y., Qiu, K., Song, Y., Gong, J., Han, J., Zheng, M., Zhou, H., and Ma, W.- Molcraft: Structure-bare drug design in continuous parameter space. arXiv preprint 2xX1::2404.2214,, 2024.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Er- mon, S., and Finn, C. Direct preference optimization: Your language model secretly a reward model. Ad- vances in Neural Information Processing Systems, 36, 2024.
Ragoza, M., Masuda, T., and Koes, D. R. Generating 3d molecules conditional on receptor binding sites with deep generative models. Chemical science, 13(9):2701-2713 2022.
Schneider, G. and Fechner, U. Computer-based de novo design of drug-like molecules. Nature Reviews Drug Discovery, 4(8):649 663, 2005.
Wallace, B., Dang, M., Rafailov, R., Zhou, I.., Lou, A., Pu- rushwalkam, S., Ermon, S., Xiong, C., Joty, S., and Naik, N. Diffusion model alignment using direct preference op- timization, 2023. URL. https::/p//.. /ett:///rxxiii org/ abs/ 2311. 12908.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824-58366 2020.
Zhang, Z. and Liu, Q. Learning subpocket prototypes for generalizable structur-based drug design. In Inter- national Conference on Machine Learning, pp. 41382- 41398. PMIR, 2023.
10

---

Enhancing Ligand Validity and Affinity in Structure-Bascl Drug Design with Multi-Reward Optimization
A. Complete Derivation of BFN-DPO
Given preference molecular dataset as D = {(p, m m') where p the target proteins and m" m' are winning-losing pair of molecules based on pre-computed rewards. With the similar argument used in Diffusion DPO, one can define an objective for BFN over the entire timesteps as:
Po (o(mi:r)
log Po(ml)r)) Prer Pre((m (m';r)
LEEN-DPO(P) -Emm, mm pDD log
BE po |m (mf)
log
(13)
Piet Per((m)rr
log Pe (e(()i))
log P
Pel (m (m
= m² pv log
3 plppp (mm" "00 Po (mmm:rmm / for (0
1o oT |mm),,
(14)
Pief (m))())
Peol (m (m:r) )
(m²)
Ps (m
1.7
0 00T 10:r-1)12- 1)p 1m
=m ml pDD log
f po ( 1mm"000 ( Pol0 1/m")
log
Peel mm
00: Pos (0.-
1.7
Po (mm Imm (%
Po (r(0mr)//mm))
Po (m'r Pel (o
log
(15)
Pref(m 1o 00: |m))
Pel (m, 10 log
pel pa(m(10(_1)
Em m pDD log
(mp Im ao)pp ( (m" pa pa(mi Jm (0) Imm)
log Prer (mg 16/-1)
+C,
Pier m 10m0))
(16)
where approximate the intractable Pa (f)m"") via Bayesian flow distribution pf and (om"i mm 01T via the sender distribution. We omit the conditioning on the proteins for brevity.
By applying Jensen's inequality and dropping constant, we have
Po (mm (0m)) (1og
log Po(xn{10(_1)
per(mm))0 )) (17)
LBUN-DPO(0) Em mn ppDD. (1,,). logo (o(" mm p0 (0
(s(m) Prer (ommm"165 (0mm))
Expanding the equation further gives
(ro -E log a (- BT (ps mm" Pp (m," (0"1) DKa. (ps(m) mm") (me (0m)))
DxL (ps(m)) m²) Ps (m)) o)-1)) + DxL. (ps (m, m²) Prer(m! | o(_])
(18)
The latent parameters 0, = 10x 0Y from which the noisy samples m, = [x,, vi] are drawn with the noise factors a₁ = [a o.), are assumed to follow distinct continuous distributions for each atom coordinates and atom types v. With the neural network & representing the predicted atom coordinates and types as [$x xy), we can reformulate each of the above KI. divergence terms
DxL. (ps(m,)) xx) p(m, (0-1))
DX (ps(x) (x: x) p(x: 00 1x-1)) + DD (ps(v, | v) p(v (0-1)
(19)
ax ||x - (0)12 + (oggss(v; vv,aY) log log Φx (0:-1)),
= 2
(20)
As result, we can show that the objective can be represented as a summation of the atom coordinates and atom type preference losses follows:
(BEN-DPO(O) mm.m" 10) -pp p.)) [ + L-
(21)

---

Enhancing Ligand Validity and Affinity in Structure-Bascl Drug Design with Multi-Reward Optimization
The atom coordinates loss is designed to directly train the model output $ (0F) approximate the clean coordinate values Xp. The concentration parameter for each time is multiplied during training:
C
E
log @ (---T
3T
a7 2
(|x"
23 2300-))111
||x"
-0 (0
(p.x" x D, 10,00, upp "pp xx" 0 upp (-14.pi))
11 11( + ||x
0 (
(22)
Similarly, the atom types loss is designed to directly train the model output Φ (0Y) to approximate the clean atom types v, which is the NM x K matrix with each row as K -ddimensional one-hot vector:
[log o (-ST ST ((Ipp((lv"",g) ,,)) - In 100--mpp--00000111111
L
-B
"'.v",'''>>>" (.T.) 0Y - p..). (pp(t P.
(Inps(lv" ,0,) 1npp(100
-(mpps(Ilv a,)) -IppR(110 (0-)
(Imps(tv',a)) at) -Inpa(
(23)
B. Implementaton Details
B.1. SE-(3) Equivariant Network
To model the interaction between ligand and protein pocket atoms, we employ an SE(3)-equivariant network, PosNet3D (Guan et al., 2022), as the foundational architecture & in the receiver distribution. A protein-ligand graph is initially constructed by performing k-nearest neighbor search on the atomic coordinates, represented as G = (V, E". At each layer, the atomic hidden states h' and coordinates x' are iteratively updated as follows:
h h
h
Σ on (dig, h1; b/, exj, jENa(1)
Ax
Σ
(x) - xz) (dig b1tt hoft e.yyt)
jENa(i)
x Ax;
1mo!
where Ni(1) represents the neighborhood of atom in graph G, while h x, and hy. x, denote the hidden states and coordinates of atoms i and j, respectively. The term d' represents the Euclidean distance between atoms and j. and eij indicates whether the interaction is between protein atoms, ligand atoms, or a combination of both. The indicator function mot ensures that only ligand atom positions are updated. The functions Φn and dz act as attention mechanisms, utilizing Iy as query and [h h/: eij] as keys and values.
For the initial layer, the atomic positions and hidden states are initialized as x = [n, xp] and h linear(@", vp, f). the final layer, the network produces an estimation x Φ* Additionally, for the discrete variable v(d), the function & applies softmax operation to yield the probability distribution:
(d) = softmax (Dr)
B.2. Data Featurization
Each protein atom is characterized by multiple features to effectively capture its structural and biochemical properties. These features include a one-hot encoded element identifier (comprising H. C, N, O, S, and Se) to specify the atomic type, a one-hot encoded amino acid type indicator of 20 dimensions to classify the amino acid identity, a binary flag to distinguish backbone atoms, and a one-hot encoded arm/scaffold region indicator that designates whether the atom is part of an arm or scaffold region based on its spatial proximity the arm prior center.
12

---

Enhancing Ligand Validity and Affinity in Structure-Bascl Drug Design with Multi-Reward Optimization
Similarly, ligand atoms are represented using a one-hot encoded element type indicator (spanning C. N, O, P. and CI) and one-hot encoded arm/scaffold indicator to differentiate between aromatic and non-aromatic atoms.
For efficient message passing within the protein-ligand complex, two dynamically constructed graphs are employed: a k-nearest neighbors (k-NN) graph linking ligand and protein atoms and a fully connected graph for ligand atoms. The edge features in the k-NN graph are derived from the outer products of the distance embeddings, which are computed via radial basis function expansion, and four-dimensional one-hot encoded edge type vector that classifies different types of atomic interactions. In contrast, the ligand graph represents ligand bonds through one-hot encoded bond type vector, which categorizes bond types as non-bonded, single, double, triple, or aromatic.
Following the existing work (Guan et al., 2023), proteins and ligands are represented through atomic coordinates alongside one-hot encoded atomic type vectors. Protein atoms are characterized by one-hot encoded vector distinguishing among 20 amino acid types, while ligand atoms are encoded using one-hot vector that identifies various elements, including H, C. N, O, F, P, S, and CL. Additionally, a binary flag is incorporated to distinguish between protein and ligand atoms. To further refine these representations, two separate single-layer multi-layer perceptrons (MLPs) are applied, each mapping the input features into 228-dimensional latent space, ensuring compact yet informative feature embedding for subsequent computational processing.
B.3. Preference Pair Generation
For each protein binding site, we initially generate ten molecular candidates from the pretrained model. The reward for each synthesized molecule is then evaluated based on user-defined reward function specific to the corresponding binding site. To construct preference data, identify lower-performin sample with a comparatively lower reward and establish preference ranking. The detailed procedure for reward computation and preference selection is provided in Section 4.2.
B.4. Model Hyperparameters
For the SE(3)-equivariant network, we utilize k-nearest neighbors (kNN) graphs with 32-nearest neighbor search to construct the graph. The model consists of nine layers, each with hidden dimension of 128, employing a 16-head attention mechanism. ReLU activation functions are used in conjunction with Layer Normalization (Ba et al., 2016) to enhance stability and performance.
Regarding the noise schedules, we set the parameters as follows: 3, = 1.5 for atom types and α 0.03 for atom coordinates. The model is trained using a discrete-time loss function over 1000 training steps.
For pretraining, adopt the same architecture MolCRAFT, which is based on Bayesian Flow Networks (Graves et al., 2023).
For finetuning, we employ the Adam optimizer with a learning rate of 0,005 and batch size of 32. Additionally, an exponential moving average (EMA) of model parameters is maintained with decay factor of 0.999. Training converges within two epochs on a single RTX 4090 GPU, requiring approximately six hours for completion. For sampling, we generate 100 sample steps using noise-reduced sampling strategy.
To ensure higher value corresponds to a better reward across different metrics, transformation function in Equa- tion (10) as normalizing each rewards.
exp(Js((())) b= exp(f
A(0)
B
if j corresponds to the Vina Dock,
G())
r' if j corresponds to strain energy,
r, if j corresponds QED.
where (0()) is the j-th reward for the i-th molecule, and f-()) is a simple transformation function to ensure each reward is reasonable. This maps each reward to the (0, 1) interval, ensuring that distinct reward signals become comparable despite inherently different numerical ranges. We set y = 0.4 in Equation (11)
13

---

Enhancing Ligand Validity and Affinity in Structure-Based Drug Design with Multi-Reward Optimization
B.5. Computation Time
We report the time required for one epoch of pretraining and fine-tuning our model, comparing the AliDiff baseline trained with DPO. To measure the computation time, use the official source code provided by the author and the code on single NVIDIA A6000 GPU.
Table 4. Time needed for training per one epoch (in seconds)
Model
Pretraining
Finetuning only
AliDiff (DPO with diffusion model)
7461.73
22,884
Ours (DPO with BFN model)
5695.40
16,354
This shows applying DPO to pre-trained model takes 16,354 seconds with our method per one epoch, compared to 22,884 seconds for AliDiff, representing 28.5% reduction. Our approach does not significantly increase overall training time; however, obtaining rewards from an external tool can be time-consuming. To address this, we preprocess and store rewards in advance, but this requires offline DPO framework. A promising future direction for SBDD is reducing reliance external tools, enabling more iterative and efficient online training loop.
14


Reviews:
[{"review_id": "75q9qwdKr9", "reviewer": "Reviewer_e8Ki", "summary": "This paper introduces a multi-reward optimization framework for structure-based drug design, addressing the challenge of generating ligand molecules with multiple desired properties like binding affinity, validity, and drug-likeness.  It fine-tunes generative models for these attributes together, using direct preference optimization for a Bayesian flow network and a reward normalization scheme. Experimental results show the method generates more realistic ligands with higher binding affinity compared to baselines, expanding the Pareto front observed in previous studies.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We deeply appreciate your efforts and positive feedback. Here, we address the comments and questions mentioned.\n\n---\n\n**W1)  Lack of computation complexity analysis against diffusion baselines. might be complex given the complexity of Bayesian Flow Networks (BFNs) and multi-reward optimization**\n\nTo address your concern regarding time complexity, we report the time required for one epoch of pretraining and fine-tuning our model, comparing it to the AliDiff baseline (which also uses DPO). To measure the computation time, we use the official source code provided by the author and run the code on a single NVIDIA A6000 GPU.\n\nRebuttal Table 1.  Time needed for training per one epoch (in seconds)\n| Model                     | Pretraining  | Finetuning only |\n|:---------------------------:|:----------------------------:|:-----------------:|\n| AliDiff (DPO with diffusion model) | 7461.73                  | 22,884      |\n| Ours (DPO with BFNs)         | 5695.40                    | 16,354       |\n\nThe Rebuttal Table 1 shows applying DPO to a pre-trained model takes 16,354 seconds with our method per one epoch, compared to 22,884 seconds for AliDiff, representing a 28.5% reduction. These results demonstrate that our model is more time-efficient than the diffusion-based baseline. Furthermore, coupled with its strong performance, this highlights the practicality of our model for Structure-Based Drug Design (SBDD) tasks. We have added this computation time in the appendix and appreciate R1's feedback in helping us improve the manuscript.\n\nOur approach does not significantly increase overall training time; however, obtaining rewards from an external tool can be time-consuming. To address this, we preprocess and store rewards in advance, but this requires an offline DPO framework. A promising future direction for SBDD is reducing reliance on external tools, enabling a more iterative and efficient online training loop.\n\n---\n\n**W2) Minor issue: figure 1&2 should provide axis indications**\n\nThank you for the editorial comment. We have added the axis indications in our revised manuscript.", "claims_and_evidence": "The main claim of this work is that the proposed multi-reward optimization framework generates more realistic ligands than baseline models. While not all evaluation metrics outperform baselines, the authors argue that their approach expands the Pareto front.", "methods_and_evaluation_criteria": "- Method: The paper applies multi-reward optimization and Direct Preference Optimization (DPO) for ligand generation in structure-based drug design. While the approach aligns with recent advances, it closely resembles existing frameworks (Kim et al., 2024) and apply the method on conditional generation of SBDD.\n\n- Evaluation: The evaluation criteria are well-structured, incorporating binding affinity, synthetic accessibility, strain energy, and drug-likeness metrics. The use of benchmark datasets like Cross-Docked ensures comparability with prior work.", "theoretical_claims": "The derivation of the loss function for Eq. 1 in Appendix A appears correct, though I did not examine it in detail.", "experimental_designs_or_analyses": "- Lack computation complexity analysis against diffusion baselines. might be complex given the complexity of Bayesian Flow Networks and multi-reward optimization", "relation_to_broader_scientific_literature": "- Structure-based Drug Design: A task to generate ligand conditional on protein targets. Previous works explore from auto-regressive models [GraphBP, Pocket2mol, etc] to diffusion-based models [TargetDiff, IPDiff, DecompDiff etc] and optimal transport based models [DecompOpt etc].\n\n- Bayesian Flow Networks [Graves et al., 2023].\n\n- Direct Preference Optimization [DPO 2024], and the multi-reward design in this work follows [Kims 2024].", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Minor issue: \n- figure 1&2 should provide axis indications", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q3cuNR6T7U", "reviewer": "Reviewer_Tr2z", "summary": "This paper introduces a multi-reward optimization framework for structure-based drug design, integrating Bayesian Flow Networks (BFNs) with Direct Preference Optimization (DPO). The method aims to simultaneously optimize ligand binding affinity, synthetic accessibility, and conformational stability. By incorporating reward normalization and uncertainty-regularized ensembles, the model expands the Pareto frontier across multiple benchmarks, achieving better trade-offs between affinity and molecular properties.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We deeply appreciate your thoughtful comments and positive feedback. Here, we address the comments and questions mentioned.\n\n---\n\n**W1) The model does not outperform baselines in ligand-protein clash detection, which may restrict practical applications.**\n\nWe acknowledge the reviewer’s concern regarding the clash detection performance. However, we would like to clarify the following points:\nThe clash metric reflects the number of colliding atoms, which inherently depends on the molecular size. As shown in the average size of main table 1 in our paper, Pocket2Mol and AR tend to generate smaller molecules on average compared to other models, which naturally results in fewer clashes. To address this discrepancy, we conduct an additional analysis where we filter the generated ligands based on the number of atoms to obtain comparable molecular sizes and report the clash performance accordingly.\n\nRebuttal Table 2. Clash performance after filtering for similar molecule sizes.\n|   Method  |  AR  | Pocket2Mol | TargetDiff | DecoompDiff | DecompOpt | IPDiff | AliDiff | MolCRAFT | Ours |\n|:---------:|----:|----------:|----------:|-----------:|---------:|------:|-------:|--------:|----:|\n| Avg. size | 17.7 |       17.7 |       18.0 |        17.7 |      17.2 |   17.5 |    17.6 |     17.8 | 17.6 |\n|   Clash (↓) | **4.46** |       6.24 |       7.50 |        6.70 |      9.10 |   6.20 |    7.20 |     5.87 |  4.50 |\n\n\nIn Rebuttal Table 2, when comparing models with comparable molecular sizes, our method demonstrates strong performance in clash metric. Notably, the only model exhibiting a better clash score, AR, performs poorly in binding affinity. Given that both low clash and high binding affinity are critical for practical applications, these results indicate that our model offers the best balance, making it the most suitable choice for real-world deployment.\n\n---\n\n**W2) The innovation of this paper is somewhat limited.**\n\nWe acknowledge that each methodological component may not be significantly novel. However, integrating them into a single system that simultaneously enhances key properties for Structure-Based Drug Design (SBDD) remains a meaningful contribution. In particular, we are the first to apply DPO to Bayesian Flow Networks (BFNs) and to leverage reward normalization for multi-reward optimization—improving binding affinity, validity, and drug-likeness at once. This approach, not previously explored in the literature, helps make SBDD more practical.\n\n---\n\n**W3) Validation is limited to the CrossDocked dataset, raising concerns about generalizability to other protein targets or real-world drug discovery scenarios. However, this is primarily due to the inherent issues and limitations of the SBDD task at its current stage, rather than a problem with the paper itself.**\n\nWe completely agree with this point. As research in SBDD continues to grow, we expect more robust and diverse benchmarks to emerge, allowing broader validation and ensuring greater generalizability to real-world scenarios.\n\n---\n\n**Q1) I'm curious about what results would be obtained if the fine-tuned model were sampled again and the fine-tuning process were repeated. Will the model continue to improve, or will its performance deteriorate due to certain reasons (such as overfitting)?**\n\nWe appreciate the insightful question, which encourages further discussion. To address it, we provide results for a model fine-tuned iteratively for two stages (i.e., multistage) with newly sampled molecules using DPO, compared to our model and the pretrained model.\n\nRebuttal Table 3. Performance comparison of the pretrained model, a multistage fine-tuned model, and our model, \n| Method |  Vina Score Med. (↓) |  SE Med. (↓) |   Clash Avg. (↓)  |  SA (↑)  |   QED (↑) | \n|:--------------|---------:|---------:|-----------:|---------:|------------:|\n|    Pretrained    |  -7.04 |7.62 |7.09 |0.69 |0.50|\n|     Multistage DPO (2 stage)    |  -7.16 |  **5.15** |   7.54 |  **0.74**|   0.54  | \n| Ours |  **-7.38** |  5.56   |  **6.69** | **0.74**  | **0.55**| \n\nAlthough the two-stage DPO achieves competitive performance, it does not surpass our model choice. Xu et al. also suggest that fine-tuning via DPO does not necessarily guarantee significant improvements [1]. Furthermore, iteratively measuring rewards using external tools at each training step is computationally expensive compared to the training process itself, which led us to adopt an offline strategy. Nevertheless, we believe that exploring an online, iterative approach with efficient sampling remains a promising direction for future research.\n\n **Reference**\n\n[1] Xu, Shusheng, et al. \"Is dpo superior to ppo for llm alignment? a comprehensive study.\" arXiv preprint arXiv:2404.10719 (2024).", "claims_and_evidence": "Yes, it seems to be fine.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes, it seems to be fine.", "experimental_designs_or_analyses": "Yes, it seems to be fine.", "relation_to_broader_scientific_literature": "Refer to the strengths and weakness.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n\n1. The integration of multi-reward optimization with BFNs addresses the limitations of single-objective approaches, balancing critical molecular attributes effectively.\n2. Extensive comparisons with state-of-the-art models and ablation studies validate the necessity of multi-reward strategies, showing improvements in binding affinity and conformational stability.\n3. The writing is clear and easy to follow.\n\nWeakness:\n\n1. Based on paper [1], which finetunes a diffusion model using DPO. The innovation of this paper is somewhat limited. \n2. The model does not outperform baselines in ligand-protein clash detection, which may restrict practical applications.\n3. Validation is limited to the CrossDocked dataset, raising concerns about generalizability to other protein targets or real-world drug discovery scenarios. However, this is primarily due to the inherent issues and limitations of the SBDD task at its current stage, rather than a problem with the paper itself.\n\n\n\n[1]. Gu, et, al. Aligning target-aware molecule diffusion models with exact energy optimization. NeurIPS 2024.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. I'm curious about what results would be obtained if the fine-tuned model were sampled again and the fine-tuning process were repeated. Will the model continue to improve, or will its performance deteriorate due to certain reasons (such as overfitting)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eAl6mdmWEC", "reviewer": "Reviewer_rt1G", "summary": "This paper proposes a multi-reward optimization framework to enhance ligand validity and binding affinity in structure-based drug design (SBDD). By integrating Direct Preference Optimization (DPO) with Bayesian Flow Networks (BFNs), the authors achieve joint optimization of multiple objectives, such as binding affinity strain energy, and QED. Experiments demonstrate that the proposed method is comparable to existing baseline models in terms of binding affinity and molecular validity, while also expanding the Pareto frontier in multi-objective optimization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## [R3] Reviewer rt1G\n\nWe sincerely appreciate your positive and constructive feedback. Below, we address the comments and questions raised. Due to the word limit, additional experimental results can be found in [anonymous pdf link](https://anonymous.4open.science/r/anonymouspdf-718B/anonymous%20rebuttal.pdf) \n\n---\n\n**Q1) What difficulties arise when using multi-reward optimization?**\n\nMulti-reward optimization can be viewed as a multi-objective optimization, which sometimes underperforms compared to single-objective optimization due to task conflicts. Therefore, we encountered difficulties in designing a reward signal that simultaneously satisfies multiple metrics required for SBDD task.\n\nTo provide more insight, we analyze the reward values of winning and losing samples using both all seven rewards and three rewards (ours) and report their average rewards. [Rebuttal Table 4](https://anonymous.4open.science/r/anonymouspdf-718B/anonymous%20rebuttal.pdf)  shows the average difference between winning and losing samples with seven rewards is much closer ($\\times 10$) than that of the three rewards. We conjecture that the small difference introduces conflicts in the fine-tuning process, making it difficult for the model to optimize all properties.\n\n---\n\n**Q2) Why choose Vina, SE, and QED instead of other metrics (such as Vina, SE, SA; Vina, SE, Clash, QED)**\n\nWe first rule out using all possible metrics (as mentioned in Q1) and run single-reward experiments (Table 3) to identify unique, non-overlapping signals. We observe that optimizing one Vina metric improves all Vina metrics, so only one metric for binding affinity is sufficient. We chose QED over SA because improving QED also boosts SA, and we exclude clash due to its lack of improvement when used alone. Additionally, Vina and SE exhibit conflicting objectives, i.e., improving binding affinity (Vina) tends to compromise structural stability. This tendency is also observed in our single-reward experiment, so we include both to balance these competing objectives. Through this manual search, we determine that Vina, SE, and QED best complement each other in multi-reward optimization, as supported by our experimental results.\n\nFor further information, we additionally conduct an experiment using the reward combination given in R3's question statement—Vina, SE, QED, and Clash—whose results are presented in [Rebuttal Table 5](https://anonymous.4open.science/r/anonymouspdf-718B/anonymous%20rebuttal.pdf). As shown in the table, this combination leads to performance improvements over the pretrained model across most metrics. However, it still underperforms compared to our final model choice.\n\n---\n\n**W1) Details in temperature parameter to ensure the rewards make sense.**\n\nWe sincerely appreciate the reviewer’s insightful question, which helped us clarify the formulation. In the revised manuscript, for clarity, we have revised the formula in line 227 as\n\n$ \\hat{r}_{i}\\^{(j)} = \\frac{\\exp\\bigl(f_j(r_i^{(j)})\\bigr)} {\\sum\\_{b=1}^B \\exp\\bigl(f_j(r_b^{(j)})\\bigr)}$, \nwhere  $f_j(\\cdot)$ is $ -r,  \\quad  \\frac{10}{r}, \\quad r \\ $ if $j$-th reward is Vina Dock, SE, and QED, respectively.\n\nThis ensures a higher value corresponds to a better reward across different metrics. \n\n---\n\n**W2) Limited Improvement in Figure 3 compared to MolCRAFT, which serves as the backbone**\n\nThe results from Figure 3 are based solely on pass/fail criteria while presenting a broader range of physical plausibility (e.g., bond lengths, angles, etc). Therefore, we believe it is important to interpret Figure 3 in conjunction with Table 1 in the main paper. While the pass rate in Figure 3 shows a modest increase, the overall physical validity of our generated molecules has significantly improved, as shown in validity metrics, SE and clash. Additionally, we highlight that SOTA models such as IPDiff and AliDiff in binding affinity exhibit lower pass rates in Figure 3. This implies that these models sacrifice physical plausibility to achieve higher binding affinity while our approach improves both key properties.\n\n---\n\n**W3) Lack of Intuitive 3D Visualization as a 3D method**\n\nWe appreciate the reviewer’s valuable suggestion regarding the intuitive presentation of our contribution. We will incorporate additional visualizations in the revised manuscript, highlighting the docking poses of our generated ligands with improved binding affinity. \n\n---\n\n**W4) $K$ in lines 103 and 242 is not specified and inconsistent.**\n\nWe recognize the inconsistency in notation. As R3 mentioned, $K$ in line 103 represents the number of atom features, whereas $K$ in line 242 denotes the number of rewards. we have revised the manuscript to clarify and correct the error.\n\n---\n\n**W5) Unnumbered formulas**\n\nTo improve readability, we have numbered the equations accordingly in the revised manuscript.\n\n---\n\n**W6) Unspecified value $\\gamma$ in 242**\n\nWe chose $\\gamma = 0.4$. It has been included in the revised manuscript.", "claims_and_evidence": "The main claims of the paper are supported by experiments and theory.", "methods_and_evaluation_criteria": "For the formula in the left column of line 227, the paper does not explain how to set the temperature parameter to ensure that the rewards make sense. \nFor example, the Vina score is negative, with smaller values being better, while the strain energy is positive, with smaller values also being better. As rewards, higher values are better for both after applying softmax.", "theoretical_claims": "1. In the right column of line 103, it is not specified what K represents. It seems to refer to the number of atomic features or atomic types. Moreover, it conflicts with the K in the formula in the left column of line 242, where K seems to represent the number of rewards. \n\n2. Most of the formulas are not numbered, which affects readability. \n\n3. The value of $\\gamma$ in the formula in line 242 is not specified in the implementation details.", "experimental_designs_or_analyses": "1. Figure 3 shows the results of the molecular validity check. Compared to MolCRAFT, which serves as the backbone, the proposed method that utilizing multi-reward optimization shows limited improvement.\n\n2. As a 3D method, and with the authors claiming to have expanded the Pareto frontier, there is a lack of intuitive presentation of the improvements. For example, compared to baselines, the docking poses of this method in the visualization are more realistic, and the Vina score is improved.", "relation_to_broader_scientific_literature": "The paper is closely related to existing literature. In the BFN section, it utilizes the backbone of the previous work, MolCRAFT [1]. The multi-reward part is inspired by Kim et al. (2024) [2], and the paper provides its own insights on combining DPO and multi-reward optimization with BFN.\n\n[1] Qu Y, Qiu K, Song Y, et al. Molcraft: Structure-based drug design in continuous parameter space[J]. arXiv preprint arXiv:2404.12141, 2024.\n\n[2] Kim K, Jeong J, An M, et al. Confidence-aware reward optimization for fine-tuning text-to-image models[J]. arXiv preprint arXiv:2404.01863, 2024.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n\nSee \"Relation To Broader Scientific Literature\"\n\n\n\nWeaknesses\n\nSee \"Methods And Evaluation Criteria\", \"Theoretical Claims\", \"Experimental Designs Or Analyses\".", "other_comments_or_suggestions": "Line 377, right column, \"DOP\" ->\"DPO\"", "questions_for_authors": "Why choose Vina, SE, and QED instead of other metrics (such as Vina, SE, SA; Vina, SE, Clash, QED)? What difficulties arise when using multi-reward optimization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a multi-reward optimization framework based on DPO for structure-based drug design. It aims to optimize multiple desired properties like binding affinity, validity, and drug-likeness.  Experimental results show the method generates more realistic ligands with higher binding affinity compared to baselines and a better Pareto front. 

Strengths:
- The integration of multi-reward optimization with BFNs addresses the limitations of single-objective approaches, balancing critical molecular attributes effectively.
- Experiments show improvement in binding affinity and conformational stability. 

Weakness:
- The paper uses existing training approach for multi-objective ligand optimization. The novelty is limited. 
- Validation is on the CrossDocked dataset, which is synthetic. 

Other concerns raised by reviewers such as computational complexity,  ligand-protein clash results, choice of metrics, and temperature parameter are addressed by authors during rebuttal.
Decision: Accept (poster)