Unified Guidance for Geometry-Conditioned Molecular Generation

We are grateful for the reviewer's insightful feedback and are pleased with the positive comments on UniGuide's novelty and effectiveness for various drug discovery tasks. We would like to clarify the remaining questions and concerns in the following.

 

The contribution and the generalizability of UniGuide are somewhat overstated; the available conditions are described in a quite general way, including not only structures and surfaces, but also densities

We appreciate the reviewer's feedback and acknowledge the ambiguity between UniGuide's contributions, as shown in Fig. 1, and the empirical evaluations presented in our experimental section. To resolve this ambiguity:

• We extended the setting from App.G on density-conditioned guidance as part of our rebuttal and explain it in more detail down below.
• We plan to incorporate this novel setting into the final part of Sec.5.1 in the updated manuscript, referring to it as density-based drug design (DBDD), a more challenging task than LBDD.

Given its potential [of guiding w.r.t. densities] to enhance the model's capabilities, it is suggested to provide additional details about this setting (e.g., task definition, condition mapping) and to reference App. G in the main body of the paper

We appreciate the reviewer's suggestion and are pleased to provide a more detailed explanation of the technical aspects of the DBDD setting.

Task Definition

• As motivated in App.G, we anticipate UniGuide to be particularly useful in scenarios where explicit information about advantageous features of the ligand is provided in the form of 3D densities. Examples of this include
  • a) volumetric densities that indicate beneficial placement of certain atom types, such as oxygen atoms [2] or
  • b) pharmacophore-like retrieval of advantageous positions for aromatic rings, as utilised in [3].
• On a technical level, the DBDD setting assumes this information to be provided as follows:
  • Instead of a reference ligand's shape, we only have the protein pocket's surface, which primarily defines exclusion zones rather than precise atom placement. Please refer to Fig.1 in the attached PDF.
  • Instead of a reference ligand's structure, we only have access to (multiple) atom-type densities that indicate preferred locations for optimal interaction with the protein.

Condition Map: Adapting UniGuide for the DBDD scenario requires minor adjustments

• The protein surface is treated like shapes in standard LBDD, defining an exclusion zone based on proximity to the surface.
• The atom densities are thresholded to reflect regions of high interest and converted to surfaces using the marching cubes algorithm [4]
• To include feature information, we employ a modified condition map similar to Eq.21 that extends the transformation from the conformation to the configuration space.
• Moreover, the number of atoms guided by each density is adjusted based on its volume, reflecting the varying influence of each density, and guidance is only applied if atoms are sufficiently close.

We have included a new figure in the attached PDF that visually demonstrates qualitative results for this setting. While our current approach represents a promising first step in tackling this task, we acknowledge the potential for further refinement. We are eager to explore future improvements within the UniGuide framework.

For FBDD [...], can the UniGuide framework also be applied to the DiffLinker model, considering it is diffusion-based backbone

While direct application of UniGuide's FBDD condition map is not possible with DiffLinker, it can be combined with the surface-based Linker Design condition map (App.F.1) to guide linker atom positions.

• DiffLinker's Modifications: Since DiffLinker fixes condition fragments in space and only learns to diffuse the linker atoms, the model becomes unsuitable for UniGuide's FBDD condition map from Sec. 4.2.
• DiffLinker's Similarity to EDM: However, DiffLinker's diffusion backbone is very similar to EDM, which we successfully combined with UniGuide for controlled generation in Linker Design, see Tab.3.
• Compatibility with Surface-Based Conditioning: Moreover, DiffLinker can still be combined with the surface-based Linker Design condition map (App.F.1) to guide the placement of linker atoms towards a specific region between the fragments.

 

We appreciate the positive feedback and are excited to move forward with these improvements. We hope that we have addressed all outstanding concerns satisfactorily and welcome further discussion.

 

[1] Hoogeboom, Emiel, et al. Equivariant diffusion for molecule generation in 3d, 2022.

[2] Zaucha, Jan, et al. Deep learning model predicts water interaction sites on the surface of proteins using limited-resolution data, 2020

[3] Zhu, Huimin, et al. A pharmacophore-guided deep learning approach for bioactive molecular generation, 2023

