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
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion model
Molecular generation
Primary area: diffusion_based_models
TL;DR: We achieve highest performance in conditional 3D molecular generation while maintaining molecular stability.

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.

Conditional Synthesis of 3D Molecules with Time Correction Sampler
Hojung Jung'*
Youngrok Park! Bongsang Kim² KAIST Al"
Laura Schmid! Se- Young Yun
Jaehyeong Jo Jinwoo Shin
Dongkyu Lee²
LG Electronics
{ghwjd7281, yr-park, yunseyoung, jinwooshinJOkaist. kc.
Abstract
Diffusion models have demonstrated remarkable success in various domains, in- cluding 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 distribu- tion. In this work, we present Time-Aware Conditional Synthesis (TACS), novel approach to conditional generation on diffusion models. It integrates adaptively controlled plug-and-pa aong-and-plyy "online" guidance into diffusion model, driving sam- ples toward the desired properties while maintaining validity and stability. A key component of our algorithm is new type of diffusion sampler, Time Correction Sampler (TCS), which used to control guidance and ensure that the generated molecules remain on the correct manifold at each reverse step of the diffusion pro- cess at the same time. Our proposed method demonstrates significant performance in conditional 3D molecular generation and offers promising approach towards inverse molecular design, potentially facilitating advancements in drug discovery, materials science, and other related fields.
Introduction
Discovering molecules with specific target properties is 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 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 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 condition-labelas 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].
the other hand, recent works [51, 18] have introduced training- guidance for controllable generation in plug-and-play manner, which we hereafter refer to 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 samples with significantly low molecular stability and validity due to the stepwise enforcement of specific conditions without considering the original distribution at each timestep.
result in generating
"Equal contribution Corresponding authors
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Mist
Deviatio Direotcor =?
Rt
Nony data marnic
31+1
Md. M dets maniial
Z1
Twneddl'y formuis conected Smodtep Ip predict moleads
Fulll deared condiiin
Tttp
X
Time precktionn Correct offcche timetop
Figure 1: (a) Overview of Time-Awar Conditional Synthesis (TACS). TACS helps generate high- quality samples that match target condition while following basic properties of the molecules. At each timestep f, online guidance is applied to push towards the desired condition. Time Predictor finds the desired timestep tp for 2y after applying the guidance. Using predicted timestep tp. Tweedie's formula is used to predict the clean molecule in.. Finally, forward process q(<x-1 (2o) is applied to proceed the next denoising step of 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 avoided by using time correction to first measure the deviated timestep /', 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 call Time Correction Sampler (TCS). TCS explicitly considers the possibility that online guidance can steer the generated sample away from the desired data manifold 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 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 main contributions 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), novel diffusion sampling technique that ensures the generated samples remain faithful to the data distribution. It includes 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 outper- forms previous state-of-the-eeee methods by producing samples that closely match the desired quantum chemical properties while maintaining data consistency.
2 Related Works
Diffusion models Diffusion models have achieved great success in variety of domains, including generation of images [12, 46), audio generation [30], videos [21, 36, 42], and point clouds [8, 37]. 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], method based on the Rose'T Fold structure prediction network [31]. Other previous literature considers different domain-speciffe 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. Recently, some of the works further advance 3D molecule generation techniques by making more reliable diffusion process [63] or fast generation [22] of 3D molecules.
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 mnergy-based model to provide conditional guidance during the inference. GeoLDM [58] utilizes 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, 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.
Adaptive inference schedule Recent efforts have explored adaptive inference schedules to enhance the fidelity of samples generated by diffusion models with timestep information. [59] proposes adjusting the noise schedule during reverse diffusion, while [60] optimizes input timesteps for more accurate reverse trajectory. TS-DPM [34] mitigates exposure bias by adaptively shifting timesteps during inference to align with the training variance. Similar to our time preidction mechanism. several works leverage auxilary network to classify time information with different motivation and purpose. DMCMC [28] introduce noise classifier for accelerating inference during earlier diffusion steps and MoreRed [61] leverages timestep prediction network to initialize or adaptively guide the reverse diffusion process using the adjusted timesteps for molecular relaxation. In contrast, our TACS is specifically designed for conditional generation, predicts timesteps at each denoising step to refine clean sample estimates via Tweedie's 's formula, ensuring alignment with the data manifold and adherence to target conditions. Further comparison is provided in Appendix E.
3 Preliminaries
Diffusion models Diffusion models [50, 20, 52] 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, Xo ~ go. until it becomes perturbed into random noise, XT ~ (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:
dx; - 3(t)x dt + ((t)dw..
(1)
2
where (()) is a pre-defined noise schedule and w, is a standard Wiener process. Then, the reverse of the forward process in Eq. (1) is also diffusion process that is modeled by the following SDE [2]:
dx, = - - -00x - V() lx,) di + v B(I)dw..
where Pt is the probability density of x, and w. is standard Wiener process with backward time flows. The reverse process in Eq. (2) can be used as a generative model when the score function Vx log p. is known. To estimate the score function from given data, neural network sg is trained to minimize the following objective [52]:
((0)- = T4,,, X(t) (x)) () Vx, log (xx) (x:)
(3)
where L~U(,, TJ, Xo is sampled from the data distribution, and À : [0, 7] - IR+ is a positive weight function.
3

---

Online guidance Recent works [17, 9, 51] leverage a conditional diffusion process where the conditional probabilityy pe (x,(c) is modeled without training on labeled pairs (x,, c). To understand this approach, observe that for the conditional generation, the reverse SDE of Eq. (2) can be rewritten as follows:
dx. -000xx V( (x)c) dt + v B(tddw.
(4)
From Bayes' rule, the conditional score Vx log (x+lc) can be decomposed as follows:
Xx, log (o, Xx, log (x,
Vx, log (c(xx)))
While unconditional V log (xe) can be approximated by the unconditional diffusion model 89 from Eq. (3), online guidance can estimate conditional score part Vx, log ((()(() without any training. Notably, DPS [9] approximates conditional score follows:
Xx, log p(c|x,) Xx, log ((c((o))
(6)
where Xg is a point estimation the final clean data using Tweedic's formula [13]:
+ Vx, log p(x,)
ā₃ e B(s)ds
at
While DPS uses a point estimate Xp. Loss Guided Diffusion (LGD) [51] uses Bayesian assumption where is a Monte Carlo estimation of g(xx) which is normal distribution with mean Xo and variance of which is hyperparameter.
For given property estimator A satisfying c A(xo), Eq. (6) can be written as follows:
1
V xx log E (xx R Vx, log
exp (-
t),
(8)
1=
where is a differentiable loss function and x are independent variables sampled from ((xo Using this approximation, we replace Vx, logp(c)xx) in Eq. (5) and conditional generation process now becomes:
dx.
3(1xx -(() (Vx. (x.) + =g(x, 0)) dt + v B())
2
where 2 is the hyperparameter for controlling online guidance strength. In this work, we use L2 loss for the differentiable loss L and set m following Song et al. [51].
3D representations of molecules In our framework, we follow EDM [23] in modeling the distribu- tion of atomic coordinates x within a linear subspace X' where the center of mass is constrained to zero. This allows for innslation-invarian modeling of molecular geometries. Further details of the zero-center-of mass subspace is provided in Appendix A.3.
4 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, 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.
4.1 Time Correction Sampler
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.

---

Algorithm Tware Conditional Synthesis (TACS)
Input: Total number of diffusion timesteps T,, online guidance strength target condition c, diffusion
model 0, time predictor Φb time-clip window size A.
1: Xx ~ (2)
2: for f- = do
3: if O'linne guidance then
4:
g(X,, t) Vx L[A(xo)) c)
Online guidance from Eq. (8)
5:
x f X1 + 8(x 1)
6:
(pred arg max (x((x)) xx))
7:
fpred t clip(tpred. A)
8:
X, Tweedie (x,, (pred)
from Eq. (11)
9:
Xt-1 t forward Xg. 1)
>> from Eq. (1)
else
11:
X1-1 rx)
+ one reverse step by diffusion model
end if
13: end for
Time Predictor During conditional generation process of diffusion model, sample follows the reverse SDE as in Eq. (4). However, due to error accumulation during the generation process [33, 6]. sample at timestep & 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 the sample's current position.
To achieve this, we train a neural network, a time predictor, to estimate the proper timestep of given noised data. Specifically, given a random data point with unknown timestep, time predictor models how likely x ~ Pt for each timestep [o. T]. For training, we parameterize a time predictor 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 used as follows:
(10) where timestep is sampled from the uniform distribution U(,, T), Xg chosen from the data distribution, x, is constructed from Eq. (I), and (x:); is the t-th component of the model output P for given input X,. 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.
(Ø) - E, [log (x:))).
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 sample X; at time / during the reverse diffusion process, we use the corrected timestep /pred predicted by time predictor, instead of the current timestep I, to estimate the final sample &n as follows:
X1 (1 al so Ne, (prd))
Tweedie(Xr, tprea) == f(x,, tprea) =
(11)
a
From the better prediction of the final sample Xu. 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.
5

---

CFG (Baseline)
(Baseline)
0.5
0.5
CFG
OG
(Ours)
TACS (Ours)
0.4
0.4
Distance
0.3
0.3
0.2
0.2
0.1
0.1
15
(online guidance strength)
(online guidance strength)
MAE from target condition
(b) L2 distance from data distribution
Synthetic experiment on H3 dataset. TACS is robust in generating samples that (a) match the desired condition and (b) stick to the original data distribution.
Figure 2:
4.2 Unified guidance with time correction
Finally, we present TACS in Algorithm I, which is novel diffusion sampler for conditional gener- ation that integrates TCS with online guidance. For each timestep of the reverse diffusion process (Eq. 4), 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. 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 fpred often deviates significantly from the current time step Naively applying this information 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 fpod,, we set a time window so that Epeed remains in the time interval A, A], where A is hyperparameter for the window size.
5 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 H3" molecules, where the ground state energies are computed using the variational quantum eigensolver (VQE). In Section 5.2, we assess our method using QM9, 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.
5.1 Synthetic experiment with H
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 prior works [23, 5], which train auxiliary classifiers estimate each condition, this quantum computational chemistry-bu bbasd approach can leverage exact calculation of the condition for 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 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

---

Table 1: Conditional generation with target quantum properties on QM9. TACS generate samples
with lowest MAE while maintaining similar level of stability and validity as other baselines. TCS can generate molecules with high stability and validity. notation is marked for the values for the best MAE within methods with molecule stability above 80%.
Ce mETK
J (D)
(Bole")
Method
MAE
MS (%)
Valid (9)
MAE
MS (%)
Valid (%)
MAE
MS (%)
Valid (%)
U--oondd
6,879-9.005
1,613+0.003
8.98++0.020
EDM
1.072::0.005
81,0#066
90,4:0.2
1.118:00006
$0,2002
91.2=003
2.77:00.000
79.6:20.3
89.9:022
EEGSDE
1.44410.023
81.010.5
90.5200.2
0,85410.002
79.420.1
90.6.00.5
2.62:10.090
79.8.10.3
89.220.33
0G
0.18410.004
30.2104
51,926.8
26.65418.578
6.820.5
29.7=0.3
0,55::0.019
31.2::022
48.7:0.3
TCS(our))
0.904-10.011
90.310.5
0.971110.005
1.85:00.005
87.540.2
TACSS)uu))
0.659 0.659-+0.00 +00008
83.6+0.2
92,4+0.1
0.387*
83.3+0.3
91.3-0.3
1.44"+0.007
86.0+0.1
92,5+0.1
L.bound
0.040
0.043
0.090
Ar (meV)
HOMO (meV)
LLUMO (meV)
Method
MAE
MS (9)
Valid (%)
MAR
MS (5)
Valid (%)
MAR
MS (5)
Valid (%)
U-bound
146424
645+44
1457:5
EDM
673=7
818220.5
90.9=0.3
372=1
79.6=0.1
91.620.2
602:4
810=20.2
91.420.5
EEGSDE
539-5
80.1::0.4
90.5.10.3
300:55
78.7=0.7
91.2:00.2
494119
81.4:0.6
91.1=0.2
DG
95.2:14
31.3=0.77
61.9#3.4
233:88
11.5=0.4
40.8=2.6
170::1
21.1=0.5
42.3=0.9
TCS(ours)
594=4
91.9:0.4
96.0=:022
338:15
92.7=0.2
96.6:20.2
493:19
91.7=3.9
96.2=0044
TACS(ours)
332"-43
8.68
93,9
168
87.3-0.7
93.0+0.2
289 :14
82.7+0.7
91,3-002
L-bound
65
39
36
diffuusin 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, 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 L.2 distance. Further details and additional discussions are provided in Appendix B.1.
Results
5.2 Conditional generation for target quantum chemical properties
Dataset We evaluate method 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], we test 6 types of quantum chemical properties and split the dataset into 100k/18k/134 molecules for training, validation, and test. The training set is further divided into two disjoint subsets of 50k molecules each: D,, for property predictor training and D, for generative model training. Further details are provided in Appendix B.2.
Evaluation To evaluate how generate samples meet the desired target condition, a property pre- diction model Φp [47] is trained on D.. Then, MAE for K number of samples is calculated as R Ek (xp(((()) -c(0)) where x+) represents i-th generated molecule and cl0) 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 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 separate prediction model).
Results Table shows the result of conditional generation of TACS and TCS with baselines. We generate K=10$ samples for the evaluation in each experiment and the average values and standard deviations are reported across runs. For 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
7

---

TACStOurs)
BECSDE
True Timestep
Train
Test
1000
Predicted Timestep
Molecule Stability
800
0.4
Og
600
0.6
400
00
200
0.15
0.20
0.25
0.30
0.35
MAE
Actual Timestep
(a) Parcto front
(b) Time predictor performance
Figure 3: (a) Comparative analysis of molecule stability and MAE across five distinct methods, highlighting the enhanced stability our approach at comparable MAE levels. (b) Performance of the time predictor for train and test set QM9.
Table 2: Quantitative results showing similarity and stability of generated molecules compared to target structures.
Table 3:
Unconditional generation us- ing Geom-Drug dataset.
Geom-Drug
QM9
Method
AS (%)
Valid (%)
Method
Similarity
MS (%)
Data
86.5
99.9
cG-SchNet
0.499110.022
Conditional EDM
0.67110.004
ENF
TCS (Ours)
0.792:10.077
90.42
G-Schnet
TACS (Ours) 0.01)
0.694+00000
90.45
GDM
75.0
90.8
TACS (Ours) (: 0.05)
0.695-10.003
90.02
GDM-AUG
77.7
TACS (Ours) (= 0.1)
0.713-100877
EDM
81.3
92.6
EEGSDE (s 0.1)
0.547+1.002
74.07
EDM-Bridge
82,4
92.8
EEGSDE (# 0.5)
0.600+00.002
TCS(Ours)
89.6
97.6
EEGSDE (s 1.0)
0.5401:0.029
90.44
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 higher MAE compared to TACS. This highlights the importance 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 the correct data manifold at each denoising step, even if they deviate from the true time-manifor 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.
5.3 Target structure generation
We conduct experiments target structure generation with QM9 as in [5]. For evaluation, report Tanimoto similarity score [16] which captures similarity of molecular structures by molecular fingerprint and molecular stability check whether basic properties of molecules satisfied during the conditional generation process. We put additional details 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 This demonstrates TACS's ability to generalize on different tasks.

---

Table 4: Comparison between performance of TACS when using argmax function and expectation for correcting timesteps using time predictor.
G.
malk
J. J.D))
(Bahr")
Ar (meV)
FHOMO (meV)
'LUMO (meV)
Method
MAE
MS
MAE
MS
MAE
MS (5)
MAE
MS
MAE
MS (9)
MAE
MS. (%)
Argmax
0.659
83.3
0.387
83.3
1.44
86.0
332
88.8
168
87.3
289
82.7
Expectation
0,703
84.6
0.451
90.2
1.36
87.3
351
90.7
182
89.8
334
90.1
5.4
Unconditional generation on Geom-Drug dataset
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 the QM9. For fair comparisons, 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 Seoom-Dugg, For evaluation, we use atom stability (AS) and validity of generated samples. The result Table 3 shows that samples generated by TCS satisfy the highest atom stability and the validity with high margin compared to the basclines. Details of the experiments are in Appendix B.
5.5 Ablation Studies
Time prediction function We investigate how the the design of time prediction function affects the performance of TACS. Specifically. for line 6 in Alg. rather than using argmax function obtain corrected timestep épred. we choose to use expectation value by fpred = B[ø(=)]. Table 4 shows the comparison between methods six different types of quantum chemical properties. Expectation based time prediction results in molecules with higher MAE and higher molecular stability.
Online guidance strength z To analyze the effect of online guidance strength z in Eq. (9), measure MAE, molecular stability, and validity of samples generated by TACS for different values. Table 5 shows the result with target condition on éLUмO val- ues. One can observe while trade-off occurs for dif- ferent z values, performance of TACS is robust in varying Interestingly, our experiment shows that there exists an optimal value of which generates samples with the lowest MAE that is even compara-
Table 5: with target property ELUMO-
Ablation study on the OG strength
Method
MAE
MS (%)
Valid (%)
TCS (z OG
493
91.7
96.2
170
21.1
42.3
311
80.3
90.7
= = 1.0
236
74.9
86.3
2 = 0.5
288
82.7
91.3
ble to applying online guidance without time correction (OG). As expected, using 2 = 0 (TCS) generates molecules with the highest MS and validity but with the highest MAE.
Time window length A We measure how the perfor- mance of TACS varies with time winodw length A. Ta- ble 6 shows the MAE, molecular stability values for dif- ferent window sizes when the target property is a. 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.
Table 6: Results of TACS on target property a when varying the time win- dow size.
Window size (A)
MAE
MS (%)
1.481
87.25
4
1.460
86.52
6
1,460
85.73
8
1.448
86.84
In Appendix C, provide the results of further ablation studies including the effect of me sampling and effective diffusion steps for TACS. Overall, the results show that our method is robust the choice hyperparameters and generalizable different datasets and tasks.
10
1.441
86.08
12
1,442
85.17
14
1.459
82.76
16
1.530
79.10
Discussion
Exploiting quantum chemistry In Section 5.1, we demonstrate that quantum computing-based guidance can serve as 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 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 mod- els [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 interesting future direction.
Conclusion
In this work, we introduce Time-Awar Conditional Synthesis (TACS), novel approach for condi- tional 3D molecular generation using diffusion models. Our algorithm leverages 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, it can generate molecules that are close 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 quantum chemistry and diffusion models.
Limitation Although we demonstrate the effectiveness of our algorithm on multiple datasets and tasks, we use trained neural network to estimate chemical properties of each molecule for main experiments. Using exact computational chemistry-baas based methods might improve our algorithm.
Societal impacts 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 harmful substances.
Acknowledgments
This work was supported by Institute for Information & communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (No. RS-2019-11190075, Artificial Intelligence Graduate School Program (KAIST); No. RS-2022-00457882, Al Research Hub Project), the National Research Foundation of Korea NRF grant funded by the Korean government (MSIT) (No. - Stochastic Analysis and Application Research Center (SAARC)), and LG Electronics.
10

---

References
[1] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
[2] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications. 12(3):313-326, 1982.
[3] Frank Arute, Kunal Arya, Ryan Babbush, Dave Bacon, Joseph Bardin, Rami Barends, Rupak Biswas, Sergio Boixo, Fernando GSL Brandao, David A Buell, et al. Quantum supremacy using a programmable superconducting processor. Nature, 574(7779):505--55 -100 2019.
[4] Simon Axelrod and Rafael Gomer-Bombarelli. Geom, energy-annotated molecular conforma- tions for property prediction and molecular generation. Scientific Data, 9(1::185, 2022.
[5] Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy-guided sde for inverse molecular design. In The eleventh international conference on learning representations, 2022.
[6] Joe Benton, VD Bortoli, Amaud Doucet, and George Deligiannidis. Nearly d-linear convergence bounds for diffusion models via stochastic localization. 2024.
[7] Keith T Butler, Daniel W Davies, Hugh Cartwright, Olexandr Isayev, and Aron Walsh. Machine learning for molecular and materials science. Nature, 559(77155:547-555, 2018.
[8] Ruojin Cai, Guandao Yang. Hadar Averbuch-Elor Zekun Hao, Serge J. Belongie, Noah Snavely, and Bharath Hariharan. Learning gradient fields for shape generation. In ECCV, 2020.
(9] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc I. Klasky, and Jong Chul Ye. Diffu- sion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687, 2022.
[10] Jacob M. Clary, Eric B. Jones, Derek Vigil-Fowler, Christopher Chang, and Peter Graf. Explor- ing the scaling limitations of the variational quantum eigensolver with the bond dissociation of hydride diatomic molecules. International Journal of Quantum Chemistry, 123(1 11))2)777777
2023. doi: https://doi.org org/ 10.11000.1100 argg/001102//ua..22.. 1002000a..20097.
[11] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diff- dock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference on Learning Representations, 2023.
[12] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-878 779048794, 2021.
[13] Bradley Efron. Tweedic's formula and selection bias. Journal of the American Statistical Association, 106(496):1602-6611 2011.
[14] Richard Feynman. Simulating physics with computers. In Feynman and computation, pages 133-153. CRC Press, 2018.
[15] Niklas Gebauer, Michael Gastegger, and Kristof Schütt. commetry-adapteer generation of 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems, 32, 2019.
[16] Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof T Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature communications, 13(1):9733 2022.
[17) Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. Advances in Neural Information Processing Systems, 15:14715-14728. 2022.
[18] Xu Han, Caihua Shan, Yifei Shen, Can Xu, Han Yang, Xiang Li, and Dongsheng Li. Training- free multi-objective diffusion model for 3d molecule generation. In The Twelfth International Conference Learning Representations, 2023.

---

[19] Jonathan Ho and Tim Salimans. 1.Xi::2207.12598. 2022.
Classifier-free
diffusion
guidance.
arXiv preprint
[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 2020.
[21] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-8644 8646, 2022.
[22] Haokai Hong. Wanyu Lin, and Kay Chen Tan. Fast 3d molecule generation via unified geometric optimal transport. arXiv preprint arXiv:2405. 1.04405..522... 15252, 2024.
[23] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pages 8867-8887. PMLR, 2022.
[24] Frank Jensen. Introduction to computational chemistry. John wiley & sons, 2017.
[25] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35:24240-22422 24253, 2022.
[26] Jachyeong Jo. Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference on Machine Learning, pages 10362-10383. PMIR, 2022.
[27] Jachyeong Jo, Dongki Kim, and Sung Ju Hwang. Graph generation with diffusion mixture. arXiv preprint arXiv:2302.03356, 2023.
(28) Beomsu Kim and Jong Chul Ye. Denoising meme for accelerating diffusion-basee based generative models. arXiv preprint arXiv:2209.14590 2022.
[29] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
[30] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint (00::200000976.1 2020.
[31] Rohith Krishna, Jue Wang. Woody Ahem, Pascal Sturmfels, Preetham Venkatesh, Indrek Kalvet, Gyu Rie Lee, Felix biomolecular modeling and design with rosettafold all-atom. Science, page ead12528, 2024.
Ivan Anishchenko, Ian Humphreys, et al. Generalized
[32] Greg Landrum et al. Rdkit: Open-source cheminformatics software. 2016.
[33] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 35:22870-22882. 2022.
[34] Mingxiao Li, Tingyu Qu Ruicong Yao, Wei Sun, and Marie-Franciie Moens. Alleviating exposure bias in diffusion models through sampling with shifted time steps. arXiv preprint arXiv:2305.15:3 155883 2023.
[35] Sijia Liu, Pin- Yu Chen, Bhavya Kailkhura, Gaoyuan Zhang, Alfred O Hero III, and Pramod K Varshney. primer on zeroth-orde optimization in signal processing and machine learning: Principals, recent advances, and applications. IEEE Signal Processing Magazine, 37(55:43-54 2020.
[36] Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, Jianfeng Gao, et al. Sora: A review background, technology, limitations, and opportunities of large vision models. arXiv preprint 17X1v.2402.17177, 2024.
[37] Shitong Luo and Wei Hu, Diffusion probabilistic models for 3d point cloud generation. In CVPR, 2021.
12

---

[38] Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations (ICLR). 2022.
[39] Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-LLoatoo Se (3) equivariant augmented coupling flows. Advances in Neural Information Processing Systems, 36, 2024.
[40] Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Onal Ertugrul. Elucidating the exposure bias in diffusion models. arXiv preprint arXiv:2308..523 2023.
[41] Alberto Peruzzo, Jarrod McClean, Peter Shadbolt, Man-Hong Yung. Xiao-QI Qi Zhou, Peter J Love, Alán Aspuru- Guzik, and Jeremy L O'brien. A variational eigenvalue solver on a photonic quantum processor. Nature communications, 5(1):4213, 2014.
[42] Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoory Vyas, Bowen Shi, Chih-Yat Yao Ma, Ching Yao Chuang. et al. Movie gen: A cast of media foundation models. arXiv preprini Stti::2410.137220 2024.
[43] John Preskill. Quantum computing in the nisq era and beyond. Quantum, 2:79, 2018.
[44] Edward O Pyzer-Knapp. Changwon Suh, Rafael Gömez-Bombarelli, Jorge Aguilera- Iparraguirre, and Alán Aspuru-Guzik. What is high-throughput virtual screening? a perspective from organic materials discovery. Annual Review of Materials Research, 45:195-216 2015,
[45] Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[46] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[47] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[48] Peter W Shor. Algorithms for quantum computation: discrete logarithms and factoring. In Proceedings 35th annual symposium on foundations of computer science, pages 124-134. Icce, 1994.
[49] Gregory Sliwoski, Sandcepkumar Kothiwale, Jens Meiler, and Edward W Lowe. Computational methods in drug discovery. Pharmacological reviews, 66( (661):334335, 2014.
[50] Jascha Sohl-Dicke Dickstein, Eric Weiss, Niru Maheswaranathhn, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256 2265.. PMLR, 2015.
[51] Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning, pages 32483-32498. PMLR, 2023.
[52] Yang Song, Jascha Sohl-Dickstein Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.134 13456, 2020.
[53] Jules Tilly, Hongxiang Chen, Shuxiang Cao, Dario Picozzi, Kanay Setia, Ying Li, Edward Grant, Leonard Wossnig, Ivan Rungger, George H Booth, al. The variational quantum eigensolver: a review of methods and best practices. Physics Reports, 186:1-128, 2022.
[54] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang. Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. The Eleventh International Conference Learning Representations, 2023.
13

---

[55] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahem, Andrew Borst, Robert J Ragotte, Lukas Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, 62007976):1089-11000 2023.
[56] Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems, 1136633-365455 2022.
[57] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.0292. 2022.
[58] Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference on Machine Learning, pages 18592-38610 PMLR, 2023.
[59] San-Roman, R., Nachmani, E. Wolf, L. Noise estimation for generative diffusion models. ArXiv Preprint ArXiv:2104.0266 (2021)
[60] Xia, M., Shen, Y., Lei. C., Zhou, Y., Zhao, D., Yi, R., Wang, W. & Liu, Y. Towards More Accurate Diffusion Model Acceleration with A Timestep Tuner. Proceedings Of The IEEE/CVF Conference On Computer Vision And Pattern Recognition (CVPR). pp. 5736-5745 (2024,6)
[61] Kahouli, K., Hessmann, S., Müller, K., Nakajima, S., Gugler, S. & Gebauer, N. Molecular relaxation by reverse diffusion with time step prediction. ArXiv Preprint ArXiv:2404. 10935. (2024)
[62] Ye, Z., Chen, Z., Li, T., Huang, Z., Luo, W. Qi, G. Schedule On the Fly: Diffusion Time Prediction for Faster and Better Image Generation. ArXiv Preprint ArXiv:4412.01243. (2024)
[63] Cornet, F., Bartosh, G., Schmidt, M. & Naesseth, C. Equivariant neural diffusion for molecule generation. 38th Conference On Neural Information Processing Systems. (2024)
14

---

A Method details
A.1 Time Predictor
For model architecture, we use EGNN [47] with I, 7 layers, each with hidden dimension by = 192. We the same split of training / test dataset of QM9.
Time predictor training 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:
Preneeredietel ( = x, [log (X+(x))))),
(12)
where 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 e to the input Xt.
The rationale behind using cross-entropy loss rather than simple regression loss is because p(t |x) not be estimated from the point estimate if there exists intersection between support of marginal distributions pt and ps which from different timesteps and t. respectively. In other words, this implies there exists x E Rd such that Px(x) > and ps (x) simultaneously holds.
Finally, we train the time predictor within hours with NVIDIA A6000 GPUs.
A.2 Quantum online guidance
Quantum Machine Learning Quantum computing is expected to become powerful computa- tional 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 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 Figensolver Variational Quantum Figensolver (VQE) [41] is 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 I, trial wave function (x)) which is parameterized by a quantum circuit (0) is prepared to obtain ground state energy Eg with following inequality:
Eg (O)A (O))
(13)
(0) (O))
Specifically, parameters 0 in quantum circuit is iteratively optimized to minimize the following objective:
E(O)- (v) (O))
(14)
(v 0
When quantum circuit 0 expressive enough, one can see that (a,)).), where 0, is 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 VQE-based algorithm calculate exact values of quantum chemical properties instead of using classifier which is usually a neural network.
In each denoising step of difffusion model, first apply Tweedie's formula (Eq. 7) to estimate clean molecule Xg. Then we use VQE to calculate ground state energy of the estimated molecule by iteratively updating 0 for (xx, 0) (x)) (Xo))))))))) (())). After obtaining 0, (Xo) which minimizes ExR,, 0), we obtain target property value En (&)). To obtain gradient in Eq. (8), we use zeroth-order method [35] with respect to the position of atoms to obtain gradient follows:
Eo (2 + h:)) Eo h2)
Xx, log Exo P(C(22)) x -xx, Ex
(15)
2h:
15

