Conditional Synthesis of 3D Molecules with Time Correction Sampler

Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that

• Clear motivation and a novel idea to the challenging problem.
• Extensive related works for the audience.
• Necessary ablation studies.

We initially address your concerns below.

---

W1 Additional experiments on another task.

Thank you for the suggestion. In Table. R2 of the uploaded pdf, we provide the results of our methods (TCS/TACS) in target substructure generation following the setting of EEGSDE and in another dataset GEOM-DRUG (Table. R3). TACS outperforms EEGSDE in similarity scores while maintaining high stability, and achieving higher stability and validity outperforming baselines even in a more complex dataset (GEOM-DRUG). For a detailed analysis of these results, please refer to the discussion in [GR1].

---

W2 Generalizability beyond molecules / comparison with [1]

We appreciate your insightful suggestion. As detailed in [GR1], we have applied TCS to image generation on CIFAR-10 dataset and our method results in improved FID scores. Our results indicate that TCS effectively mitigates exposure bias across different domains. For instance, on CIFAR-10 dataset, applying TCS only improved the FID score from 3.353 to 2.663, while ADM-ES+TCS further reduced it to 2.179. This consistent improvement across molecular and image domains suggests that our approach addresses fundamental aspects of exposure bias in diffusion models, rather than being data-specific. Methodologically, our approach differs from Epsilon Scaling [1] in that we focus on correcting the timesteps based on predicted variance by time predictor, while [1] scales the network output directly. This allows for more precise adjustments at each step, potentially leading to better quality outputs across various domains. For a more detailed analysis of these results and their implications, please refer to [GR1].

---

W3 Correction

We sincerely appreciate your careful reading and for bringing this to our attention. We will thoroughly review and correct any typographical errors in the final version of our manuscript.

---

Q1 Online guidance for discrete variables

While our work focus on 3D molecules with continuous variables, adapting TACS for discrete variables using zero-order gradient estimation—similar to our approach with synthetic data and VQE—presents an intriguing direction for future research. We will include a discussion on this potential extension in the limitations and future work section of the manuscript.

---

Q2 Optimizing number of atoms

Exploring variable atom numbers is indeed a promising direction for future research, which we would like to mention in the final manuscript. TACS could potentially be extended to handle variable atom numbers, possibly by incorporating a learnable atom number prediction step in the generation process.

[1] Ning, M., Li, M., Su, J., Salah, A.A. and Ertugrul, I.O., 2023. Elucidating the exposure bias in diffusion models. arXiv preprint arXiv:2308.15321. [2] 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).

Dear reviewer ByqK

We are pleased if our response has addressed all of your concerns. If you have any further questions, please let us know.

Your careful review and insightful comments on our paper are greatly appreciated. Thank you once again for putting in your valuable time and effort.

Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that

• Offers useful insights in reducing the exposure bias of diffusion models.
• Comprehensive analysis and ablation studies which could help future works.
• Demonstrates the possibility of quantum computing.

We initially address your concerns below.

---

Q1 Size of Time Window

We appreciate your feedback and we agree that sensitive analysis on time window size can be an important point. We used $\Delta = 10$ for all of the experimental results reported in our original manuscript and the value is chosen from selecting the best performed MAE values. Here, we provide the experimental results for target property $\alpha$ in Table R6 of the uploaded pdf. The results indicate that performance is relatively stable for $\Delta$ between 2 and 16.

We appreciate this constructive suggestion and we will include the detailed ablation study on window size in our final manuscript for a clearer understanding of our method's behavior.

---

Q2 Online guidance using VQE.

In each denoising step of diffusion model, we first apply Tweedie’s formula to estimate clean molecule and use VQE calculates ground state energy of this predicted molecule by updating $\theta$ for $E(x,\theta) = \langle\psi(\theta)|H(x)|\psi(\theta)\rangle$. Then we use the zeroth-order method with respect to the position of atoms $x$ to obtain gradient $\nabla_x E(x, \theta)$. This gradient is used as an online guidance for our synthetic experiment. We will add this information in the Appendix in our final version of the paper.