[4] Lorensen, William E., and Harvey E. Cline. Marching cubes: A high resolution 3D surface construction algorithm, 1998

We thank the reviewer for their constructive feedback regarding our manuscript, particularly regarding the presentation and interpretation of UniGuide. We believe we can address these concerns effectively, as detailed in the following answers. 

 

Only highlight the best diffusion-based approach 

• We are happy to revise the updated manuscript's presentation style as the reviewer suggested.
• We want to emphasise that we never intended to mislead the reader by highlighting the results in a way that favours UniGuide. We are confident in the benefits provided by UniGuide and agree that they should be presented without ambiguity.

UniGuide seems not so competitive in LBDD and Linker design in FBDD tasks

• Clarifications on LBDD performance
  • We stress that "on-par" performance for UniGuide compared to specialised approaches is a significant achievement, as our method imposes no constraints or additional training. The general response provides clarifications on LBDD performance and the reasoning behind the experimental settings. We highlight that UniGuide achieves very high shape similarity and low graph similarity, even though it does not rely directly on the reference structure (as in SQUID).
• Clarifications on FBDD (Linker design) performance
  • We agree that UniGuide performs on par with DiffLinker overall in the linker design task and does not strictly outperform it. Yet, we want to emphasise again that we simply guide the sampling of an EDM [1] model that was not optimised for linker design (unlike DiffLinker, which is a conditional model limited to the linker design task only). In Fig.4 and App. F.2, we showcase how UniGuide can be flexibly adapted to more tasks such as scaffolding or fragment growing.
  • Despite not being trained explicitly for linker design, UniGuide(EDM) performs competitively with DiffLinker: Our method generates the most diverse linkers (Uniqueness) and successfully recovered nearly 60% (Recovery) of the actual linkers from the test set. While baseline methods may achieve higher recovery rates, they produce less diverse solutions. Moreover, UniGuide (EDM) generates more complex structures as measured by the number of rings. However, this also leads to a slightly worse SA score as these molecules are not as easily synthesisable.
  • Another advantage of UniGuide not relying on extra training for the linker design task is that it is agnostic to the fragmentation procedure used to obtain the condition fragments. This means that UniGuide will generalise to unseen fragments if the underlying molecule fits within the training distribution.
  • We realise that the original manuscript lacked clarity on these aspects, which may lead to misunderstandings. In the updated version, we will better contextualise the FBDD results.

Additional baselines

• SBDD:
  • We include DecompDiff and IPDiff in the updated version of the manuscript as requested by the reviewer. Furthermore, we are happy to include additional Vina metrics to indicate the pose quality better. Please refer to Tab.1 in the attached PDF. We have also provided additional experimental details and comparisons to different baselines in the general response. 
• FBDD:
  • We thank the reviewer for pointing us towards the LinkerNet work. Like DiffLinker, LinkerNet is a conditional model that is additionally able to model fragments with more degrees of freedom, i.e. their pose and centre of mass. Since this is also reflected in their experimental setup (LinkerNet can predict the condition fragment poses, the baselines randomly rotate them), it is not directly comparable to the experiments conducted in Table 3. However, this is an exciting addition to the linker design task, which can be readily implemented with minor adjustments to the condition maps presented by UniGuide. We will investigate this addition for inclusion in the updated version of the manuscript.

Q: Connection to retrieval and molecular translation

• Comparison to molecular translation: While UniGuide's primary focus is novel drug design with geometric constraints, its adaptability allows for future exploration of molecular translation in 3D, similar to existing 2D approaches [3,4]. This could include, for example, optimising a molecule's properties while preserving its 3D shape using the surface condition map. However, this would necessitate property regressors and the combination of UniGuide with classifier guidance.
• Comparison to retrieval-based methods: Retrieval-based methods, such as shape-based Virtual Screening (VS),  involve retrieving existing data based on predefined criteria or queries, whereas UniGuide focuses on generating new samples based on a predefined condition without involving existing data. Our comparisons for LBDD show that we surpass VS in terms of 3D and 2D similarity (Tab.1). We are curious to explore the inclusion of retrieval-based methods into the condition map in the future.

 

We hope that we have adequately addressed the concerns raised and believe the reviewer's feedback has significantly enhanced the presentation of our results and the experimental section. We look forward to further discussion.

 

