Unlocking the Power of Gradient Guidance for Structure-Based Molecule Optimization

Q19: Distribution shift

Yes, Figure 3 primarily shows the general effect of guidance. The purpose here is to show that the guidance mechanism not only modifies the distribution but also surpasses the reference distribution within a reasonable range, thereby addressing and overcoming limitations in data distribution quality to a certain extent. For a more detailed comparison, please refer to Figure 11-17 in Appendix E.1.

Q20: Vina Dock results

We appreciate the emphasis on Vina Dock results. While we agree that they provide valuable insights, we chose Vina Score and Vina Min as indicators of binding affinity in ablation studies, consistent with prior SBDD conventions. These metrics offer meaningful approximations while being significantly faster to compute.

Below are the Vina Dock results together with Success Rate:

For statistical significance, we report the p-values of pairwise T-test between B.C. (k=130) and Vanilla (k=1), Vanilla MC (k=1), Full B.C. (k=200), respectively. It can be seen that our proposed strategy shows statistically significant improvement w.r.t. Vanilla (p<<0.05), Vanilla MC (p<<0.05 except for QED, which is not included as an objective in our implementation). Compared with Full B.C., it's true that Vina affinities remain close between the two and shows no significant difference, but the improvement in molecular properties such as QED and SA is clear (p<0.05), showing a better consistency across continuous and discrete modalities and thereby resulting in a higher Success Rate (51.8% vs. 44.0%).

Q21: More about energy proxies

In order to make a fair comparison, we stick to the same setting as our gradient-based counterpart TAGMol, where each energy proxy is indeed as large as TargetDiff backbone (2.8M parameters). Additionally, we are actively exploring if reducing the model size of energy proxy has a notable impact on performance.

For sampling of energy proxies, we report the overall time in Q10.

For accuracy of the time-dependent energy proxies, we report the Mean Absolute Error (MAE) in the table below, averaged over 10 evenly divided time intervals.

For training of the energy functions, $\theta$ is defined the same way as backbone MolCRAFT, where the parameter space is defined by $\sigma_1=0.03, \beta_1=1.5$, as described in Appendix C, with the transformation in Eq. 4 in MolCRAFT.

Q22: Typos in paper

We have corrected the typo, and we'd like to thank reviewer F2z6 again for the time and efforts devoted to reviewing our paper. Sincerely, we welcome any further questions and hope the reviewer is willing to reassess our paper based on the above answers.

We thank reviewer F2z6 for providing such detailed feedback that helps us improve our manuscript, and we are happy to discuss it further and hopefully clarify some possible misunderstanding.

Q1: Missing citations

Thanks for pointing this out, and we've added the missing citations.

Q2: Equivalence of BFN and Diffusion

We'd like to point out that BFNs are not equivalent to DMs. Although BFN-Solver [1] tried to mathematically unify BFN and diffusion via SDE, its conclusion is based on the reduction of graphical models (from Fig. 2a to 2b in GeoBFN [4]), i.e., marginalizing out the $y$ in BFN. By doing so, uncertainty, also arguably the essence of Bayesian inference, is not fully reflected in such a reduction. For example Eq. (9) in [1] does not correctly reflect the original loss for continuous data in Eq. (101) in [2], the latter with a different loss weighting of $\frac{1-\gamma(t)}{\gamma(t)}$, which is Var(x) in Eq. (83) and closely related to uncertainty. It's worth noting that [1] completely inverts the BFN and makes it look just like diffusion model by starting from a "forward process" that seemingly comes out of nowhere, which actually is not a prerequisite for BFN in which what needs to be specified is only the prior distribution and Bayesian update rule instead.

