TurboHopp: Accelerated Molecule Scaffold Hopping with Consistency Models

First, we thank the reviewer for the great insights for authors to consider. As stated in your limitations, we acknowledged that comparison with de-novo generative molecules were important. However, since direct comparison would be unfair, we had to build repurposed versions of de-novo models which we plan to additionally release. Please refer to the global rebuttal section for further details.

Q1. Regarding inference speed and generation quality (Table 1 main, Table 3,5,6,7,8,9 as additional reference)

A1. In order to explain the relationship between speed and quality, authors extended previous Table 1 that compares DiffHopp, TurboHopp and variations. As shown in Table 1, fewer steps in TurboHopp leads to overall poor quality compared to versions with more steps. Among the 3 variants of TurboHopp, TurboHopp100 seemed to have the best efficiency regarding sample quality as well as generation time. Table 1 as well as Table 3 show that our model is efficient with comparable metrics compared to diffusion-based models despite the reduce in time . However, in Table 5,6,7,8,9, where we compared geometric properties with existing conditional diffusion models, we see some geometric properties remain a property to improve. In future research if we train the model to learn bond distributions, this may resolve this issue.

Q2. Novelty issue

A2. We apologize if we confused the reviewer. Novelty is capped at one and the digits after the plus-minus sign is the standard deviation.

Q3. Consistency models and improvements

A3. We did not train using consistency distillation but trained a consistency model from scratch using consistency training [3]. We referred to improved techniques used for consistency training as well and found that MSE loss instead of psuedo-huber loss is better and that training without skip connection in the model is more stable [4] . In addition, we had to change the original consistency model so that it is suitable for conditional molecule generation. It had to take account conditions(protein, functional groups etc.) as well as expand the model to be multimodal and SE(3) equivariant. Overall,we believe our framework can be broadly adapted to many of the molecule generative diffusion models existing and hopefuily improve sampling efficiency.

Q4. Advantages over de-novo conditional generative diffusion models(Algorithm 1, Table 3)

A4. Scaffold-Hopping aims at finding novel scaffolds that connect key functional groups (interacting with protein residues) while de-novo generative models focus on building the whole molecule given a protein pocket. Both models are, SBDD models, and they both aim on building molecules that have great potential in finding novel, potent molecules. However, because of the lack of 3D-binding complex data, it is a realistically hard task for even advanced models to learn the vast chemical space of potent molecules. Therefore, by conditioning the functional groups with potential interactions, we can reduce the chemical space the model should learn, which leads to more efficient learning and sampling(Figure 1, paper). This is additionally powered by the speed boost of our consistency model in allowing 1) faster, efficient generation which is often requested in the pharmaceutical domain and 2) faster optimization in certain molecular metrics which opens possibilities of applications in human-in-the-loop model optimization.

Model structure-wise, straight comparison between de-novo SBDD models and scaffold-hopping models is unfair since for input conditions, scaffold-hopping models have an extra ligand condition to consider during generation. This may be more harder for the model to learn since there's an extra constraint, but on the flip-side also mean less chemical space for the model to learn. We do, however, compare inpainting variations of de-novo conditional generative diffusion models which we have additionally built for [1] and [2]. Scaffold hopping can be seen as an inpainting task, where models trained on generating de-novo molecules treat the scaffold as a missing piece that needs to be filled in, while the known functional groups of the molecule are provided as context. In a gist, as shown in Table 3, we have better outcomes in overall molecular metrics. Also, inpainting variations of SBDD models are too slow to practically use compared to our model. For more information, please refer to the global rebuttal.

Q5. Questions regarding RL performance(Table 1,2)

A5. The authors agree that there is a significant impact of reinforcement learning. However, we do not believe that surpassing DiffHopp primarily relies on RL. In particular, by leveraging the consistency model’s ability to give a number of high quality predictions within significantly less steps, we can use metric sampling to achieve competitive results(Table 1). However, we also showed that the inclusion of RL provided significant gains over diffusion based models (Table 2). To counteract reward hacking in standard docking tasks, we followed the methods of [5], using a combination of QED, synthesizability and docking scores to mitigate this issue (Equation 11). In future research, we plan to create a more robust reward function aligned to better tackle reward hacking further. While it is true that to some extent we are overoptimizing for these metrics, we found that many metrics excluded from the reward function were able to maintain quality. This hints that the molecules were not being overoptimized. See the response for reviewer Cq9Y for more information