---

Q3 Table 1 correction.

Thank you for the feedback. Our original intention was to highlight the best performance in conditioning (MAE) with molecular stability above 80%, aligning with our research goal of achieving desired quantum properties while maintaining molecular stability and validity. This is because molecules with lower stability are unlikely to exist in the real world (please refer to Figure 3, second row, in our manuscript for an example). However, we admit that this can cause the confusion of the readers and thus decided to change the notations accordingly.

In the updated table, we use bold for the best overall performance and color highlighting for the best performance with stability above 80%. We believe that this better represents the trade-offs in our method and baselines, while emphasizing the importance of molecular stability in practical applications.

---

Q4 Comparison with Time Shift Sampler in [1]

Regarding the comparison with [1], we would like to clarify the key differences between TS [1] and our TCS / TACS:

• Our method directly selects timesteps based on a trained time predictor, which has demonstrated robustness in our experiments. In contrast, [1] selects timesteps by calculating variance of image pixels at each step and matching the noise level from the predefined noise schedule.

• In [1], the corrected timestep is used directly for the start of the next step, while our approach maintains the predefined timestep after accurately estimate clean sample by corrected time step. By taking every single diffusion step while carefully using predicted time, TCS can generate samples closer to the target distribution.

To further validate our approach, we provide additional experimental results comparing TS-DDPM [1] and TACS on the QM9 dataset with step size 10 In Table A.3 below. The consistent improvements across various quantum properties underscore the robustness of our approach.

[Table A.3] Comparison between TS ([1]) and TCS / TACS

We will include this comparison in our final manuscript..

---

Q5 Addition and corrections of Appendix.

We appreciate your careful review, which has helped improve the clarity and completeness of our paper. In our final version, we will change the links and add information accordingly. Please refer to [GR2] for more details on these additions.

---

Q6 Possible relationship between start of OG and Figure 4.

Thank you for the insightful comment. We also think there might be a connection between starting time of online guidance and the convergence of predicted time to the ground truth values as in Fig.4 of the main paper. We suspect this happens because applying online guidance can be especially effective after the sample converges to some degree as pointed out in [2]. Investigating the exact relationship between this convergence and where to start OG would be an interesting future direction.

[1] Li et al., Alleviating exposure bias in diffusion models through sampling with shifted time steps, arXiv 2023 [2] Han, Xu, et al. "Training-free Multi-objective Diffusion Model for 3D Molecule Generation." The Twelfth International Conference on Learning Representations. 2023.

Dear Reviewer 1ecU,

We are glad to hear that our response properly addressed your concerns. Your careful review and insightful comments on our paper are greatly appreciated. Thank you once again for putting in your valuable time and effort.

Thank you for the constructive and helpful comments. We have addressed your comments below.

W1 Generalizability of TCS

As detailed in [GR1], we have applied TCS to image generation using the CIFAR-10 dataset, demonstrating significant improvements in FID scores. These results show that our method generalizes well beyond molecular data. We believe that our algorithm can be applicable in other domains as well since the design of TCS and TACS does not include any domain-specific knowledge. It would be an excellent direction for future work to investigate our method in various domains.

Q1 Experiments on other datasets As mentioned in [GR1], we have conducted experiments on the Geom-Drug dataset for molecule generation. Our method shows improved performance in generating molecules with molecular stability compared to the previous baselines like in EEGSDE. This further demonstrates TACS's ability to generalize to more complex molecular datasets and tasks. We will add these results to the final version of our manuscript.

Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that our work

• Introduces a promising approach to the exposure bias
• Show improvements compared to baselines with desired quantum properties while maintaining stability and validity.
• Well-written and clearly outlines the methodology and motivation.

We initially address your concerns below.

---

W1 Evaluation on other benchmarks