A further note is that BFN and diffusion models fundamentally differ in both their training and sampling behaviors. While the original BFN [2] uses a conjugate prior modeled as a Gaussian noisy channel, this does not imply that BFNs are inherently limited to this kind of training. This is one possible instantiation but does not constrain the broader applicability of BFNs. Thus the so-called unified perspective might only be applicable to the overlapping cases with diffusion and BFNs (the Gaussian-distributed parameter), but it is not immediately clear if it still holds when the conjugate prior is something other than Gaussian, for example the von Mises distribution in [3]. As for sampling, the BFN Solver proposed by [1] was found to have issues of diversity in molecule generation [5].

Therefore, we sincerely hope that reviewer F2z6 could adopt a forward-thinking perspective of BFN. BFN represents a generalizable framework rather than a single, narrow implementation, and the foundational elements (e.g. input, output, sender, receiver distributions) do not prescribe any specific implementation. A limited view of its core concept is reductive and might fail to reflect the flexibility inherent in the framework's design.

Q3: Unclear contributions

Our contributions lie in both the methodology (deriving gradient guidance within BFNs and proposing a generalized sampling strategy) and in implementing and evaluating SBMO applications. As for reviewer's concerns over the former, we would like to clarify by highlighting the significance of both the guidance within BFN (below) and the backward correction strategy (in Q4). Given the answer to Q2, the best form of exerting gradient guidance remains open within the novel framework of BFN. For example, the guidance can be formalized over either latent $y$, the parameter $\theta$, or the sample $x$. While we admit that there is a connection between our choice of guidance over Gaussian distributed variables and diffusions as suggested by Remark 4.2, what such guidance amounts to from the fresh view of BFN remains unexplored. We make such analysis available and show how the gradient redistributes the probability through guided Bayesian update, offering a unique perspective in understanding BFNs, and we believe contextualizing the guidance within BFN is certainly a non-trivial contribution. Furthermore, simply reducing BFN to diffusions and claiming trivialness overlooks the distinct advantages of BFN and discourages deeper exploration of this new generative model class.

Q4: Novelty of backward correction strategy

Rather than a slight modification, the backward correction strategy effectively bridges the vanilla Bayesian update from [2] and the variance reduction strategy in [6] by generalizing the concept of update and offering an interpolation. The paper also analyzes and explains why [6] achieves higher sample quality than [1] and explores this trade-off further. Refer to Q8 for more details.

[1] Xue et al. Unifying Bayesian Flow Networks and Diffusion Models through Stochastic Differential Equations.

[2] Graves et al. Bayesian Flow Networks.

[3] https://openreview.net/forum?id=Lz0XW99tE0

[4] Song et al. Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks.

[5] https://github.com/ML-GSAI/BFN-Solver/issues/2

[6] Qu et al. MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space.

Q11: Questions regarding beam search

Our method does not employ a default beam search strategy (refer to Line 13 in Table 1). We compare it against both generative baselines and gradient-based methods without this strategy. For baselines such as DecompOpt and RGA, which often generate more samples and select the top-K results, we perform a similar "best-of-N" operation (shown in Line 14 of Table 1) for direct comparison. This ensures a fair evaluation across methods.

Q12: More optimization baselines

Thank you for suggesting the inclusion of AutoGrow4 as a powerful baseline. We have included the result in the revised manuscript for a more comprehensive comparison.

Q13: Questionable reproducibility

We understand the importance of reproducibility and are committed to ensuring transparency. While the submitted code includes scripts for training classifiers, we have also uploaded trained model checkpoints in the supplementary materials to further support verification.

Q14: Information about the size of the reference molecule during sampling

We appreciate your careful review and this important question. To clarify, we do not use ground-truth information about the size of the reference molecule during sampling. Instead, we follow the procedures used in TargetDiff and MolCRAFT, sampling molecule sizes according to a prior distribution derived from the pocket size vs. ligand size in the training data.

Q15: Weakness of DecompOpt

Thank you for the feedback. Our intent was to emphasize the potential advantage that atom-level gradient guidance offers over molecule-level oracle functions in certain scenarios. Additionally, DecompOpt's inference time is the slowest among all methods compared (see our response to Q10), which makes it less suitable for large-scale molecular datasets where computational efficiency is critical.