[1] Hoogeboom et al. Equivariant diffusion for molecule generation in 3d

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

[3] Jin et al. Hierarchical generation of molecular graphs using structural motifs

[4] Jin et al. Junction tree variational autoencoder for molecular graph generation

We thank the reviewer for their valuable feedback and will address their concerns below.

 

Superiority of UniGuide and insufficient experimental support

UniGuide's appeal builds upon multiple aspects:

• No Training: UniGuide is based on guidance, and as such, it does not require any additional training but achieves the desired sampling behaviour at inference time
• The unification aspect of UniGuide: Unconditional models can readily be adapted to new settings via UniGuide. This adaptation effectively reduces to the definition of a condition map C, which can be non-differentiable. To support this, we guide an unconditional EDM model for the LBDD task (Tab.1) as well as the Linker Design task (Tab.3) and show that we outperform specialised approaches.
• The special case of $S=Z$ (not the general case!) can be addressed using techniques based on inpainting. However, UniGuide leads to better or equal performance, making it favourable over inpainting. This is supported across our experiments, such as in Tab.2 or App.F.2.
• Extensive experimental support: We benchmark UniGuide against specialised non-diffusion and diffusion-based models (including conditional and guidance-based alternatives) across tasks involving geometric conditions. Our results prove that UniGuide provides better or on-par performance, indicating that controlling unconditional models with UniGuide is effective, simple, and easily extended to novel settings. Yet, we acknowledge the need for a more explicit discussion and refer the reviewer to our general response for further clarification on SBDD, LBDD (where UniGuide excels), and FBDD (see also the response to reviewer YA7u).

Shape-based generation via technique akin to validity guidance

• We agree that the proposed loss used for validity guidance [1] could be adjusted to the LBDD case, and we thank the reviewer for pointing this out.
• Yet, UniGuide is more general and decouples the surface computation (input to the condition map) and gradient computation, which is consistently applied to the L2 loss between the clean estimate and the target condition.  
• This decoupling has multiple implications:
  • Due to the loss formulation, the condition map does not have to be differentiable, making it agnostic to how the surface points are computed and, therefore, more flexible.
  • At the same time, the condition map follows an explicit geometric intuition, which makes it easily adjustable to novel scenarios. This aspect is especially well reflected in the new experiments discussed in our general response. We guide towards a volume reflecting the symmetry requirement to generate symmetric proteins.
• In our general response, we discuss incorporating drift terms into UniGuide (akin to validity guidance) for SBDD. As part of this addition, we will conduct experiments to compare an adjustment of validity guidance with our surface condition map. We will share the results as soon as they are available.

Limitation: Reliance on the base model

• While UniGuide's performance is inherently tied to the performance of the underlying base model (as noted in our limitations), this connection also has a positive effect as it encourages research on unconditional generation. Our experiments show that UniGuide can directly translate improvements to better task-specific performance. UniGuide also facilitates the application of these models to diverse downstream tasks, broadening their utility.

Eq10 does not make much sense

• We apologise for the confusion caused by the poor presentation of Eq.10.
• In the current manuscript, we missed to specify that the condition $c$ must lie in the same space as the samples $z\_t$ (the configuration space). Only this assumption makes the Gaussian approximation possible. To clarify this, we will add further explanations near Eq.10.
• We would like to highlight that we discuss and lift this assumption on the condition $c$ in our method section (Sec.4). Note that the condition map $C$ serves as a transformation that takes the source condition $s$ and the clean approximation of $z\_t$ as inputs, and outputs a suitable target condition $c$ that lies in the configuration space and can be used directly in the guidance loss presented in Eq.14/15.

Eq.17 does not match the definition of C

• In Eq.17, we define the condition map as $C(s, \hat{z}^{\mathcal{A}}\_t)=C(\tilde{z}, \hat{z}^{\mathcal{A}}\_t)$, where $\tilde{z}\in Z$ is a configuration that specifies $m<N$ nodes (for $m=N$ there is no point in guidance). 
• In order to match the dimension of the clean data estimate $\hat{z}\_t$ with $\tilde{z}$, we subset $\hat{z}\_t$ to $m$ nodes as indicated by the superscript $\mathcal{A}$ in the condition map.
• For example, for SBDD, the source condition is a protein with $m= N^{P}<N$ nodes.
• While we believe Eq.17 and its definition are consistent, we would appreciate further clarification if we have overlooked any aspects.

