Conditional Synthesis of 3D Molecules with Time Correction Sampler

Summary

This paper focuses on diffusion-model-based molecular inverse design (i.e., conditional molecular generation), and proposes a novel approach to address the inconsistency between target distribution and that after online guidance. Specifically, a time predictor is trained to predict the time of the manifold that the sample lies on, and then a correction is done to ensure that the generated molecules remain on the correct manifold by Tweedie's formula. The proposed method shows satisfactory results on the QM9 dataset.

Strengths

- The motivation is clear and the idea is novel. The inconsistency is indeed a challenge in online guidance of diffusion models. TCS and TCAS solve this problem elegantly. - The authors have introduced extensive related work to help the audience to better understand the literature. - The authors provided necessary ablation studies in the appendix.

Weaknesses

- The experiments are not sufficient. More experiments (like designing molecules with given substructures in EEGSDE). - The idea of TCS is good, and an experiment beyond molecules may further demonstrate the generalizability of the proposed approach. For example, a comparison with [1] is recommended, which the authors have also mentioned when introducing the exposure bias in diffusion models. Besides, I think the accuracy of the time predictor is highly dependent on the data itself. It is suspected that this approach may not work in other fields. - The presentation of the paper needs to be improved. There are some obvious typos (e.g., line 63). This is just a kind reminder and does not impact my rating of this paper. [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.

Questions

See the above weaknesses. Besides, - as this paper focuses on the online guidance of the diffusion models for continuous variables. What about the version of TCS and TACS for the discrete variables, if applicable? - EEGSDE and some other related works that optimize the properties of generated molecules (e.g., [1]) cannot optimize the number of atoms, which is also important in molecule design. Do the authors have some idea about this? (This is only a related question and does not impact my rating.) [1] 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).

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have discussed the limitations in terms of how the signal of guidance is computed. However, I wonder whether the use of computational quantum chemistry methods works as the authors mentioned. My concerns lie on the differentiability required by guidance.

Summary

This study utilizes a predicted time estimator to correct the data manifold during the guided generation of diffusion models for molecules, to mitigate the discrepancy between the forward and reverse distribution. The authors show that by adjusting the noised sample according to predicted time, as opposed to relying on the pre-defined time schedule, the guided generative process remains on the correct manifold, leading to higher generation quality for several conditional generation tasks.

Strengths

This study offers useful insights into reducing the exposure bias of diffusion models using time correction. It also performs comprehensive analysis and ablation studies on the proposed framework. The results could help future works on diffusion models for 3D molecules in enhancing the generative quality with property guidance. In addition, the authors demonstrate the possibility of using quantum computing (instead of data-driven classifier, which is the common practice) for the online guidance.

Weaknesses

Though the proposed method and most results are solid, some important results and experimental details seem to be missing or incorrect. See “Questions”. I'm willing to adjust my score if the questions are properly addressed.

Questions

Major: 1\. Line 173: the value of time window size is an important hyperparameter, but the value is not provided. Furthermore, the impact of the time window size choice should be evaluated in the ablation. 2\. Appendix C seems incomplete. How is the function incorporated into the guidance? 3\. Some highlighted values in Table 1 are not the best value, and some values seem abnormally high or low. Please confirm. 4\. A previous work [1] directly searches the time window for a better match of the time-dependent variance. How does the proposed method compare to it? [1] Li, Mingxiao, et al. "Alleviating exposure bias in diffusion models through sampling with shifted time steps." arXiv preprint arXiv:2305.15583 (2023). Minor: 1\. Several references to the Appendix needs to be fixed, e.g.: Line 264 should be B.4; Line 149: the said comparison is not in the Appendix. Line 248: the reference is broken and the said result is not in the Appendix. 2\. According to Fig 4, the corrected time becomes very close to the actual time after 400. Appendix B.4 shows the molecule generation quality is the highest if the time correction and OG starts at 400. Is there any possible relation between these observations?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have properly addressed the limitations.

Summary

