Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Authors: Keyue Qiu
Yuxuan Song
Jie Yu
Hongbo Ma
Ziyao Cao
Zhilong Zhang
Yushuai Wu
Mingyue Zheng
Hao Zhou
Wei-Ying Ma
Venue: ICML 2025 poster (2025)
Keywords: molecule optimization
structure-based drug design
structure-based molecule optimization
bayesian flow network
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: Gradient-based BFN for targeted drug design

Abstract: Structure-based molecule optimization (SBMO) aims to optimize molecules with both continuous coordinates and discrete types against protein targets.
A promising direction is to exert gradient guidance on generative models given its remarkable success in images, but it is challenging to guide discrete data and risks inconsistencies between modalities.
To this end, we leverage a continuous and differentiable space derived through Bayesian inference, presenting Molecule Joint Optimization (MolJO), the gradient-based SBMO framework that facilitates joint guidance signals across different modalities while preserving SE(3)-equivariance.
We introduce a novel backward correction strategy that optimizes within a sliding window of the past histories, allowing for a seamless trade-off between explore-and-exploit during optimization.
MolJO achieves state-of-the-art performance on CrossDocked2020 benchmark (Success Rate 51.3\%, Vina Dock -9.05 and SA 0.78), more than 4x improvement in Success Rate compared to the gradient-based counterpart, and 2x ``Me-Better'' Ratio as much as 3D baselines. 
Furthermore, we extend MolJO to a wide range of optimization settings, including multi-objective optimization and challenging tasks in drug design such as R-group optimization and scaffold hopping, further underscoring its versatility.
Code is available at https://github.com/AlgoMole/MolCRAFT.