---

where k is a hyperparameter for multipoint estimate when using the zeroth-order optimization.
A.3 Properties of the Zero-Center-ff-Maee Subspace
O} be the subspace of IR Mxx where the center of mass is zero. Here we discuss some properties of X that are used the TACS framework. First, note that X is a linear subspace of R MMxx x
Let X' == {x E Mx33 Σ
with dimension (M 1) x 3. There exists an isometric isomorphism (M + X, i..., linear bijective map that preserves distances: || ||||| for all & E R (M-1)x3 Intuitively, & allows to map between the lower-dimensional space (M and the constrained subspace X' without distortion. I E X' is random variable with probability density g(x), then the corresponding density of & & (x)
is given by ((&) (())
Similarly, a conditional density g(xly) on X' can be written as ((x1)) = q (((())(((())) In practice, computations involving probability densities on A' can be performed in R2 (MM 33 and mapped back X using & as needed. This allows TACS 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].
A.4 Classfier-free guidance
For conditional generation, we need conditional score Vx log Px ic) where e is our target condition. Classifier-free guidance (CFG) [19] replaces conditional score with combination of unconditional score and conditional score. Here, diffusion model trained with combination of unlabeled sample (Xo: 0) and labeled samples (xo. c). The reverse diffusion process (Eq. 2) in CFG changes follows:
dx; = -00x0 3()
(~ (x.) + (1 V> (x,1())))))) ai+ ((()d()))))))))))))))))))) (16)
Here, w a conditional weight which controls the strength of conditional guidance. When w 1, the process converges to the unconditional generation and with larger w, can put more weights on conditional score. In the experiments, we use w = 0 following [23, 5].
B Experimental details
B.1 Synthetic experiment with Hg*
Here, we provide experimental results on our synthetic experiment. Our results show that TACS is robust through different online guidance strength and outperforms naive mixture of CFG and OG.
Table 7: MAE with target condition
Table 8: L2 distance from data distribution
Z
CFG+0G
TACS
Z
TACS CFG+0G
0
0.4823
0.0817
0
0.0386
0.0294
0.0961
0.0530
0.0653
0.0305
2 0.1069
0.0438
2
0. 0818
0.0339
5
0.1107
0.0377
5
0,6553
0.0383
10
0.1430
0.0529
10
3.5112
0.1820
25
0.5427
0.3089
25
18.8983
8.2466
Geometric Optimization is known that H.*" molecule has a ground state energy around - H..44Ha. Assuming we don't have any prior knowledge of this information, we try to generate molecules with conditioning the target ground state energy -a.HH.. Applying TACS proves to be strong in this case also, which implies that our method can robustly guide the molecule without destroying distances.
16

