Paper: Molecule Relaxation by Reverse Diffusion with Time Step Prediction (ICLR-2024)

Reviewer comment (Novelty):
I believe that the idea of using diffusion in this paper is quite straightforward and not very innovative, so I find this aspect lacking in novelty.

Author response:
We respectfully disagree with the reviewer regarding a lack of novelty in our approach. The prediction of the diffusion time step is a novel contribution and applying diffusion as denoising to perform relaxation was not done before. Furthermore, the diffusion time predictor  is not “loosely connected” to the relaxation but a prerequisite for applying the diffusion model framework to relaxation: It enables the reverse-diffusion from arbitrarily noisy inputs, e.g. molecules that are close to their stable conformation. This is because diffusion models use the time step $t$ as an input, and therefore we must determine  from which $t$ the reverse diffusion should start for a given unstable molecule. We have added Fig. 2a to show how the predicted timestep relates to the RMSD of the unstable molecules (i.e. their distance to the ground-truth stable structure). One can see that the time step prediction enables the model to start relaxation at different noise levels, depending on the configuration of the unstable structure. 

Additionally, our analysis show that adaptive scheduling based on the diffusion time prediction is beneficial as it allows the model to flexibly adjust the process (i.e. jump back and forth in time) which leads to significantly improved performance of the two flexible models (MoreRed-JT and MoreRed-AS) over the model where only the starting point is predicted (MoreRed-ITP). We added a figure to show some example trajectories of predicted time steps over relaxation that visualize this flexibility (see Fig. 2b and Section 4.1). This is also where the applicability of the time step prediction to other diffusion-based frameworks comes from: We show that having a flexible reverse diffusion process might not only be beneficial when starting from non-Gaussian input such as unstable molecules but also for data generation from scratch, e.g. when sampling molecules or images from Gaussian noise (see Tables 1 and 2).

Therefore, the time step predictor is both an integral part of MoreRed and a promising contribution to other fields. We argue that our approach is an innovative method for relaxation of molecules. In our evaluation on the QM7-X dataset, we demonstrate  its usefulness in a thorough study where we show that it is a viable alternative to ML force field models that does not require the expensive training data with unstable structures with force labels, but still learns to map onto the data manifold of stable molecules.

Reviewer comment (Novelty):
The proposed "Diffusion Time Step Prediction" seems to have a loose connection with the main task of molecule relaxation studied in the paper. The authors also mentioned that it appears to be a generic technique that could be applied in other domains like image generation.

Author response:
We respectfully disagree with the reviewer regarding a lack of novelty in our approach. The prediction of the diffusion time step is a novel contribution and applying diffusion as denoising to perform relaxation was not done before. Furthermore, the diffusion time predictor  is not “loosely connected” to the relaxation but a prerequisite for applying the diffusion model framework to relaxation: It enables the reverse-diffusion from arbitrarily noisy inputs, e.g. molecules that are close to their stable conformation. This is because diffusion models use the time step $t$ as an input, and therefore we must determine  from which $t$ the reverse diffusion should start for a given unstable molecule. We have added Fig. 2a to show how the predicted timestep relates to the RMSD of the unstable molecules (i.e. their distance to the ground-truth stable structure). One can see that the time step prediction enables the model to start relaxation at different noise levels, depending on the configuration of the unstable structure. 

Additionally, our analysis show that adaptive scheduling based on the diffusion time prediction is beneficial as it allows the model to flexibly adjust the process (i.e. jump back and forth in time) which leads to significantly improved performance of the two flexible models (MoreRed-JT and MoreRed-AS) over the model where only the starting point is predicted (MoreRed-ITP). We added a figure to show some example trajectories of predicted time steps over relaxation that visualize this flexibility (see Fig. 2b and Section 4.1). This is also where the applicability of the time step prediction to other diffusion-based frameworks comes from: We show that having a flexible reverse diffusion process might not only be beneficial when starting from non-Gaussian input such as unstable molecules but also for data generation from scratch, e.g. when sampling molecules or images from Gaussian noise (see Tables 1 and 2).