Empower Structure- Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Keyue Qiu 12 Yuxuan Song 12 Jie Yu? Hongbo Ma² Ziyao Cao 12
Zhilong Zhang
Yushuai Wu 1 Mingyue Zheng² Hao Zhou' 1 Wei-Ying Ma
Abstract
1. Introduction
Structure-basee drug design (SBDD) plays critical role in drug discovery by identifying three-dimensionall (3D) molecules that are favorable for protein targets (Isert et al., 2023). While recent SBDD focuses on the initial identifi- cation of potential drug candidates, these compounds must undergo series of further modifications for optimized prop- erties, a process that is both complex and time-consuming (Hughes et al., 201 1). Therefore, structure-bases molecule optimization (SBMO) has garnered increasing interest in real-world drug design (Zhou et al., 2024a), emphasizing the practical need to optimize for specific therapeutic criteria.
Structure- based molecule optimization (SBMO) aims to optimize molecules with both continu- ous coordinates and discrete types against pro- tein targets. A promising direction is to exert gradient guidance on generative models given their remarkable success in images, but it is chal- lenging to guide discrete data and risks incon- sistencies between modalities. To this end, we leverage a continuous and differentiable space derived through Bayesian inference, present- ing Molecule Joint Optimization (MolJO), the gradient-based SBMO framework that facilitates joint guidance signals across different modali- ties while preserving SEL())equivvariance. We introduce a novel backward correction strategy that optimizes within a sliding window of the past, allowing for a seamless trade-off between exploration and exploitation during optimiza- tion. MoJJO achieves state-of-the-art performance on CrossDocked2020 benchmark (Success Rate 51.3%. Vina Dock -9,05 and SA 0.78), more than 4x improvement in Success Rate compared to the gradient-based counterpart, and 2x "Meeerrter" Ratio as high 3D baselines. Furthermore, we extend MolJO to a wide range of settings, includ- ing multi-objective optimization and challeng- ing tasks in drug design such as R-group opti- mization and scaffold hopping, further underscor- ing its versatility. Code is available at https: //git trithan huub. CoomAlggMole haMMMMoeMMMMMRRM.. Mc1CRAFT.
Concretely, SBMO can be viewed as a more advanced task within the broader scope of general SBDD, requiring pre- cise control over molecular properties while navigating the chemical space. Specifically, SBMO addresses two key as- pects: (1) SBMO prioritizes targeted molecular property enhancement according to expert- seperf--p objectives, while generative models for SBDD focus primarily on maximizing the likelihood of data without special emphasis on property improvement (Luo et al., 2021; Peng et al., 2022). Therefore, these models can only produce outputs similar their train- ing data, limiting the ability improve molecular properties. (2) SBMO is capable of optimizing existing compounds with 3D structural awareness, addressing a critical gap left by previous molecule optimization methods with 1D SMILES or 2D graph representations (Bilodeau et al., 2022; Fu et al., 2022), and allowing for a more nuanced control. The focus on structure makes SBMO particularly suited for key design tasks, such as R-group optimization and scaffold hopping-
Recent works have explored SBMO through evolutionary- based resampling and gradient-based methods. For instance, DiftSBDD (Schneuing et al., 2022) uses a gradient-free evo- lutionary sampling method, and DecompOpt (Zhou et al., 2024a) further introduces a fragment--condttonee 3D gen- erative model for resampling. These approaches rely on iterative, computationally expensive oracle calls to select top-of-N candidates. One generalized and orthogonal solu- tion is gradient guidance, eliminating the need for oracle simulations while being flexible enough to be incorporated into strong-pertt poonngpnrooringggg generative models in plug-and and--lay
'Inttittte for Al Industry Research (AIR), Tsinghua University DDpartment of Computer Science and Technology. Tsinghua Uni- versity 'Shanghai Institute of Materia Medica, Chinese Academy of Sciences "Peking University. Correspondence to: Hao Zhou <zhouhao air.singhuua.cdu.cm.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Empower Structure-Basse Bascd Molecule Optimization with Gradient Guided Bayesian Flow Networks
A Structure-Based Molecule Optimization
C Joint Gradient Guidance over Continuous-iscrete Data
(1) Enhance property
(2) Optimize existing structure
discrete
Reference
GED: 0.53 6AO.11 5578
y ~
05
f
Guided
8
fly)
apl
00
OED Vnn
continuous
B
"Me-Ber Better" Ratio of Improved Molecules
Vina Score
Vina Min
10gm(0
Va, (ogp(8,10..)) (/0,
Gradient Guidance
Pocket2Mol
Vina Dock
TargetDiff
CRAFT
DecompDiff
D Backward Correction
DecompOpt
d
TAGMo
QED
AIl
Figure 1. Overview. A. Structure- molecule optimization, including (1) guiding molecule design by expert-specified objectives, (2) optimizing existing compounds the structure space. B. Study on the ratio of "me-better" molecules (with improved propertics), where all other baselines fall short in the overall improvement. C. Overall illustration of MolJO, utilizing joint gradient signals over continuous-disrrete data, where the distributions of 0 (for continuous x and discrete v) are taken from true guided trajectories. Graphical model of our proposed backward correction strategy, keeping a sliding window of size k.
fashion, demonstrated in a wide range of challenging real- world applications including image synthesis (Dhariwal & Nichol, 2021; Epstein et al., 2023).
In this paper, we address the multi-modality challenge for gradient guidance by leveraging a continuous and differ- ential space, representing an aggregation of noisy samples from the data space derived through Bayesian inference (Graves et al., 2023). We design MolJO (Molecule Joint Optimization), principled, end-to-end differentiable frame- work that enables gradient-based optimization of continuous and discrete variables. We introduce novel sampling strat- egy called backward correction, enhancing the alignment of gradients over different steps. By maintaining a sliding window of past history for optimization, the backward cor- rection strategy enforces explicit dependency on the past, effectively alleviating the issue of inconsistencies. More- over, it balances the exploration of molecular space with the exploitation of better-aligned guidance signals, offering flexible trade-off.
However, current gradient- based methods have not fully realized their potential in SBMO, for they have historically suffered from the continuous-diserete challenges: (1) it is non-trivial to guide discrete variable within probabilistic generative process. More specifically, standard gradient guidance is designed for continuous variables that follow Gaussian distributions, making them not directly applicable to molecular data that involve discrete atom types. Methods attempting to adapt gradient guidance discrete data often resort approximating these variables continuous, either by adding Gaussian noise (Bao et al., 2022) or by assuming that classifiers follow a Gaussian distribution (Vignac et al., 2023). Unfortunately, these approximations can lead to sub- optimal results, as they do not accurately reflect the discrete nature (Kong et al., 2023). (2) Gradient guidance might introduce inconsistencies between modalities. For instance, TAGMol (Dorna al., 2024) formalizes guidance exclu- sively over continuous coordinates, resulting in a disconnect between the discrete and continuous modalities. This may explain why TAGMol struggles to optimize overall molecu- lar properties as shown in Fig. 1, despite its improvement in Vina affinities. By solely guiding the continuous coordi- nates, TAGMol enhances spatial protein-ligand interactions but fails to optimize e.g. synthesizeability, which depends more molecular topology, especially discrete atom types.
Our main contributions are summarized as follows.
We propose MolJO, the joint gradient-based method for SBMO that establishes the guidance over molecules, offering better controllability and effectively integrat- ing gradient guidance for continuous-diserte variables within a unified framework.
We design a novel backward correction strategy for ef- fective optimization. By keeping a sliding window and correcting the past given the current optimized version, we achieve better-aligned gradients and facilitate a flex- ible trade-off between exploration and exploitation.
2

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
MolJO achieves the best Vina Dock of -9005, SA of 0.78 and Success Rate of 51.3%, and "Me-Better" Ra- tio of improved molecules that is 2x as much as other 3D baselines. We generalize MolJO to various needs including R-group optimization and scaffold hopping, highlighting its versatility.
only implicitly and belatedly in the generative process, and such imbalanced guidance would probably result in subopti- mal performance for lack of joint optimization.
3. Preliminary
3D Protein-Ligand Representation. A protein binding
2. Related Work
site p (xp, vp) represented as point cloud of Np atoms with coordinates Xp (xp.............. pxp....,xpp = R N and Kp-dimensional atom features Vp NNP) m = (XM, V.r) contains NM atoms, where (i) E R is the atomic coordinate and v RR the atom type. For brevity, the subscript for molecules "M and the pocket condition p are omitted unless necessary.
Pocket-Aware Molecule Generation. Pocket-aware gen- erative models aim to learn a conditional distribution over the protein-ligand complex data. Initial approaches adopt ID SMILES or 2D graph representation (Bjerrum & Threlfall, 2017; Gómez-Bombarelli et al., 2018), and recent research has shifted its focus towards 3D molecule generation in order better capture interatomic interac- tions. Early atom-based autoregressive models (Luo al., 2021; Peng et al., 2022; Liu et al., 2022) enforce an atom ordering to generate molecules atom-by-atom Fragment- based methods (Powers et al., 2022; Zhang et al., 2023; Lin et al., 2023) alleviate the issue of ordering by decom- posing molecules into motifs instead of atom-level genera- tion, but they risk more severe error accumulation and thus generally require post-processing or multi-stage treatment. Non-auttoregressive methods based on diffusions (Schneu- ing et al., 2022; Guan et al., 2022: 2023) and BFNs (Qu et al., 2024) target full-atom generation for enhanced per- formance and efficiency. However, the needs of optimizing certain properties and modifying existing compounds are not adequately addressed in the scope of previous methods, limiting their usefulness in drug design.
[vp..... p......, vp") E RRNKKr, Similarly, ligand molecule
Bayesian Flow Networks (BFNs). We briefly introduce how BFN views the generative modeling as message ex- change between a sender and receiver, with more details in Appendix A. The sender distribution (y X: xxa) builds upon the accuracy level a applied to data and defines the noised y. The varying noise levels constitute the schedule (()) f (('))di', similar to that in diffusion models.
A key motivation for BFN is that the transmission ought to be continuous and smooth, therefore it does not directly operate on the noisy latent as diffusions, but on the struc- tured Bayesian posterior 0 given noisy latents instead. The receiver holds a prior belief Oo, and updates the belief upon observed y, yielding the Bayesian update distribution:
E pu (0.1 0:-1,x; ((1
8(0. ((0)-1,)....)) (1)
Gradient- Molecule Optimization. Inspired by classifier guidance for diffusions (Dhariwal & Nichol, 2021), pioncering approaches are committed to adapting the guid- ance method to handle the complicated molecular geome- tries in the setting of pocket-unaware generation. EEGSDE (Bao et al., 2022) derives an equivariant framework for con- tinuous diffusion, and MUDM (Han et al., 2023) further explores time- independenn property functions for guidance. As they enforce continuous diffusion process for discrete variables, these methods are not applicable in advanced molecular modeling (Guan et al., 2022; 2023) where dis- crete data are processed by a discrete diffusion, for it is unnatural to apply progressive Gaussian noise that drives Categorical data away from the simplex. DiGress (Vignac et al., 2023) proposes classifier guidance for discrete diffu- sion of molecular graphs, yet it additionally assumes that the probability of classifier follows a Gaussian, which is ungrounded and often a problematic approximation. Based on the continuous-discred diffuuion for SBDD, TAGMol (Dorna et al., 2024) retains the guidance only for continuous coordinates, because there lacks a proper way to propagate the gradient over discrete types. The discrete part is affected
where 5(-) is Dirac distribution, and Bayesian update func- tion h is derived through Bayesian inference.
Intuitively, BFN aims to predict the clean sample given ag- gregated 0, i.e. conditioning on all previous latents. 0 is fed into neural network & to estimate the distribution of clean datapoint &, L.e. the outpul distribution po 4 t)). The receiver distribution is obtained by marginalizing out x:
PR(Y: 0,-1;t;,
ps(y, xx; (2)
(o (x +
The training objective is to minimize the KL-divergence between sender and receiver distributions:
(x) Ih (0.00 18x0))
E
Dps / x; pn(y: 01-1, t, (3)
4. Method
In this section, we introduce MolJO that guides the distribu- tion over 0 utilizing aggregated information from previous
3

---

Empower Prructure--ascce Molecule Optimization with Gradient Guided Bayesian Flow Networks
latents. Though different from guided diffusions that operate on noisy latent y, this guidance aligns with our generative process informed by 0. By focusing on lower-variance 0, we can effectively steer the clean samples towards desirable direction, ensuring a smooth gradient flow.
through Gaussian-distriuuted latent y and ensures that the discrete variables are still on the probability simplex without relying assumptions (Vignac et al., 2023) or relaxations (Bao et al., 2022: Han et al., 2023), and (2) alleviates the inconsistencies between modalities (Dorna et al., 2024) by joint gradient signals.
Notation. Following Kong et al. (2024) and denoting the guided distribution x as product of experts (Hinton, 2002) modulated by energy function E that predicts cer- tain property, we have π (0.(0,-1) X (0 0 (0.), where & is the pretrained network for BFN. P5 (0+) = exp - E f:)] is the unnormalized Boltzmann distribu- tion corresponding to the time-dependent energy function.
Guiding 0t for Continuous x. For continuous coordi- nates x E Nx3 it natural to adopt a Gaussian sender distribution y" N(x, or II). With a prior belief 0g 0.
we have the Bayesian update function for posterior 0 given noisy y" in Graves et al. (2023):
110 y", ...) = 0 1i-1Pi-1 + y
y",,,)
(6)
Overview. As illustrated in Fig. 1, we introduce MolJO as follows: in Sec. 4.1, we propose the concept of gradient guidance over the multi-modality molecule space, derive the form of guided transition kernel π (0:(0;-1) via first-order Taylor expansion, and explain the underlying manipulations of distributions the guidance corresponds to. In Sec. 4.2, we present generalized advanced sampling strategy termed backward correction for Ps, which allows for a flexible trade- off between explore- and-eeeoo by maintaining sliding window of past histories. We empirically demonstrate our strategy helps optimize consistency across steps, ultimately improving the overall performance.
Pi
with 0; 3 1. =1+ 3 given the schedule 5 01 21/n 1 for positive 0, and steps.
Remark 4.2. In the continuous domain, guidance over 0" is analogous guided diffusions, since guiding 0" corre- sponds to guiding noisy latent y" using the uncertainty- adjusted gradient ( [2)) )a'gg gy= that accounts for how changes in the sample space yd propagate to parameter space 0":
0- = 0 + o ge=
0 1- + (8* + 0o Pi
go. )as
a.
Pi
4.1. Equivariant Guidance for multi-modality Molecular Data
1 + (5* + o gy )ai
Ps
In this section, we derive the detailed guidance over for molecule m = (x, v) with N atoms, where x E RNN33 rep- resent continuous atom coordinates and v E (1...... K}N for K discrete atom types, and thus 0 == [0 0%). latent y = y") for the continuous and discrete modality.
Guiding 0e for Diserete v. For N-dimensional discrete types v E {l...., KJN, the noisy latent represents the counts of each type among K types, where we have y" ~ N(y") W(y"la(Ke) (a(Ke - 1), aKI), ey = [e...).... RKN C,0y = 6,60 E R with Kronecker delta func- tion 5. Further explanation is left to Appendix B.
E
Guidance over Multi-Modalities. To steer the sampling process towards near-optimal samples, we utilize the score Vo log Pe(()) (0) as a gradient-based property guidance, for which we have the following proposition (proof in Ap- pendix C.1), followed by details for each modality.
k 0 as a posterior belief is updated from the prior 0j
exp(y")0:
h(0) as) ΣK (Xp
(8)
Proposition 4.1. Suppose 0r ~ NO o"I) and K: A N(y 6, I by definition of BFN generative process, we can approximate the guided transition kernel (0. 0--1):
where the redundant o; = 3 - B with 3r = ( given a positive hyperparameter 3y.
Remark 4.3. In the discrete domain, guiding all latents y" amounts to reweight of the Categorical distribution for 0% changing the probability of each class in accordance with the gradient. Take an extreme case to illustrate, where gy is filled with one-hot vectors 5n:
0- ~ N + +0 "g=g o"I)
(4)
y. ~ N(y% + o" "ggyy,""'
(5)
where gradient ga o E(O, (2) 0=
gay
-Vg- (0,,t() By" = go By
20t
exp (0) 1)x
(0)) E. exp(yr)(0' (exp(o") exp(5"))
The guidance is formalized over both continuous coordi- nates and discrete types, and differs from previous guided diffusion for molecules in that (1) it guides the discrete data
= (0,)) 1+C < (0))
99)

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
for all k + d, where C = [exp(o") - 10 > 0 as
the property of additive accuracy once pe follows certain form as described by Graves al. (2023), the one-step backward correction can be derived as follows:
D exp(y?"]
the variance o" 0. It is obvious that the guidance lowers the probability for all classes but the favored d, redistributing the mass for discrete data in a more structured way than diffusion counterparts.
Pa(O ( 0 1, 02-2)
E (R 11 c.
E
PU (0: 1 (()
Dlpe (0.
10, 1. 2::000
Equivariance. Our proposed guided sampling that utilizes joint gradient signals is still equivariant as shown in the proposition below, with the proof in Appendix C.2.
E (0
PU
(0: 0 8
; + 0:))
(11)
where the original - po (x; [ 22 fz-1)) that has been used to update 0:-1 from 0.-2 at previous step, is now replaced by the optimized x;. By iteratively tracing back, we arrive at the k 1 step corrected estimation of ps:
Proposition 4.4. The guided sampling process preserves SE(3))equivariance when & is SE(3)- -qqiivaann if the energy function E(O,p.. p, 1) is also -qqui)-quivariaat and the complex is shifted to the space where the protein's 's Center of Mass (CoM) is zero.
Ps (0m.)mn- On-k)
E po ((0n))1)))))))
Pu (0, &ni im - K++11
Σ
a,)
(12)
4.2. Bayesian Update With Backward Correction
Algorithm 1 Gradient Guided Sampling of MolJO
Plugging Eq. 6 and 8 together with the sender distributions defined above into the right hand side according to Eq, I, yields the form of the backward corrected Bayesian update
Require:
9 ()), network (,),, schedules (t),
number of sample steps n, back correction steps k,
number of atom types K, energy function E(O, t),
A A38n + 0n A8 Pr I
guidance scale
Pu (0% (0 k. xn)) = =N(
1: Initialize belief 0 := 10, 0-] + 10, K.
Pn
2: i = 1 to n do
(13)
3: [t,t_ + [=1 m0, ا n
Pu (00 vn)
E y~N (Kev, e.. -1) KI)
4:
[x, p))
5:
[go- t --
oz E(O, p, t), o- E0-E(OO,pp) p,
8 8(0% exp(y) 0
6:
(P.P.:: A8" A8" + [1 + (t), +
(14)
Di=! exp(y:)
(()) (t) - 8r [t]
7: Retrieve a 8_, from the past
where i = 18, is drawn from the output distribution [m (0,-1,, p)) given pocket p. A 8%
8: Sample 0" according to Eq. 4 and 13
9: Sample y" and update 0" according to Eq. 5 and 14
&h and A,3" 3 3h are obtained from correspond- ing accuracy schedules. The concept of sliding window unifies different sam- pling strategies proposed by Graves et al. (2023)
10: end for
11: [x. ē,, + $ 1.p)
12: V arg (.,)
3.3
13: return 8, V]
5-50
Gookure Smarahy
8.5
b-102
(k=1) and Qu et al. (2024) (k=n). To under-
8.5
Here we propose general backward correction sampling strategy inspired from the optimization perspective, and analyze its effect on aligning the gradients. Recall that from Eq. we can aggregate 0, from previous latents:
0.4
stand its effect, we visu- alize the cosine similar- ity of gradients at each step w.r.t. the previous step in Fig. 2. By chang- ing the size i of sliding window, it succeeds in balancing sample quality (explore) and optimiza- tion efficiency (exploit), where it first focuses on exploring the molecular space with rapidly changing structures and gradients, and then exploits better-aligned guidance signals over gradually refined structures.
202
150
Generitie Froceor
p0(0; 0.-1)= 1
E Px ((.).....))
pu (0, 0 a;);)
Figure 2. Gradient cosine similar- ity, where & denotes the back- ward correction window size, For << k < 200, the similarity be- fore timestep A is omitted for it overlaps with k 200, i.e. cover- ing all the past.
(10)
Backward correction aims at "correcting the past to further optimize". Since we obtain optimized 0; from the guided kernel T (0, 00-1), there will be an optimized version of X = x for the next step. By backward correcting the Bayesian update distribution Pt given the optimized x", we are able to reinforce the current best possible parameter 0, instead of building on the suboptimal history. By utilizing
5

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
In practice, we employ the gradient scale 8 temperature
generated molecules that pass certain criteria (Vina Dock <
parameter, equivalent to (0. exp (ops=sslll F t)]- We 80 0² 0*) to stabi-
1818, QED 0.25, SA 0.59) following Guan et (2022).
further bypass the derivative Ty lize the gradient flow. The general sampling procedure is summarized in Algorithm 1.
(2) Molecular properties, including drug-likeness (QED) and synthesizability score (SA). (3) Metrics for sample dis- tribution, such as diversity (Div). A more comprehensive set of metrics are detailed in Appendix F.
5. Experiments
5.2. Unconstrained Optimization
5.1. Experimental Setup
In this section, we demonstrate the ability of our framework to improve molecular properties in both single and multi- objective optimization. We sample 100 molecules for each protein and evaluate MolJO in optimizing binding affin- ity and molecular properties. For additional evaluation of molecular conformation besides optimization performance, please see Appendix G.
We conduct two sets of experiments for structure-based molecule optimization (SBMO), although the constrained setting seems within the scope of unconstrained one, it is biologically meaningful and more practical in rational drug design, and further showcases the flexibility of our method.
Task. For molecule E EM where M denotes the set of molecules, there are oracles (m) .M- + R for property i, each with desired threshold 6; E R. MolJO is capable of different levels of controllability: (1) unconstrained op- timization, where we identify a set of molecules such that {m E MM| (m) 2 6, Vij. i.e. the goal is to optimize a number of objectives. (2) constrained optimization, where we aim to find set of molecules that contain specific sub- structures such that {m E M| a, (m) 2 6;, m, Vi).
MolJO effectively enhances molecular property w.r.t. generative models. The optimized distribution greatly improves upon the original generated distribution, as shown in Table I (row 14 vs. row 9).
MolJO outperforms gradient- based method with 4x higher Success Rate. As shown in Table 1, our model achieves sfate-of-the-an in affinity-related metrics while be- ing highly drug-like, with the best Success Rate of 51.3%, a four-fold improvement over TAGMol (row 14 vs. row 13).
Dataset. Following previous SBDD works (Luo et al., 2021), we utilize CrossDocked2020 (Francoeur al., 2020) to train and test our model, and adopt the same processing that filters out poses with RMSD 1 and clusters proteins based on 30% sequence identity, yielding 100,000 training poses and 100 test proteins.
MolJO has more potential than oracle-based baselines if equipped with oracles. RGA (Fu et al., 2022) and DecompOpt (Zhou et al., 2024a) show satisfactory Success Rate, enjoying the advantage of oracle-based screening at some expense of diversity. while AutoGrow4 (Spiegel & Durrant, 2020) falls short in QED, yielding a suboptimal Success Rate. Given the same concentration use of Z-score (Zhou et al., 2024a), we report variant of MolJO with top-of-N, selecting a tenth of top scoring molecules and showing that it is more effective than oracle-based methods once in a similar setting. Moreover, the higher diversity of DecompOpt and MolJO suggests the superiority of 3D structure-aware generative models over 2D optimization baselines (row 15 vs. row 10-12).
Baselines. We divide all baselines into the following: (1)
Generative models (Gen), including AR (Luo al., 2021),
GraphBP (Liu et al., 2022), Pocket2Mol (Peng et al., 2022),
FLAG (Zhang al., 2023), DifISBDD (Schneuing et al., 2022), TargetDiff (Guan et al., 2022), DecompDiff (Guan et al., 2023), IPDiff (Huang et al., 2024) and MoICRAFT (Qu et al., 2024). (2) Oracle-based optimization (Oracle) that rely on docking simulation in each round, such as Au- toGrow4 (Spiegel & Durrant, 2020), RGA (Fu et al., 2022), and DecompOpt (Zhou et al., 2024a), (3) Gradient-guided (Grad) TAGMol (Doma et al., 2024). Detailed descriptions of baselines are left in Appendix F.
MolJO is 2 as effective in proposing "me-better" candi- dates. For gradient- method TAGMol (Dorna et al., 2024), although produces seemingly promising high affin- ity binders, they come the expense of sacrificed molecular properties like QED and SA, demonstrating the suboptimal control of coordinate-oni, guidance signals. Notably, the ratio of all-better samples is below 17% for all other base- lines, and MolJO is twice as effective (39.8%) in generating feasible drug candidates that pass this criteria (Fig. 1).
Metrics. We employ the commonly used metrics as fol- lows: (1) Affinity metries calculated by Autodock Vina (Eberhardt et al., 2021), in which Vina Score calculates the raw energy of the given molecular pose residing in the pocket, Vina Min conducts a quick local energy minimiza- tion and scores the minimized pose, and Vina Dock per- forms a relatively longer search for optimal pose to calculate the lowest energy. Success Rate measures the percentage of
6

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Table Summary of different properties of reference molecules and generated molecules by our model and other baselines, where G+O
denotes equipping our method with top-of-NV in oracle simulations. Additional baselines and results can be found in Appendix F.I. (1)
(J) denotes a largrr / smaller number is better. Top results are highlighted with bold text and underlined text.
Vina Score (4)
Vina Min ())
Vina Dock ())
QED (t)
SA (*)
Div (1)
Success
Methods
Avg.
Med.
Avg.
Med.
Avg.
Med.
Avg.
Avg.
Avg.
Rate (1)
Reference
-6.36
-6.46
-6.71
-6.49
-7.45
-7.26
0.48
0.73
25.0%
①
AR
-5.75
-5.64
-6.18
-5.88
-6.75
-6.62
0.51
0.63
0.70
6.9%
GraphBP
4.80
-4.70
0.43
0.49
0.79
0.1%
① Pocket2Mol
-5.14
4.70
-6.42
-5.82
-7.15
-6.79
0.57
0.76
0.69
24.4%
FLAG
45.85
36.52
9.71
-2.43
-4.84
-5.56
0.61
0.63
0.70
1.8%
Gen
3
DiffSBDD
-1.44
-4.91
-4.52
-5.84
-7.14
-7.30
0.47
0.58
0.73
7.9%
TurgetDiff
-5.47
-6.30
-6.64
-6.83
-7.80
-7.91
0.48
0.58
0.72
10.5%
DecompDin
-5.19
-5.27
-6.03
-6.00
-7.03
-7.16
0.66
0.73
14,99
IPDiff
-6.41
-7.01
-7.45
-7.48
-8.57
-8,51
0.52
0.59
0.74
16.5%
MolCRAFT
-6.55
-6.95
-7.21
-7.14
-7.67
-7.82
0.50
0.67
0.70
26.8%
AutoGrow4
-8.99
-9.00
0.46
0.76
0.47
14.3%
Oraclle
T RGA
-8.01
-8.17
0.57
0.71
0.41
46.2%
G DecompOpt
5.75
-5.97
-6.58
66.70
-7.63
-8.02
0.56
0.73
0.63
39.4%
Grad
G TAGMol
-7.02
-7.77
-7.95
-8.07
-8.59
-8.69
0.55
0.56
0.69
11.1%
MoUOO
-7.52
-8.02
-8.33
-8.34
-9.05
-9.13
0.78
0.66
51.3%
G +
O
MoUJO (N=10)
-8.54
-8.81
-9.48
-9.09
-10.50
-10.14
0.67
0.79
0.61
70.3%
Table 2. Constrained optimization results, where Redesign means g-group optimization with fragments of the same size redesigned,
Growing means fragment growing into larger size, Hopping means scaffold hopping.
Vina Score (t)
Vina Min (L)
Vina Dock (4)
QED (t)
SA (1)
Connected Avg. (?)
Success
Methods
Avg.
Med.
Avg.
Med.
Avg.
Med.
Avg.
Avg.
Rate (1)
Reference
-6.36
-6.46
-6.71
-6.49
-7.45
-7.26
0.48
0.73
100%
25.0%
TargetDiff
-6.14
.6.21
-6.79
-6.58
-7.70
-7.61
0.50
0.64
85.5%
18.9%
TAGMol
-6.60
-6.66
-7.10
-6.80
-7.63
-7.76
0.53
0.62
87 .7.0%
19.2%
Redesign
MoICRAFT
-6.63
-6.70
-7.12
-6.91
-7.79
-7.72
0.49
0.67
96.7% %
22.7%
MolOO
-7.13
-7.28
-7.62
-7.39
-8.16
-8.20
0.57
0.68
95.1%
29.0%
TirgetDiff
-6.73
-7.29
-7.60
-7.67
-8.89
-8.79
0.39
0.52
11.2%
TAGMol
-7.00
-7.70
-8.08
-7.81
-8.92
-8.78
0.47
0.53
78.7%
11.8%
Growing
MolCRAFT
-6.96
-7.47
-7.86
-7.73
-8.80
-8.65
0.44
0.59
91.7%
19.9%
MoUOO
-8.08
-8.35
-8.79
-8.58
-9.21
-9.45
0.53
0.62
93.2%
32.7%
TargeDDiff
-5.72
-5.78
-6.00
-5.83
-6.31
-6.66
0.39
0.65
63.3%
6.2%
TAGMol
-6.17
-6.10
-6.46
-6.07
-7.19
-6.80
0.44
0.62
68.7%
6.9%
Hopping
MolCRAFT
-6.31
-6.17
-6.58
-6.40
-7.25
-7.15
0.42
0.67
89.9%
14.6%
MollO
-6.86
-6.50
-7.13
-6.70
-7.67
-7.58
0.46
0.68
90.5%
23.6%
MoLJO excels even in optimizing large OOD molecules. Note that for fair comparison, we restrict the size of gener- ated molecules by reference molecules so that both genera- tive models and optimization methods navigate the similar chemical space,
Table 3. Properties of molecules with a larger average size, where Vina stands for Vina Dock Avg., .. SR for Success Rate.
Methods
Vina
QED
SA
SR
Size
Reference
-7.45
0.48
0.73
25.0%
22.8
observe clear correlation between properties and sizes in Fig. 4. For model variants capable of exploring larger number of atoms, we report the results in Table 3 with sizes, where MolJO consistently outperforms other baselines, demonstrating its robustness. detailed discussion can found in Appendix E.
DecompDiff
-8.39
0.45
0.61
24.5%
29.4
Decompppt
-9.01
0.48
0.65
52.5%
32,9
MolCRAFT
-9.25
0.62
36.6%
29.4
MoJOO
-10.53
0.50
0.72
64.2%
30.0
cal use cases including R-group optimization and scaffold hopping, achieved by infilling (details in Appendix D.2).
5.3. Constrained Optimization
MolJO captures the complex environment of infilling, Table 2 shows that our method generates valid connected molecules and captures the complicated chemical environ- ment with better molecular properties than all basclines, showcasing its potential for lead optimization. for
Constrained optimization seeks to optimize the input ref- erence molecules for enhanced properties while retaining specific structures. We generalize our framework with such structural control and show its potential for pharmaceuti-

---

Empower Setucture-Bascee Molecule Optimization with Gradient Guided Bayesian Flow Networks
Reference
g-group design results
Reference
Scaffold hopping results
2PC8
1A2G
Mina Score -5,45, QED 0.34, SA 0.86
Mina Score +8 96, GED O. 0.84, SA..25
Vina Score 40, OED p SA
Vina Scone -5.33, QED q SA 0.94
2AZY
2E24
B
Vina Score 9.93, QED 0.54, SA.0.58
Vina Score -11.26, QED 0,64, SA 0.88
Vingl Séore 5.39 QED 0.11. SA
Siore -$ $.98, Q.9 SA
stacking
cation
Hydrogen bond
Sals bridge
Figure 3. Visualization of the binding modes of the reference molecule (carhons in green) and the optimized molecule (in eyan) within the protein pocket (PDB ID: 2PC8, 2AZY, IA2G, 2E24). The molecules and key residues (in blue) are shown in stick, while the protein's main chain drawn cartoon (in gray). Dashed lines various colors indicate different types of non- non-bonding interactions. Left: R-group optimization results. Right: scaffold hopping results.
Table 4. Performances of no correction (Vanilla), SDE and backward correction strategy (B. (B.C.) without and with gradient guidance. Positive numbers in green show the relative improvement, while non positive numbers in black indicate performance gain.
Vina (1)
Min (1)
Dock (1)
QED (t)
SA (?)
Grad Sampling
Avg.
Med.
Avg.
Med.
Avg.
Med.
Avg.
Avg.
Vanilla
-5.23
-5.81
-6.30
-6.17
-7.37
-7.31
0.46
0.62
SDE
6.62
-7.08
-7.31
-7.24
-8.22
-8.32
0.51
-6.50
-7,03
-7.14
-7.95
-7.87
Vanilla
-5.47 (14.6%)
-5889 (x1.4%)
.6.29 (-0.6%)
6.31 (12.3%)
-7.49 (12.2%)
-7.46 (12.1%)
0.46 (10.0%)
0.62 (10.0%)
SDR
-7.11 (-7.4%)
-7.53 (46.3%)
-7.76 (46. 196
-7.73 (16.8%
-8.39 (42,1%)
-8.66 (14.1%)
0.50 (-1.9%)
0.68 (14.69)
B.C.
-7.52 (+15.76)
-8.06 (-15.19%)
-8.34 (+1666))
-8.40 (+17.89)
-9.11 (+14.6%8
.9.25 (+17.5%)
0.56 (+14.3%)
0.77 (+11.6%)
diffusion baselines, they generate fewer valid connected molecules especially in the challenging case with scaffold hopping, with diffusion baselines lower than 70% valid- ity, and proves to be less effective in proposing feasible candidates, with Success Rate < 20%.
different modalities, we ablate the gradient for coordinates or types. As shown in Table 11, utilizing gradients guide both data modalities is consistently better than applying singlc-modality gradient only. For affinities, optimizing co- ordinates is effective in improving the spatial interactions, while for drug-like properties, guidance over atom types plays a crucial role. This underscores the significance of de- riving appropriate guidance form jointly, and again supports our finding that a single coordinate guidance as in TAGMol is insufficient and yields suboptimal results.
Optimized molecules form more key interactions for binding. The visualization for constrained optimization is shown in Fig. 3. It can be seen that the optimized molecules establish more key interactions with the protein pockets, thus binding more tightly to the active sites. For example, the optimized molecule for 2PC8 retains the key interaction formed by its scaffold, with R- R-group grown deeper inside the pocket, forming another ग- stackings.
Backward correction boosts both the unguided sampling and the effect of guidance. We denote sampling 0, ac- cording to Eq. 10 Vanilla for point estimate of ~ ps in Eq. 1, advanced SDE proposed by Xue et al. (2024) with classifier guidance, and for B.C. set k 130 as backward correction steps. Table shows that our method of correcting the past yields better results with guidance. Note that for vanilla case, the gradient guidance does not work as much probably due to the suboptimal history, and SDE-based classifier guidance may have suffered from dis- cretization errors, while correcting a sufficient number of past steps shows consistent boosts.
5.4. Ablation Studies
We conduct ablation studies to thoroughly validate our de- sign. More details are left to Appendix F.2. For all the 100 test proteins, sample 10 molecules each.
Joint guidance is consistently better than single-min modality guidance. To validate our choice of joint guidance over
8

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
6. Conclusion
with recurrent neural networks (rnns). arXiv preprint 2017.
We present MolJO, the joint gradient-based SE(3)- equivariant framework within Bayesian Flow Networks to solve the structure-based molecule optimization problems, which only requires differentiable energy functions instead of expensive oracle simulations. The general framework further equips gradient- bradeett-based optimization method with backward correction strategy, offering flexible trade-off be- tween exploration and exploitation. Experiments show that MolJO is able to improve the binding affinity of molecules by establishing more key interactions and enhance drug- likeness and synthesizability, achieving state-of-the toe-the-art per- formance on CrossDocked2020 (Success Rate 51.3%, Vina Dock -9.05 and SA 0.78), together with 4x improvement compared to gradient-based counterpart and 2 "Me-Better" Ratio as other 3D baselines.
Cheng, X., Zhou, X., Yang, Y., Bao, Y., and Gu, Q. Decom- posed direct preference optimization for structure-based drug design, 2025. URL https: /openrevie...
net/ http//oommmmdbbyyykk/o witmmiiidbbbllrrrrmmmm
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances neural information processing systems, 34:8780-877 8794, 2021.
Dorna, V., Subhalingam, D., Kolluru, K., Tuli, S., Singh, M., Singal, S., Krishnan, N. M. A., and Ranu, S. Tagmol: Target-aware gradient-guided molecule generation. arXiv preprint arXiv:2406.0/64 2024.
Du, Y., Jamasb, A. R., Guo, J., Fu, T., Harris, C., Wang. Y., Duan, C., Liò, P., Schwaller, P., and Blundell, T. L. Ma- chine learning-aiin aided generative molecular design. Nature Machine Intelligence, pp. 1-16, 2024.
Impact Statement
This work is aimed at facilitating structure-based molecule optimziation (SBMO) for drug discovery pipeline. The positive societal impacts include effective design of viable drug candidates. While there is a minimal risk of misuse for generating harmful substances, such risks are mitigated by the need for significant laboratory resources and ethical conduct.
Eberhardt, J., Autodock vina 1.2. O: New docking methods, expanded force field, and python bindings. Journal of chemical information and modeling, 6189)388,,)3888, 2021.
D., Tillack, A. F., and Forli, S.
Epstein, D., Jabri, A., Poole, B., Efros, A., and Holynski, A. Diffusion self geff-guidance for controllable image generation. Advances in Neural Information Processing Systems, 36: 16222-16239, 2023.
Acknowledgments
This work is supported by the Natural Science Founda- tion of China (Grant No. 62376133) and sponsored by Beijing Nova Program (20240484682) and the Wuxi Re- search Institute of Applied Technologies, Tsinghua Univer- sity 220242001120).
Francoeur, P. G., Masuda, T., Sunseri, J., Jia, A., lovanisci, R. B., Snyder, I., and Koes, D. R. Three- stree-dimmesiooal convolutional neural networks and a cross-docked data set for structure-bard based drug design. Journal of chemical information and modeling, 60(9):4200-44) -2155, 2020.
Fu, T., Gao, W., Coley, C., and Sun, J. Reinforced genetic algorithm for structure-based based drug design. Advances in Neural Information Processing Systems, 35:12325- 12338, 2022.
References
Bao, F., Zhao, M., Hao, Z., P., C., and Zhu, J. Equiv- ariant energy-guided sde for inverse molecular design. In The eleventh international conference on learning repre- sentations, 2022.
Hombzz-ommbrreli, R., Wei, J. N., Duvenaud, D., Hernández-Lobato J. M., Sánchez-Lengelingg B., She- berla, D., Aguilera- Iparraguirre, J., Hirzel, T. D.. Adams, Aguilera-Iparrauungg
Bengio, E., Jain, M.. Korablyov, M.. Precup, D., and Ben- gio, Y. Flow network based generative models for non- iterative diverse candidate generation. Advances in Neu- ral Information Processing Systems, 34:27381-27394, 2021.
P., and Aspuru-Guzit A. Automatic chemical de- sign using a data-driven continuous representation of molecules. ACS central science, 4(2):268-276 2018.
Graves, A., Srivastava, R. K., Atkinson, T. and Gomez, F. Bayesian flow networks. arXiv preprint arXiv:2308.00031 2023.
Bilodeau, C., Jin, W., Jaakkola, T., Barzilay, R., and Jensen, K. F. Generative models for molecular discovery: Recent advances and challenges. Wiley Interdisciplinary Reviews: Computational Molecular Science, 2255)e1608. 2022.
Guan, J., Qian, W. W., Peng, X., Su, Y., Peng. J., and Ma, J. 3d equivariant diffusion for target-aware molecule genera- tion and affinity prediction. In The Eleventh International Conference on Learning Representations, 2022.
Bjerrum, E. J. and Threlfall, R. Molecular generation

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Guan, J., Zhou, X., Yang, Y., Bao, Y., Peng, J., Ma, J., Liu, Q.. Wang, L., and Gu, Q. DecompDiff: Diffu- sion models with decomposed priors for structure-based drug design. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 11827-11848 1846. PMLR, 23-29 Jul 2023. Utthttps: 0202/guan22 html.
generation. In Krause, A., Brunskill, E., Cho, K., En gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 17391-17408. PMIR, 23-29 Jul 2023, htts:: httpttttttttptptooottoee http//prooeeedngneeee//e/e mlr. press/ v202 2202/kong23b.. html.
Kong, L.., Du, Y., Mu, W., Neklyudov, K., De Bortol, V., Wang. H., Wu, D., Ferber, A., Ma, Y.-A., Gomes, C. P., al. Diffusion models as constrained samplers for op- timization with unknown constraints. arXiv preprint arXiv:2402.18012 2024.
/ //t:://prrro http///oceedddee ngs. mlr ....press////////
Halgren, T. A., Murphy, R. B., Friesner, R. A., Beard, H. S., Frye, L.. I.., Pollard, W. T., and Banks, J. L.. Glide: new approach for rapid, accurate docking and scoring. 2. enrichment factors in database screening. Journal of medicinal chemistry, 47(7):1550-175) 2004.
Lin, H., Huang, Y., Liu, M., Li, X., Ji, S., and Li, S. Z. Diffbp: Generative diffusion of 3d molecules for target protein binding, 2022.
Han, X., Shan, C., Shen, Y., Xu, C., Yang. H., Li, X., and
Lin, H., Huang, Y., Zhang, O., Liu, Y., Wu, L., Li, S., Chen, Z., and Li, S. Z. Functional- group-----d- diffu- sion for pocket-specific molecule generation and elabo- ration. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 34603- 34626. Curran Associates, Inc., 2023, URL https: / peexiv.orge abs http//aab//33...3 221. 13769.
Li, D. Training-fre free multi-objeccive diffusion model for 3d molecule generation. In The Twelfth International Conference on Learning Representations, 2023.
Harris, C., Didi, K., Jamasb, A. R., Joshi, C. K., Mathis, S. V., Lio, P., and Blundell, T. Benchmarking generated poses: How rational is structure- -baeee drug design with generative models? arXiv preprint 0784:3330..07133, 2023.
Lin, H., Huang, Y., Zhang, O., Wu, L., Li, S., Chen, Z., and Li, S. Z. Functional-grruup-baae based diffusion for pocket- specific molecule generation and claboration, 2024.
Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8):1771- 1800, 2002. doi: 10. 10.11620899776622661888111
Lin, H., Zhao, G., Zhang, O., Huang, Y., Wu, L., Tan, C., Liu, Z., Gao, Z., and Li, S. Z. CBGBench: Fill in the blank of protein- molecule complex binding graph. In The Thirteenth International Conference on Learning Representations, 2025. URL. WtRt:/tps:///pt /pw.nevie...eeeee net/forum?ia wtw.orrmmmmmimmooo////222
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-awand 3d molecule diffusion models. In The Twelfth International Confer- ence on Learning Representations, 2024.
Lipinski, C. A., Lombardo, F., Dominy, B. W., and Feeney, P. J. Experimental and computational approaches to es- timate solubility and permeability in drug discovery and development settings. Advanced drug delivery reviews, 23(1-3):3-22 1997.
Hughes, J. P., Rees, S., Kalindjian, S. B., and Philpott, K. L. Principles of early drug discovery. British journal of pharmacology, (62(6)):239-112)) 2011.
Isert, C., Atz, K., and Schneider, G. Structure-based drug design with geometric deep learning. Cur-
Liu, M., Luo, Y., Uchino, K., Maruhashi, K., and Ji, S. Generating 3d molecules for target protein binding. In International Conference on Machine Learning, 2022.
rent Opinion in Structural Biology, 79:102548, April 2023. ISSN 0959440X. 102548. URL UBLLhtpss//Lnniiingub.but com/ http://erreeeee/l///094444022220
doi: 10. s0.1016/j.bbi.2233..
Luo, S., Guan, J., Ma, J., and Peng. J. A 3D Gen- erative Model for Structure- Based Drug Design. Ad- vances in Neurul Information Processing Systems, 34: 6229-6239, 2021. URL http: http::::/axxvvvv.ro abs/
Jin, W. Barzilay, R., and Jaakkola, T. Junction tree vari- ational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323- 2332. PMLR, 2018
Nigam, A., Friederich, P., Krenn, M., and Aspuru-Guzik, A. Augmenting genetic algorithms with deep neu- ral networks for exploring the chemical space. In In- ternational Conference on Learning Representations,
Kong, L., Cui, J., Sun, H., Zhuang, Y., Prakash, B. A., and Zhang, C. Autoregressive diffusion model for graph
10

---

Empower Structure- Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
2020. URL httph:://ppeeereeeeeee..// 1m=HIIy
Sun, F., Zhan, Z.. Guo, H., Zhang, M., and Tang, Graphvf: Controllable protein-specifie 3d molecule generation with variational flow, 2023. URL pettps://arxiv..
Olivecrona, M.. Blaschke, T., Engkvist, O., and Chen, H.
Molecular de-novo design through deep reinforcement learning. Journal of cheminformatics, 911):1-14, 2017.
Vignac, C., Krawezuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffusion for graph generation. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / / weweenrvieww.eteee
Peng, X., Luo, S., Guan, J., Xie, Q., Peng. J., and Ma, J. Pocket2Mol: Efficient molecular sampling based on 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 of Proceedings of Machine Learning Research, pp. 17644-17655 PMLR, 17-23 Jul 2022. URL Uttpstts: htwpatttpppppppeeeeeeeeeeeee /ptp//peeeeingssseoee mlr. press/ v1 v162/peng222 html.
wit.orrmmiideeee Nu86WX.
Wu, H., Song, Y., Gong, J., Cao, Z., Ouyang, Y., Zhang, J., Zhou, H., Ma, W.-Y., and Liu, J. A periodic bayesian flow for material generation. In The Thirteenth In- lernational Conference on Learning Representations, 2025. URL https: http//ppppeeeeeee//e/eeeeeoeeeep net/ forum? /ttp:/uttt///////ee/et/ettt id-120XW99tBC.
Pinheiro, P. O., Jamasb, A., Mahmood, O., Sresht, V., and Saremi, S. Structure-based drug design by denoising voxel grids. In ICML,, 2024.
Xie, Y., Shi, C.,, Zhou, H., Yang, Y., Zhang, W., Yu, Y., and Li, L. (MARS): Markov molecular sampling for multi- objective drug discovery. In International Conference on Learning Representations, 2021. URL https:// openreview net. / forum? id=.Hdd-eHbHb....
Polykovskiy, D., Zhebrak, A., Sanchez- B., Golo- vanov, S., Tatanov, O., Belyaev, S., Kurbanov, R., Arta- monov, A., Aladinskiy, V., Veselov, M., et al. Molecular sets (moses): benchmarking platform for molecular gen- eration models. Frontiers in pharmacology, 11:565644, 2020.
Xue, K., Zhou, Y. Nie. S., Min, X., Zhang, X., ZHOU, J., and Li, C. Unifying bayesian flow networks and diffusion models through stochastic differential equa- tions. Forty-first International Conference on Machine Learning, 2024. URL https: httpt::t/ppe/eeeeeeeeeee /wwpnrevee.eeeeee/oo//e/eeor net/ forum 1id-11dd111 j 111444011.
Powers, A. S., Yu, H. H., Suriana, P. A., and Dror, R. O. Fragment-bar based ligand generation guided by geometric deep learning on protein-ligand structures. In ICIR2022 Machine Learning for Drug Discovery, net forum?
Zhang, Z., Min, Y., Zheng, S., and Liu, Q. Molecule Gener- ation For Target Protein Binding with Structural Motifs. 2023.
2022. URL bettpa://openreeeee id=1 -
Qu, Y. Qiu, K., Song, Y., Gong, J., Han, J., Zheng, M., Zhou, H., and Ma, W.-Y. MolCRAFT: Structure-based drug design in continuous parameter space. In Forty- first International Conference on Machine Learning, 2024. URL https: http:///pppeeeeeeeeeee hwtp:/wwwwwwe.eeeeeeeee net/forum? forum? id=KaAQu5rNU1.
Zhou, X., Cheng, X., Yang, Y., Bao, Y., Wang, L.,, and Gu, Q. Decompopt:
Controllable and decomposed diffusion models for structure- molecular optimization. arXiv preprint arXiv:2403.13829 2024a.
Zhou, X., Wang, L., and Zhou, Y. Stabilizing policy gradi- ents for stochastic differential equations via consistency with perturbation process. In Forty-first International Conference on Machine Learning, 2024b. URL https: / /ppenrevi net/ / forum? d=ytz2na7oDB.
Schneuing, A., Du, Y., Harris, C., Jamasb, A., Igashov, I., Du, W., Blundell, T., Lió, P., Gomes, C., Welling, M.. Bronstein, M., and Correia, B. Structure-based Drug Design with Equivariant Diffusion Models, Oc- tober 2022. URL http: : arxiv. wawiiioorrg/aaae 13695. arXiv:2210.13695 [cs, b-bi
Song. Y., Gong, J., Qu, Y., Zhou, H., Zheng, M., Liu, J., and Ma, W.Y Unified generative modeling of 3d molecules via bayesian flow networks. arXiv preprint Mkki:22003.154411 2024.
Spiegel, J. O. and Durrant, J. D. Autogrow4: an open- source genetic algorithm for de novo drug design and lead optimization. Journal of cheminformatics, 12:1-16, 2020.

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
A. Overview of Bayesian Flow Networks
In this section, we provide further explanation of Bayesian Flow Networks (BFNs) that are designed to model the generation of data through process of message exchange between a "sender" and a "receiver" (Graves et al., 2023). The fundamental elements include the sender distribution for the sender, and input distribution, output distribution, receiver distribution for the receiver.
This process is framed within the context of Bayesian inference, where the sender distribution is factorized distribution py//|m, ax) that introduces noise to each dimension of the data m and sends it to the receiver. The receiver observes y, has access to the noisy channel with accuracy a at timestep f, and compares it with its own receiver distribution pa(y]6 (,, p; p) based on its current belief of the parameters 0, the timestep and any conditional input such as protein pocket p.
The generative process for the receiver begins with prior distribution (referred to as input distribution pr (m (m) IL= Pr (m) ((())) that is also factorized for N-dimensional data) that defines its initial belief about the data. For continuous data, the prior can be chosen as Gaussian (Song al., 2024; Qu al., 2024), or another distribution such as von Mises distribution (Wu al., 2025), while for discrete data, the prior is modeled as a uniform categorical distribution. Then, the receiver uses its belief 0 with the help of neural network model the inter-dependency among dimensions and compute the output distribution Po Po(ii] 0 p; t), which represents its estimate of the possible reconstructions of the original data m, and is used to construct the receiver distribution (Eq. 2).
The Bayesian update function in Eq. defines how the prior belief 0g is updated to the conjugate posterior 0,. Ideally, the update requires aggregating all possible noisy latents y from the sender distribution ps. However, during the actual generative process, only the receiver distribution PR available, which leaves an exposure bias, and different approximations determine different forms of mapping [ly to the posterior, showcasing the flexibility in the design space of BFN.
Through the iterative communication between sender and receiver, the receiver progressively updates its belief of the underlying parameters, and training is achieved by minimizing the divergence between the sender and receiver distributions (Eq. 3). This is analogous to the way Bayesian inference works in parameter estimation: as more noisy data y is observed, the receiver's 's posterior belief about the data m becomes increasingly accurate, which implies the reconstruction would be made easier. It is shown that BFN for Gaussian priors can be sampled from the view of SDE (Xue et al., 2024), but this type of sampling displays distinct behaviors empirically, possibly due to discretization errors.
B. Sender Distribution for Discrete Data
The continuous parameter 0" for discrete types v is updated by observed noisy y". Here we briefly introduce how to configure the sender so that y" follows a Gaussian as well. For detailed derivation, we refer the readers to Graves et al. (2023).
While true discrete data can be viewed as a sharp one-hot distribution, it can be further relaxed by a factor w E [0, 1] into a Categorical distribution defined by the probability p(k((d) (d)); w) = - + wd for from 1 to K along the d-th dimension, where 8 is the Kronecker delta function.
Instead of focusing on the density sampling from it once, note that the counts of observing each class in m independent draws follow a multinomial distribution, namely e e Multi(m,p). Dropping the superscripts, Graves et al. (2023) derives the following conclusions:
Proposition B.1. When the number of experiments m is large enough, the frequency approximates its density for class indexed k, i.e. limm-roe Co. = p)), following the law of large numbers. Furthermore, by the central limit theorem, m.
it follows that c-mp ~ ~N I) when m - oc.
mmp(1-p)
Proposition B.2. Denoting yk = (c - x Ing & with g = 1 + T and ps (uke: = lim p(u (v:)) with a = mm2 holds from the change of variables that ps(yn/v; = -Na( - aK).
Thus, naturally follows that such noisy y" N(y")a(Kee - 1),
12

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
C. Proofs
C.1. Proof of Guided Bayesian Update Distribution
Lemma. If random vector X has probability density fxx x (x E)e where c is constant vector with the same dimension as X, then X ~ NO + Ec, ))
Proof. We obtain the proof by completing the square as shown below.
log f(x) =0 7(x 15-0% 10-00)=-~- -0-) + c²x
=C' (x 0" -Ec) 'x -0 -0--Ec) -Ec)
(15)
where C' and C" are constant scalars.
Proposition (4.1). Assuming Of ~ NV 7) and ye ~ (xy σ) by definition of the generative process of BFN, we can approximately sample 0 y. from the guided transition kernel ग =0,00 -1) according to Eq. 4 and 5.
Proof. Under the definition of R,, with the parameters 0, | 0,] referred to as [0; (v.") and slight abuse of notation, we have
y x ps (0 (0-))p
(16)
where some parentheses and protein pocket condition p have been omitted for brevity.
Eq. 13 and 14 guarantee that 0f NO y Ny o²). Plugging (=(O,, yy) x e -00 into Eq. 16, we get
π (0, y. 11) & N (0 o N (y Iy σ²) B(o yf.i))
(17)
With 4, fixed, perform first-order Taylor expansion to E y" (44 at (0 11 Y
E ~ 1, go T O 0
By lyF y;
(18)
where gradient gø. Vo. E(0, 0=0. -
t.) (:)) := Substitute into Eq. 17:
x -0,, y & N 00 o²) N (y o" of
(19)
Eq. 19, together with the lemma above, leads Proposition 4.1.
C.2. Proof of Equivariance
Proposition (4.4). The guided sampling process preserves SE(3)--quivariance when is SE(3)-equivariant, if the energy function E p, f) is also parameterized with -E(3)-equivariant neural network, and the complex is shifted to the space where the protein's 's Center of Mass (CoM) is zero.
Proof. Following Schneuing al. (2022), once the complex is moved so that the pocket is centered at the origin (i.e, zero CoM), translation equivariance becomes irrelevant and only O(3) eqquivarraee needs to be satisfied.
For any orthogonal matrix R. E 3 such that R R. I, it is easy to see that the prior 0 = is 0(3)-invariant. Given that & ~ (x (00, p, 1)) and the equivariance of Φ, it suffices to prove the invariant likelihood for the transition kernel.
Given the parameterization of the pretrained energy function E(O, p. t) is SEgi)-equivariant, then the gradient gø- (0)
- ( p. f4) is also equivariant according Bao al. (2022).
Without loss of generality, we consider the guided transition density for i S k. which simplifies to
x ( 0 I.
=N n (0 1 y m1(11 7)gg- (0 P. 4- m1 I)
13

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
def B(t))
where 143(6))
Then we can prove that it is 0())-invrriant:
T =RO; -O Rp)
(RO
Rp) + 7(11 7/ go
- (RO ( 1. v. Rp, t.1)
I)
( (RO Rp) + +7((1 Rgo 1. 71 7
(equivariance of ge=)
RR 1. 7((1 % Rgo
-.- 7/ 7/ I
(equivariance of 4)
-N T
x1 7: g)-
P, ( (111~x))
(equivariance of isotropic Gaussian)
(0 0 1.
also applies to cases where i k. as we can recurrently view the starting point of backward corrected history 07_, as the new O(3)-invariant prior Of and iteratively make the above derivation.
D. Implementation Details
D.I. Model Details
Backbone. Our BFN backbone follows that of MoICRAFT (Qu et al., 2024), and we conduct optimization during sampling on the pretrained checkpoint without finetuning.
Training Property Regressors. To enable differentiable oracle function, additionally train the energy function based on the molecules and their properties (Vina Score, QED, SA) in CrossDocked dataset (Francoeur et al., 2020) by minimizing the squared loss for property e over the data distribution Pdata:
L Epases E p.t) c/²
(20)
where the Bayesian posterior is derived by Graves et al. (2023) as:
0r No0 7: + 7: - 71 - I))
(21)
0. 8(0% softmax(y)))
(22)
for y" N(y" (t)(t) 1), (t, is the Dirac delta distribution.
The input parameters to all energy functions belong to the parameter space defined by B1 = 1.5 for atom types, 0; 0.03 for atom coordinates, n = 1000 discrete steps. The energy network is parameterized with the same model architecture as TargetDiff (Guan et al., 2022), i.e. kNN graphs with & = 32, N = layers with 128 hidden dimension, 16-headed attention, and the same featurization, i.e. protein atoms (H. C, N. O, S, Se) and ligand atoms (C, N, O, F, P, S, Cl). For training, the Adam optimizer is adopted with learning rate 0.0005, batch size is set to The training takes less than 8 hours on single RTX 3090 and converges within epochs.
Sampling, To sample via guided Bayesian flow, we set the sample steps to 200, and the guidance scale to 50, For the combination of different objectives, we simply take an average of different gradients.
D.2. Task Details
R-group optimization. Lead optimization cases involve retaining the scaffold while redesigning the remaining R-groups, usually when the scaffold forms desirable interactions with the protein and anchors the binding mode, and the remaining parts need further modifications to secure this pattern and enhance binding affinity. Following Polykovskiy et al. (2020), we employ RDKit for fragmentation and atom annotation with R-group or Bemis-Murch scaffold.
14

---

scaffolding Empower Saructure-Bascl Molecule Optimization with Gradient Guided Bayesian Flow Networks
QED
SA
Vina Score
Ro-0.32
1.0
0.8
-6
0.6
-8
0.6
-10-
122
-14
0.4
0.3
-18
3
13
18
23
28
33 #Atoms
38
4.3
48
53
58
13
18
23
28
38 *Atoms
33
48
53
58
3
13
18
23
28
33
38
43
48
53 58
Figure 4. Distribution of molecular properties (QED, SA, Vina Score) over the number of atoms for CrossDocked2020, For each size, the mean and error bars are shown in boxplot.
Table Molecular properties under different sizes, where Ref Size denotes 23 atoms average, and Large Size around 30 atoms. Top results are highlighted with bold text for each size category.
Vina Score (L) Avg.
Vina Min (U)
Vina Dock (U)
QED (1)
SA (1)
Div (f)
Success
Methods
Med.
Avg.
Med.
Avg.
Med.
Avg.
Avg.
Avg.
Rate (t)
DecompDiff
-5.19
-5.27
-6.03
-6.00
-7.03
-7.16
0.51
0.66
0.73
14.9%
Ref Size
DecompOpt
-5.75
-5.97
-6.58
-6,70
-7.63
-8.02
0.56
0.73
0.63
39,4%
MolCRAFT
-6.59
-7.04
-7.27
-7.26
-7.92
-8.01
0.50
0.69
0.72
26.0%
MoJO
-7.52
-8.02
-8.33
-8.34
-9.05
-9.13
0.56
0.78
0.66
51.3%
DecompDiff
5566
-6.04
-7.04
-7.09
-8.39
-8.43
0.45
0.61
0.68
24.5%
Large
DecompOpt
-5.87
-6,81
-7.35
-7.72
-8.98
-9.01
0.48
0.65
0.60
52.5%
Size
MolCRAFT
-6.61
-8.14
-8.14
-8.42
-9.25
-9.20
0.46
0.62
0.61
36.6%
MolJO
-7.93
-9.26
-9.47
-9.73
-10.53
-10.48
0.50
0.72
0.57
64.2%
Scaffold hopping. Unlike R-group design, scaffold hopping involves redesigning the scaffold for given molecule while keeping core functional groups, for example, overcome the patent protection for a known drug molecule while retaining pharmaceutical activity. This is technically need to fulfill are generally larger than those in R-group design, and the hopping is usually subject more chemical constraints. We construct scaffold hopping as dual problem to R-group optimization using annotation, although does not need to be so.
challenging task for generative models, because the missing parts they
E. Effect of Molecular Size on Properties
The size of molecules is found have a notable impact on molecular properties, including Vina affinities (Qu et al., 2024). We quantify the relationship and plot the distribution of molecular properties w.r.t. the number of atoms with the Pearson correlation coefficient in Fig. 4. It is not surprising to see non-negligible correlation between properties and molecular sizes, since the sizes of molecules typically constrain their accessible chemical space. To ensure fair comparison, we adhere to the molecular space with similar size to the reference. For further comparison among different model variants, we report the molecular properties under different sizes in Table Results show that our method consistently achieves the highest success rate, demonstrating its robust optimization ability even an Out-of-Distribee (OOD) scenario.
15

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
F. Full Optimization Results
Baselines.
We provide detailed description all baselines here:
AR (Luo et al., 2021) uses MCMC sampling to reconstruct molecule atom-by-atom given voxel-wise densities.
GraphBP (Liu et al., 2022) is autoregressive atom-based model that uses normalizing flow and encodes the context to preserve 3D geometric equivariance.
Pocket2Mol (Peng et al., 2022) generates one atom and its bond a time via an E()-evuuirriian network. predicts frontier atoms expand, alleviating the efficiency problem in sampling.
FLAG (Zhang et al., 2023) a fragment-base based model that assembles the generated fragments using predicted coordinates and torsion angles.
DiftSBDD (Schneuing et al., 2022) constructs an equivariant continuous diffusion for full-atom generation given pocket information, and applies Gaussian noise to both continuous atom coordinates and discrete atom types.
TifgetDiff (Guan et al., 2022) adopts continuous-discrete diffusion approach that treats each modality via corre- sponding diffusion process, achieving better performance than continuous diffusion such as DiftSBDD.
DicompDiff (Guan al., 2023) decomposes the molecules into contact arms and linking scaffolds, and utilizes such chemical priors in the diffusion process.
IPDiff (Huang al., 2024) pretrains an affinity predictor, and utilizes this predictor to extract features that augment the conditioning of the diffusion generative process.
MoCCRAFT (Qu et al., 2024) employs Bayesian Flow Networks for molecular design with an advanced sampling strategy, showing notable improvement upon diffuusion counterparts.
AutoGrow4 (Spiegel & Durrant, 2020) is an evolutionary algorithm that uses genetic algorithm to optimize 1D SMILES with docking simulation. Starting from the initial seed molecule, AutoGrow4 iteratively conducts mutations and crossovers, then makes oracle calls for docking feedback, and retains the top-scoring molecules in the end.
RGA (Fu et al., 2022) is built on top of AutoGrow4, and utilizes pocket-aware RL-trained policy to suppress its random walking behavior in traversing the molecular space.
DecompOpt (Zhou et al., 2024a) trains conditional generative model on decomposed fragments and the binding pocket, following the style of DecompDiff. The optimization is done by iteratively resampling in the 3D diffusion latent space given the top K arms ranked by oracle functions updated fragment condition input.
TAGMol (Dorna al., 2024) exerts gradient-based property guidance on the pretrained TargetDiff' backbone, and the gradient is enabled only in the continuous diffusion process for coordinates.
Metrics. Besides the common evaluation metrics such as binding affinities calculated by Autodock Vina (Eberhardt et al., 2021) and QED, SA calculated by RDKit, elaborate other metries follows:
Diversity measures the diversity of generated molecules for each binding site, Following SBDD convention (Luo et al., 2021), is based on Tanimoto similarity over Morgan fingerprints, and averaged across 100 test proteins.
Connected Ratio is the ratio of complete molecules overall, i.e. with only one connected component.
Lipinski enumerates the Lipinski rule of five (Lipinski et al., 1997) and checks how many satisfied. These rules typically seen as empirical reference that helps to predict whether the molecule is likely to be orally bioavailable.
Key Interaction, i.e. key non-covalent interactions formed between molecules and protein binding sites as an in- depth measure for binding modes, including æ interactions, hydrogen bonds (donor and acceptor), salt bridges and hydrophobic interactions calculated by Schrödinger Glide (Halgren al., 2004).
Strain Energy measures the internal energy of generated poses, indicating pose quality (Harris al., 2023).
16

---

Empower Saructure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Steric Clash calculates the number of clashes between generated ligand and protein surface, where clashing means the distance of ligand and protein atoms is within certain threshold. This reveals the stability of complex some extend, yet it does not strictly mean violation of physical constraints, since the protein is not overly rigid and might also undergo spatial rearrangement upon binding, as noted by Harris et al. (2023).
Redocking RMSD reports the percentage of molecules with an RMSD between generated and Vina redocked poses lying within the range of 2A, which suggests the binding mode remains consistent after redocking.
F.1. Molecule Optimization
Overall Distributions.
We additionally report the property distributions for SA, QED and Vina Score in Fig. 12, 13, 14, respectively, demonstrating the efficacy of our proposed method in optimizing a number of objectives for "me-better" drug candidates. We additionally report in Table the error bars as 95% confidence intervals for our main results (Table 1).
Table 6. Main results with error bars as 95% confidence intervals.
Model
Vina Score (U)
Vina Min (l)
Vina Dock (L)
QED (1)
SA (f)
Reference
-6.362 0.615
-6.707 0.491
-7,450 0.456
0.476 0.040
0.728 # 0.027
-5,754 = 0.066
-6.180 0.049
-6,746 # 0.082
0.509 = 0.004
= 0.003
Pocket2Mol
-5.139 = 0.063
-6.415 + 0.058
-7.152 # 0.097
0.573 + 0.003
0.756 + 0.002
TargetDiff'
-5,466 = 0.172
-6.643 +:2 0.102
-7,802 = 0.075
0,480 +: 0.004
0.585 = 0.003
FLAG
45.978 = 0.778
0.552
-5.237 0.142
0.609 + 0.003
0.626 + 0.002
DecompDiff
-5.190 0.060
6.035
-7.033 0.073
0.505 0.004
0.661 0.003
IPDiff
6..417
-8.572 0.072 0.088
0.519 0.004
0.595
MoICRAFT
-6,587 0.122
-7.265 0.070
-7.924 # 0.097
0.504 0.004
0.003
DecompOpt
-4,899 + 0.415
0.210
-8,425 # 0.528
0.429 + 0.011
0.625 0.006
TAGMol
-7.019 0.175
-7,951 + 0.088
-8.588
0.135
0.553 a 0.004
0.003
MoUOO
-7.516 0.136
-8.326 0.078
-9.048 0.083
0.556 0.003
0.775
QED (+)
(+)
Lipinski (+)
Vina Score (-)
0.0
0.8
0.9
-25
5.0
0.6
0.7
-7.5
0.5
12.5
15.00
0.3
11.55
0.0
Ref
Gen
Ref
Opt
Ref
Gen
Opt
Ref
Gen
Figure Distribution shift from test set (Ref), backhone without guidance (Gen) to guided MolJO (Opt).
Affinity Analysis. We present the tail distribution of Vina affinities Table 7, demonstrating that our method not only excels in optimizing overall performance shown in Fig. 5, but also enhances the quality of the best possible binders.
To better understand the enhanced binding affinites, we further analyze the distribution of non-covalent interactions that are known to play an important role in stabilizing protein-ligand complexes. Fig. 6 demonstrates that the improved affinity results are achieved by forming a greater number hydrophobic interactions, more hydrogen bond acceptors and æ interactions.
Combination of Objectives. In Table 8, we report the results for an exhausted combination of different objectives under the unconstrained setting, where 1000 molecules are sampled in total. can be seen that combining two objectives yields nearly the best optimized performances for each objective, with the choice of Affinity + SA even displaying improvement in QED. However, from the QED + SA setting, we observe a negative impact on binding affinity. possible that
17

---

Empower Structure- Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Denor
Hydrophoblc
Piinteraction
Salt
105
14
62-
Outs
Dur
9us
Infernns
Oar
Oata
Figure 6. Non-covalent interaction distributions of reference and optimized molecules.
Table Tail distribution of Vina affinitics.
Vina Score 5%
Vina Min 5%
Reference
-9.98
-9.93
-10.62
AR
10.05
-00.3
-0056
Packet2Mol
-10.47
-11.77
-12.36
TargetDiff'
-11.10
-11.57
-11.89
DecompDiff
-10.04
-10.96
-11.77
IPDiff
-12.98
-13.40
-13.63
MoICRAFT
-12.14
-12.34
-12.58
DecompOpt
100.78
-11.70
-12.73
TAGMol
-13.15
-13.50
-13.67
MolIO
133.99
-13.90
-14.18
too high a requirement of QED and SA further constrains the chemical space for drug candidates, limiting the types of potential interactions with protein surfaces. When it comes to all objectives, MolJO achieves balanced optimization results, f.e. satisfactory QED and SA comparable to single objective optimization the combination of two, and enhanced affinities compared with the results without affinity optimization, though slightly inferior to the best possible affinity optimization results. This might stem from QED SA problems described above, suggesting careful handling of these two objectives. In this regard, we simply choose Affinity + SA objectives in all our main experiments for a clear demonstration of our optimization ability.
For better understanding of the correlation between objectives, we plot the pairwise relationships for the molecules in the training set in Fig. 7, and calculate the Spearman's 's rank coefficient of correlation p. The Spearman is 0.41 between SA and QED, and it is reasonable to such positive correlation between SA and QED, since these are both indicators of drug-liken likensik with certain focus and thus alternative to some extent. This aligns with our findings that adopting the Affinity + SA objectives can also benefit QED, and justifies our choice of optimization objectives in this sense. Moreover, although there is also a slightly positive correlation between Vina Score and SA (p = 0.33), meaning that it is nontrivial to simultaneously optimize both properties, our method succeeds finding the best balanced combination of properties, demonstrating the superiority joint optimziation compared with TAGMol.
Table 8. Combinations of different objectives. Top-2 results highlighted in bold and underlined, respectively.
Vina Score ())
Vina Min (L)
Objective
QED (t)
SA (t)
Counected (?)
Avg.
Med.
Avg.
Med.
Affinity
-7.74
-7.96
-8.21
-8.19
0.52
0.68
0.87
QED
-6.84
-7.32
-7.54
-7.65
0.66
0.70
0.99
SA
-6.25
-7.24
-7.48
-7.65
0.57
0.78
0.97
QED+SA
-6.55
-7.23
-7.52
0.65
0.74
0.99
Affinity+QED
-7.46
-8.04
-8.18
-8.20
0.64
0.67
0.98
ASA
-7.08
-7.88
-8.05
-8.21
0.57
0.75
0.97
All
-7.09
-7.47
-7.79
-7.76
0.62
0.73
0.98
18

---

Empower Setucturr-Basce Molecule Optimization with Gradient Guided Bayesian Flow Networks
0.8
0.6
QED
0.4
0.2
0.0
1.0
0.8
SA
0.6
0.4
Vina Score
-10
-15
0.0
0.5
0.25
0.50
0.75
1.00
115
-10
-5
QED
SA
Vina Score
Figure 7. Pairwise correlation of different properties. On the diagonal are histograms showing single property distributions on Cross- Docked2020.
19

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Top-of-N Comparison. We have made the same top-of-N selection (N 10) for picking the top 1/10 from the generated molecules for each baseline. Specifically, for each of the 100 test proteins, a tenth out of roughly 100 molecules selected based on z-Score reranking, where 5- norm (Vina) + nOD)(QED) 1.5 norm(SA), and norm denotes zero mean and unit variance normalization.
This generally shows what the "concentrated space" for desirable drug-like candidates looks like for generative models. As shown Table 9, our method displays the best Success Rate top-of-N evaluations, indicating better optimization efficiency. Though IPDiff also displays superior Vina affinities, comes the expense of low SA (0.62) and shows only moderate Success Rate, as there observed to be a slightly positive correlation between Vina Score and SA (p 0.33) in Fig. meaning that it is nontrivial to simultaneously optimize both properties. MolJO does better job in finding the best balanced combination of properties, demonstrating the superiority of gradient-based joint optimziation.
Table 9. Top-of-N (N 10) performances for baselines.
Score ())
Vina Min (.)
Vina Dock (4)
Method
QED
SA (f)
Div
Success Rate (1)
Avg.
Med.
Avg.
Med.
Avg.
Med.
AR
-6.71
-6.35
-7.12
-6.63
-7.81
-7.33
0.64
0.70
0.60
19.1%
Pocket2Mol
-5.00
-5.39
-7.18
-6.50
-8.32
-7.78
0.67
0.84
0.59
40.5%
FLAG
50.37
43.14
6.27
-3.38
-6.57
-6.47
0.74
0.78
0.71
9.6%
TargetDiff
-7.06
-7.57
-8.10
-9.31
-9.18
0.65
0.67
32.6%
DecompDiff
-5.78
-5.82
-6.73
-6.57
-8.07
-8.03
0.74
0.61
32.1%
MoICRAFT
-7.54
-7.89
-8.40
-8.13
-9.36
-9.05
0.65
0.77
0.63
55.0%
IPDiff
-8.15
-8.67
-9.36
-9.27
-10.65
-10.17
0.60
0.62
0.69
34.6%
MolIO
-8.54
-8.81
-9.48
-9.09
-10.50
-10.14
0.67
0.79
0.61
70.3%
Additional Baselines. For more comprehensive comparison, have added the results of DiffBP (Lin et al., 2022), D3FG (Lin et al., 2024) and VoxBind (Pinheiro et al., 2024) from recently proposed benchmark CBGBench (Lin et al., 2025) and concurrent work DecompDPO (Cheng et al., 2025). It can be seen that our MolJO maintains superiority in optimizing the overall properties, reflected by its highest Success Rate.
Table 10. Comparison with additional baselines, where the results for DifBBP, D3FG and VoxBind are calculated based on the samples released by CBGBench, and results for DecompDPO follow the numbers reported by the authors.
Vina Score (4) Avg.
Vina Min (.) Avg.
Vina Dock (L) Avg.
Method
QED (t)
SA (f)
Div (t)
Success Rate (1)
Med.
Med.
Med.
DiffBP
-7.34
0.47
0.59
D3FG
-2599
-6.78
0.49
0.66
VoxBind
-6.16
-6.21
-6.22
-6.73
-7.68
-7.59
0.54
0.65
21.4%
DecompDPO
-6.10
-7.22
-7,93
-8.16
-9.26
-9.23
0.48
0.64
0.62
36.2%
MolJO
-7.52
-8.02
-8.33
-8.34
-9.05
-9.13
0.56
0.78
0.66
51.3%
F.2. Ablation Studies
Efffec of Joint Guidance. Table 11 shows the effectiveness of joint guidance over coordinates types. Utilizing gradients guide both data modalities is consistently better than applying single gradient only, since the energy landscape of molecular system is function of both the atom coordinates and the types. Lack of direct control over either modality can lead to suboptimal performance due to not efficiently exploring the chemical space where certain atomic types naturally pair with specific spatial arrangements. Specifically, it can be seen that for affinities, the optimization is closely related to coordinates, while for drug-like properties, simply propagating gradients over coordinates displays no improvement at all. This validates our choice of finding appropriate guidance form jointly, and a single coordinate guidance would be insufficient for generating desirable molecules.
Effect of Backward Correction. We conduct ablation studies regarding the proposed backward correction strategy. w/o correction denotes sampling 0, according to Eq. 10. Fig. 8 shows that increasing the steps & in Eq. 12 that have been corrected backward boosts the optimization performance once sufficient past steps are corrected for optimization.
20