Q16: Meaning of gradient guidance

Our formulation of gradient guidance is based on the product of experts (Q9), applicable to both classifier-based and classifier-free guidance. In classifier-free scenarios, the energy function can be viewed as the negative log likelihood estimated by a conditional model for given properties. In practice, our approach resembles classifier guidance since we train property predictors to serve as the energy function. However, since there are no explicit "labels" as conditions in our optimization setting for controllable generation, referring to the energy function as a "classifier" may not be fully appropriate.

Q17: Challenge of applying Gaussian noise to discrete data

Our claims in line 87 primarily refer to existing diffusion-based guidance methods for molecular graph generation that operate with continuous Gaussian transition kernels or discrete diffusion over categorical features. These methods can yield suboptimal results due to two key challenges:

1. Continuous Representations often rely on dequantization methods that may be less effective in generating discrete features, as noted by DiGress [1] and [2]. This is why they are sometimes perceived as "unnatural" (line 138).
2. Discrete Representations, though better capturing the inherent discreteness of data, propagating continuous gradients to discrete variables is not inherently valid [3], which can limit their effectiveness.

In contrast, our framework bridges these gaps by leveraging a structured parameter space derived through Bayesian inference. This enables the continuous gradient guidance of discrete variables, combining the strengths of both continuous and discrete optimization approaches for joint, synchronized optimization.

Furthermore, structure-based molecule optimization (SBMO) requires such joint optimization of continuous and discrete variables. This need is illustrated by the suboptimal performance of existing approaches, such as TAGMol, in Figure 1. Our proposed MolJO framework offers a unified solution that handles both modalities efficiently.

[1] Vignac et al. Digress: Discrete denoising diffusion for graph generation.

[2] Kong et al. Autoregressive Diffusion Model for Graph Generation.

[3] Murata et al. G2D2: Gradient-guided Discrete Diffusion for image inverse problem solving.

Q18: Inconsistency issue

Inconsistency might arise across different modalities (continuous-discrete) and different timesteps (within the generative process). For the former, it is empirically demonstrated by the performance boost, and for the latter it is shown in Figure 2 (see also our response to Q8).

Dear Reviewer F2z6,

We appreciate your insightful feedback, which has been instrumental in improving our manuscript.

Please let us know if we have adequately addressed your concerns. Your further comments are most welcome. Thank you again for your invaluable contribution to our research.

Warm Regards,

The Authors

We are happy to further discuss with reviewer F2z6 and update the results.

Q1: Regular SDE

Yes, we followed [1] and borrowed the code for SDE Euler update for both continuous and discrete data. The sampling is done by discretizing the following form of SDE: $\mathrm{d}{\bf{x}} = [F(t){\bf{x}} - G(t)^2 \nabla_{{\bf{x}}}\log p_t({\bf{x}})] \mathrm{d}t + G(t)\mathrm{d}\bar{{\bf{w}}}$ where the $\bf{x}$ corresponds to $\bf{\mu}$ for continuous data, and $\bf{y}$ for discrete data. The estimated score function is given as $s({\bf{\mu}}(t), t)=-\frac{1}{\sqrt{\gamma(t)(1-\gamma(t))}}\epsilon({\bf{\mu}}(t), t), \quad s({\bf{y}}(t), t)=-\frac{{\bf{y}}(t)}{K\beta(t)}+\epsilon_s({\bf{y}}(t), t)-\frac{1}{K}$

A few changes have been made since our backbone MolCRAFT is not trained to predict $\epsilon$, but $x_0$ instead, thus we recalculated the coefficient for updating continuous $\mu$ by $\hat{x}$ and empirically verified the modification.