Therefore, the time step predictor is both an integral part of MoreRed and a promising contribution to other fields. We argue that our approach is an innovative method for relaxation of molecules. In our evaluation on the QM7-X dataset, we demonstrate  its usefulness in a thorough study where we show that it is a viable alternative to ML force field models that does not require the expensive training data with unstable structures with force labels, but still learns to map onto the data manifold of stable molecules.

Reviewer comment (Novelty):
If this paper solely relies on diffusion to address the molecule relaxation task, I think it lacks significant innovation.

Author response:
We respectfully disagree with the reviewer regarding a lack of novelty in our approach. The prediction of the diffusion time step is a novel contribution and applying diffusion as denoising to perform relaxation was not done before. Furthermore, the diffusion time predictor  is not “loosely connected” to the relaxation but a prerequisite for applying the diffusion model framework to relaxation: It enables the reverse-diffusion from arbitrarily noisy inputs, e.g. molecules that are close to their stable conformation. This is because diffusion models use the time step $t$ as an input, and therefore we must determine  from which $t$ the reverse diffusion should start for a given unstable molecule. We have added Fig. 2a to show how the predicted timestep relates to the RMSD of the unstable molecules (i.e. their distance to the ground-truth stable structure). One can see that the time step prediction enables the model to start relaxation at different noise levels, depending on the configuration of the unstable structure. 

Additionally, our analysis show that adaptive scheduling based on the diffusion time prediction is beneficial as it allows the model to flexibly adjust the process (i.e. jump back and forth in time) which leads to significantly improved performance of the two flexible models (MoreRed-JT and MoreRed-AS) over the model where only the starting point is predicted (MoreRed-ITP). We added a figure to show some example trajectories of predicted time steps over relaxation that visualize this flexibility (see Fig. 2b and Section 4.1). This is also where the applicability of the time step prediction to other diffusion-based frameworks comes from: We show that having a flexible reverse diffusion process might not only be beneficial when starting from non-Gaussian input such as unstable molecules but also for data generation from scratch, e.g. when sampling molecules or images from Gaussian noise (see Tables 1 and 2).

Therefore, the time step predictor is both an integral part of MoreRed and a promising contribution to other fields. We argue that our approach is an innovative method for relaxation of molecules. In our evaluation on the QM7-X dataset, we demonstrate  its usefulness in a thorough study where we show that it is a viable alternative to ML force field models that does not require the expensive training data with unstable structures with force labels, but still learns to map onto the data manifold of stable molecules.

Reviewer comment (Novelty):
The "time step prediction" aspect doesn't seem closely related to the main task.

Author response:
We respectfully disagree with the reviewer regarding a lack of novelty in our approach. The prediction of the diffusion time step is a novel contribution and applying diffusion as denoising to perform relaxation was not done before. Furthermore, the diffusion time predictor  is not “loosely connected” to the relaxation but a prerequisite for applying the diffusion model framework to relaxation: It enables the reverse-diffusion from arbitrarily noisy inputs, e.g. molecules that are close to their stable conformation. This is because diffusion models use the time step $t$ as an input, and therefore we must determine  from which $t$ the reverse diffusion should start for a given unstable molecule. We have added Fig. 2a to show how the predicted timestep relates to the RMSD of the unstable molecules (i.e. their distance to the ground-truth stable structure). One can see that the time step prediction enables the model to start relaxation at different noise levels, depending on the configuration of the unstable structure. 

Additionally, our analysis show that adaptive scheduling based on the diffusion time prediction is beneficial as it allows the model to flexibly adjust the process (i.e. jump back and forth in time) which leads to significantly improved performance of the two flexible models (MoreRed-JT and MoreRed-AS) over the model where only the starting point is predicted (MoreRed-ITP). We added a figure to show some example trajectories of predicted time steps over relaxation that visualize this flexibility (see Fig. 2b and Section 4.1). This is also where the applicability of the time step prediction to other diffusion-based frameworks comes from: We show that having a flexible reverse diffusion process might not only be beneficial when starting from non-Gaussian input such as unstable molecules but also for data generation from scratch, e.g. when sampling molecules or images from Gaussian noise (see Tables 1 and 2).