---

Empower Garucture-Bascl Molecule Optimization with Gradient Guided Bayesian Flow Networks
Table 11. Ablation studies of joint optimization for atom types and coordinates, where w/o type means the gradient is disabled for types.
Top 2 results are highlighted with bold and underlined text.
Vina Score (1) Avg.
Vina Min (L)
Objective
Methods
QED (1)
SA (t)
Med.
Avg.
Med.
Ours
-7.74
-7.96
-8.21
-8.19
0.52
0.68
Affinity
w/o type
-7.13
-7.58
-7.82
-7.80
0.50
0.66
w/o coord
-6.61
-7.23
-7.42
-7.53
0.52
0.71
Ours
-6.84
-7.32
-7.54
-7.65
0.66
0.70
QED
w/o type
-6.41
-7.03
-7.20
-7.26
0.52
0.67
w/o coord
-6.50
-7.20
-7.44
-7.42
0.65
0.70
Ours
-6.25
-7.24
-7.48
-7.65
0.57
0.78
SA
w/o type
-6.29
-6.85
-7.07
-7.09
0.51
0.70
w/o coord
-6.71
-7.22
-7.60
-7.60
0.57
0.77
It can be inferred that sampling Pa (0, 1, up until O, results in a chain of parameters (0r, | mod k), and 0: ~ (o, 1. when i I (n mod k).
[n/k]] where T; = iik-nn
Smaller number & of corrected steps moves the starting point or, closer to 0g and sees more updates along the chain. We observe that when & small, the sampling process tends to suffer from error accumulation instead of error correction due stochasticity. Once l larger than 50, the process is better balanced in exploiting the shorteut (i.e. interval k)) and exploring the stochasticity to reduce approximation errors via few updates (i.e. [n/k]). The final k is set to 130, while strategy is robust within the range h E (50, 200].
QED (+)
(A(1)
Completeness (1)
Vina Score (4)
0.575
0.58
-6.6
0.550
0.75
0.97
-6.0
0.70
0.96
0.95
0.85
0.84
-7.6
0.60
0.425
0.53
-7.8
0.400
0.55
100
100
200
100
200
200
Cotrection Steps
Steps
Cornection Steps
Correction Steps
Correction
Figure 8. Ablation study of backward correction. Correction Step on the x-axis means the length of history k, and w/o Correction means vanilla update (k = 1) with Monte-Cart Carlo estimate of y.
Additionally, we have conducted pairwise t-tests comparing our guided Backward Correction (A: 130) approach against both Vanilla (k: 1) and SDE guidance. The results in Table 12 show statistically significant improvements < 0.05, N 1000) for our proposed strategy.
Table 12, P-values for pairwise t-tests.
p-value
Vina Score
Vina Min
Vina Dock
QED
Ours Vanilla
2.63E-13 2.55E-19
3.31B-31 6.48E-19
2.79E-35 7.84E-4
8. 80100-115
1,98E-26
Ours vs. SDE
2.10E-50
1.82E-12
Effect of Scales. We conduct a grid search of guidance scales, and report the full results of ablation studies on different guiding scales within the range 10.1, 1, 10, 20, 50, 100} for different objectives (Affinity, QED, SA) in Table 13, where 10
21