For classifier guidance, initially we simply subtract the estimated score function by the scaled gradient $s\nabla_{\bf{x}}\log p_E({\bf{x}})$, and we apologize for possible miscalculation of the scale of gradient. We have updated the results in the table below, with a clearer comparison between SDE-based and BFN-based sampling. It shows that SDE-based guidance requires 5x sampling time (i.e. the same discrete-time steps as in training) while being slightly inferior in optimizing both objectives (Affinity + SA).

Q2: Beam search results for generative models

Thanks for your interest. We have made the same top-of-N selection based on z-score for picking the top 1/10 out of roughly 100 generated molecules w.r.t. each of the 100 test proteins (a tenth of 10,000 in all), and provided the result in the table below. This generally shows what the "concentrated space" for desirable drug-like candidates looks like for generative models. Supplementing it with Table 1, we found that the performance of DecompOpt is comparable to a top-of-10 version of DecompDiff, and ours to that of MolCRAFT in terms of success rate. Considering that the sampling time of DecompDiff (~309s) and MolCRAFT (~22s), DecompOpt (~11714s) is significantly heavier than ours (~146s), and does not offer additional benefits in computational efficiency compared to a top-of-10 DecompDiff.

Thanks for the thorough review, and especially for the general advice of introducing the necessary components of BFN in more detail. We have added a more intuitive explanation in Appendix A, and addressed the remaining questions below. Please feel free to raise any further questions, which will definitely help us improve our manuscript.

Q1: BFN notations

[BFN in preliminaries.] “The receiver holds a prior belief $\theta_0$, and updates…” Exactly what does $\theta$ represent? Are these pointwise parameters? Is there a distribution defined over them?

Yes, that's correct. We denote $\mathbf{\theta}=[\theta^{(1)},…,\theta^{(N)}]$ as the parameters of an N-dimensional factorized distribution, where the input distribution is defined by $p(m|\theta) = \prod_{d=1}^N p(m^{(d)} | \theta^{(d)})$. This formulation allows $\theta$ to encapsulate intrinsic parameters across different dimensions and modalities.

Specifically, these parameters can be pointwise vectors in $\mathbb{R}^3$ for N continuous atom coordinates (representing positions in 3D space) and on the probability simplex in $\mathbb{R}^K$ for N discrete atom types (defining categorical distributions over K possible types).

Together, these parameters define a factorized distribution over N atoms. This structure enables joint optimization and parameterization of molecular structures across different modalities within the Bayesian Flow Network framework.

[BFN in preliminaries.] Eq (1) shows a distribution over $\theta_i$, but eq. (2) shows $\theta$ as a deterministic variable (given observations y_0, …, y_i). I understand Eq (1) is the posterior given a datapoint $x$, integrating out potential observations, but Eq (2) represents the deterministic updates we get for some given observations at different noise levels?

Thanks for pointing this out. It's true that Eq. 1 represents the posterior distribution over $\theta_i$, given a data point $x$. In this formulation, we integrate out potential observations $y$ over the sender distribution $p_S$ to model the uncertainty of the latent state.

On the other hand, Eq. 2 represents deterministic updates of $\theta$ when given specific observations of $y$ at different noise levels. This deterministic update reflects how the receiver distribution $p_R$ evolves as it processes known observations, resulting in a more precise latent state estimation for $\theta$ under specific conditions.

During training, since the sender distribution is known, we can integrate out $y$ within $p_S$. However, during sampling, we only work with the receiver distribution $p_R$. This leads to different forms of the function mapping $f$ that governs updates, as discussed in the additional discussion provided in Appendix A.

Q2: Reason for guidance over $\theta$

The parameters $\theta$ are fed through a NN to model output distribution over clean data. Why is it sensible to apply guidance over theta (if they are connected to the clean samples through a NN, and the energy is typically defined over clean samples)? Coming back to the question above, what do these latents represent?

Our approach applies guidance over $\theta$ because the energy function is defined on the intermediate latents. This enables us to guide the aggregated latent states within the generative trajectory, which are then used to predict clean datapoints, rather than directly working with clean samples. By doing so, we ensure consistent updates that reflect both observed noise and the desired target properties.