Therefore, the time step predictor is both an integral part of MoreRed and a promising contribution to other fields. We argue that our approach is an innovative method for relaxation of molecules. In our evaluation on the QM7-X dataset, we demonstrate  its usefulness in a thorough study where we show that it is a viable alternative to ML force field models that does not require the expensive training data with unstable structures with force labels, but still learns to map onto the data manifold of stable molecules.

Reviewer comment (Experiments):
The application of the proposed method needs to be described so that the impact of this method is not clear. I would like to see how the proposed molecule relaxation can be used in chemistry or biology.

Author response:
Thank you for pointing this out. We updated the Introduction to reflect the possible impact of our framework more clearly. Our main target application is to find stable molecule structures by relaxing unstable structures proposed by cheap methods, e.g., with empirical MMFF, and our MoreRed can complement the well-established force field (FF) models in the domain where unstable structures with force labels for training FF models are not available or hard to obtain, e.g. when using very accurate but costly QM methods or when working with large structures. This could, for example, have a huge impact on crystal structure search, where large databases like MaterialsProject [2] exist, but only provide stable structures with labels. In such cases, all FF approaches cannot be used due to the lack of unstable structures. In our studies, we thoroughly evaluate our method on the QM7-X dataset as it allows us to train both MoreRed and an ML force field for a proper comparison. Our main contribution is that we propose a novel approach to perform structure optimization based on learning a pseudo potential energy surface using solely stable structures. While achieving promising results in systems in the scale of QM7-X opens a new promising path, extending diffusion models to handle periodic boundary conditions or very large molecules still is an open challenge in itself, and thus remains a task for future work [3].

Reviewer comment (Evaluation):
The proposed method is 10 times slower than the structure optimization using a force field model.

Author response:
This is partially true. We reported an average inference time per step for a batch of 128 structures of 0.03 s for the force field and 0.05 s for the MoreRed variants and a median (not mean because the distribution of the number of steps is highly skewed) number of relaxation steps until convergence of 118 for the FF model, 53, 219, and 1000 for MoreRed-ITP, MoreRed-JT, and MoreRed-AS, respectively. Hence, it is true that using the best performing variant MoreRed-AS takes longer in the median. However, the simplest variant that only predicts the starting $t$ (MoreRed-ITP) requires less steps than the ML force field in median while still outperforming it with regard to the RMSD ratio (see Figure 3a), and the MoreRed-JT variant, which outperforms the force field model and requires only two times more steps in median. On top of that, we find that all MoreRed variants can easily be parallelized during sampling and run the whole denoising on the GPU, while parallelization of the force field with the L-BFGS optimizer can be non-trivial. The parallelized optimization of 20k structures in our experiments took about 3 hours for MoreRed-AS on NVIDIA A100 while the sequential optimization with the force field model took about 49 hours utilizing the same hardware requirements and using L-BFGS from the ASE library, according to the log files. Because of the added overhead of loading from and to the GPU and converting to and from ASE classes after each relaxation iteration using L-BFGS in ASE, the running time increases further. We will give a more detailed theoretical calculation on the computation time for structure optimization in the questions section below as asked.

Reviewer comment (Evaluation):
It needs to be explained why the RMSD ratio, the RMSD after relaxation divided by the RMSD of the unstable initial structure, is used for the comparison. I am worried that the proposed method can be only accurate for initial structures near their relaxed structures.

Author response:
We agree that this might be confusing. We decided to use the RMSD ratio as it is a concise way to express whether the relaxation brought the structure of interest closer to its stable state or not (cases with ratio > 1 have not succeeded in getting closer to the stable state and we refer to these as outliers or failing cases). Besides, when reporting only the RMSD, it could be hard to interpret because it is not a normalized measure and it can happen that one model shows less RMSD level but fails in more cases as one can see in the higher outlier percentage for the force field model than MoreRed variants. However, in order to make our results more transparent and convincing, we added several plots in Fig. A6 which shows the RMSD error (not the ratio) and updated the manuscript to discuss these results in Appendix A.4.3. It can be seen that MoreRed performs well on both structures initially close to or far away from their relaxed structure. Only the simplest variant, which predicts only the starting time step and applies reverse diffusion with fixed scheduling (MoreRed-ITP), shows impaired performance for unstable QM7-X structures with comparatively larger RMSD. This can be explained by the restricted flexibility of this variant. The adaptive scheduling of the other two variants allows them to dynamically steer the relaxation and push the sample closer to the target distribution by moving back and forth in time to account for its errors (see newly added Fig. 2b with trajectories of predicted time steps). Hence, these two variants are clearly more successful in relaxing the test structures with large RMSD. This reinforces the significance of the time step prediction for our approach and we have updated Section 4.1 to better reflect this.