---

B.2 Experiments on QM9
Dataset details The QM9 dataset [45] a widely- benchmark 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.
Table 9: Quantum chemical properties
Property
Description
Measure of molecule's is ability to form instantaneous dipoles in an external electric field.
Polarizability (o)
Energy gap between the HOMO and LUMO orbitals, indicating
HOMO-LUMO gap (Ac)
electronic excitation energies and chemical reactivity.
Energy of the highest occupied molecular orbital, related to
HOMO energy (éhOмO)
ization potential and donor reactivity.
Energy the lowest unoccupied molecular orbital, related to
LUMO energy (el.UMo)
electron affinity and acceptor reactivity.
Measure of the separation of positive and negative charges in a molecule, indicating polarity.
Dipole moment (μ)
Measure of how much the temperature of a molecule changes when it absorbs or releases heat.
Heat capacity (C.)
Performance metrics of the metries that used in our experiments check whether generated samples satisfy basic molecular properties.
To evaluate the quality of generated molecules, Table 10 shows descriptions
Table 10: Performance metrics
Property
Description
Proportion of generated molecules that are chemically valid, as determined by RDKit. A molecule is valid if RDKit can parse it without encountering invalid valences.
Validity (Valid)
Percentage of atoms within generated molecules that possess cor- rect valencies. Bond types predicted based atom distances and types, and validated against thresholds. An atom is stable if its total bond count matches the expected valency for its atomic number.
Atom Stability (AS)
Percentage of generated molecules where all constituent atoms are stable.
Molecule Stability (MS)
Baselines We compare our method against Equivariant Diffusion Models (EDM) [23] which learn a rotationally equivariant denoising process for property-conditione generation and Equivariant Energy Guided SDE (EEGSDE) [5], which guides generation with 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, and evaluating & the modified dataset. "L.-Bound" is a lower bound estimation of the predictive capability. This value is obtained by assessing the loss of &. on D,, providing reference point for the minimum achievable performance. If proposed model, denoted
17

