Conditional Synthesis of 3D Molecules with Time Correction Sampler

Authors: hojung  jung
youngrok  park
laura  schmid
jaehyeong  jo
dongkyu  lee
bongsang  kim
se-young  yun
jinwoo  shin
kaist  ai
l g electronics
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:07:23

Abstract: Diffusion models have demonstrated remarkable success in various domains, including molecular generation. However, conditional molecular generation remains a fundamental challenge due to an intrinsic trade-off between targeting specific chemical properties and generating meaningful samples from the data distribution. In this work, we present Time-Aware Conditional Synthesis (TACS), a novel approach to conditional generation on diffusion models. It integrates adaptively controlled plug-and-play "online" guidance into a diffusion model, driving samples toward the desired properties while maintaining validity and stability. A key component of our algorithm is our new type of diffusion sampler, Time Correction Sampler (TCS), which is used to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion process at the same time. Our proposed method demonstrates significant performance in conditional 3D molecular generation and offers a promising approach towards inverse molecular design, potentially facilitating advancements in drug discovery, materials science, and other related fields.

Main paper: 
 introduction: 
 Discovering molecules with specific target properties is a fundamental challenge in modern chemistry, with significant implications for various domains such as drug discovery and materials science [49, 44, 7] . While diffusion models have shown great success in the generation of real-world molecules [55, 31] , their primary goal is often simply to generate realistic molecules without considering specific properties, which can lead to producing molecules with undesirable chemical properties.Existing works address this issue by leveraging controllable diffusion frameworks to generate molecules with desired properties [23, 5] . One approach is to use classifier guidance [54] , which utilizes auxiliary trained classifiers to guide the diffusion process [5] . An alternative is to use classifier-free guidance (CFG) [19] , which directly trains the diffusion models on condition-labeled data. While both approaches can generate stable molecules, they struggle to generate truly desirable molecules due to the complex structures and the discrete nature of atomic features [24] .On the other hand, recent works [51, 18] have introduced training-free guidance for controllable generation in a plug-and-play manner, which we hereafter refer to as online guidance (OG). This approach can directly estimate the conditional score with unconditional diffusion model. However, our analysis shows that applying online guidance into the molecular generation can result in generating samples with significantly low molecular stability and validity due to the stepwise enforcement of specific conditions without considering the original distribution at each timestep. . TACS helps generate highquality samples that match target condition while following basic properties of the molecules. At each timestep t, online guidance is applied to push x t towards the desired condition. Time Predictor finds the desired timestep t p for x t after applying the guidance. Using predicted timestep t p , Tweedie's formula is used to predict the clean molecule xt 0 . Finally, forward process q(x t-1 |x 0 ) is applied to proceed to the next denoising step of t -1. (b) Motivation for TACS. Applying online guidance (⃗ g, purple) can shift the generated samples away from the correct data manifold corresponding to the current timestep. This undesirable deviation (red) can be avoided by using time correction to first measure the deviated timestep t ′ , then adjusting the guidance to get corrected guidance vector ( ⃗ g * , green), which keeps the generated samples stay on the correct data manifold.To address this gap, we propose Time-Aware Conditional Synthesis (TACS), a novel framework for generating 3D molecules. TACS utilizes the online guidance in tandem with a novel diffusion sampling technique that we call Time Correction Sampler (TCS). TCS explicitly considers the possibility that online guidance can steer the generated sample away from the desired data manifold at each step of the diffusion model's denoising process, resulting in an 'effective timestep' that does not match the correct timestep during the generation. TCS then corrects this mismatch between the timesteps, thereby effectively preventing samples from deviating from the target distribution while ensuring they satisfy the desired conditions. By combining online guidance with TCS and integrating them into a diffusion model, TACS allows generated samples to strike a balance between approaching the target property and remaining faithful to the target distribution throughout the denoising process.To the best of our knowledge, TACS is the first diffusion framework that simultaneously addresses inverse molecular design and data consistency, two critical objectives that often conflict.We summarize our main contributions as follows:• We propose Time-Aware Conditional Synthesis (TACS), a new framework for diffusion model that utilizes adaptively corrected online guidance during the generation.• We introduce Time Correction Sampler (TCS), a novel diffusion sampling technique that ensures the generated samples remain faithful to the data distribution. It includes a time predictor, an equivariant graph neural network that accurately estimates the correct data manifold during inference by predicting the time information of the generation process.• Through extensive experiments on a 3D molecule dataset, we demonstrate that TACS outperforms previous state-of-the-art methods by producing samples that closely match the desired quantum chemical properties while maintaining data consistency. 