---

Empower Structure-Bascl Molecule Optimization with Gradient Guided Bayesian Flow Networks
molecules sampled for each of the 100 test proteins.
For binding affinity, the optimization performance steadily improves with increasing scales, but the ratio of complete molecules significantly decreases when the scale is greater than 50.
For QED and SA, MolJO achieves best results when the scale is around 20 and 50.
In order to maintain the comparability with molecules without guidance, we stick scale range where the connected ratio remains acceptable, and therefore set the guidance scale 50 for all our experiments.
Table 13, Full ablation studies on different guiding scales for different objectives. Top-1 values are highlighted in bold.
Vina Score (L)
Vina Min (L)
Objective
Scale
QED (t)
SA (t)
Connected (t)
Avg.
Med.
Avg.
Med.
0.1
-6.28
-6.98
-7.17
-7.25
0.50
0.70
0.96
-6,24
-7.01
-7.27
-7.29
0.50
0.69
0.96
10
-6.69
-7.46
-7.46
-7.67
0.51
0.70
0.97
Affinity
-7.03
-7.87
-7.84
-8.08
0.70
0.98
50
-7.64
-8.38
-8.39
-8.64
0.53
0.68
0.90
-9.33
-9.55
-9.87
-9.85
0.55
0.63
0.55
0.1
-6.03
-6.22
-7.10
-7.19
0.51
0.70
0.97
-6.24
-7.09
-7.11
-7.31
0.56
0.71
0.96
-6.12
-7.07
-7.29
-7.41
0.66
0.71
0.98
QED
20
-6.33
-7.23
-7.34
-7.64
0.66
0.69
0.98
50
-6.84
-7.32
-7.54
-7.65
0.66
0.70
0.99
LO0
-6.25
-6.83
-7.02
-7.10
0.62
0.60
0.95
0.1
-6.30
-7.11
-7.19
-7.29
0.50
0.70
0.96
-6.17
-7.16
-7.66
-7.37
0.52
0.73
0.97
-5.87
-7.23
-7.39
-7.72
0.57
0.78
0.98
SA
20
-6.14
-7.24
-7.49
-7.72
0.56
0.79
0.98
-6.38
-7.29
-7.86
-7.77
0.54
0.79
0.99
100
-6.08
-7.36
-7.51
-7.85
0.54
0.78
0.98
G. Evaluation of Molecular Conformation
PoseCheck Analysis. To measure the quality of generated ligand poses, further employ PoseCheck (Harris et al., 2023) to calculate the Strain Energy (Energy) of molecular conformations and Steric Clashes (Clash) w.r.t. the protein atoms in Fig. 9 and 10, respectively.
Our proposed MoJJO not only significantly outperforms the other optimization baselines in both Energy and Clash, but also shows competitive results with strong-perine poongppeformmig generative models, in which Pocker2Mol achieves lower strain energy via generating structures with fewer rotatable bonds as noted by Harris et (2023), and fragment-based model FLAG directly incorporates rigid fragments in its generation. As for clashes, we achieve the best results in non-autoregressive methods.
Notably, IPDiff ranks the least in Strain Energy and displays severely strained structures despite its strong performance in binding affinities. This arguably suggests that directing utlizing pretrained binding affinity predictor as feature extractor might result in spurious correlated features, even harming the molecule generation.
RMSD Distribution. We report the ratio of redocking RMSD below 2Å between generated poses and Vina docked poses reveal the agreement of binding mode. Due to issues of poses generate by Autodock, not pose pairs are available for calculating cenmmetty-correcter RMSD, where are faithfully evaluated. As shown in Fig. 11, the optimization methods all display tendency towards generating few outliers, which might be attributed to the somewhat out-of disttii-iitiiittonnnn (OOD) nature of optimization that seeks to shift the original distribution. Among all, DecompOpt generates the most severe outliers with RMSD as high as 160.7 Å, and its unsatifactory performance is also suggested by the lowest ratio RMSD 2Å (24.3%), while for gradient- based TAGMol and our method, it only has negligible impact and the ratio is generally more favorable.
report the non-corrected RMSD instead to make sure that all samples
22