In 5.1 Ligand-based drug design, why do the authors not compare their model with conditional EDM? 

In our general response, we clarify our reasoning for the LBDD evaluation and hope that our answer explains why we did not compare UniGuide (EDM) to conditional EDM. 

How could UniGuide (shapeMol), a conditional model with soft constraints, outperform Cond-Shape2Mol?

• We believe that what is meant by "Cond-Shape2Mol" maps to the conditional ShapeMol [2]. We followed this suggestion and added UniGuide (ShapeMol) to our comparisons; see Tab.3 (PDF).

For FBDD, add inpainting comparison

• We added this in Tab.2 in the PDF and would refer to our general response for additional details on this setting.

 

Again, we thank the reviewer for their constructive comments, which have significantly improved our work.

 

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

[2] Chen et al. Shape-conditioned 3d molecule generation via equivariant diffusion models

We would like to thank the reviewer for their valuable feedback and will address the raised concerns in our answer below.

 

However, the paper's contribution is hard to evaluate. Though the author claims a new general framework for Uniguide, the form takes exactly as a gradient guidance form for diffusion models which has been explored in the previous literature and following ups [1]

• We acknowledge that UniGuide builds upon the existing framework of gradient-based guidance for diffusion models. However, our primary goal was to develop a method capable of conditioning on general geometric information, a capability absent in previous work. We achieve this through two key innovations: self-guidance and the condition map.

Key Differentiators:

• Condition Map: UniGuide leverages a novel condition map to incorporate diverse geometric information, such as protein structures, molecular fragments, and molecular surfaces, as guidance conditions. This is a significant departure from previous work like [1, 2, 3], which are limited to conditioning only on specific (quantum) properties. 

• Self-guidance: UniGuide is a self-guiding approach that modifies the reverse process of the diffusion model without relying on additional networks to guide the generation, similar to prior work [5]. This highlights the flexibility of UniGuide as it does not require additional training (neither training a specialised conditional diffusion model nor training additional networks for guidance). This is unlike prior work [1] that requires training additional regressors to guide the generation.

• Property-Based Conditioning: In principle, UniGuide can also perform (quantum) property-controlled generation. As this task focuses on global graph properties, it would require additional regressors. 

  • EEGSDE [1], on the one hand, computes the guidance loss between a condition property value $\mathbf{c}$ and the output of a property prediction network that is finetuned over different noise levels of the diffusion process $\mathbf{m}_{\theta}(\mathbf{z}_t,t)$. 
  • In contrast, UniGuide uses a property regressor trained on clean samples, leveraging a clean approximation ( $\mathbf{\hat{z}}_t$ ) of the noisy data ( $\mathbf{z}_t$ ) for guidance. This approach, explored in [2], differs from our focus on self-guidance. While previous studies [1,2,3] explore property-based conditioning with various conditioning mechanisms, our work investigates self-guidance without relying on external regressors or classifiers, and its broad application across various drug discovery tasks.

• Multiple Conditioning: We further indicate that, unlike prior works [1,2], UniGuide can support conditioning both on geometric conditions (scaffolds, proteins or shapes) and global graph properties (quantum properties), purely during inference. However, this requires additional property regressors [2].

• Motivation for new drug discovery strategies: Inspired by the performance on the LBDD task, we further motivate the applicability of UniGuide for the generation of molecules given molecular densities, see App.G and our detailed answer to Reviewer D4nw. We further provide a qualitative example in Fig.1 in the attached PDF.

UniGuide is Not a "Direct Application": Consequently, we can confidently state that our work does not merely constitute a "direct application" of gradient-based guidance. UniGuide introduces fundamental innovations that enable geometric conditioning, unlocking novel applications and expanding the capabilities of controlled molecule generation.

The important related works are missing 

• The reason why we omitted the discussion on property-based conditioning is that the current self-guidance formulation of UniGuide does not support drug discovery tasks involving global graph properties. However, it is possible to combine UniGuide with classifier-guidance to enhance the generated molecules' quality for different downstream drug discovery applications.
• We acknowledge that the omission of this discussion may have led to some confusion. We will incorporate the suggested prior works [1, 2] into the related works section and describe how they differ from UniGuide.

 