The latent $y$ can be viewed as the observed noisy representations of clean data, and $\theta$ a posterior belief updated by noisy observations, therefore can be understood as aggregated $y$. For more intuitive discussion of how it is derived through Bayesian inference, please refer to Appendix A.

Q3: Guidance notations

[Proposition 4.1.] Is $\mu_\phi$ an output of the NN $\Phi(\theta_{i-1})$? Also, in the definition of $\theta$ (line 201) you use $\theta=[\mu, z]$, but Eq (6) uses $y$, stating “recall $z=f(y)$”. I suppose this comes from Eq (2)? Are the $y$ noisy observations? If so, why is Eq (6) sampling $y$? (Should it be sampling $z$ which is part of $\theta$?)

Yes, $\mu_\phi$ is indeed one of the outputs of the neural network, as it parameterizes the distribution over clean coordinates.

Regarding the use of $\theta=[\mu, z]$ and subsequent references to $y$ in Eq. 6, and $y$ do represent noisy observations drawn during training and guidance stages. The function $z=f(y)$ defines a mapping that links these noisy observations to their discrete state representations. The reason why Eq 6. samples $y$ is actually related with numerical stability: for continuous variables, the noisy samples aggregate linearly into $\mu$ with a constant multiplier, as demonstrated in Remark 4.2. But for discrete variables, the gradient w.r.t. $z$ needs to be multiplied by $[\frac{\partial h}{\partial y}]^{-1}=\frac{1}{z_i(1-z_i)}$ that actually differs along each dimension, and due to the fact that $z$ lies on the probability simplex and will be increasingly sharp towards the end of generation (i.e. Approaching a one-hot), such a gradient would be highly unstable.

Q4: Gaussian distribution in Proposition

[Proposition 4.1.] Where are the original Gaussians from line 207 coming from (in-line equations, just before Eq (5))? Why are those the correct “unguided kernels” $p_\phi$?

Sorry for the confusion caused. Proposition 4.1 actually assumes that the unguided distributions are Gaussian, while they are then fulfilled in Eq. 14, 15. We have revised the proposition to make it more clear.

Q5: Sample generation

After reading section 4.1, it is unclear to me how samples are actually generated by the model without using backward correction.

We apologize for any confusion regarding sample generation. A more detailed overview of the Bayesian Flow Network (BFN), including an intuitive description of its generation process without backward correction, has been provided in Appendix A for clarity.

Q6: More about sender distribution for discrete data

[Line 233.] This line uses the notation $e_v$ without definition. $e_v$ is defined later in line 240. I’d suggest to introduce variables before using them in equations. That same paragraph states “Surprising as it may seem, this is mathematically grounded…” in which way?

Thank you for pointing out the use of the notation $e_v$. We have revised the draft to introduce it earlier, ensuring clarity and consistency. The statement for mathematical ground highlights how the formulation of sender distribution for discrete data is constructed with theoretical rigor, as opposed to heuristics that tried to relax the discrete variable into some continuous space. To illustrate, continuous diffusion applies Gaussian noise to the one-hot encoding, while BFN constructs a "multi-hot" y that follows a multinomial distribution stemming from the categorical. Details can be found in Appendix B.

Q7: Higher-order approximation of guided distributions

[Prop 4.1.] “it suffices to sample guided … [guided Gaussians]”. These expressions are based on a 1st order Taylor expansion. Would these become increasingly exact as the updates become smaller? (Related to the discretization used?) I think the approximation used here could be briefly discussed in the main text, as it is claimed before, in line 192, that the analytic expressions for the guided kernels are derived.