[3] Song, Yang, et al. "Consistency models." arXiv preprint arXiv:2303.01469 (2023).

[4] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models." arXiv preprint arXiv:2310.14189 (2023).

[5] Ghugare, Raj, et al. "Searching for high-value molecules using reinforcement learning and transformers." arXiv preprint arXiv:2310.02902 (2023).

We first thank the reviewer for the constructive feedbacks regarding additional baselines/metrics for the authors to consider. Please refer to global rebuttal as well as PDF attached!

Q1. Other scaffold-hopping baselines (Table 3,4)

A1. Thank you for adding valuable information regarding baseline models. Regarding the first two models the reviewer mentioned: Scaffold-GVAE[1], GraphGMVAE[2] are both VAE models built for scaffold hopping. DecompOpt[3] is a recently published diffusion-based model for molecular optimization and has scaffold-hopping related experiments with results. Baselines regarding Decompopt are in the global rebuttal section.

1) Comparison with VAE-based models (Table 4)

Unfortunately, most VAE based scaffold-hopping models are either SMILEs-based which makes comparison with our model difficult. Also, the codes/data for most of them are missing. The authors of ScaffoldGVAE[1] provide a trained model, but the vocabulary necessary for tokenization, encoding, and decoding is missing. Consequently, we couldn't use the model with the provided checkpoint alone. To properly implement ScaffoldGVAE, we would need to generate a new vocabulary and retrain the model using the 1.9 million ChEMBL compounds mentioned in their paper. Given the limited time frame for the review process, we found it challenging to experiment with multiple baselines that require such extensive setup and training. To ensure fairness in terms of training data, we decided to train ScaffoldGVAE on the PDBbind dataset, which is the same dataset used for our model. This approach allowed us to evaluate ScaffoldGVAE under comparable conditions to our model, maintaining consistency in the data used across experiments. Our approach involved pretraining on the training data followed by target-specific fine-tuning for evaluation on the test set. The results are as follows:

To define connectivity for 2D molecules, we considered a molecule to have connectivity if the generated scaffold could be successfully substituted for the original core and form bonds with the remaining functional groups. Despite the high QED values, most scaffolds generated failed to connect with functional groups, leading to low validity and connectivity. Furthermore, most of the scaffolds had low similarity with the reference scaffold. This is presumably due to the fact that the SMILEs based models do not consider 3D pocket structure during generation. In addition most VAE based scaffold-hopping models require tailoring on specific targets. Our model's performance can be attributed to its ability to utilize 3D structural information of target proteins. This allows for generation of molecules even for pockets with few ligands known to bind .

2) Comparison with other SBDD diffusion models (including Decompopt) (Table 3)

Please refer to the global rebuttal section.

Q2. Related works regarding further research on RL-applied molecular science tasks. (Table 2)

A2 Authors thank the reviewer for highlighting these works which we will cite. We do note that these works are different from ours in that [4] proposes to use a critic function which we found to be superfluous for consistency models. Further, we highlight that they focus on a single reward metric, while we create an array of target metrics to ensure that no metrics are sacrificed. Also during a single optimization task, we optimize our model towards multiple targets while [4] focuses on single target. Moreover, we differ from [5] since [5] is related to molecule docking tasks while we are covering a different problem regarding molecule generation. Nevertheless, we agree these works are important related research to cite and will include these works.

Q3. Metric Based Sampling (Table 1)

A3 Metric-based sampling is a technique that selects the best result towards the end of the sample trajectory, rather than selecting the final output. Figure 8 in Appendix D of our paper illustrates an example between conventional end-of-trajectory sampling and metric-based sampling. We extended Table 1 to ablate the effectiveness of this method, which shows a slight improvement. This is further demonstrated by enhancements observed in DiffHopp (Table 1: DiffHopp_scored vs. DiffHopp), indicating that it can yield higher quality samples. However, in DiffHopp, it significantly reduces sampling efficiency, as it necessitates more frequent evaluations of the generated products towards the end(100 sec vs 440 sec)

Q4. Geometric property metrics evaluation (Table 5,6,7,8,9)