Reviewer comment (Theory):
No explanation why the diffusion time step prediction is also helpful for unconditional generation from normal distribution prior.

Author response:
Thank you for bringing this up. As mentioned in our previous point, the MoreRed variants with adaptive scheduling achieve the best results. This comes from their flexibility in running longer/shorter relaxation with adapted time steps. They can account for errors resulting from the stochasticity of the noise prediction by going back in time or skip to lower time steps if a relaxation step brings them close to the data manifold (as can be seen in the newly added Fig. 2b and Fig. A8 left with trajectories of predicted time steps). We also added a more detailed general response 2 above for all reviewers. Besides, we have added an explanation to Section 4.3 in the paper and in details in the Appendix A.4.5 about  why the time step predictor is also promising for unconditional data generation.

Reviewer comment (Writing):
What is the application of molecule relaxation without energy/force prediction? I guess you have some assumptions about applying the proposed method, but it is not described in the main text.

Author response:
As mentioned in our first point above, our target application is to find stable molecule structures by relaxing unstable structures proposed by cheap methods. Users of course want to know the physical property of the obtained stable structures.  To this end, one can for example use QM methods applied only to the stable molecules, which is much cheaper than those applied throughout the whole relaxation process.  Alternatively, one can train an additional output head or separate model to predict physical properties, which can  also be  trained only on stable structures.

Reviewer comment (Experiments):
Is it possible to output the energy or force of stable structures? It extends the applications of this method, such as crystal structure predictions. If the energy prediction head is added to the proposed method, Matbench Discovery[1] can be a good platform for evaluating the proposed method.

Author response:
Yes it is possible to add an energy prediction head to the backbone representation, similar to the prediction of the time step, which are both invariant quantities. Note that we are using the same backbone used to train the force field model, which predicts the energy as well. However the Matbench Discovery dataset can currently not be used for evaluation, because our method focuses on the optimization of molecules, not materials and on the accuracy of the optimized structure. Extending diffusion models to handle periodic boundary conditions still is an open challenge in itself [3] and remains as a task for future work that can open potential paths.

Reviewer comment (Evaluation):
In the second paragraph of page 8, the authors mention that adding their synthetic noise on stable molecules causes force field methods to not be able to optimize the molecules to their stable starting point, whereas their diffusion model is able to perform this. However, this is not a valid benchmark, since the force field methods are supposed to predict forces on physically plausible molecules whereas the diffusion process that the authors have designed makes no guarantees about the plausibility of the molecules. This experiment can be reframed to show a weakness of the proposed method, in which unphysical starting points still become "stable", which should not be the case.

Author response:
It is correct that our method exhibits a strong robustness with respect to the starting point of the relaxation. It is trained to converge to physically plausible molecules regardless of the physicality of the input. We believe that this robustness is a significant advantage, not a weakness, since it expands the space of the input where the method can be applied. In our assumed applications, the goal is to relax structures that are proposed by some cheap empirical force field. A robust model would allow one to freely choose the cheap method, e.g. using different implementations from different chemoinformatics packages. However, classical ML force fields tend to fail as soon as the starting point stems from a different method than used when assembling the training data set. This is one of their main weaknesses and it occurs because the input is out of distribution for the ML model (regardless of the physicality of the input). Therefore, we believe that our proposal of a more robust approach is a valuable contribution.