---

as, surpasses the performance of Naive (U-Bound), it indicates successful incorporation of conditional property information into generated molecules. Similarly, outperforming the L-Bound demonstrates the model's capacity to incorporate structural features beyond atom count and capture the intricacies of molecular properties. These baselines establish upper and lower bounds for evaluating mean absolute enror (MAE) metrics in conditional generation tasks.
Diffusion model training For a fair comparison with EDM and EEGSDE [23, 5], we adopt their training settings, using model checkpoints provided in the EEGSDE code: https: h/tpss//// //gi ggtbb.......... com/ gracezhao / The diffusion model is trained for 2000 epochs with a batch size 64, learning rate of 0.0001, Adam optimizer, and exponential moving average (EMA) with decay rate of 0.9999. During evaluation, we generate molecules by first sampling the number of atoms M ~ M)) and the property value с ~ p(c] M)). Here p(M) is the distribution of molecule sizes in the training data, and p(c) M) is the conditional distribution of the property given the molecule size. Then generate molecule conditioned M and using the learned reverse process.
Hyperparameters for TACS We analyze different hyperparameter settings for all six quan- tum chemical properties (a, Ac, CHOMO- CLUMO- μ. C.) for the result in the Table We vary the TCS starting timestep fres E f200, 400, 600, 800), online guidance starting timestep los E (200, 400, 600, 800), online guidance ending timestep log E {10, 20, 30), gradient clipping threshold {oo, I, 0.13 for Eq. (8), and guidance strength 2 E {1.5, 1.0, 0.5} in Eq. (9). Except for cLUMO (optimal with 0.5, fog 0), we find Ircs 600, toG 600, toG 20, and
K consistently achieve low MAE with molecular stability above 80%.
B.3
Additional Experiment Details
B.3.1
Experiments on Geom-Drug
Dataset Details Geom- -oou- dataset [4] contains approximately 450K molecules with up 181 atoms and an average of 44.4 atoms per molecule. Following [5], we split the dataset into train- ing/validation/tesst 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].
Time Predictor Training
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.
B.3.2 Experiments Target Structure Generation
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 Metries 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, and S, be the sets of bits that are set to I in the fingerprints of generated and target molecules respectively. The Tanimoto similarity is defined as |S, SnSts///SSS U. where | denotes the number of elements in set. We evaluate the similarity on 10,000 generated samples.
Baselines For both experiments, we directly compare with baseline results reported [22]. For Geom- 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.
18