A4 Authors compared geometric properties including those related to bond (Please refer to Table 5,6,8,9) and ring distributions(Please refer to Table 7). Results show that our model has closer bond length/angle/atom-atom length distributions to the reference molecules compared to TargetDiff, but poorer results compared in bond/torsion angles. In all aspects, DecompDiff was outstanding largely because it learns the distribution of bonds. For ring distributions, results show that our model is capable of generating similar ring types compared to the reference. In future research, we plan to design our model to learn bond properties , and we expect better results regarding geometric properties.

Thank you for your feedback! We wish to clarify some details concerning our results. Due to the unavailability of Decompopt's code implementation, performance metrics come from published literature, precluding a direct computational comparison (the scaffold masking method might be different, number of optimization runs are not shown, docking methods may be different). Consequently, the values presented for Decompopt should be considered as reference values rather than for direct comparisons. Furthermore, our variant, TurboHopp 50-RL, not only achieved higher docking scores exceeding that of Decompopt and reference but also demonstrated a significantly faster generation speed compared to other SOTA models. We kindly request the reviewer to consider the efficiency of generation regarding inference time and quality!

We sincerely appreciate the reviewer's positive feedback on our model. Thank you for highlighting the typos and mathematical notations to fix; we will ensure they are corrected in the final draft.

Q1. Concerns regarding reward hacking (Table 2)

A1.

As the reviewer mentioned, authors agree that it is true that generators are prone to learning blind spots in the reward function, and this might indeed affect how useful it is in real-world properties. SMILEs-based RL algorithms exploit their state and action spaces to design chemically trivial molecules with exceptionally high docking scores. To counteract reward hacking in standard docking tasks, we follow the methods of [2], using a combination of QED, synthesizability and docking scores to mitigate this issue (Equation 11). In future research, we plan to create a more robust reward function aligned to better tackle reward hacking further. To this end, KL regularization to the original model usually fixes this problem (a common practice in LLM alignment with RLHF) and could be used if other metrics / real world fabrication showed issues. Other methods [3] have also been developed for these issues and likely would be extendable to consistency models. As shown below (please refer to Table 2 of PDF in global rebuttal), TurboHopp RL maintains other metrics other than docking score and also has high diversity compared to TurboHopp. Furthermore, TurboHoppRL trained on PDBBind test sets had competent generation quality on CrossDocked test sets, without training on CrossDocked test sets, exceeding the reference docking score, which indicates that our model generalizes well to new data. Contrary to other works related to RL for diffusion [4] which optimize towards a single protein pocket, this may be because we optimize on an array of target metrics to ensure that overfitting does not occur.

[2] Ghugare, Raj, et al. "Searching for high-value molecules using reinforcement learning and transformers." arXiv preprint arXiv:2310.02902 (2023).

[3] Uehara, Masatoshi, et al. "Fine-tuning of continuous-time diffusion models as entropy-regularized control." arXiv preprint arXiv:2402.15194 (2024).

[4] Zhou, Xiangxin, Liang Wang, and Yichi Zhou. "Stabilizing Policy Gradients for Stochastic Differential Equations via Consistency with Perturbation Process." arXiv preprint arXiv:2403.04154 (2024).

Thank you for answering the raised question.

I have read the other reviews and the authors' responses. On the one hand, the other reviewers seemed to raise valid points, e.g., a lack of baselines. On the other hand, the authors added information in this regard during the rebuttal. Assuming that the issues with respect to baselines and related work are solved (I hope reviewer VN4j will comment on this), I'd like to stick to my score because pocket-conditioned 3D molecular scaffold hopping is interesting to the community and the manuscript is of high quality.

Dear Reviewer Cq9Y,

Thank you for your positive feedback, your support and constructive engagement with our work!

Best regards,

Authors

For all reviewers, thank you for your time and effort for having interest in our work. We additionally add the results of TurboHoppRL-50 trained on CrossDocked.

Through RLCM(RL for consistency models) training, we were able to increase all metrics (Validity, Connectivity, Diversity, Novelty, QED, Docking Scores) except synthesizability exceeding those of the reference set especially for docking scores. Please refer to the novelty/diversity of generated molecules regarding model overfitting.

If you have any issues to discuss about please leave any comments!