For our application of relaxing structures obtained from cheap methods, we do not see any problem if MoreRed relaxes unphysical starting points to stable states. It would be helpful if the reviewer would further explain the situations in which this behavior is undesirable, so that we can discuss this point in the paper. If desired, one might add sanity checks to identify relaxations of unphysical starting points. For example, MoreRed is trained to map unphysical systems to physically plausible, stable systems. Accordingly, the RMSD of atom positions before and after relaxation will be comparatively large if a relaxation starts from an unphysical system, since it will involve a lot of movement of the atoms. Such sanity checks are cheap to compute and also often employed when running relaxations with DFT or with ML force fields (note that these methods also have no guarantee not to converge when the starting point is  unphysical).

Reviewer comment (Evaluation):
Similar to the criticism above, I don't believe that the time step prediction performance as reported in section 4.1 is relevant as independent Gaussian noise on the atoms is somewhat easier to predict than stable vs unstable molecules. Since QM7-X contains unstable structures as well, a better approach would be to see the correlation between the time step predictions of the network versus the ground truth RMSDs.

Author response:
Thank you for this idea. We added Fig. 2a where we examine the correlation between the RMSD of unstable structures in QM7-X and the predicted initial time step by our model.  It mainly shows two insights: There is a clear correlation between the RMSD and the predicted time step and the model always predicts time steps greater than zero for these unstable structures. There are few outliers in MoreRed-AS/-ITP where the model predicts larger time steps for structures with intermediate RMSD. However, we examine the performance of all three MoreRed variants with respect to the RMSD before and after relaxation in the newly added Fig. A6c and their performance is not impaired for structures with intermediate RMSD.

Reviewer comment (Evaluation):
The test runs should be done with a few different random seeds so that we see how much of a discrepancy exists between runs for this method, as it is statistical in nature.

Author response:
We agree, and have added Figure A6b to our paper to report the standard deviations across 5 test runs with different random seeds. Overall, the observed deviations are low, and do not affect our conclusions.

Reviewer comment (Novelty):
How does the time step prediction compare to non-fixed variance approaches such as [1] and [2]? It seems very similar to the approach taken there, why not mention how your work is different?

Author response:
In [1,2], the conditional variance of $x_{t-1}$ given $x_t$ is learned for improving the test likelihood and data generation quality when the reverse diffusion steps are reduced.  Although we also use an additional model to predict the diffusion time step, which correlates to the noise component of $x_t$, our approach is fundamentally different from [1,2].  First, the conditional variance corresponds to the noise amplitude that a single reverse step should ADD after the noise removal, which corresponds to the last term in Eq.(3) in our paper, while our diffusion time predictor implicitly measures how much noise the whole remaining reverse steps should REMOVE from the current sample, i.e., $\|x_t - x_0\|^2$.  Apparently, the conditional variance prediction in [1,2] cannot be used to predict the diffusion time for setting the initial diffusion time for molecule relaxation nor for adaptive scheduling.  Second, [1,2] always use fixed scheduling, even if it is different from the training scheduling, while we use adaptive scheduling by using the diffusion time predictor.

Reviewer comment (Theory):
Dataset scaling is frequently mentioned as a benefit to this method, although it seems that with larger datasets, the probability that some of the modes of the mixture of Gaussians you use as your prior distribution would get mixed. That is, once you add some noise to some training structure \( y_i = x_i + \epsilon \), then you might have the issue that for some \( j \neq i \), \( ||y_i - x_i|| > ||y_i - x_j|| \). Is this a problem?