---

C Additional experiments
C.1 When to apply TCS and OG
We ablate when start the time-corrected sampling (TCS) and online guidance (OG) during the reverse diffusion process. The notation brcs and foc indicates we apply TCS after timestep frcs and OG after the timestep foo in the reverse diffusion process. We report results for property CLUMO 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 start applying TACS after later step l 400, generated samples have higher MAE, MS, and Validity, Interstingly, if apply TACS only in the later part (after 200), MAE, MS, and validity decreases again. We leave further investigation on this phenomenon and explanation for future works.
In our experiments, use trcs toG 600 our default setting.
Table 11: Ablation study when to start time-corrected sampling (TCS) and online guidance (OG) during the reverse diffusion process for the target property ELUMO. We use 2 = 1 for the experiment. The best value in each column is bolded.
[TCS: loG]
MAE
MS (%)
Valid (%)
800, 800
236
74.9
86.3
6600, 600
236
74.9
86.2
400, 400
360
86.6
93.3
200, 200
248
72.1
84.5
C.2 Number of MC samples
Additional experiments conducted the effect number of MC samples m in Eq. (8). LGD [51] estimates assuming (xx,) xx)) is normal distribution with mean &o (Eq. 7) and variance o² which 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, test how the performance of TACS is affected by number of MC samples Table 13 shows that performance of TACS is robust number of MC samples but we did not observe any performance increase with the number of MC samples in [18].
Table 12: MC effect with varying 5 MC samples used and the target property is
MAE
MS (%)
0.0001
1.506
86.06
0.0005
1.390
84.33
0.001
1.501
86.92
0.005
1.395
86.35
0.01
1.464
85.10
0.05
1.801
85.04
0.1
2.186
82.69
0.3
2.936
75.67
Table 13; Varying number of MC samples with o = 0.005. Target property is a.
# Samples
MAE
MS (%)
1 1.440
86.10
5
1.395
86.35
10
1.505
82.21
15
1.545
83.04
20
1.468
86.76
19