Thanks for the suggestion! We have added the approximation used in line 192, mentioning that the guided form is based on a first-order Taylor expansion, and indeed, higher-order terms would theoretically improve the approximation, yet second-order terms would require additional time complexity such as computing a Hessian matrix, and often lead to instability during sampling, which we will leave for future work. Here we briefly list a possible solution for sampling $(\mu_i, y_i) \sim \mathcal{N} \left(\begin{pmatrix} \mu_i^* \\\ y_i^*\end{pmatrix}, \Sigma_{\pi} \right)$, where the covariance matrix $\Sigma_{\pi} = \begin{pmatrix} \sigma + H_{\mu\mu}^{-1} & H_{\mu y}^{-1} \\\ H_{\mu y}^{-1} & \sigma' + H_{yy}^{-1} \end{pmatrix}$, and $H_{\mu\mu}$ ($H_{yy}$) is the Hessian matrix of second derivatives with respect to $\mu$ ($y$), $H_{\mu y}$ is the cross-derivatives.

Q8: Definition of guided $\pi$

[Detail in Line 187.] The definition of $\pi(\theta_i | \theta_{i-1})$ should have a $\propto$ instead of $=$? If $p_\phi$ and $p_E$ are both normalized, their product is not necessarily normalized. For instance, consider the product of two Gaussian densities, the resulting thing is not normalized.

Yes. Thank you for pointing out the typo in the definition. We have corrected this in our revised manuscript.

Dear Reviewer LDEf,

We appreciate your insightful feedback, which has been instrumental in improving our manuscript.

Please let us know if we have adequately addressed your concerns. Your further comments are most welcome. Thank you again for your invaluable contribution to our research.

Warm Regards,

The Authors

Thanks for the detailed feedback! While we appreciate the reviewer’s recognition of the importance of equipping molecular generative models with conditional generation and optimization capabilities, we would like to first clarify that MolJO is not based on diffusion or flow models. Instead, our method leverages Bayesian Flow Networks (BFNs), which represent a distinct paradigm for generative modeling, and enable more structured and flexible updates, particularly suited for joint continuous-discrete optimization task. We believe that our approach offers a valuable perspective on this challenge, and we have added a more detailed introduction of BFN in Appendix A.

W1: Missing literature of molecule optimization

Thanks for the reference to the review paper, which helps us better understand the literature of molecule optimization. We have added the related works in Appendix I.

W2.1: Clarification of the scope

Some claims in the paper are not appropriate, previous structure-based drug design models also optimize both the atom type and coordinate [2, 3] (gradient-based algorithm and evolutionary algorithm). I am not sure if it is appropriate to call it the first proof-of-concept for gradient-based optimization of continuous-discrete variables because (1) the scope is very narrow, (2) other papers have worked on it [2, 3]

Thanks for raising this important question. We acknowledge that previous models have also made significant strides and addressed the optimization of both atom types and coordinates. However, our method introduces a unique approach within the novel generative framework of BFN through synchronized, joint gradient optimization within a unified differentiable space, which distinguishes it from prior models, such as those described in references [2] and [3], in tackling the challenge of optimizing both atom types and coordinates.

• Evolutionary-based DiffSBDD [3] relies on iterative search for optimized properties, selecting top-K molecules to seed the new population. It uses supervision signals only implicitly, and we aim to open the door to directly utilizing the atom-level gradient signals.
• Gradient-based algorithm [2] builds on GDSS [4] for graph generation but does not incorporate SE(3)-equivariance in its modeling. This approach treats the graph features as continuous variables, which means that it dequantizes node vectors and adjacency matrices using uniform or Gaussian noise and subsequently learns Gaussian diffusion models over these dequantized variables. While such dequantization and Gaussian diffusions can capture some aspects of graph structures, they face inherent challenges with compatibility between continuous Gaussian diffusion processes and discrete data, making it difficult for the generation of consistent and valid graph structures [6].

Moreover, when it comes to optimizing both coordinates and types, more sophisticated designs are required to ensure the consistency across different modalities, such as those nuanced noise schedules and multiple stages in MolDiff [5], thus it can be expected that potential suboptimality could arise from indirect guidance approaches.