As detailed in General Response-[GR1], we provide additional experiments on Geom-Drug and CIFAR-10 which demonstrate TACS's ability to generalize beyond QM9, addressing molecules with more complex structures (Geom-Drug) and even extending to image generation (CIFAR-10).

[Table A.1] TCS on CIFAR 10

---

W2 Error bars and statistical significance

Thank you for pointing out the missing error bars. We have included them in Table.R1 of the uploaded pdf and will add in our final manuscript.

---

W3 Missing information in Appendix A.

We appreciate your careful review. In our final version, we will expand on several points as follows:

R3-1 Comparison with DMCMC [3] DMCMC uses a classifier to predict noise levels for MCMC on the product space of data. In contrast, our time predictor directly predicts timesteps for correction in diffusion sampling itself. Moreover, while the purpose of noise prediction in [3] is for fast sampling, our work use time predictor to accurately produce samples from the desired data distribution.

R3-2 Impact of MC numbers We will also include detailed analysis on MC numbers (line 248), demonstrating TACS's robust performance across different sample sizes. We will also include our analysis of experiments conducted on the effect of MC sample numbers. Unlike [1], where increasing samples from 5 to 10 more than doubled performance at increased computational cost, our method shows robust performance across different sample numbers.

For details on these additions, we kindly refer to [GR2].

R3-3 MAE distribution given the threshold. We will also clarify our discussion on MAE and molecular stability (MS) thresholds. The revised text will emphasize TACS's lower MAE compared to methods achieving MS above 80%, and its performance even for samples below this threshold.

---

W4 Calibrate the claim that “TACS outperforms competitive baseline methods across all metrics", since online guidance refers to a published baseline model [1].

We will revise our claim as "TACS outperforms competitive baselines when considering MAE, stability, and validity collectively," which more accurately reflects our results. The updated Table. R1 of the uploaded pdf in the supplementary material supports this claim.

---

Q1 Stability values reported in Table 1.

Stability values denote the molecule stability.

Q2 Online guidance results in Table 1.

Online guidance yields higher MAE for some properties such as $\mu$ compared to that reported in [1], and we suspect this stems from the sensitivity of online guidance strength.

First, data distribution of $\mu$ in the QM9 training dataset is sharp as can be seen in Figure.6-(a) in [2] the distribution of $\mu$.

This in turn, can make online guidance sensitive to its guidance strength, where we sometimes observe diverging gradients when naively applying online guidance.

Additionally, we provide additional experimental results for impact of online guidance strength on $\mu$ to support our claim.

We will include the above result in our final version of the paper for the better understanding of the reader.

Q3 How to obtain predicted time from the time predictor

Thank you for an interesting question. We provide experimental results on expectation-based $\mathbb{E}[{\phi(\boldsymbol{x})}]$ and argmax-based time prediction in the table below. The result shows that the performance remains robust in general with some trade-offs, where argmax-based prediction achieves slightly better MAE while the expectation-based approach maintains higher molecular stability.

[Table A.2] Different type of time prediction. Above is MAE for target condition and below is molecular stability of generated samples.

---

Q4 Is the time predictor only trained on samples from the forward process?

Yes, we trained the time predictor with the samples with the forward process as described in Line 142. Our work used a time predictor for reducing the gap between the marginal distribution of the forward and reverse processes, and thereby it provides good guidance to the desired data distribution only when trained with noisy samples from the forward process.

---

[1] Han, Xu, et al. "Training-free Multi-objective Diffusion Model for 3D Molecule Generation." The Twelfth International Conference on Learning Representations. 2023. [2] Bao, Fan, et al. "Equivariant energy-guided sde for inverse molecular design." The eleventh international conference on learning representations. 2022. [3] Kim, Beomsu, and Jong Chul Ye. "Denoising mcmc for accelerating diffusion-based generative models." arXiv preprint arXiv:2209.14593 (2022).