We hope that the we were able to appropriately adress reviewer's concerns and are looking forward to a constructive discussion.

 

[1] Bao et al. Equivariant energy-guided sde for inverse molecular design

[2] Han et al. Training-free Multi-objective Diffusion Model for 3D Molecule Generation

[3] Hoogeboom et al. Equivariant diffusion for molecule generation in 3d

[4] Dhariwal et al. Diffusion models beat gans on image synthesis

[5] Song et al. Loss-guided diffusion models for plug-and-play controllable generation

We thank the reviewer for the feedback and positive evaluation of our method! We provide detailed answers in the following. 

 

Additional details on training required for UniGuide’s generality

• We highlight that UniGuide does not require extra training as it controls the generation of pretrained unconditional diffusion models during inference using guidance. 
• Therefore, Uniguide can be directly applied to a suitable (unconditional) model, i.e. a model that matches the configuration space required for the task and is trained on a fitting dataset (e.g. ZINC). We discussed this in the limitations, but we will ensure that this aspect is clearly conveyed in the updated manuscript.
• Drug design tasks have varying requirements. SBDD and FBDD need protein information, while LBDD and linker design do not. Therefore, the underlying model to which we apply UniGuide varies with the required configuration: We use the unconditional protein-ligand model from [1] for SBDD and FBDD and the molecular generative models ShapeMol or EDM for LBDD and linker design.
• In summary, UniGuide’s benefit is additive: It does not require a specialised model (no extra training), and the same model can be applied to multiple tasks.

Comparison with simpler conditioning mechanisms? 

• Controlling the generation of diffusion models using guidance in combination with additional classifiers/regressors is common for the image and text domains [7]. In contrast, UniGuide is a self-guiding [6] approach (Eq.12) that does not require additional models to control the generation, unlike [7]. Another key difference is our focus on the molecular domain, enabling unified guidance while respecting the requirements of molecular structures: (1) geometric conditions and (2) equivariance.
• The condition map is the central element in achieving this. It maps from the space of source conditions S to the configuration space Z (the diffusion model’s output space), realising (1), and maintaining equivariant updates (2) when satisfying Theorem 4.1.
• We compare to (i) conditionally trained diffusion models and (ii) an inpainting-inspired [1] technique for SBDD and FBDD and show that UniGuide is generally favourable, see e.g. Tab.2. 
• For LBDD, UniGuide is the first method to successfully tackle this task purely at inference time using the surface condition map.

Q: Computational complexity? 

• Inference: 
  • The computation cost associated with UniGuide results from computing the guidance gradients at inference time, see Eq.15.  We report runtime comparisons in App.E.2. UniGuide is slightly slower than the inpainting-inspired method [1]. At the same time, the DiffSBDD-cond model provides a baseline for pure sampling without guidance as the computational cost for controlled generation was already invested upfront for the conditional training (which is much higher).
• Training:
  • Details on the unconditional training of EDM and ShapeMol[U] are provided in App.C and App.D.1, respectively. For the SBDD experiments, we use DiffSBDD checkpoints, as provided by the authors [1]; see App.E.2.

Q:Relation with latent diffusion models? Could one encode geometric conditions using a latent encoder?

• UniGuide can readily be applied to guide a latent diffusion model, for example, GeoLDM [4] instead of EDM [5].
• In the special case when the source condition is a molecular structure ($S=Z$), one could use the encoder of the latent diffusion model to map the condition to the latent space and guide with UniGuide there.
• In the general case of $S \neq Z$, a possible solution is to leverage the decoder D of the respective latent diffusion model and keep the condition map in the data space. That is, instead of computing $C(s, \hat z_t)$, one estimates the clean latent code and applies the decoder afterwards: $C(s, D(\hat z_t))$.

Q: Does UniGuide extend to other tasks? How does it compare to [2-3]?

UniGuide also applies to the domains presented in [2-3]. We did not include them in the current version of the manuscript as they focus on docking and protein generation only, while we focus on drug discovery centred around small molecules.