Unlike prior works that typically treat discrete variables as continuous or optimize the entire molecule via resampling, our work uniquely emphasizes explicit, synchronized, joint gradient-guided optimization across continuous and discrete variables with SE(3)-equivariance. This represents a fresh perspective compared to existing methods, and ensures smooth and consistent information flow across modalities.

(For better readability, the index for reference [2, 3] remains consistent with the original question.)

[2] Lee et al. Exploring chemical space with score-based out-of-distribution generation.

[3] Schneuing et al. Structure-based drug design with equivariant diffusion models.

[4] Jo et al. Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations.

[5] Peng et al. MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation.

[6] Vignac et al. DiGress: Discrete Denoising diffusion models for graph generation.

Thank you for writing the rebuttal, I have a couple of more comments/questions

First of all, I agree with what you said, but every paper needs to be at least somehow different to be published, but not every paper claims to be the first of its kind because of its difference or novelty. The idea is essentially, by given so many relevant works in this field, I suggest removing the phrase "first proof of concept", I think it makes more sense to be used if very few works have been done in the field or if it provides a highly novel concept which potentially people have not thought it could work or nothing close to it worked before ("proof of concept").

Re the word "gradient", I am not too strongly against the word, as long as you can find a good root in the derivative-free and discrete optimization literature so it doesn't confuse people.

Re time-dependent energy function, the reason why I ask is that to achieve exact guidance, it always requires a time-dependent energy function, in the diffusion model literature, it is usually evaluated by sampling from p(x_0|x_t), \log \sum exp(-energy function). I am trying to understand if you did it this way or you simply obtain it through the same function by evaluate on different x_t and t which is often inexact.

Thank you for this thoughtful suggestion. We agree that the phrase may imply a broader claim than intended, and our goal was to highlight the unique integration of gradient guidance for joint continuous-discrete optimization within Bayesian Flow Network (BFN). We recognize that emphasizing our contributions in terms of methodological advances and practical implications rather than positioning it as the "first" will better reflect its value, and we have revised the draft accordingly.

Regarding the use of "gradient", we agree that ensuring clarity is crucial, especially when applying gradient-based methods to discrete variables. We will be careful in using "gradient" in this context.

Regarding time-dependent energy function, we trained a predictor $E(\theta_t, t)$ directly over the parameter space (the belief as aggregated noisy samples) to estimate the $\nabla_{\theta_t} p_E(\theta_t)$, given $\theta_t$ at different accuracy levels. Then, we obtain the guidance term from the same function when sampling $\theta_t$, which proves to be effective.

Our guidance operates in a training-based manner, essentially different from the training-free methods that estimate the time-dependent term through a loss function defined over clean data space, which is challenging and usually requires multiple model evaluations for the true posterior mean $\hat{x}_0$ (i.e. sampling $p(x_0|x_t)$ many times) [1] or multiple back propagations with the log-mean-exp-negative operation [2]. One such training-free guided molecular diffusion [3] based on DPS [1] requires turning off the guidance during the early "chaotic stage" and applying the MC sampling and the time-travel technique to improve $\hat{x}_0$ estimates, while our guided Bayesian update provides uncertainty-aware guidance with reduced variance without needing such adjustments.

We believe that our work offers a promising foundation for further exploration of guidance within BFN, and we look forward to more exciting discoveries in this emerging class of generative models.

[1] Diffusion Posterior Sampling for General Noisy Inverse Problems.

[2] Loss-guided diffusion models for plug-and-play controllable generation.

[3] Training-free Multi-objective Diffusion Model for 3D Molecule Generation.

Q2: How does the backward correction window size impact performance versus computational cost, and what determines the optimal balance?

Thank you for raising this question, which helps us in clarifying our proposed method further. The size of backward correction window actually does not have a notable impact on computational cost, because we've derived the simulation-free solution for the backward corrected Bayesian update (Eq. 14, 15), thus the update does not require traversing the window of size k $(i-k, i-k+1, \dots, i-1)$, but instead jumps from $i-k$ toward the next step $I$.