This paper presents a framework for generating 3D molecules called Time-Aware Conditional Synthesis TACS. The proposed approach uses conditional generation with adaptively controlled plug-and-play online guidance into a diffusion model to drive samples toward the desired properties while maintaining validity and stability. To prevent generated samples deviating from the data distribution during the conditional generation process authors introduce a Time Correction Sampler to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion process. Authors compare their TACS results with Equivariant Diffusion Models EDM and Equivariant Energy Guided Stochastic Differential Equations EEGSDE.

Strengths

This paper presents a framework for generating 3D molecules called Time-Aware Conditional Synthesis TACS. The proposed approach uses conditional generation with adaptively controlled plug-and-play online guidance into a diffusion model to drive samples toward the desired properties while maintaining validity and stability. To prevent generated samples deviating from the data distribution during the conditional generation process authors introduce a Time Correction Sampler to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion process. Authors compare their TACS results with Equivariant Diffusion Models EDM and Equivariant Energy Guided Stochastic Differential Equations EEGSDE.

Weaknesses

How to efficiently use the Time Corrected Sampler and whether this method improves the performance in other domains such as in image generation.

Questions

Have the authors experimented TACS approach on other datasets?

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

In this study, authors have used a trained neural network to estimate chemical properties of each molecule. Using an exact computational chemistry-based method could improve the guidance.

Summary

This paper proposes Time-Aware Conditional Synthesis (TACS), a method that aims to improve the robustness of property-conditioned diffusion models for 3D molecule generation. The key idea is to mitigate the exposure bias of the conditional denoising process by training a time prediction model that matches samples to the most likely marginal distribution of the forward process before applying online guidance via Tweedie's formula. Experiments on a synthetic dataset and QM9 show that TACS generates valid samples that match the conditioning label more closely than alternative methods.

Strengths

* The paper introduces a promising approach to keep the generated samples aligned with the marginal distributions of the forward process and address the problem of exposure drift in conditional molecule generation. * The paper demonstrates that TACS performs better than several well-established baselines on the QM9 dataset, showing improvements in generating molecules with desired quantum properties while maintaining stability and validity. * The paper is well-written and clearly outlines the methodology and motivation behind the method.

Weaknesses

* The method is only compared to well-established baselines on a single dataset (QM9). It would be good to evaluate the model on at least one other benchmark, to ensure that the results indicate a general trend and are not specific to the very distinct data distribution of QM9. * None of the quantitative empirical results in Section 5 include error bars or measures of statistical significance. * The main text contains multiple references to Appendix A.1, which I assume are incorrect links. I could not find the comparison with relevant work in [2] (referenced in line 149) or any details on the MAE distribution of samples below and above an 80% stability threshold (referenced in lines 227-229). Furthermore, the model performance analysis for $m>1$ MCMC samples referred to in line 248 seems to be missing from the Appendix. * Minor Point: The results in Table 1 and Figure 4 show that online guidance is often able to generate samples with much lower property MAEs at the expense of stability and validity. This tradeoff is discussed in the Ablation Studies paragraph, but it would be good to calibrate the claim that TACS "outperforms competitive baseline methods across all metrics", since online guidance refers to a published baseline model [1].

Questions

* Do the stability values reported in Table 1 refer to atom or molecule stability? * The online guidance results in Table 1 are sometimes much worse than those reported in [1], especially for $\mu(D)$, $\epsilon_\text{HOMO}$ and $\epsilon_\text{HOMO}$. Do you know what could cause that? * The paper mentions having to use a classification rather than a regression model to estimate the time step $t_\text{pred}$ because $p(t\vert\mathbf{x})$ cannot be estimated from a point estimate (line 147). However, Algorithm 1 then discards the information about the full distribution by taking the maximum likelihood estimate (I assume that's what $\operatorname{argmax}\phi(\mathbf{x}')$ in line 6 means). Would it be better to use $\mathbb{E}_{p(t\vert\mathbf{x}')}[t]$ instead? * Is the time predictor is only trained on 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] Kim, Beomsu, and Jong Chul Ye. "Denoising mcmc for accelerating diffusion-based generative models." arXiv preprint arXiv:2209.14593 (2022).

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors list the reliance on a potentially flawed predictive model for sample ranking as the main limitation. It would be good to also discuss any limitations of the method itself.