---

Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
Overall Conformation Quality and Validity. The overall results in Table 14 show that gradient-based method actually improves upon the conformation stability of backbone in terms of energy and clash, demonstrating its ability faithfully model the chemical environment of protein-ligand complexes, while DecompOpt generates heavily strained structures similar to DecompDiff, and TAGMol ends up with even worse energy than its backbone TargetDiff. Moreover, from the perspective of validity reflected by Connected Ratio, the optimization efficiency of RGA and DecompOpt is relatively low as suggested by the ratio of successfully optimized molecules.
Ring Size. For comprehensive understanding the effect of property guidance, we additionally report the distribution of ring sizes in Table 15, showing that the gradient-based based property guidance generally favors more rings, but our result still lies within reasonable range, and even improves upon the ratio of 4-membered rings.
Table 14, Summary of conformation stability results. Energy, Clash are calculated by PoseCheck. Connected the ratio successfully generated valid and connected molecules.
Energy Med. (L)
Clash Avg. (1)
RMSD < 2Â (t)
Connected (t)
Reference
114
5,46
34.0%
100%
AR
608
4.18
36.5%
93.5%
Pocket2Mol
186
6.22
31.3%
96.3%
FLAG
396
40.83
8.2%
97.1%
TargetDiff
1208
10.67
31.0%
90.4%
DifompDiff
983
14.23
25.1%
72.0%
IPDiff
5861
10.31
17.9%
90.1 %
MoICRAFT
196
6.91
42.4%
96.7%
RGA
52.2%
DecompOpt
861
16.6
24.3%
2.64%
TAGMol
2058
7.41
37.2%
92.0%
MolJO
163
6.72
43.5%
97.3%
1.0
Reference
0.8
Pocket2Mol
0.6
TargetDiff
DecompDir
IPDiff
0.4
MoiCRAFT
DecompOpt
TAGMol
0.0
10a
10a
10a
10a
Strain Energy (kcal/mol)
Figure Cummulative density function (CDF) strain energy distributions of generated molecules and reference molecules.
H. Inference Time
We report the time cost in Table 16 for optimization baselines in the table below, which is calculated the time for sampling batch of 5 molecules on a single NVIDIA RTX 3090 GPU, averaged over 10 randomly selected test proteins.
'Connected ratio of DecompOpt and RGA is calculated based the optimization results of all rounds provided by the authors. For each N rounds, a total of molecules ought to he generated, thus we divide the total number by N k n.
23