Q3: Can this joint optimization approach extend beyond the three basic properties (Affinity, QED, SA) to handle more complex molecular properties relevant to drug discovery?

Yes, certainly. As shown in the table below, we conduct an experiment that solely optimizes Lipinski's Rule of Five (RO5), and our guidance method succeeds in optimizing molecular properties such as QED, SA without compromising Vina Affinities, even though they are not listed as optimization objectives.

We admit that for more complex molecular properties, such as ADMET and pharmacokinetics, richer datasets that contain such information will be needed in order to build a good data-driven differentiable proxy.

We appreciate the constructive feedback from reviewer C1Hy as well as the efforts in reviewing our paper. We hope that the concerns have been addressed below, and we welcome any further questions.

W1: Presentation

Thanks for the advice! We've added more explanations of BFN in Appendix A. Briefly speaking, the major difference between BFN and discrete diffusion / GFlowNets lies in that the latter two operate in a discrete state space instead of a continuous space of noisy observations $y$. For BFN over discrete data, there is no need to utilize a transition matrix (as in discrete diffusion) or a (state, action) pair (as in GFlowNets) to navigate the probability simplex, enjoying the benefit of smoother transformation.

W2: Limited Objective Scope

We admit that other biological objectives are yet to be tested within MolJO framework, but an extension would not be difficult. We've added the Lipinski's Rule of Five (RO5) objective. Please see Q3 below. If there are specific objectives that reviewer C1Hy would like to see, we will definitely give it a try in our revision.

As for scaling to $k$ objectives, it requires $k$ gradients that either come from neural networks to make the predictions, or some other calculations, so in the worst case the time scales linearly to the number of objectives. However, due to the fact that different objectives are usually not equally hard and so will the network parameters and calculation time, the actual time cost depends on the implementation.

W3: Computational Analysis Gaps

We've added the comparison of computation time for optimization baselines in Appendix H, which is calculated as the time for sampling a batch of 5 molecules on a single NVIDIA RTX 3090 GPU, averaged over 10 randomly selected test proteins.

W4: Hyperparameter Sensitivity

Thanks for bringing this up, which inspires us to investigate key hyperparameters like guidance scale and correction window size further.

For guidance scale, we found that the gradient norm is generally between 0.02 and 0.06, which might explain why scale = 50 appears optimal in the previous study. We conduct a new experiment and adopt the L2 gradient normalization that automatically adjusts the guidance scale and thus does not require manual tuning. The updated result is shown in the table below, and it can be seen that our guidance method remains robust to the normalized gradient scale.

For correction window size, it can be seen from Figure 9, Appendix F.2 that the performance of backward correction strategy remains robust when the window size k > 50, i.e., correcting a sufficient part (more than 1/4) of the optimization history. This brings consistent performance boost.

Q1: How reliable is the oracle for real-world drug discovery? Is the cost of calling the oracle a concern?

This is a great question. The reliability of using oracles for real-world drug discovery, particularly for Vina scores as affinity, remains a topic of discussion. The most reliable method is definitely wet-lab experiments, yet it is often infeasible in the early-stages. While AutoDock Vina may not be perfect, it is one of the most widely used open-source docking tool, offering a balance between speed and accuracy for in-silico calculation [1]. Thus, the oracle is of considerable value for benchmarking molecular properties for drug discovery.

As for the cost, Vina Dock usually takes tens of seconds per molecule (exhaustiveness=16), and therefore, in order to call the oracles and evaluate 10K molecules in the main experiment, it requires more than 10 hours to run on a single CPU. This becomes increasingly costly when the optimization methods like RGA typically run several rounds before arriving at the final outputs, making it a CPU-heavy process.

[1] Su et al. Comparative Assessment of Scoring Functions: The CASF-2016 Update.