related works: 
 Diffusion models Diffusion models have achieved great success in a variety of domains, including generation of images [12, 46] , audio generation [30] , videos [21, 36, 42] , and point clouds [8, 37] . A particular highlight in the success of diffusion models is their potential to generate molecules that can form the basis of new, previously unseen, medical compounds. Multiple approaches have been explored to achieve this. For instance, graph diffusion such as GeoDiff [57] , GDSS [26] , and Di-Gress [54] can generate graph structures that correspond to molecular candidates. Additionally, some approaches incorporate chemical knowledge tailored to specific applications, such as RFdiffusion for protein design [55] , a method based on the RoseTTAFold structure prediction network [31] . Other previous literature considers different domain-specific applications, such as diffusion for molecular docking [11] , or molecular conformer generation [25] . Most relevant to our work, diffusion models have also shown promising results in synthesizing 3D molecules [23, 58, 27] , generating stable and valid 3D structures.Conditional molecular generation Deep generative models [38, 23, 58] have made considerable progress in synthesizing 3D molecules with specific properties. Specifically, conditional diffusion models [23, 5, 58, 18] have achieved noticeable improvements in synthesizing realistic molecules. EDM [23] trains separate conditional diffusion models for each type of chemical condition, while EEGSDE [5] trains an additional energy-based model to provide conditional guidance during the inference. GeoLDM [58] utilizes a latent diffusion model [46] to run the diffusion process in the latent space. MuDM [18] applies online guidance to simultaneously target multiple properties. However, existing methods either produce unstable and invalid molecular structures or are unable to accurately meet the target conditions. To overcome these limitations, we propose TACS, a novel framework which ensures the generative process remains faithful to the learned marginal distributions in each timestep while effectively guiding the samples to meet the desired quantum chemical properties. 
preliminaries: 
 Diffusion models Diffusion models [50, 20, 52] are a type of generative model that learn to reverse a multi-step forward noising process applied to the given data. In the forward process, noise is gradually injected into the ground truth data, x 0 ∼ p 0 , until it becomes perturbed into random noise, x T ∼ N (0, I), where T is the total number of diffusion steps. We follow the Variance Preserving stochastic differential equation (VP-SDE) [52, 20] where the forward process is modeled by the following SDE:EQUATIONwhere β(t) is a pre-defined noise schedule and w t is a standard Wiener process. Then, the reverse of the forward process in Eq. ( 1) is also a diffusion process that is modeled by the following SDE [2] :EQUATIONwhere p t is the probability density of x t and wt is a standard Wiener process with backward time flows. The reverse process in Eq. ( 2) can be used as a generative model when the score function ∇ x log p t (x) is known. To estimate the score function from given data, a neural network s θ is trained to minimize the following objective [52] :EQUATIONwhere t ∼ U [0, T ], x 0 is sampled from the data distribution, and λ : [0, T ] → R + is a positive weight function.Online guidance Recent works [17, 9, 51] leverage a conditional diffusion process where the conditional probability p t (x t |c) is modeled without training on labeled pairs (x t , c). To understand this approach, observe that for the conditional generation, the reverse SDE of Eq. ( 2) can be rewritten as follows:EQUATIONFrom Bayes' rule, the conditional score ∇ x log p t (x t |c) can be decomposed as follows:EQUATIONWhile unconditional score ∇ xt log p t (x t ) can be approximated by the unconditional diffusion model s θ from Eq. ( 3), online guidance can estimate conditional score part ∇ xt log p(c|x t ) without any training. Notably, DPS [9] approximates conditional score as follows:EQUATIONwhere x0 is a point estimation of the final clean data using Tweedie's formula [13] :EQUATIONWhile DPS uses a point estimate of x 0 , Loss Guided Diffusion (LGD) [51] uses Bayesian assumption where x 0 is a Monte Carlo estimation of q(x 0 |x t ), which is a normal distribution with mean x0 and variance σ 2 t which is a hyperparameter. For a given property estimator A satisfying c = A(x 0 ), Eq. ( 6) can be written as follows:EQUATION)where L is a differentiable loss function and x i 0 are independent variables sampled from q(x 0 |x t ). Using this approximation, we can replace ∇ xt log p(c|x t ) in Eq. ( 5) and conditional generation process now becomes:EQUATION)where z is the hyperparameter for controlling online guidance strength. In this work, we use L 2 loss for the differentiable loss L and set m = 1 following Song et al. [51] .In our framework, we follow EDM [23] in modeling the distribution of atomic coordinates x within a linear subspace X , where the center of mass is constrained to zero. This allows for translation-invariant modeling of molecular geometries. Further details of the zero-center-of-mass subspace is provided in Appendix A.3.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]Answer: [No] Justification: We provide experimental details and will provide code after it is polished.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