• However, we checked the mentioned papers and would like to share some ideas and results on how the presented tasks can be accomplished with UniGuide:
  • In the context of docking [2], UniGuide could be used to condition on information about the protein complex. Contact information, for example, could be leveraged by guiding to specific positions of the C-a atoms. Coarser positional information, cf. Fig.G.1 [2], could be done in a similar fashion to the shape-based generation presented for LBDD.
  • Compared with the settings from [3], UniGuide can be used as an alternative to generate symmetric proteins. We adopted this idea and chose the C8 symmetry for this PoC. In particular, we designed a pie-like volume element that resembles the required angles and guided the protein generation with the surface condition map from Eq.21. We present the results of this experiment in Fig.2 in the pdf.
• Additionally, as prompted by reviewer D4nw, we expand the exploration of Density-based Drug Design in Fig. 1 of the PDF.

 

We hope that our answers sufficiently clarify the raised questions. We are thankful for the reviewer’s suggestion to investigate the protein settings and are excited about our results for generating symmetric proteins. 

 

[1] Schneuing et al. Structure-based Drug Design with Equivariant Diffusion Models

[2] McPartlon et al. LATENTDOCK: Protein-Protein Docking with Latent Diffusion

[3] Watson et al. De novo design of protein structures and function with RFdiffusion

[4] Xu et al. Geometric latent diffusion models for 3d molecule generation

[5] Hoogeboom et al. Equivariant diffusion for molecule generation in 3d

[6] Song et al. Loss-guided diffusion models for plug-and-play controllable generation

[7] Bansal et al. Universal guidance for diffusion models

Follow-Up on Comparison of UniGuide with validity guidance

 

We want to follow up on the “Validity Guidance” discussion (ZcJd), as mentioned in [1]. We have conducted experiments using the surface loss computation and self-guidance for the LBDD task. We selected the hyperparameters $\sigma$ and $\gamma$ in the surface loss computation to achieve a high DICE score between the implicitly defined surface ($S=\gamma$) and the meshes UniGuide utilises for LBDD ($\sigma=1$, $\gamma=2$, DICE > 0.8). Our surface calculations are performed using the Open Drug Discovery Toolkit (ODDT), which assigns specific radii to individual atom types and employs the marching cubes algorithm to generate meshes [2].

We performed several runs around the above hyperparameter specifications. The runs were similar in performance and we report the best result below:

Although Validity Guidance for LBDD yields low graph similarity, the shape similarity remains suboptimal compared to UniGuide. Additionally, we frequently encounter numerical instability for this guidance term, an issue not present with UniGuide’s formulation of LBDD.

We believe this ablation will provide a valuable comparison for the final version of our manuscript and hope we have addressed the reviewer's remaining concerns.

 

[1] Guan, J., Zhou, X., Yang, Y., Bao, Y., Peng, J., Ma, J., ... & Gu, Q. (2023). DecompDiff: Diffusion Models with Decomposed Priors for Structure-Based Drug Design. In International Conference on Machine Learning.

[2] Lorensen, William E., and Harvey E. Cline. Marching cubes: A high resolution 3D surface construction algorithm, 1998

   

Follow-Up on SBDD

 

We want to follow up on the discussion regarding baselines for SBDD (ZcJd & YA7u), specifically IPDiff [1] and DecompDiff [2]. As previously mentioned, incorporating a form of "Clash Drift" into UniGuide, particularly within the condition map, is straightforward and helps prevent the generated ligand from colliding with the pocket. We have implemented this adjustment, and our initial results indicate that similar to IPDiff and DecompDiff, it improves the overall VINA scores. Specifically, UniGuide's VINA score, compared to the results reported in Table 1 of the rebuttal PDF, improved significantly from −5.074 to -5.81. These initial results are promising and we are confident that adjustments of UniGuide's guidance parameters will lead to further improvments.

We again appreciate the reviewers' suggestions to include [1, 2] and believe that these additional results underscore UniGuide's flexibility to readily incorporate various conditions, enabling the comparison of underlying models across different levels. We hope that these findings address all the reviewers' remaining concerns, and we are looking forward to discussing these additional results.

 

[1] Huang et al. Protein-ligand interaction prior for binding-aware 3d molecule diffusion models

[2] Guan et al. DecompDiff: diffusion models with decomposed priors for SBDD