Author response:
The mixture components get overlapped when the training samples are dense or the noise level is large, and this degrades the diffusion time prediction performance for samples with large noise, as seen in Fig.A1 right (which was Fig.2 in the original submission).  However, the prediction for low noise samples is barely affected,  unless the dimension $D$ of the data space is very small.  Assume that there are two training molecules $x_a, x_b$ with the distance $r = \|x_a-x_b\|$, and consider Gaussian blurred samples of $x_a$ with the standard deviation equal to $\sigma = r$.  Although, in this situaion, $x_b$ lies in the high density shell of the Gaussian blurry of $x_a$ (the bump in the left plot in Fig A.1 left in the revision), the Gaussian blurred samples of $x_a$ are uniformly distributed in the still high $(D-1)$-dimensional shell, and therefore the probability that the Guassian blurry produces a sample close to $x_b$ is extremely low.  On the other hand, the diffusion time predictor is trained with a lot of samples from the neighborhood of $x_b$ as low noise samples from $x_b$ because the density is high.  Accordingly, the models are trained so that they recognize the molecules close to $x_b$ as low noise samples.  This intuition can be mathematically confirmed by computing the density ratio between the two mixture components around $x_b$, i.e., $Gauss_D(x_b + \epsilon; x_b, \delta^2) / Gauss_D(x_b + \epsilon; x_a, r^2)$ for $\|\epsilon\| ~ \delta << r$, which is extremely high when $D$ is not very small. We added this explanation in the discussion in Appendix A.2.1.

Reviewer comment (Experiments):
For table 1, there are many methods that generate molecules on QM9, why not report some of those results as well?

Author response:
We added comparisons with E-NF [1], G-SchNet [2], EDM [3] and GeoLDM [4]  in Appendix A.4.5 and Table A1.  However, please note that the baselines and our MoreRed solve different problems. The baselines, which learn $p(X,Z)$ from the QM9 dataset, generate both atom types $Z$ and positions $X$, while our MoreRed, which learns $p(X|Z)$, generate atom positions given the atom types. We restricted our scope to the atom position generation because our main goal is to perform molecule relaxation.

**Citations:**

[1] Garcia Satorras, Victor, et al. "E (n) equivariant normalizing flows." Advances in Neural Information Processing Systems 34 (2021): 4181-4192.

[2] Gebauer, Niklas, Michael Gastegger, and Kristof Schütt. "Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules." Advances in neural information processing systems 32 (2019).

[3] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

[4] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." International Conference on Machine Learning. PMLR, 2023.

Reviewer comment (Experiments):
**Regarding the evaluation settings**: As stated in Section 4.2, the authors use the QM7-X dataset to evaluate the molecule relaxation performance. The unstable molecules in this dataset were generated by sampling from the Boltzmann distribution of the stable molecules. The confidence of such data generation procedure is in doubt: (1) In real-world applications, given a molecule, we usually use either random conformation sampling or cheap conformation optimization (e.g., empirical MMFF) to obtain initial molecular structures. The authors should further clarify whether these generated unstable molecules in QM7-X match the real-world settings. (2) Beyond the small organic molecules, molecular relaxation is widely used to investigate the equilibrium state of molecular systems like the adsorbate-catalyst complex or protein-ligand complex. There exist large-scale benchmarks that are more related to real-world settings like Open Catalyst Project (IS2RE and IS2RS tasks). It would be more convincing to verify the proposed methods on these challenging tasks.

Author response:
We agree that the distribution of the initial unstable molecules in the real-world setting does not follow the Boltzmann distribution, and in this sense, QM7-X is maybe not the optimal data set for evaluation.  However, due to the high computational cost of generating datasets for chemical structures, only a limited number of these datasets is available to use. For a fair comparison of our method to machine learning force fields on the task of structure optimization, we require a dataset that provides a large number of stable/equilibrium structures with various compositions, as well as corresponding unstable/non-equilibrium variants of these with labels for energy and forces. Although the QM7-X dataset has some disadvantages, it is to the best of our knowledge the only molecule dataset fulfilling these requirements. For this reason it is often used to test the performance of machine learning force field models [6-8], which were proposed for the same real-world applications. Besides, it is common practice to utilize datasets in the scale of QM7-X like QM9, QM7 and MD17 to validate novel machine learning force fields [7-14]. 

The stable structures of QM7-X are generated in exactly the same way you described. First,  a set of sufficiently different initial 3D structures were obtained using the MMFF94 force field for each SMILE string to cover many conformations. Then, all structures were subsequently optimized with accurate DFT simulations. The unstable structures were generated by displacing each molecular structure along a linear combination of normal mode coordinates computed at the DFTB3+MBD level within the harmonic approximation. Based on the equipartition theorem from classical statistical mechanics, the normal modes for the displacements were chosen in such a way that the energy difference between the unstable and stable structures follow a Boltzmann distribution. In their paper, the authors of the QM7-X dataset argue theoretically and provide analysis showing that the resulting structures are physically relevant  and cover the potential energy surface.