time-aware conditional synthesis: 
 In this section, we propose our framework, Time-Aware Conditional Synthesis (TACS). Section 4.1 presents the key component of TACS: the Time Correction Sampler, a novel sampling technique that leverages corrected time information during the generation process. Section 4.2 introduces the overall framework of TACS, which accurately integrates online guidance into the diffusion process using our sampling technique to generate stable and valid molecules that meet the target conditions.Diffusion models are known to have an inherent bias, referred to as exposure bias, where the marginal distributions of the forward process do not match the learned marginal distributions of the backward process [40] . To mitigate exposure bias in the diffusion models, we propose the Time Correction Sampler (TCS), which consists of two parts: time prediction and time correction.Time Predictor During conditional generation process of diffusion model, a sample follows the reverse SDE as in Eq. ( 4). However, due to error accumulation during the generation process [33, 6] , a sample at timestep t of the reverse process of diffusion may not accurately reflect the true marginal distribution at timestep t. This discrepancy between forward and reverse process can especially increased when applying online guidance for every denoising steps and consequently lead to the generation of molecules with low stability and validity. We aim to mitigate this issue by correcting the time information based on the sample's current position. end if 13: end for by equivariant graph neural network (EGNN) ϕ. Then, cross-entropy loss between the one hot embedding of timestep vector and the logit vector for the model output is used as follows:EQUATIONwhere timestep t is sampled from the uniform distribution U [0, T ], x 0 is chosen from the data distribution, x t is constructed from Eq. ( 1), and pϕ (x t ) t is the t-th component of the model output pϕ for a given input x t . We empirically evaluate the performance of the time predictor on the QM9 train and test datasets. Forward noise, corresponding to each true timestep, is added to the data, and the time predictor estimates the true timestep, with accuracy measured accordingly. As shown in 3b, the predictor struggles in the white noise region, but achieves near-perfect accuracy after timestep 400.Time correction TCS works by first modifying Tweedie's formula (Eq. 7) using the estimated timestep from time predictor to utilize the information of the proper timestep estimated by time predictor during the denoising diffusion process.Specifically, for a sample xt at time t during the reverse diffusion process, we use the corrected timestep t pred predicted by time predictor, instead of the current timestep t, to estimate the final sample x0 as follows:EQUATIONFrom the better prediction of the final sample x0 , we perturb it back to the next timestep t-1 using the forward process using Eq. (1).Intuitively, TCS encourages the generated samples to adhere more closely to the proper data manifolds at each denoising step. The iterative correction of the time through the generative process ensures the final sample lie on the correct data distribution. Consequently, for 3D molecular generation, TCS helps synthesizing molecules with higher molecular stability and validity, both of which are crucial for generating realistic and useful molecules.Finally, we present TACS in Algorithm 1, which is a novel diffusion sampler for conditional generation that integrates TCS with online guidance. For each timestep of the reverse diffusion process (Eq. 4), we apply online guidance (Eq. 8) to guide the process toward satisfying the target condition. Subsequently, TCS is applied to correct the deviation from the proper marginal distribution induced by online guidance. This ensures the samples follow the correct marginal distribution during the generative process, as shown in Fig. 1 . In Section 5, we experimentally validate that TACS is capable of generating valid and stable molecules satisfying the target condition. Time clipping During the generation, we empirically observe that the predicted timestep t pred often deviates significantly from the current time step t. Naively applying this information to TCS can be problematic as the dramatic changes in timestep may skip crucial steps, resulting in unstable or invalid molecules. To prevent large deviation of t pred , we set a time window so that t pred remains in the time interval [t -∆, t + ∆], where ∆ is a hyperparameter for the window size. 