---

Empower Structure-Basce -aated Molecule Optimization with Gradient Guided Bayesian Flow Networks
20.0
20.0
17.5
17.5
15.0
15.0
12.5
12.5
Clash
10.0
10.0
7.5
7.5
5.0
5.0
2.5
0.0
0.0
Reference
A8
Pocket2Mol
FLAG
TargetDiff Decompoiff
IPD
MolGRAFT Decompogt
TAGMal
Ours
Figure 10. Box plot for clash distributions of generated molecules and reference molecules.
14
00
odd
12
12
10
10
RMSD (A)
8
8-
TargetDiff
DecompOpt
Figure Boxplot for RMSD distributions generated molecules and reference molecules.
1.0
0.8
0.8
0.7
0.7
3 0.6
0.6
0.3
0.3
0.2
0.2
PocketzMel
TargetDiff DecompDiff
(Pbiff
MolcHAFT DecompOpt
Figure 12. Violin plot for SA distributions of generated molecules and reference molecules.
24

---

Empower Srtucture-Bascee Molecule Optimization with Gradient Guided Bayesian Flow Networks
0.8
0.6
OED
0.4
0.2
0.0
PacketZMol
TargetDiff DacompDitt
iPDIt
MolCRAFT DecompOpt
Figure 13. Violin plot QED distributions of generated molecules reference molecules.
0.0
0.0
-2.5
2.5
-.00
5.00
Vina Score
-7.5
7.5
-10.0
10.0
122.5
12.5
-15.0
15.0
-17.5
17.5
-20.0
-20.0
Reference
AR
Packet2Mal
FLÁG
TargelDiff DecompDif
IPDiff
MolCRAFT DecompOpt TAGMal
Ours
Figure 14. Violin plot for Vina Score distributions of generated molecules and reference molecules.
0.0
-2.5
-2.5
-5.0
5.0
Vina Min
-7.5
7.5
-10.0
10.0
-12.5
12.5
-15.0
-17.5
17.5
-00.0
20.0
TargetDiff Decompoif
(Dccomppp
Figure Violin plot for Vina Min distributions generated molecules and reference molecules.
25

---

Empower Structure-Basse Beree Molecule Optimization with Gradient Guided Bayesian Flow Networks
0.0
0.0
-2.5
-2.
5.00
5.0
Vina Dock
-7.5
7.5
-10.0
10.0
-12.5
12.5
-15.0
15.0
-17.5
17.5
-20.
-20.0
Reference
PocketZMol
FLAG
ToraetDiff DecompDiff
IPDIff
MolCRAFT DecompCpt. TAGMol
Figure 16. Violin plot for Vina Dock distributions of generated molecules and reference molecules.
Table 15. Proportion (%) of different ring sizes in reference and generated ring structured molecules, where 3-Ring denotes three- membered rings and the like.
#Rings Avg.
3-Ring
4-Ring
5-Ring
6-Ring
Reference
2.8
4.0
49.0
84.0
Train
3.0
3.8
0.6
56.1
90.9
AR
3.2
50.8
0.8
35.8
71.9
Pocket2Mol
3.0
0.3
0.1
38.0
88.6
FLAG
39.9
84.7
TargetDiff
3.1
0.0
7.3
57.0
76.1
DecompDill'
3.4
9.0
11.4
64.0
83.3
IPDiff
3.4
0.0
51.0
83.7
MolCRAFT
3.0
0.0
0.6
47.0
85.1
DecompOpt
3.7
6.8
11.8
61.4
89.8
TAGMol
4.0
0.0
8.5
62.5
82.6
MolJO (Af))
3.6
92.5
MoJJO (QED)
3.7
0.0
0.5
58.0
96.1
MoJJO (SA)
3.8
0.0
0.2
37.0
97.8
MoJJO (Aff+SA)
3.9
0.0
0.1
37.0
98.1
MoJJO (AII)
3.6
0.0
0.3
44.4
97.6
Table Inference time cost optimization baselines, error indicating standard deviation across randomly selected proteins.
Model
Ours
TAGMol
DecompOpt
RGA
AutoGrow4
Time (s) 146 + 11
667 +
69
11714 1115
458 + 43
2586 +
360
26

---

Empower Saructure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks
I. More Related Works
General Molecule Optimization. As an alternative to target-aware generative modeling of 3D molecules, the optimization methods are goal-directed, obtain desired ligands usually by searching in the drug-like chemical space guided by property signals (Bilodeau et al., 2022: Sun al., 2023; Du al., 2024). General optimization algorithms were originally designed for ligand-based drug design (LBDD) and optimize common molecule-specifee properties such LogP and QED (Olivecrona et al., 2017; Jin et al., 2018; Nigam et al., 2020; Spiegel & Durrant, 2020; Xie et al., 2021; Bengio et al., 2021), but could be extended to structure- based drug design (SBDD) given docking oracles. However, since most early attempts did not take protein structures into consideration thus were essentially not target-aware, it means that they need to be separately trained on the fly for each protein target when applied to poeket-specific scenarios. RGA (Fu et al., 2022) explicitly models the protein pocket the design process, overcoming the transferability problem of previous methods. DiffAC (Zhou et al., 2024b) utilizes policy gradients for SDEs to fine-tune pretrained diffusion models given affinity signal, demonstrating the potential of RL method but limited to Vina Score optimization only.
Constraint Molecule Optimization. Real-world lead optimization typically requires retaining specific substructures to preserve critical molecular interactions properties. Recently, CBGBench (Lin et al., 2025) introduces generative graph completion framework for systematic evaluation, including tasks such as linker design, fragment growing, side chain decoration, and scaffold hopping, showing the potential for rigorous benchmarking of structure based molecular optimization under predefined substructural constraints.
27


Reviews:
[{"review_id": "4LDkfyCya4", "reviewer": "Reviewer_LNs4", "summary": "The paper proposes MoIJO, a gradient-guided framework for SBMO. The key contributions are:\n\nJoint gradient guidance over both continuous (coordinates) and discrete (atom types) modalities via Bayesian Flow Networks, avoiding modality inconsistency.\n\nBackward correction strategy with a sliding window to balance exploration-exploitation.\n\nSE(3)-equivariance preservation through energy function design.\n\nExperiments on CrossDocked2020 show SOTA results. The method also generalizes to multi-objective optimization and drug design tasks.", "questions": "", "limitations": "", "rating": 4, "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 sincerely appreciate the reviewer's thorough reading and insightful questions, which have helped us identify areas for improved clarity and presentation. We shall address each point in our responses below as well as our revised manuscript, and we welcome further questions.\n\n**Q1: Smoothness of Energy Landscapes**\n\nThanks for raising the insightful question about the limitations of first-order Taylor expansion in our approach. We have followed established guided diffusion methodology [1] which assumes that with sufficiently small step sizes, the changes in the energy landscape between consecutive steps remain modest. We acknowledge that this approach inherently assumes the energy landscape is relatively smooth in the local region of expansion, which may not hold universally for complex energy functions. \n\nAs the reviewer correctly pointed out, higher-order terms could potentially provide more accurate gradient estimates, which is an active area of research in guided diffusion models [2]. However, higher-order methods introduce computational overhead and potential numerical instability during sampling, which we will leave for future work. While the empirical results demonstrate the practical utility of our approach despite this approximation, we plan to further develop adaptive guidance schemes that adjust based on estimated local curvature. \n\n[1] Diffusion Models Beat GANs on Image Synthesis.\n\n[2] Inner Classifier-Free Guidance and Its Taylor Expansion for Diffusion Models.\n\n**Q2: BFN Derivation Clarity**\n\nWe will enhance our presentation in the revision with a more detailed, step-by-step explanation of the Bayesian Flow Network derivation. \n\n**Q3: Schematic Diagram of Backward Correction**\n\nThank you for this suggestion. In Figure 1D, we have a schematic diagram illustrating the backward correction process, and we will continue to make it clearer in our revision.\n\n**Q4: Statistical Significance**\n\nWe acknowledge the importance of statistical validation. We have conducted pairwise t-tests comparing our guided Backward Correction against both Vanilla and guided SDE. The results show statistically significant improvements (p<0.05), and we will add it to our revision. \n\n||vs. Vanilla|vs. SDE|\n|---|---|---|\n|Vina Score|2.63E-13|2.55E-19|\n|Vina Min|3.31E-31|6.48E-19|\n|Vina Dock|2.79E-35|7.84E-4|\n|SA|8.10E-115|2.10E-50|\n|QED|1.98E-26|1.82E-12|\n\n**Q5: Molecule Size Bias**\n\nThe reviewer makes an excellent point regarding molecule size bias, which is a crucial issue for SBDD. We apologize for the confusion caused by Table 4, and we would like to clarify that rather than solving the size bias issue directly, our primary goal in size-controlled experiments was to demonstrate that our guidance approach effectively improves molecules across the size ranges, and consistently outperforms baselines across different molecular sizes.\n\nFor Table 4, we will improve our presentation by calculating optimal scores separately for each size (Ref & Large), which will better highlight that our improvements stem from enhanced molecular quality rather than size bias.\n\n**Q6: SE(3)-equivariance with Real-world Pockets**\n\nThe reviewer correctly identifies a limitation for unknown protein pockets, where reference positions are required to clip and obtain the pocket region. We acknowledge this as a limitation of current SBDD methods and certainly worth exploration. For known protein pockets, we simply subtract the centroid to ensure the center of mass is at the origin. \n\n**Q7: RMSD Ratio Without Symmetry Correction**\n\nThank you for this important observation. Here we followed the calculation from MolCRAFT for fair comparison with existing methods. This approach calculates symmetry-corrected RMSD between Vina Docked PDBQT and generated molecules when possible, but falls back to non-symmetry-corrected values in cases where symmetry correction cannot be applied. We appreciate the reviewer highlighting this important consideration for accurate assessment of molecular pose consistency, and we will clarify the detail in our revision and investigate the issue further in the future.\n\n**Q8: Strain Energy Higher Than Ref**\n\nThank you for this important observation. Although our results are generally within reasonable ranges for computational generation, we agree there is room for improvement. Future work could incorporate additional chemical prior knowledge to further reduce strain energy. We appreciate the reviewer highlighting this point, as it identifies an important direction for continued refinement of our approach.\n\n**Q9: Missing Citation to GraphVF**\n\nWe sincerely thank the reviewer for bringing GraphVF to our attention. This work indeed makes valuable contributions by encoding different modalities in latent space for joint coordinate-type optimization. We agree that the unified space represents a promising direction. We will update our manuscript to include this important reference and add relevant discussion on how it relates to our work.", "claims_and_evidence": "1 The experiments lack statistical significance tests (e.g., p-values), making the improvement questionable.", "methods_and_evaluation_criteria": "1 The molecule size bias is not fully addressed. Larger molecules inherently have better Vina scores (Fig. 5), but MoIJO’s superiority on size-controlled subsets (Table 4) is only briefly discussed.", "theoretical_claims": "1 The Taylor expansion (Eq. 18) assumes E(θ,t) is locally linear, which may not hold for complex energy functions. The approximation error is unquantified.\n\n2 The SE(3)-equivariance proof assumes protein CoM is zero, but how do you keep this for real-world pockets?", "experimental_designs_or_analyses": "1 The reported RMSD ratio is based on non-symmetry-corrected values (Page 21), which may overestimate pose consistency.\n\n2 While MoIJO outperforms baselines in strain energy (Table 9), the absolute energy values (163 kcal/mol) are still higher than reference molecules (114 kcal/mol).", "relation_to_broader_scientific_literature": "The work builds on BFNs and gradient-guided diffusion.", "essential_references_not_discussed": "GraphVF[1] Combines SE(3)-flows and GNNs for joint coordinate-type optimization.\n[1] Sun, F., Zhan, Z., Guo, H., Zhang, M., & Tang, J. (2023). GraphVF: Controllable Protein-Specific 3D Molecule Generation with Variational Flow. arXiv [q-Bio.BM]. Retrieved from http://arxiv.org/abs/2304.12825", "other_strengths_and_weaknesses": "The BFN derivation (Appendix A) is overly condensed; a step-by-step example would improve readability.", "other_comments_or_suggestions": "Suggestion: Add a schematic diagram of the backward correction process", "questions_for_authors": "How does the Taylor expansion in Proposition 4.1 handle highly non-convex energy landscapes? Would higher-order terms significantly affect guidance?", "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": "5VH18gqmIZ", "reviewer": "Reviewer_MxvP", "summary": "This paper proposes MolJO, a framework that jointly guides continuous 3d coordinates and discrete atom types of 3d molecules based on the geometry of the target protein pocket and one or more molecular property classifiers. The paper also proposes a backward correction strategy that corrects parameters of Bayesian update distribution based on the current optimized sample and parameters from previous timepoints. It is shown that MolJO outperforms existing methods for generating molecules with better docking scores and molecular properties such as QED, SA, notably methods that do not involve gradient-based guidance or do not do guidance over discrete atom types.\n\n### update after rebuttal\n\nThank you to the authors for providing explanations and running additional experiments for the Top-of-N comparison and error bars. It would also be good to include a Top-of-N comparison for TagMol in future versions.\n\nFor Table 3, does w/o guidance mean no guidance over both atom types and coordinates? If so, authors should include an ablation row where guidance is done only over coordinates and not atom types (equivalent to TagMol if I understand correctly), so that we can observe the impact of only turning on guidance over atom types (current row 4 of table 3), since this is a key argument being made by authors (and also TagMol + BC). Also, it doesn't seem like there is a huge difference between the docking scores of rows 1 and 4, which measures the impact of guidance--a p-value would be helpful here. Do authors have any intuition for why BC and guidance work synergistically?\n\nI think the method is interesting but the clarity of the work can be significantly improved, I will keep my score as-is.", "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 sincerely thank the reviewer for the careful reading and insightful feedback that helps improve the clarity and completeness of our work. \n## Questions\n**Q1: Table Value Inconsistencies**\n\nWe apologize for the confusion. The values in Table 1 and 3 were obtained from different runs. For Table 3 (ablation studies), we sampled 10 molecules per pocket instead of 100 per pocket. We will clarify this in our revision.\n\n**Q2: Fig (Table) 3 Clarification**\n\nThanks for bringing up this issue that helps improve our clarity. All models use the same size specifications from DecompDiff. The differences result from some methods having failed molecule generations (e.g., invalid, not connected). We apologize for the confusion and will correct this labeling (from \"Figure\" to \"Table\") in our revision.\n\nFor guided SDE, we followed Xue et al. and perform guidance by modifying the conditional score following common practices.\n\n**Q3: \"Me-better\" Definition**\n\nThank you for noting this oversight. \"Me-better\" refers to molecules with improved specific properties over existing compounds—a term borrowed from traditional drug discovery [1]. We will add the definition in the revision.\n\n[1] Me-Better Drug Design Based on Nevirapine and Mechanism of Molecular Interactions with Y188C Mutant HIV-1 Reverse Transcriptase. https://pubmed.ncbi.nlm.nih.gov/36364174/\n\n**Q4: Effect of Joint Guidance and Backward Correction**\n\nThank you for this insightful question. The ablation study is actually presented in Table 3, which directly addresses this question by isolating the contributions of each component. We apologize for the confusion caused by our initial abbreviation \"w/ (guidance)\" (MolJO-based) and \"w/o (guidance)\" (MolCRAFT-based) that can be misinterpreted as w/ or w/o (BC). We will improve the clarity in our revision.\n\nAs the reviewer suggested, we evaluated \"w/ guidance, Vanilla\" (row 4) that corresponds to MolJO at k=1 (guidance without BC), and \"w/o guidance, BC\" (row 3) that corresponds to MolCRAFT with BC. Table 3 shows that both components contribute to performance gains. Notably, the combined improvement exceeds the sum of individual improvements, suggesting these components work synergistically.\n\n**Q5: Top-of-N Comparison**\n\nThanks for the great suggestion. We will add top-of-N comparisons for all methods in our revision, which generally shows what the \"concentrated space\" for desirable drug-like candidates looks like for generative models. Our method shows the best Success Rate (70.3%), indicating better optimization efficiency. \n\n|Method|Success Rate|Vina Score Avg|Vina Min Avg|Vina Dock Avg|QED|SA|Div|\n|---|---|---|---|---|---|---|---|\n|AR|19.1%|-6.71|-7.12|-7.81|0.64|0.7|0.6|\n|Pocket2Mol|40.5%|-5.8|-7.18|-8.32|0.67|0.84|0.59|\n|FLAG|9.6%|50.37|6.27|-6.57|0.74|0.78|0.71|\n|TargetDiff|32.6%|-7.06|-8.1|-9.31|0.64|0.65|0.67|\n|DecompDiff|32.1%|-5.78|-6.73|-8.07|0.61|0.74|0.61|\n|MolCRAFT|55.0%|-7.54|-8.4|-9.36|0.65|0.77|0.63|\n|IPDiff|34.6%|-8.15|-9.36|-10.65|0.6|0.62|0.69|\n\n## Additional Clarifications\n**Q6: Tables 1-2 Component**\n\nThank you for highlighting this potential source of confusion. Yes, Tables 1-2 results include backward correction. To clarify, our contribution over TAGMol is two-fold: (1) We derived joint guidance over atom types and coordinates within the BFN framework. (2) We proposed the BC sampling algorithm that further improves optimization performance. As described in Q4, the results in Table 3 demonstrate that both contributions are significant. In our revised manuscript, we will make these distinctions clearer to avoid confusion.\n\n**Q7: Backward Correction Explanation & Novelty**\n\nWe acknowledge that \"correcting the past\" is potentially confusing terminology. The correction applies to timesteps [n-k, n), while preserving parameters from [0, n-k).\n\nAs for its novelty, we develop a more flexible sampling strategy, establishing a sliding window approach that generalizes previous methods and explores a more nuanced control of variance. We appreciate the reviewer connecting this to variance reduction techniques, and we will investigate this connection further in our future work.\n\n**Q8: Error Bars**\n\nThanks for the advice! We report the error bars as 95% confidence intervals for our main result in Table 1, and will add it to the Appendix.\n\n||Vina Score|Vina Min|Vina Dock|QED|SA|\n|---|---|---|---|---|---|\n|AR|0.066|0.049|0.082|0.004|0.003|\n|Pocket2Mol|0.063|0.058|0.097|0.003|0.002|\n|TargetDiff|0.172|0.102|0.075|0.004|0.003|\n|FLAG|0.778|0.525|0.142|0.003|0.002|\n|DecompDiff|0.060|0.048|0.073|0.004|0.003|\n|IPDiff|0.141|0.088|0.072|0.004|0.003|\n|MolCRAFT|0.122|0.070|0.097|0.004|0.003|\n|DecompOpt|0.415|0.210|0.528|0.011|0.006|\n|TAGMol|0.175|0.088|0.135|0.004|0.003|\n|Ours|0.136|0.078|0.083|0.003|0.003|\n\n**Q9: Presentation & Typos**\n\nWe thank the reviewer for pointing these out, and we will revise our manuscript as requested.", "claims_and_evidence": "The paper claims that gradient guidance is needed over discrete atom types in the setting of SBMO, since optimizing for molecular properties requires knowledge of the molecular topology, hence existing methods such as TAGMol suffer because they only do guidance over atom coordinates. MolJO proposes two novelties: gradient-based guidance over discrete atom types and backward correction. \n\n- One thing that is not clear to me: do the values in Tables 1-2 incorporate backward correction? The difference between TAGMol and MolJO in Table 1 (lines 13-14) is about the same or smaller than the difference between MolJO with BC vs MolJO without BC in Table 3 and Figure 8 in terms of docking score. Since the authors claim that their improvement over TAGMol is due to guidance over atom types, it would be helpful to clarify this. \n\n- The backward correction section is a bit hard to follow -- what exactly do authors mean by \"correcting the past\" if parameters from past timesteps (n-k) are being used to update p_U at the current n-th timestep? Also, is the goal of the method essentially variance reduction like in [1]? \n\n- The guidance component based on molecular properties is very important since this is one of the key problems authors are tackling but implementation of property-guidance is in the appendix. In my opinion this should be in the main paper. \n\n[1] https://link.springer.com/article/10.1007/s10107-016-1030-6", "methods_and_evaluation_criteria": "Authors demonstrate their method to improve molecule generation on an established benchmark on both constrained and unconstrained optimization. Authors evaluate their method on three types of molecular properties (affinity, QED, SA) and their combinations. The proposal of gradient-based guidance over atom types for optimizing molecular properties is cool jointly with continuous coordinates is interesting.", "theoretical_claims": "I have gone through the proofs in the main paper and they seem reasonable. Equations (11) and (12) are a bit hard to read, it might help to put the RHS all on one line instead of splitting it.", "experimental_designs_or_analyses": "The experiments generally make sense and authors compare with a lot of baselines in Table 1 which is good. There are several things that I would like to clarify.\n\n- Why don't the tables contain error bar?\n- Why did authors only report top-10 for MolJO and not other methods?\n- For Figure 3, how did authors restrict the size of generated molecules? The table seems to have different sizes for each method (also should this be called a Figure or a Table?)\n- For the ablation in Table 3, should the values in the last row  (row 6) match line 14 in Table 1? If not, what are the difference between Table 3, row 6 and Table 1, row 14?\n- For Table 3, it's confusing to me why the SDE from Xue et al. is called \"SDE with classifier guidance\", but then authors do an ablation of methods without and without guidance. Can authors clarify what they mean by SDE classifier guidance?", "relation_to_broader_scientific_literature": "The paper tackles an important problem in the scientific community, which is generation of molecules based on the structure of a protein target site, while also optimizing for specific molecular properties, since molecule often have to meet several properties.", "essential_references_not_discussed": "To my knowledge, the work is not missing essential references.", "other_strengths_and_weaknesses": "Strengths:\n- The combination of gradient-based guidance over discrete and continuous data types is an interesting application of BFNs.\n- The empirical gains from the proposed backward-correction are notable based on the ablation study.\n- Authors compare with a large amount of baselines and show gains on several metrics.\n\nWeaknesses:\n- The novelty of the backward correction sampling is not clear since authors write that Qu et al. implement the sampling for when k=n, and in Fig 8 write that the \"strategy is robust within the range k \\in (50, 200]\" and n=200, so the need for a variable window size is not obvious to me since (if I understand correctly), all samples are generated over 200 steps?\n- Some experimental information is unclear or not described in the main paper (listed above; importantly, information on property guidance).", "other_comments_or_suggestions": "- page 6: \"fourthfold\" --> \"four-fold\"?", "questions_for_authors": "[1] Why are the values different between Table 3 row 6 and Table 1 row 14? Based on the writing it seems like these should be the same settings; if they are different, then authors should clarify this in the writing.\n[2] Can authors clarify experimental information like: how they restrict the size of generated molecules in Fig. 3; how classifier guidance SDE differs from \"without and with guidance\" settings in Table 3?\n[3] What are \"me-better\" molecules? I couldn't find a definition in the paper.\n[4] Are the gains in Table 1 rows 13 vs 14 only due to guidance over atom types, or also due to BC? Authors should revise the claims if the latter is the case, or ideally provide an ablation for this (run MolJO at k=1 or MolCRAFT with BC).  \n[5] Why don't authors do a top-of-N comparison for other methods?", "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": "XPdOCN1q88", "reviewer": "Reviewer_fSvs", "summary": "This paper propose a gradient-based molecule optimization framework for the SBDD task, which in experiment achieves state-of-the-art performance on CrossDocked2020 benchmark. Besides, it extend MolJO to a wide range of optimization settings, including multi-objective optimization and challenging tasks in drug design such as R-group optimization and scaffold hopping, further underscoring its versatility.", "questions": "", "limitations": "", "rating": 4, "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 sincerely appreciate the reviewer's thorough reading and insightful feedback, which have helped us improve clarity and presentation. We shall address each point in our responses below, and we welcome further questions.\n\n## Questions\n**Q1: Similarities and Differences between Gaussian BFN and Guided Diffusion**\n\nThank you for this insightful observation. We explain the fundamental differences and structural similarities between Gaussian BFN and guided diffusion as follows:\n- Key Differences: Our approach guides in parameter space ($θ$) rather than in data space ($y$). In the continuous case, $θ$ exhibits lower input variance and therefore provides more informative signals for guiding the generative process toward the desired output. Guided diffusion typically steers the sampling process in the data space directly, while our BFN-based approach performs inference in the parameter space. Such parameter-space guidance arguably connects deeper to the final target properties due to lower input variance, which allows for more explicit Bayesian modeling of uncertainty and more reliable property optimization. \n- Similarities: Both approaches use gradient information to steer the generative process toward desired outputs, operating for Gaussian distributed variables.\n\n**Q2: Regarding the Guided Term's Energy Function E**\n\nWe thank the reviewer for suggesting this clarification. While we included this discussion in the Appendix, we agree it should be introduced earlier for clarity and will revise accordingly in revision.\n\nThe energy function forms part of a Boltzmann distribution over parameters $θ$. Although the target property is explicitly defined only at $t=0$, Bayesian inference defines the parameter space $θ$ for any timestep $i$. This allows us to associate property values with every $\\theta_i$ throughout the generative process, enabling prediction of time-dependent properties during intermediate optimization stages.\n\nWe train a predictor $E(\\theta_t, t)$ directly over the parameter space to estimate $\\nabla_{\\theta_t} p_E(\\theta_t)$ given $\\theta_t$ at different accuracy levels. This approach proves effective as it provides a consistent guidance signal throughout the sampling process.\n\n**Q3: Connection to Energy-Based Compositional Diffusion**\n\nWe appreciate the connection to [6], which introduced how to interpret diffusion models from the perspective of energy-based model (EBM), and thus applies the additivity from EBM to diffusions.\n\nIn our approach, the energy function can be viewed as the negative log-likelihood estimated by a conditional model for given properties. In practice, our method resembles classifier guidance since we train property predictors to serve as the energy function. However, our framework doesn't require explicit \"labels\" as conditions for controllable generation.\n\n**Q4: Preservation of Key Interactions**\n\nWe appreciate the reviewer's expertise in drug design and the important point raised about preserving key pharmacophoric interactions. We fully agree that in real-world drug optimization, certain critical molecular substructures forming intermolecular interactions (e.g., π-π stacking, hydrogen bonding) should be preserved rather than arbitrarily modified.\n\nOur MolJO framework actually addresses this concern by design. The framework allows for precise control over which molecular substructures should be modified and which should be preserved, while our energy function can guide the redesign of remaining substructures. This better aligns with practical applications such as scaffold hopping.\n\n## Other Suggestions\n**Q5: Additional Related Work and Baselines**\n\nWe thank the reviewer for suggesting these important methods for comparison. We agree that including these baselines will enhance the completeness of our evaluation. We commit to including comprehensive comparisons with these methods in our revision.\n\nWe have cited the results from CBGBench (DiffBP, D3FG) and borrowed the samples to calculate the median Vina affinities and Success Rate based on the statistics available for VoxBind. For DecompDPO, we directly cite the numbers from their paper. It can be seen that our MolJO maintains superiority in optimizing the overall properties, reflected by its highest Success Rate (51.3%).\n\n||Success Rate|Vina Score Avg|Vina Score Med|Vina Min Avg|Vina Min Med|Vina Dock Avg|Vina Dock Med|QED|SA|Div|\n|---|---|---|---|---|---|---|---|---|---|---|\n|DiffBP|-|-|-|-|-|-7.34|-|0.47|0.59|-|\n|D3FG|-|-|-|-2.59|-|-6.78|-|0.49|0.66|-|\n|VoxBind|21.4%|-6.16|-6.21|-6.82|-6.73|-7.68|-7.59|0.54|0.65|-|\n|DecompDPO|36.2%|-6.10|-7.22|-7.93|-8.16|-9.26|-9.23|0.48|0.64|0.62|\n\n**Q6: CBGBench for Constraint Optimization**\n\nWe appreciate the reviewer highlighting CBGBench. This is indeed an excellent work that sets up extensive experimental design for structure-based molecule optimization. We will add a thorough discussion of CBGBench, particularly in the context of constraint optimization tasks.", "claims_and_evidence": "Yes, most of the claims are clear and supported by the evidence.", "methods_and_evaluation_criteria": "The method is mainly operated in the probability parameter space instead of the raw data space with the paradigm of BFN. The application is wide and appropriate for the task.", "theoretical_claims": "The theory and proposition in the paper follows BFN and others such as GeoBFN and MolCRAFT. The formulation of update of parameters is correct.", "experimental_designs_or_analyses": "The experiment is complete, while some of the methods are missing. We hope several baseline should be included for more comprehensive comparison, such as VoxBind [1], DiffBP [2] and D3FG [3], in which the metrics can be found in the recently proposed benchmark [4]. Besides, for the optimization methods, can DecompDPO be a competitor? If the comparison is hard to conduct, please explain why.\n\nFinally, in some specific scenarios like fragment growing and scaffold hopping, CBGBench [4] also take these tasks into consideration, please discuss related work and appropriately include evaluated baselines in the benchmark to demonstrate the superiority of the proposed method.\n\n[1] https://arxiv.org/abs/2405.03961\n[2] https://pubs.rsc.org/en/content/articlelanding/2025/sc/d4sc05894a\n[3] https://arxiv.org/abs/2306.13769\n[4] https://arxiv.org/abs/2406.10840", "relation_to_broader_scientific_literature": "The concept related to probabilistic models has been mentioned in BFN. The task related to SBDD has been partially adopted in MolCraft. However, I have not previously observed molecular optimization based on BFN in prior work.", "essential_references_not_discussed": "To enhance the completeness of the paper in the Pocket-Aware Molecule Generation section, it is essential to include VoxBind and DiffBP to provide a more comprehensive overview of recent advancements in the field. Additionally, in flexible SBDD, the recently proposed FlexSBDD[5] represents a state-of-the-art (SOTA) approach in SBDD-related drug design. It is recommended to include it in the Related Work section.\n\nFor the Gradient-Based Molecule Optimization section, additional optimization methods, such as DecompDPO and various 2D-based approaches, should be incorporated. This will facilitate a smoother introduction for readers unfamiliar with molecule optimization and help contextualize the proposed approach within the broader landscape.\n\nIn the experimental evaluation, particularly in the tasks of scaffold hopping and fragment growing, incorporating CBGBench for reference would be beneficial. This would not only provide a standardized benchmark for assessing performance but also help clarify the significance of these tasks in molecular optimization.\n\n[5] https://arxiv.org/abs/2409.19645", "other_strengths_and_weaknesses": "All of them are listed and mentioned.", "other_comments_or_suggestions": "In conclusion, I suggest that：\n- Add relevant related work and baselines, such as DiffBP, D3FG, and VoxBind, as these methods have been evaluated on previous benchmarks with established metrics. This will enhance the completeness of the paper.\n- Include a discussion and reference to CBGBench in the constraint optimization section.", "questions_for_authors": "- Although the optimization in this work is based on BFN, in the case of Gaussian BFN (where coordinates are treated as variables), the optimization process is structurally similar to guided diffusion. Could you explain the similarities and differences between the two in terms of their formulation and underlying principles？\n- The guided term’s energy  E —how is it obtained? It would be helpful to introduce this in the beginning of Methods for clarity.\n- Moreover, this guidance mechanism bears some resemblance to energy-based compositional diffusion [6], as it can be seen as a superposition of sampling across two energy landscapes. However, since this work is BFN-based (modeling in parameter space), its physical interpretation is less explicit. Could this aspect be further discussed?\n- In Claim 2 of Section 5.3: “Optimized molecules form more key interactions for binding,” while this supports the idea that interaction expansion allows the model to explore a broader chemical space, in many real-world drug optimization tasks, certain lead compounds possess specific key interactions (e.g., π-π stacking or hydrogen bonding). During optimization, it is often desirable to preserve these critical pharmacophoric interactions rather than modify them arbitrarily. For example, in cases like 1A2G and 2PC8, the optimized molecules retain the original interactions while expanding upon them, which enhances their practical applicability. How do the authors perceive this issue in the context of their model’s optimization strategy?\n\n**If my suggestions can be adopted and the questions I raised can be clarified, I will consider appropriately increasing my rating.**\n\n\n[6] https://arxiv.org/pdf/2206.01714", "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": "VeJoRY8WyU", "reviewer": "Reviewer_Kj8q", "summary": "In this paper, the authors propose a method that leverages gradient guidance in the context of structure-based drug design. In particular, they augment MolCraft (that uses Bayesian Flow Networks as generative model for structure-conditioned ligand design) to be compatible with gradient guidance according to some (learned) energy function. The guidance is applied on both continuous (atom coordinates) and discrete (atom types) variables.  The authors also propose a a backward correction strategy for more effective optimization. They show good results on CrossDocked2020 benchmark and on sub-structure conditioned generation.\n\n## Update after rebuttal\nI thank the authors for their rebuttal. However my main concerns were not really addressed, mainly lack of novelty (yet another classifier guidance to a generative model very similar to diffusion) and irrelevant experiments (only trained on crossdocked on properties that are totally irrelevant). If the authors argue that this is a general framework for classifier guidance on BFN, results should have been shown on other datasets/tasks (molecular data or not) and other properties. One model trained in one dataset (known to be vey flawed) on irrelevant properties does not  show empirically that this is a \"general method\".\nTherefore, I will keep my rating.", "questions": "", "limitations": "", "rating": 2, "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 sincerely appreciate the reviewer's thorough reading and insightful feedback, which have helped us identify areas for improved clarity and presentation. We shall address each point in our responses below, and we welcome further questions.\n\n**Q1: Explanation of BFN Fundamentals**\n\nWe thank the reviewer for highlighting the need for improved clarity. We will enhance the paper's readability by providing a better overview and more detailed explanations of our guidance approach.\n\nFor experimental details of guidance, our gradient consists of pretrained MolCRAFT and plug-and-play energy functions. We describe the guided sampling in Algorithm 1 and Appendix D.1, and we will make them clearer.\n\n**Q2: Novelty**\n\nThanks for raising this important question. While the work of Xue et al. establishes connections between BFNs and SDEs, our contribution goes beyond merely showing the theoretical feasibility. We derive a principled approach to gradient guidance specifically within the BFN framework, rather than reducing BFNs to SDEs and applying existing techniques.\n\nOur contributions lie in both the methodology (deriving gradient guidance within BFNs and proposing a generalized sampling strategy) and in implementing and evaluating SBMO applications. (1) Methodologically, we show how gradient works through guided Bayesian update, offering a unique perspective as distinct from discretizing an SDE, and we believe contextualizing the guidance within BFN is a novel contribution. (2) Practically, the empirical results in Table 3 also show MolJO's distinct advantage compared to simply reducing BFN to SDE.\n\n**Q3: Limited Dataset Evaluation**\n\nWe appreciate this concern. We add the evaluation on PoseBusters V2 test set (180 out of 384 complexes, after excluding those with sequence identity > 30% or with non-standard residues) as a held-out test. We also report the PoseBusters passing rate (PB-Valid), showing that MolJO's improvements are not dataset-specific or the result of overfitting.\n\n||PB-Valid|RMSD < 2|PB-Valid & RMSD < 2|Vina Score Avg|Vina Score Med|Vina Min Avg|Vina Min Med|Vina Dock Avg|Vina Dock Med|SA|QED|Connected|Success Rate|\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|MolCRAFT|56.4%|43.6%|26.5%|-6.93|-6.95|-7.18|-7.14|-7.77|-7.76|0.65|0.45|93.9%|20.6%|\n|MolJO|68.0%|53.1%|36.3%|-7.74|-7.73|-8.16|-8.09|-8.66|-8.69|0.74|0.57|95.7%|43.5%|\n\n**Q4: Hyperparameter Choices and Potential Overfitting**\n\nAs mentioned in Q3, the consistent performance on both datasets confirms that MolJO generalizes well beyond the specific choices for CrossDock. Moreover, they remain robust across reasonable ranges. For property predictors or training / sampling, we used the same architecture and the same BFN hyperparameters ($\\beta_1, \\sigma_1$) as MolCRAFT without specifically tuning for the task. \n\n**Q5: PoseCheck Metrics in Main Paper**\n\nWe agree that the PoseCheck metrics should be presented in the main text rather than the Appendix for better accessibility, and we will revise our manuscript accordingly. \n\nThough not directly incorporating strain energy as an objective, MolJO indeed improves the conformation quality as suggested by Figure 9, Appendix G, where our CDF is consistently above that of MolCRAFT. Furthermore, the evaluation on PoseBusters (see Q3) reveals notable improvements in PB-Valid. \n\n**Q6: Baseline Comparisons**\n\nWe appreciate this excellent point. Our ablation studies in Section 5.4, Table 3 indeed provide direct comparisons, showing how each component contributes to performance improvements. Furthermore, we're expanding our comparison to include additional optimization baselines such as DecompDPO, where MolJO maintains competitive. Please refer to Q5 in Reviewer fSvs.\n\n**Q7: Relevance of Optimized Properties**\n\nWe appreciate the reviewer's expertise in drug discovery. Our primary contribution is a general optimization framework whose efficacy can be validated through these in-silico metrics. Notably, we've observed that improvements correlate with enhanced key interactions and PB-Valid that are relevant to drug design. \n\n**Q8: Computational Overhead**\n\nCompared to MolCRAFT, MolJO requires approximately 2× or longer inference time in our updated study. The computational overhead depends on the complexity of the energy functions employed. MolCRAFT generally takes ~22s, previous MolJO with 9-layered energy proxies took ~146s, and we have experimented with 4-layered proxies that take ~45s. MolJO can be further accelerated with an efficient strategy where gradient is applied only at selected timesteps rather than at every step [2], which we leave for future work. We thank the reviewer for motivating this analysis, as it led to more efficient implementations.\n\n[1] A Periodic Bayesian Flow for Material Generation. \n\n[2] Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models.", "claims_and_evidence": "- The paper is not very well written an difficult to follow. It would be helpful to have better general overview of BFN/MolCraft models on the main paper for more clarity. The paper also misses a lot of experimental details on how guidance is done, making the understanding of experimental results a bit tricky.\n- The main contribution of the paper is to show that it is possible to do gradient guidance with Bayesian Flow Networks. This is not particularly surprising, given the relation between BFNs and diffusion models/flow matching (Xue et al. ICML24). \n- Moreover, the authors show results on only a single dataset, that has been highly overfitted in the last few years. It would be nice to see results of guidance in the context of BFN in other datasets/tasks to show that the propose model really works in practice  (either molecular datasets or even other modalities like images, other molecules, language, etc).\n- A lot of design choices are needed to be made (guidance temperature, the backward correction, the property predictors, training/sampling hyperparameters etc). It's not trivial to conclude if this approach only works because it has been overly finetuned for the dataset or if this approach would work in other settings.\n- It seems to me that the hyperparmeters of the model (and there are many of them) have bee tuned on the test set of crossdocked.", "methods_and_evaluation_criteria": "- The proposed method (gradient guidance on top of MolCraft) makes sense. \n- The authors show results on a single dataset (CrossDocked2020), which has been highly studied/overfitted in the last few years. \n- It is well known that this dataset nor the metrics (eg docking score, QED, SA, etc) are very relevant for actual drug design. Some other metrics (like those from PoseCheck/PoseBuster papers) provide a bit more insight into quality of the molecules. I think the results from PoseCheck should be displayed on the main document instead of on the appendix. From Table 9, we can see that MolJO has similar Posecheck metrics as in MolCraft, which hints to the fact that the gradient guidance does not improve the quality of the generated conformations.", "theoretical_claims": "The theoretical claims seem coherent, but I did not went through the details", "experimental_designs_or_analyses": "- I feel that comparing the proposed method with other approaches that do not do guidance is not very informative. The most relevant comparison to do is between the proposed approach and MolCraft (since this is a version of the model w/o gradient guidance). \n- It would also be nice if we would see the improvement of gradient guidance in tasks/datasets other than CrossDocked. This could be other molecule datasets (either conditioned on target pocket or not) or other modalities (proteins, images, or anything else where BFN has been applied).\n- The properties optimized in this paper are not relevant for drug discovery and it is difficult to say if any contribution on the paper would actually reflect any improvement on real use-cases.\n- With respect to inference time: What is the computational overload of the proposed gradient guidance? How does the inference time compares with MolCraft?", "relation_to_broader_scientific_literature": "This works proposes to do gradient guidance on top of Bayesian Flow Networks (BFN). In particular, they build on top of MolCraft, a BFN-based generative model for pocket-conditioned ligand generation. Structure-based drug design an important problem on drug discovery, however, the dataset used by ML practitioners---as well as the metrics used to measure performance on this dataset---are know to not be very useful in practice.", "essential_references_not_discussed": "N/a", "other_strengths_and_weaknesses": "See above.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "See above.", "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": {"ethical_review_concerns": "n/a"}}]

Author rebuttal:

Meta-review: This study addresses the problem of structure-based molecular optimization. In this setting, the task involves optimizing over discrete objects, which makes it challenging to apply molecule generation methods based on gradient guidance. This paper proposes a method to overcome this challenge. All reviewers agreed on the practical importance of the problem addressed in the paper, and the effectiveness of the proposed method is clearly demonstrated through experiments on real benchmark datasets. During the rebuttal period, the authors provided additional experimental results, which successfully resolved most of the reviewers' concerns.
Decision: Accept (poster)