---

C.3 Results on Novelty, Uniqueness
We report additional metrics on the novelty, uniqueness 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 effective making generated molecules that stick to the original data distribution while satisfying to meet the target condition.
Table 14: Novelty and uniqueness of generated sample for the target property ÉLUMO-
Method
Novelty (%)
Uniqueness (%)
EDM
84.5
99.9
EEGSDE
84.8
84.8
TACS
71.6
99.8
D Mathematical Derivations
In our derivation of equivariance properties of the time predictor, we closely adhere to the formal procedures outlined in [23, 57, 47
Definition 1 (E(3) Equivariance). A function f RR - RNxd is E(3)-eqqivariant for any orthogonal matrix R R333 and translation vector R33
(vR R f(X),
(17)
where X E Nx3 and 1 E IR is the all-ones vector.
Definition 2 (Permutation Invariance). function g Rgd + Ed is commutaaionninvariion for any permutation matrix P E {0, 10,1)
9(HH (H),
(18)
where H E Nxd
Proposition 1. Let f: RN -> Ndd be an E(3) {(3)-equivariaan function and g: R - Re. be ionnutation-ivvaiiann function. Then, the composition h g Nx3 Rd is invariant to E(3) transformations, i.e.,
MRX + v1 = h(X).
(19)
Proof. For any orthogonal matrix R p333 and translation vector 1R3
1(RX + v1 = =[( + v1
= f(x)) (E(3) equivariance of
=9(((X)) (Permutation invariance of g)
=x(X).
(20)
L
Theorem 1 (Time Predictor Equivariance). Let G (V.E) be a graph representing a molecule, where V (1....,N] is the of nodes (atoms) and C V is the set of edges (bonds). Let X, E N denote the atomic coordinates and с pd be condition vector diffusion timestep t. (Xt, c)) parameterized by a composition of an
Consider a time predictor pt c) = softmax(f (Xt,
Eq3)--auuvariant graph neural network EGNN, Nx3 x Xxd,, x Rd + IR xx.
permutation-
invariant readout function p Nxd - d and a multilayer perceptron R BT, i.e. i..,
fo(X, (p((ENN (X, H,,c))),
(21)
where H, E R Ndd, are node features and T. is the total number of diffusion timesteps. Then, the time predictor pe Po (t(X,, c) is invariant to E(3) transformations, i.e.. for any orthogonal matrix R 3x3 and translation vector E 3
(o(t +
v1 T1TTC cc) (st)) VE .T.
(22)
20

---

Proof. The E(3) equivvarur graph neural network EGNN, satisfies [47]:
OGNN, (RX, (RX, + v1", H,, c) R EGNN, (X,, H,,c).
(23)
By the permutation invariance of and the invariance of y to orthogonal transformations, have:
sR + vI .c) V(((EGNN, (RX; v1 c),) =>(ppR EGNN, (X,, H,, c)))
=V(OEEGNN, (X, c),
(24)
Consequently,
PX - v1 .c) = softmax (fo (RX, + v1 .c))
softmax(fooe (X.,
c))
(25)
E
Comparison with other related works
While 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 classifier to predict noise levels of the given data during the reverse diffusion process which 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 diffusion sampling itself. Moreover, while the purpose of noise prediction in DMCMC is for fast sampling, 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.
Table 15: Comparison between TACS and Time shift sampler conditional molecular generation.
TS-DDPM
TACS (ours)
Method
MAE
MS (%)
MAE
MS (%)
C 1.066
74.89
0.659
83.6
P
1.166
73.55
0.387
83.3
a
2.777
75.20
1.44
86.0
((C)
665.3
82.72
332
88.8
HOMO
371.8
72.74
168
87.3
LUMO
607.6
74.98
289
82.7
21

---

Visualization of generated molecules
Conditional generation with target quantum chemical property
Figure 4: Conditional generation of target property QM9. Visualization of molecules generated by TCS (top), online guidance (middle), and TACS (bottom).
Conditional generation with target structure
EEGSDE
0.5
TCS
TAGS
0.1
Figure 5: Visualization of how generated molecules align with target structures.
Unconditional generation on Geom-Drug
Drug-EDM
Drug-TCS
Figure 6: Selection of samples generated by the denoising process of EDM and TCS in Geom- -oon-
22

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the experiments with QM9 datset, we demonstrate algorithm has superior performance compared to other baselines.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We put limitations of work in Section 7.
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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
23

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: In Appendix B, we provide experimental details for all of our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
24

---

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.
Please see the NeurIPS code and data submission guidelines (https://nips.cr cc/
Coblli/guuddes/ooeSuubbaee iono for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //n /nips. wwtppcccpcc/cceeccccccccccccclct httpeiguuideeeeeeetniiieeeeeee
/Coo//eoomissssseeee crriissssoppiiyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
25

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 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).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eurrpp.pp/ee cc/ cccppblicc/etteeccceeeeeeeeeee/t htt.ai/eeehhccuuieeelleieeeeee
Answer: [Yes]
Justification:
We follow Code of Ethics during the research.
Guidelines:
The answer NA means that the authors have 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., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 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 new materials. Naturally, there small risk of potentially malicious or unintended use (i.c. 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 why the paper does not address societal impact.
26

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We properly credit the original owners of assets used in the paper.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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 package should be provided. For popular datasets, paperswi twwode..ooo..... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
27

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "piPrGKcaiK", "reviewer": "Reviewer_VutT", "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.", "questions": "* Do the stability values reported in Table 1 refer to atom or molecule stability?\n* 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?\n* 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? \n* Is the time predictor is only trained on samples from the forward process?\n\n---\n\n[1] Han, Xu, et al. \"Training-free Multi-objective Diffusion Model for 3D Molecule Generation.\" The Twelfth International Conference on Learning Representations. 2023.\n[2] Kim, Beomsu, and Jong Chul Ye. \"Denoising mcmc for accelerating diffusion-based generative models.\" arXiv preprint arXiv:2209.14593 (2022).", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n* 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. \n* 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.\n* None of the quantitative empirical results in Section 5 include error bars or measures of statistical significance.\n* 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.\n* 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].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that our work\n\n- Introduces a promising approach to the exposure bias\n- Show improvements compared to baselines with desired quantum properties while maintaining stability and validity.\n- Well-written and clearly outlines the methodology and motivation.\n\nWe initially address your concerns below.\n\n---\n\n**W1** Evaluation on other benchmarks\n\nAs 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).\n\n[Table A.1] TCS on CIFAR 10\n\n|     | ADM   | ADM-ES |ADM-TCS(Ours)| ADM-ES+TCS(Ours)|\n|-----|-------|--------|-------------|-----------------|\n| FID | 3.353 | 2.213  | 2.663       |     2.179       |\n|     |       |        |             |                 |\n\n---\n\n**W2**  Error bars and statistical significance\n\nThank 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.\n\n---\n\n**W3** Missing information in Appendix A.\n\nWe appreciate your careful review. In our final version, we will expand on several points as follows:\n\n**R3-1** Comparison with DMCMC [3]\nDMCMC 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. \n\n**R3-2** Impact of MC numbers\nWe 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. \n\nFor details on these additions, we kindly refer to [GR2].\n\n**R3-3** MAE distribution given the threshold.\nWe 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. \n\n---\n\n\n**W4** Calibrate the claim that “TACS outperforms competitive baseline methods across all metrics\", since online guidance refers to a published baseline model [1].\n\nWe 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.\n\n\n---\n\n**Q1**  Stability values reported in Table 1.\n\nStability values denote the molecule stability.\n\n\n\n**Q2** Online guidance results in Table 1. \n\nOnline 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.\n\nFirst, 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$. \n\nThis in turn, can make online guidance sensitive to its guidance strength, where we sometimes observe diverging gradients when naively applying online guidance. \n\nAdditionally, we provide additional experimental results for impact of online guidance strength on $\\mu$ to support our claim.\n\nWe will include the above result in our final version of the paper for the better understanding of the reader.\n\n**Q3** How to obtain predicted time from the time predictor\n\nThank 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.\n\n[Table A.2] Different type of time prediction. Above is MAE for target condition and below is molecular stability of generated samples.\n| Method     | Cv     | $\\mu$      | $\\alpha$   | $\\Delta(\\epsilon)$   | $\\epsilon_{HOMO}$ | $\\epsilon_{LUMO}$ |\n|------------|--------|--------|--------|-------|-------|-------|\n| Expectation| 0.7032 | 0.4511 | 1.559  | 351.7 | 181.8 | 334.3 |\n| Argmax     | 0.659  | 0.387  | 1.44   | 332   | 168   | 289.0 |\n\n| Method     | Cv     | $\\mu$      | $\\alpha$   | $\\Delta(\\epsilon)$    | $\\epsilon_{HOMO}$ | $\\epsilon_{LUMO}$ |\n|------------|--------|--------|--------|-------|-------|-------|\n| Expectation| 84.67 | 90.18 | 87.3  | 90.71 | 89.84 | 90.11 |\n| Argmax     | 83.3  | 83.3  | 86.0  | 88.8  | 87.3  | 82.7  |\n\n---\n\n**Q4**  Is the time predictor only trained on samples from the forward process?\n\nYes, 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.\n\n---\n\n[1] Han, Xu, et al. \"Training-free Multi-objective Diffusion Model for 3D Molecule Generation.\" The Twelfth International Conference on Learning Representations. 2023.\n[2] Bao, Fan, et al. \"Equivariant energy-guided sde for inverse molecular design.\" The eleventh international conference on learning representations. 2022.\n[3] Kim, Beomsu, and Jong Chul Ye. \"Denoising mcmc for accelerating diffusion-based generative models.\" arXiv preprint arXiv:2209.14593 (2022).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "op5Fes0Y8T", "reviewer": "Reviewer_9Yep", "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.", "questions": "Have the authors experimented TACS approach on other datasets?", "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.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive and helpful comments. We have addressed your comments below.\n\n**W1**  Generalizability of TCS\n\nAs 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. \n\n**Q1** Experiments on other datasets\nAs 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SG7ynlH5B3", "reviewer": "Reviewer_1ecU", "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.", "questions": "Major:\n1\\. 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. \n\n2\\. Appendix C seems incomplete. How is the function incorporated into the guidance?\n\n3\\. Some highlighted values in Table 1 are not the best value, and some values seem abnormally high or low. Please confirm.\n\n4\\. 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?\n\n[1] Li, Mingxiao, et al. \"Alleviating exposure bias in diffusion models through sampling with shifted time steps.\" arXiv preprint arXiv:2305.15583 (2023).\n\nMinor:\n\n1\\. Several references to the Appendix needs to be fixed, e.g.:\n\nLine 264 should be B.4; \n\nLine 149: the said comparison is not in the Appendix. \n\nLine 248: the reference is broken and the said result is not in the Appendix. \n\n2\\. 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": "The authors have properly addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that \n\n- Offers useful insights in reducing the exposure bias of diffusion models.\n- Comprehensive analysis and ablation studies which could help future works. \n- Demonstrates the possibility of quantum computing.\n\n\nWe initially address your concerns below.\n\n---\n\n**Q1** Size of Time Window\n\nWe 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.\n\nWe 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.\n\n---\n\n**Q2** Online guidance using VQE.\n\nIn 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.\n\n---\n\n**Q3**  Table 1 correction.\n\nThank 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.\n\nIn 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.\n\n\n--- \n\n**Q4** Comparison with Time Shift Sampler in [1]\n\nRegarding the comparison with [1], we would like to clarify the key differences between TS [1] and our TCS / TACS:\n\n- 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.\n\n- 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.\n\nTo 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.\n\n[Table A.3] Comparison between TS ([1]) and TCS / TACS\n\n| Property | TS MAE | TS Stab. (%) | TACS(Ours) MAE | TACS(Ours) Stab. |\n|----------|--------|--------------|----------------|-------------------|\n| Cv       | 1.066  | 74.89        | 0.659          | 83.6              |\n| μ        | 1.166  | 73.55        | 0.387          | 83.3              |\n| α        | 2.777  | 75.2         | 1.44           | 86.0              |\n| Δε       | 665.3  | 82.72        | 332            | 88.8              |\n| HOMO     | 371.8  | 72.74        | 168            | 87.3              |\n| LUMO     | 607.6  | 74.98        | 289            | 82.7              |\n\n\nWe will include this comparison in our final manuscript..\n\n---\n\n**Q5**  Addition and corrections of Appendix.\n\nWe 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. \n\n---\n\n**Q6**  Possible relationship between start of OG and Figure 4.\n\nThank 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.\n\n[1] Li et al., Alleviating exposure bias in diffusion models through sampling with shifted time steps, arXiv 2023\n[2] Han, Xu, et al. \"Training-free Multi-objective Diffusion Model for 3D Molecule Generation.\" The Twelfth International Conference on Learning Representations. 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PQrtyPbsa8", "reviewer": "Reviewer_ByqK", "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.", "questions": "See the above weaknesses. \n\nBesides, \n- 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?\n- 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.)\n\n[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": "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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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. \n- The authors have introduced extensive related work to help the audience to better understand the literature.  \n- 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). \n- 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.\n- 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.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive and helpful comments. We have addressed your comments below. We appreciate your positive comments that \n\n- Clear motivation and a novel idea to the challenging problem.\n- Extensive related works for the audience.\n- Necessary ablation studies.\n\n\nWe initially address your concerns below.\n\n---\n\n**W1**  Additional experiments on another task.\n\nThank 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].\n\n---\n\n\n**W2**  Generalizability beyond molecules / comparison with [1]\n\nWe 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.\nOur 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.\nMethodologically, 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].\n\n---\n\n**W3**  Correction\n\nWe 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.\n\n---\n\n**Q1** Online guidance for discrete variables\n\nWhile 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. \n\n\n---\n\n**Q2** Optimizing number of atoms\n\nExploring 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.\n\n[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.\n[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).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
General response

Dear reviewers and AC,

We sincerely appreciate your valuable time and effort in reviewing our manuscript. Your insightful feedback has been instrumental in improving our work.

Our paper introduces Time-Aware Conditional Synthesis (TACS), a novel approach for conditional 3D molecular generation using diffusion models. To our best knowledge, this is the first work to systematically achieve the generation of molecules that simultaneously meet desired conditions, stability, and validity. As highlighted by multiple reviewers, TACS offers a principled and effective method (VutT, qUzm, 9Yep, 1ecU), unlike existing works of conditional molecular generation that focus solely on desired properties. Our approach has been thoroughly validated through extensive quantitative and qualitative experiments (VutT, qUzm, 1ecU) and presented comprehensively (VutT, qUzm, ByqK).

We have carefully considered the multiple concerns raised by reviewers and have addressed them comprehensively as follows. We kindly request that you review our responses below, as well as the attached supplementary PDF file. 

---

**[GR1] Generalizability Beyond QM9 Dataset (VutT, 9Yep, ByqK)**

To demonstrate the broad applicability of TACS, we conducted additional experiments on:

1. Generating molecules with target structures on QM9 dataset (Different task).
2. The Geom-Drug dataset for molecule generation, demonstrating performance beyond QM9's distinct distribution.(Bigger dataset)
3. Image generation in CIFAR-10 dataset, showcasing generalizability to a different domain.(Other domain)

For different task, we conduct additional experiments on target structure generation as in [1].
We tested our algorithms (TACS) on QM9 dataset and showed that TACS achieves a Tanimoto similarity score of 0.792±0.077 with 90.42% of molecular stability, outperforming the reported and reproduced values of [1] with high margin. This demonstrates TACS's ability to generalize to different tasks. We kindly refer the reader to see the Table. R2 in our uploaded pdf file for more details.

We also tested Time Correction Sampler(TCS) on the Geom-Drug dataset on unconditional generation of 3D molecules. As can be seen in the Table. R3 in the attached pdf file, the result shows improved performance in generating molecules with target structures.

Finally, we tested our algorithm on the image dataset. On CIFAR-10 dataset our Time Correction Sampler (TCS) achieves improved FID scores (3.353 to 2.661) without any fine-tuning of the hyperparameters.The improvement in FID score shows the effectiveness of our time correction approach in mitigating exposure bias, a problem common to diffusion models across various domains.


We believe these results altogether demonstrates that our method generalizes well beyond molecular data, indicate possible applications in other tasks, datasets, and domains.

---

**[GR2] Additional Analyses and Clarifications** (VutT, 1ecU): We appreciate the reviewers' thorough examination of our work. In our final version, we will expand on several points to offer a more comprehensive analysis:

(Line 149) Regarding the DMCMC [2] comparison, we will add a more detailed discussion in the Appendix. This will include key differences such as: DMCMC uses a classifier to predict noise levels for MCMC on the product space of data and noise, focusing on earlier generation stages. In contrast, our time predictor directly predicts timesteps for correction in diffusion sampling, aiming to produce samples closer to the desired data distribution.

(Line 248) We will also include our analysis of experiments conducted on the effect of MC sample numbers. Unlike [3], where increasing samples from 5 to 10 more than doubled performance at increased computational cost, our method shows robust performance across different sample numbers. As shown in Table. R4 of the supplementary material:

| Num. Samples | MAE   | Stab. (%) |
|--------------|-------|-----------|
| 1            | 1.44  | 86.0      |
| 5            | 1.395 | 86.35     |
| 10           | 1.505 | 82.21     |
| 15           | 1.545 | 83.04     |
| 20           | 1.468 | 86.76     |

These results  indicate that TACS maintains consistent performance in terms of both MAE and molecular stability across different numbers of MC samples. This stability eliminates the need for higher m values, avoiding extra computational costs while maintaining accuracy.

All these modifications and additional analyses will be reflected in our final manuscript accordingly.

---

**Reference**


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

Meta-review: 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.
Decision: Accept (poster)