experiments: 
 In this section, we present comprehensive experiments to evaluate the performance of TACS and demonstrate its effectiveness in generating 3D molecular structures with specific properties while maintaining stability and validity. In Section 5.1, we present synthetic experiment with H + 3 molecules, where the ground state energies are computed using the variational quantum eigensolver (VQE). In Section 5.2, we assess our method using QM9, a standard dataset in quantum chemistry that includes molecular properties and atom coordinates. We compare our approach against several state-of-the-art baselines and provide a detailed analysis of the results.Quantum online guidance We present quantum online guidance as a modified version of online guidance where quantum machine learning algorithm for the property estimation of generating molecules. Specifically, we use VQE [53] when calculating the ground state energy of generated molecules. Contrary to prior works [23, 5] , which train auxiliary classifiers to estimate each condition, this quantum computational chemistry-based approach can leverage exact calculation of the condition for a given estimate. This, in turn, is expected to generate molecules with accurate target ground state energies. A detailed explanation of this approach is provided in Appendix A.2.Setup We first construct synthetic H + 3 as follows. For each molecule, a hydrogen atom is placed uniformly at random within the 3-D unit sphere. We then augment our sample by rotating each molecule randomly to satisfy the equivariance property [23] . Then the ground state energy of each molecule is measured with VQE in order to provide conditional labels. Finally, we train unconditional diffusion model and CFG-based conditional diffusion model with the constructed data. For evaluation, we measure the ground state energy and calculate MAE (Mean Absolute Error) with the target energy. Also, we measure the average L2 distance between position of each atom and its projection to the unit sphere.The results in Figure 2 indicate that while online guidance correctly guides the samples to the target condition, it can lead to the samples deviated from the original data distribution. In contrast, molecules sampled from TACS can satisfy both low MAE and low L2 distance. Further details and additional discussions are provided in Appendix B.1. Dataset We evaluate our method on QM9 dataset [45] , which contains about 134k molecules with up to 9 heavy atoms of (C, N, O, F), each labeled with 12 quantum chemical properties. Following previous works [1, 23] Evaluation To evaluate how generate samples meet the desired target condition, a property prediction model ϕ p [47] is trained on D a . Then, MAE for K number of samples is calculated as1 K K i=1 |ϕ p (x (i) ) -c (i) |, where x (i) represents i-th generated molecule and c (i) is corresponding target quantum chemical properties. Molecular stability (MS) and validity (Valid) [23] are used to measure how generated samples satisfy basic chemical properties. Details of the evaluation metrics are provided in Appendix B.2.Baselines We use Equivariant Diffusion Models (EDM) [23] and Equivariant Energy Guided SDE (EEGSDE) [5] for the baselines. Following [23] , we put additional baselines including "Naive Upper-Bound" (randomly shuffled property labels), "#Atoms" (properties predicted by atom count), and "L-Bound" (lower bound on MAE using a separate prediction model).Results Table 1 shows the result of conditional generation of TACS and TCS with baselines. We generate K=10 4 samples for the evaluation in each experiment and the average values and standard deviations are reported across 5 runs. For all of the quantum chemical properties, TACS achieves lower MAE while maintaining comparable or higher molecular stability (MS) and validity compared to other baselines. Notably, when comparing with the baseline methods that maintains the MS above 80%, the MAE of TACS is significantly lower than the baselines. The result demonstrates the effectiveness of TACS in balancing the objectives of generating molecules with desired properties and ensuring their structural validity.TCS achieves high validity and atom stability and molecular stability, surpassing the unconditional generation performance of the baselines, but with a higher MAE compared to TACS. This highlights the importance of online guidance for precise property targeting. However, applying only online guidance yields samples with low MAE but suffers from reduced validity and stability, occasionally failing to generate valid molecules due to numerical instability. This shows the ability of TACS which places the samples on the correct data manifold at each denoising step, even if they deviate from the true time-manifold due to the online guidance. Finally, we put the results of different methods with 9 different runs for each method and plot the MAE and MS in Figure 3a . The result clearly shows that TACS and TCS reaches closer to the Pareto front of satisfying molecular stability and the target condition together.We conduct experiments on target structure generation with QM9 as in [5] . For evaluation, we report Tanimoto similarity score [16] which captures similarity of molecular structures by molecular fingerprint and molecular stability to check whether basic properties of molecules are satisfied during the conditional generation process. We put additional details of the experiment in Appendix B.3.2.Results Table 2 shows that TACS significantly outperforms baseline methods both in Tanimoto similarity and molecular stability. Interestingly, performance of TACS is robust in the online guidance strength z. This demonstrates TACS's ability to generalize on different tasks.To verify the scalability of time correction sampler, we use Geom-Drug [4] as our dataset. Geom-Drug consists of much larger and complicate molecules compared to the QM9. For fair comparisons, we follow [23, 5] to split the dataset for training, validation and test set include 554k, 70k, and 70k samples respectively. We test the performance of TCS on unconditional generation of 3D molecules when using Geom-Drug. For evaluation, we use atom stability (AS) and validity of generated samples. The result in Table 3 shows that samples generated by TCS satisfy the highest atom stability and the validity with high margin compared to the baselines. Details of the experiments are in Appendix B.3.1. Online guidance strength z To analyze the effect of online guidance strength z in Eq. ( 9), we measure MAE, molecular stability, and validity of samples generated by TACS for different z values. Table 5 shows the result with target condition on ϵ LUMO values. One can observe while trade-off occurs for different z values, performance of TACS is robust in varying z. Interestingly, our experiment shows that there exists an optimal value of z which generates samples with the lowest MAE that is even comparable to applying online guidance without time correction (OG). As expected, using z = 0 (TCS) generates molecules with the highest MS and validity but with the highest MAE. Time window length ∆ We measure how the performance of TACS varies with time winodw length ∆. Table 6 shows the MAE, molecular stability values for different window sizes when the target property is α. The result shows that the performance of TACS is robust when using moderate window size but decreases when window size becomes larger than certain point. We set ∆ = 10 for other experiments.In Appendix C, we provide the results of further ablation studies including the effect of mc sampling and effective diffusion steps for TACS. Overall, the results show that our method is robust in the choice of hyperparameters and generalizable to different datasets and tasks. 