The distribution of the unstable structures in QM7-X is indeed different from the real-world distribution, where the candidates are obtained with cheap optimization methods. However, both give physically feasible unstable structures that cover the potential energy surface with non-Gaussian noise, and therefore we expect that the results are still indicative for the realistic scenario. In the revision, we added a discussion on the choice and shortcomings of QM7-X (see Appendix A.4.1).

While large scale benchmarks like the Open Catalyst Project would provide a challenging task to test structure optimization methods, it is a special use case of molecular relaxation, which requires the model to handle structures with periodic boundary conditions, and therefore is out of the scope of our contribution. Extending diffusion models and generative models for atomistic systems to handle periodic boundary conditions is an interesting open challenge in itself [17] and we leave it as future work.

Reviewer comment (Experiments):
**Regarding the generality of the proposed methods**: As a general framework, the proposed MoreRed can use different equivariant backbone models. Given the rich literature of equivariant networks, it would be more convincing to verify the generality of MoreRed with different architectures.

Author response:
Thank you for bringing this up. We agree that it is a valuable experiment to strengthen our work. We conducted a new set of experiments using SO3net [16] as a backbone representation in our framework. This representation incorporates spherical harmonics to handle SO(3)-equivariance, distinguishing it from the PaiNN architecture.

Our findings are summarized in Appendix A4.4 and Fig. A7. Overall, our MoreRed approach works similarly well using this different equivariant neural network backbone and still outperforms the ML force field in terms of  the structure accuracy. The general performance with SO3net is slightly worse than with PaiNN and there is a slight difference between the performance of MoreRed-AS and MoreRed-JT. This likely comes from the fact that, due to the time limitation by the rebuttal deadline, hyperparameter tuning was not feasible and we had to use half the number of parameters employed in PaiNN to speed-up the training and relaxation experiments. Despite this constraint, we maintained fairness by using identical model hyperparameters for both MoreRed and the force field. Moreover, we would like to mention that these constraints were due to the long training time of the force field model, which uses 100 times more data than MoreRed and needs more than 7 days to converge when using large models with optimal parameters. In contrast, MoreRed requires solely stable structures for training and converges in less than two days. This contrast reinforces the efficiency of our approach.

Reviewer comment (Experiments):
**Regarding the quality of unconditional generation**: I do not quite grasp how the diffusion time step predictor improve the unconditional generation quality. Moreover, the experiments also lack strong and advanced baselines [3,4,5].

Author response:
Regarding the improved unconditional generation with the time step prediction, we answered the question in the general answer 2 above, because it was asked by most of the reviewers and we realized that we need to elaborate on it in more detail. Hence, we kindly ask you to refer to the general answer 2 above, where we hope we will be able to clearly address your concerns.

Regarding the “lack of strong and advanced baselines”, we want to clarify that the cited work is not providing baselines for the experiments that we conduct. We sample atom positions from scratch given a composition of atoms as input. [5], on the other hand, generates both atom positions and types, which is different from our setup. Nevertheless, we have added as Table A1  the results by [5] and other established generative models that sample both positions and types.  We elaborate the different experimental setups in Appendix A.4.5. [3, 4] have been proposed for conformer search, the task of sampling diverse conformations given a molecular graph as input. It is related (and already included in our related work section) but a distinct task with different experimental setup and metrics for evaluation. 

We would like to note that the experiments in Section 4.3 were intended to evaluate  how much adaptive schedulingwith our time step predictor improves the traditional fixed schedule under the same condition, and not to demonstrate state-of-the-art performance. However, we thank the reviewer for pointing us into this direction, as evaluating the time step predictor in a state-of-the-art diffusion model for conformer search is an interesting additional experiment that we will explore. Unfortunately, it was infeasible for us to conduct this new kind of experiment (with new data and code base for the model) in the short time span of the rebuttal period.