discussion: 
 Exploiting quantum chemistry In Section 5.1, we demonstrate that quantum computing-based guidance can serve as an accurate property predictor for the online guidance. Currently, in the absence of a non-noisy quantum computer, scaling up this exact guidance to the QM9 dataset is close to impossible due to compounding noise [43, 10] . However, future fault-tolerant quantum technology is expected to provide quantum advantage in calculating chemical properties. This can be incorporated into our algorithm when using online guidance and therefore, further improvements of our TACS are on the horizon.Connection to other fields Recent works point out the exposure bias exists for diffusion models [40] , where there is a mismatch between forward and reverse process. Our experiments in 5.2 indicate that Time Correction Sampler can provide a solution to the exposure bias problem during the sampling process in diffusion models. Moreover, since time predictor can gauge this mismatch during inference, one might leverage this information for future works.Another direction is applying our algorithm to matching [56, 27, 39] frameworks. Contrary to diffusion models, these matching models can start with arbitrary distributions and directly learn vector fields. We expect time predictor to be also effective with these types of algorithms. Investigating the connection between our algorithm and matching models will be an interesting future direction.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: In Section 5 and Appendix B, we provide details of our experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 In this work, we introduce Time-Aware Conditional Synthesis (TACS), a novel approach for conditional 3D molecular generation using diffusion models. Our algorithm leverages a Time Correction Sampler (TCS) in combination with online guidance to ensure that generated samples remain on the correct data manifold during the reverse diffusion process. Our experimental results clearly demonstrate the advantage of our algorithm, as it can generate molecules that are close to the target conditions while also being stable and valid. This can be seen as a significant step towards precise and reliable molecular generation.Despite multiple advantages, several open questions remain. For example, how can we more efficiently use the Time Correction Sampler, or more generally, whether this method improves performance in other domains such as in image generation. We expect that our work will open various opportunities across different domains, such as quantum chemistry and diffusion models.Limitation Although we demonstrate the effectiveness of our algorithm on multiple datasets and tasks, we use a trained neural network to estimate chemical properties of each molecule for main experiments. Using exact computational chemistry-based methods might improve our algorithm.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We verified our algorithm with enough number of iterations and with different types of chemical properties.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We describe the computer resources used in the Appendix A and relevant references for the model checkpoints in Appendix B.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We include a Societal Impact statement in Section 7. We note that de novo molecular generation can be of great importance for the discovery of new drug compounds as well as new materials. Naturally, there is a small risk of potentially malicious or unintended use (i.e. generating harmful compounds).Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Molecular generation has less risk for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We properly credit the original owners of assets used in the paper. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: We believe that our framework can assist in drug discovery, which requires synthesizing stable and valid molecules that satisfy target conditions. However, our work could unfortunately be misused to generate toxic or harmful substances.For model architecture, we use EGNN [47] with L = 7 layers, each with hidden dimension h f = 192. We use the same split of training / test dataset of QM9.For completeness, we restate the training of the time predictor (Eq. 10) here. The time predictor is trained for minimizing the cross-entropy loss between the predicted logits p and one-hot vector of the forward timestep t:EQUATION)where t is uniformly drawn from the interval [0, T ].For conditional diffusion model, we train separate time predictor for each of the target property by concatenating conditional information c to the input x t .The rationale behind using cross-entropy loss rather than simple regression loss is because p(t|x) can not be estimated from the point estimate if there exists intersection between support of marginal distributions p t and p s which is from different timesteps s and t, respectively. In other words, this implies there exists x ∈ R d such that p t (x) > 0 and p s (x) > 0 simultaneously holds.Finally, we train the time predictor within 24 hours with 4 NVIDIA A6000 GPUs.Quantum Machine Learning Quantum computing is expected to become a powerful computational tool in the future [14] . While at its early stage, various of quantum machine learning algorithms are proved to have advantage over classical methods [3] . In computational chemistry, these advantages are indeed expected to have huge potential since classically intractable computations like finding ground state for big molecules are expected to become feasible with exponential speed-ups [48] of quantum machines.Variational Quantum Eigensolver Variational Quantum Eigensolver (VQE) [41] is a near term quantum machine learning algorithm which leverages variational principle to obtain the lowest energy of a molecule with given Hamiltonian. For the given Hamiltonian Ĥ, trial wave function |ψ(θ)⟩ which is parameterized by a quantum circuit (θ) is prepared to obtain ground state energy E 0 with following inequality:EQUATIONSpecifically, parameters θ in quantum circuit is iteratively optimized to minimize the following objective:EQUATION)When quantum circuit θ is expressive enough, one can see that |ψ(θ ⋆ )⟩, where θ ⋆ is a minimizer of Eq. ( 14), gives the ground state of the given system. For more comprehensive review with its potential advantages, one may refer to [53] .Quantum online guidance Quantum online guidance is a type of online guidance algorithm where we use VQE-based algorithm to calculate exact values of quantum chemical properties instead of using classifier which is usually a neural network.In each denoising step of diffusion model, we first apply Tweedie's formula (Eq. 7) to estimate clean molecule x0 . Then we use VQE to calculate ground state energy of the estimated molecule by iteratively updating θ for E(x 0 , θ) = ⟨ψ(θ)|H(x 0 )|ψ(θ)⟩. After obtaining θ ⋆ (x 0 ) which minimizes E(x 0 , θ), we obtain target property value E 0 (x 0 ). To obtain gradient in Eq. ( 8), we use zeroth-order method [35] with respect to the position of atoms to obtain gradient as follows:EQUATION)where k is a hyperparameter for multipoint estimate when using the zeroth-order optimization.Let X = {x ∈ R M ×3 : 1 M M i=1 x i = 0} be the subspace of R M ×3where the center of mass is zero. Here we discuss some properties of X that are used in the TACS framework. First, note that X is a linear subspace of R M ×3 with dimension (M -1) × 3. There exists an isometric isomorphism ϕ : R (M -1)×3 → X , i.e., a linear bijective map that preserves distances: ∥ϕ(x)∥ = ∥x∥ for all x ∈ R (M -1)×3 . Intuitively, ϕ allows us to map between the lower-dimensional space R (M -1)×3 and the constrained subspace X without distortion. If x ∈ X is a random variable with probability density q(x), then the corresponding density of x = ϕ -1 (x) in R (M -1)×3 is given by q(x) = q(ϕ(x)). Similarly, a conditional density q(x|y) on X can be written as q(x|ŷ) = q(ϕ(x)|ϕ(ŷ)). In practice, computations involving probability densities on X can be performed in R (M -1)×3 and mapped back to X using ϕ as needed. This allows TACS to efficiently learn and sample from distributions on molecular geometries while preserving translation invariance. For further mathematical details on subspaces defined by center-of-mass constraints, we refer the reader to [29] and [47] .For conditional generation, we need conditional score ∇ x log p t (x t |c) where c is our target condition. Classifier-free guidance (CFG) [19] replaces conditional score with combination of unconditional score and conditional score. Here, diffusion model is trained with combination of unlabeled sample (x 0 , ∅) and labeled samples (x 0 , c). The reverse diffusion process (Eq. 2) in CFG changes as follows:EQUATIONHere, w is a conditional weight which controls the strength of conditional guidance. When w = -1, the process converges to the unconditional generation and with larger w, one can put more weights on conditional score. In the experiments, we use w = 0 following [23, 5] .Here, we provide experimental results on our synthetic experiment. Our results show that TACS is robust through different online guidance strength z and outperforms naive mixture of CFG and OG. Geometric Optimization It is known that H + 3 molecule has a ground state energy around -1.34Ha. Assuming we don't have any prior knowledge of this information, we try to generate molecules with conditioning on the target ground state energy -2.0Ha. Applying TACS proves to be strong in this case also, which implies that our method can robustly guide the molecule without destroying distances.Dataset details The QM9 dataset [45] is a widely-used benchmark in computational chemistry and machine learning. It contains 134k stable small organic molecules with up to 9 heavy atoms (C, O, N, F) and up to 29 atoms including hydrogen. This size constraint allows the molecules to be exhaustively enumerated and have their quantum properties calculated accurately using density functional theory (DFT). Each molecule in QM9 is specified by its Cartesian coordinates (in Angstroms) of all atoms at equilibrium geometry, along with 12 associated properties calculated from quantum mechanical simulations.QM9 molecular properties In Table 9 , we describe six quantum chemical properties that are used for conditional generation in our experiments. Performance metrics To evaluate the quality of generated molecules, Table 10 shows descriptions of the metrics that are used in our experiments to check whether generated samples satisfy basic molecular properties. Baselines We compare our method against Equivariant Diffusion Models (EDM) [23] which learn a rotationally equivariant denoising process for property-conditioned generation and Equivariant Energy Guided SDE (EEGSDE) [5] , which guides generation with a learned time-dependent energy function.Additional baselines Two baselines are employed following the approach outlined by [23] . To measure "Naive (U-Bound)", we disrupt any inherent correlation between molecules and properties by shuffling the property labels in D b and evaluating ϕ c on the modified dataset. "L-Bound" is a lower bound estimation of the predictive capability. This value is obtained by assessing the loss of ϕ c on D b , providing a reference point for the minimum achievable performance. 8), and guidance strength z ∈ {1.5, 1.0, 0.5} in Eq. ( 9). Except for ϵ LUMO (optimal with z = 0.5, tOG = 0), we find t TCS = 600, t OG = 600, tOG = 20, z = 1, and κ = 1 consistently achieve low MAE with molecular stability above 80%.Dataset Details Geom-Drug dataset [4] contains approximately 450K molecules with up to 181 atoms and an average of 44.4 atoms per molecule. Following [5] , we split the dataset into training/validation/test sets of 554k/70k/70k samples respectively.Performance Metrics For each experiment, we generate 10,000 molecular samples for evaluation. As in QM9 experiments, we measure atom stability (AS) and validity (Valid) using RdKit [32] .For Geom-Drug, we employ TCS with 4 EGNN layers and 256 hidden features. The time predictor ϕ is trained unconditionally with the same EGNN architecture as QM9 for 10 epochs. During generation, TCS starts from timestep 600 and utilizes a window size of 10 for time correction.Training Details For molecular fingerprint-based structure generation, we follow the evaluation protocol in [5] . The diffusion model architecture remains consistent with our QM9 experiments, using EGNN with 256 hidden features and 4 layers, trained for 10 epochs.Performance Metrics For evaluation, we use Tanimoto similarity score [16] which measures the structural similarity between generated molecules and target structures through molecular fingerprint comparison. Specifically, let S g and S t be the sets of bits that are set to 1 in the fingerprints of generated and target molecules respectively. The Tanimoto similarity is defined as|S g ∩ S t |/|S g ∪ S t |,where | • | denotes the number of elements in a set. We evaluate the similarity on 10,000 generated samples.Baselines For both experiments, we directly compare with baseline results reported in [22] . For Geom-Drug unconditional generation, these include ENF [47] , G-Schnet [15] , GDM variants [57] , EDM [23] , and EDM-Bridge [56] . For target structure generation, we compare against G-SchNet [15], GDM, GDM-AUG [57] , Conditional EDM [23] , and EEGSDE [5] with various guidance scales.We ablate when to start the time-corrected sampling (TCS) and online guidance (OG) during the reverse diffusion process. The notation t TCS and t OG indicates we apply TCS after timestep t TCS and OG after the timestep t OG in the reverse diffusion process. We report results for property ϵ LUMO in Table 11 . The result shows that applying TACS from early steps (t = 600, 800) generates samples best satisfying the target condition but with less molecular stability and validity. In contrast, when we start applying TACS after later step t = 400, generated samples have higher MAE, MS, and Validity. Interstingly, if we apply TACS only in the later part (after t = 200), MAE, MS, and validity decreases again. We leave further investigation on this phenomenon and explanation for future works.In our experiments, we use t TCS = t OG = 600 as our default setting.Table 11 : Ablation study on when to start time-corrected sampling (TCS) and online guidance (OG) during the reverse diffusion process for the target property ϵ LUMO . We use z = 1 for the experiment. The best value in each column is bolded.[t Additional experiments are conducted on the effect of number of MC samples m in Eq. ( 8).LGD [51] estimates x 0 assuming q(x 0 |x t ) is a normal distribution with mean x0 (Eq. 7) and variance σ 2 which is a hyperparameter.First, we investigate the effect of varying the variance σ in MC sampling. Table 12 shows that the result is robust in the small values of σ but when σ is larger than some point, quality of generated samples decreases (higher MAE and lower MS).Next, we test how the performance of TACS is affected by number of MC samples m. Table 13 shows that performance of TACS is robust in number of MC samples but we did not observe any performance increase with the number of MC samples as in [18] . We report additional metrics on the novelty, uniqueness of generated molecules in Table 14 following previous literature [5, 23, 58] . Novelty measures the percentage of generated molecules not seen in the training set. Uniqueness measures the proportion of non-isomorphic graphs within valid molecules. Higher values indicate better quality for both metrics. The result shows that TACS generates molecules with decreased novelty. This shows that TACS is effective in making generated molecules that stick to the original data distribution while satisfying to meet the target condition. In our derivation of equivariance properties of the time predictor, we closely adhere to the formal procedures outlined in [23, 57, 47] .EQUATIONwhere X ∈ R N ×3 and 1 ∈ R N is the all-ones vector. Definition 2 (Permutation Invariance). A function g : R N ×d → R d is permutation-invariant if for any permutation matrix P ∈ {0, 1} N ×N ,EQUATIONwhere H ∈ R N ×d . Proposition 1. Let f : R N ×3 → R N ×d be an E(3)-equivariant function and g : R N ×d → R d be a permutation-invariant function. Then, the compositionEQUATIONProof. For any orthogonal matrix R ∈ R 3×3 and translation vector v ∈ R 3 ,EQUATIONTheorem 1 (Time Predictor Equivariance). Let G = (V, E) be a graph representing a molecule, where V = {1, . . . , N } is the set of nodes (atoms) and E ⊆ V × V is the set of edges (bonds). Let X t ∈ R N ×3 denote the atomic coordinates and c ∈ R d be a condition vector at diffusion timestep t. Consider a time predictor p ϕ (t|X t , c) = softmax(f ϕ (X t , c)) parameterized by a composition of an E(3)-equivariant graph neural network EGNN ϕ : R N ×3 × R N ×d h × R d → R N ×d ′ , a permutationinvariant readout function ρ : R N ×d ′ → R d ′ , and a multilayer perceptron ψ : R d ′ → R T , i.e.,EQUATIONwhere H t ∈ R N ×d h are node features and T is the total number of diffusion timesteps. Then, the time predictor p ϕ (t|X t , c) is invariant to E(3) transformations, i.e., for any orthogonal matrix R ∈ R 3×3 and translation vector v ∈ R 3 ,EQUATION)Proof. The E(3)-equivariant graph neural network EGNN ϕ satisfies [47] :EQUATIONBy the permutation invariance of ρ and the invariance of ψ to orthogonal transformations, we have:EQUATIONConsequently,EQUATIONWhile with different motivations and methods, here, we list some of the relevant works and compare their algorithms with ours.Comparison with DMCMC DMCMC [28] trains a classifier to predict noise levels of the given data during the reverse diffusion process which is similar to our time predictor. However, they use the classifier to estimate the current noise state when conducting MCMC on the product space of the data and the noise. In contrast, our time predictor directly predicts timesteps for correction in diffusion sampling itself. Moreover, while the purpose of noise prediction in DMCMC is for fast sampling, our work use time predictor to accurately produce samples from the desired data distribution.Comparison with TS-DPM Time-Shift Sampler [34] targets to reduce exposure bias targets similar approach of fixing timesteps during the inference as our time correction method. However, while our method directly selects timesteps based on the time predictor, which has demonstrated robustness in our experiments [34] selects timesteps by calculating variance of image pixels at each step and matching the noise level from the predefined noise schedule which is often inaccurate and expensive. Moreover corrected timestep in [34] 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 (line 9 in Algorithm 1). By taking every single diffusion step while carefully using predicted time, TACS / TCS can generate samples closer to the target distribution.To further validate our approach, we provide additional experimental results comparing TS-DDPM [34] and TACS on the QM9 dataset with step size 10. The result in Table 15 shows consistent improvements across various quantum properties which shows the robustness of our approach. Conditional generation with target structure Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 


Decision: The paper introduces time-correction sampling, an approach to ensure guided samples stay close to the data manifold. Briefly, after a guidance step, a completely denoised step is predicted using Tweedie's formular, this is then noised again using the forward noising process. Rather than using the diffusion time directly, a new diffusion time is predicted by another model, to correct for exposure bias, e.g. mismatch between marginals of the forward and backwards diffusion models. The idea is interesting, and the reviewers all agree with this. The authors address the reviewers' concerns, and while the focus here is on molecular application, I envision this approach having an impact broadly, e.g. in engineering applications. My primary criticism is that I would have loved to see a simple low dimensional illustration of the method to illustrate the method does what the authors claim. While the results suggest this, there is no formal or experimental proof of the approach.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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).
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: How to efficiently use the Time Corrected Sampler and whether this method improves the performance in other domains such as in image generation.
questions: questions: Have the authors experimented TACS approach on other datasets?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: The authors have properly addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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).
limitations: 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.

Comments: comment: 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.
Dear reviewer VutT

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.
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.
Thanks for your response.

I will keep my current score. It would be promising to introduce the optimization of the number of atoms and extend this framework to the case of discrete variables.
I would like to thank the authors for the detailed response. I believe that the additional experimental results for the GEOM-DRUG and CIFAR-10 datasets corroborate the reported performance gains on the QM9 dataset and strengthen the experimental section of the paper. I will update my score accordingly.
I appreciate the authors' detailed response and extensive experiments. I will raise my score to 7.