Straight-Line Diffusion Model for Efficient 3D Molecular Generation

Thank you for the detailed and thoughtful feedback, which has greatly helped improve the clarity of our work. We carefully incorporated these explanations in the revised version. We also sincerely appreciate your recognition of our novel contribution.

We believe the main concern may stem from misunderstandings, and we hope our responses will clarify the key points. If so, we would be truly grateful if you would consider increasing the score.

Lack a direct comparison with UniGEM

Reply: We believe there may be misunderstandings of the experimental results. Contrary to the reviewer’s impression, our work does include direct comparisons with UniGEM. Specifically:

• The EDM(U) baseline used in Table 3 and Figure 7 corresponds to the main model proposed in the UniGEM paper, which uses EDM as diffusion model.
• Our results demonstrate that under the UniGEM framework, SLDM achieves significant improvements in sampling efficiency over UniGEM. For example the mol_stable:
  • T=1000: unigem 89.8 v.s. SLDM 95.42
  • T=50: unigem 85.73 v.s. SLDM 93.37
  • T=10: unigem 23.56 v.s. SLDM 87.46

We intentionally did not include UniGEM in Table 1, because Table 1 is focused on comparing different diffusion models (e.g. BFN, DDPM/EDM, Flow matching), while UniGEM is not a new generative model, but rather a change in the generation paradigm—i.e., decoupling the generation of coordinates and atom types. In fact, UniGEM can be combined with EDM, BFN, or other diffusion models. But if you believe it would improve clarity, we are happy to include UniGEM in Table 1 as well.

Why a straight-line diffusion process is particularly suitable for molecular data (lines 41-45) ? Small perturbations would still introduce noise and potentially generate invalid molecular conformations.

Reply: We believe there might be some misunderstanding of the role of noise scheduling and our statements. For any diffusion-based model, the forward (noising) process inevitably drives the signal-to-noise ratio (SNR) toward zero, resulting in fully corrupted, uninformative structures regardless of whether the trajectory is straight or not.

Our contribution is not to avoid this endpoint, but to control the rate at which information is destroyed. In particular, SLDM employs a slower decay of SNR, meaning the molecular structure is not immediately and completely perturbed at early steps. Instead, noise is introduced more gradually, preserving meaningful spatial relationships longer and enabling a more uniform and gradual reconstruction in the reverse process. This contrasts with EDM schedules that overly disrupt structure early, potentially making recovery harder.

We also note that the brief explanation between lines 41–45 is expanded in detail in Section 3.2 for the benefits of such a schedule in the molecular setting.

Incorporate the novelty metric? SLDM may struggle to explore the whole distributions

Reply: We thank the reviewer for raising this important point regarding novelty metric in molecular generation.

1. On novelty evaluation in QM9: As discussed in the EDM paper [1], novelty metrics are not meaningful on QM9 dataset because it is an exhaustive enumeration of small organic molecules within a narrow chemical space. EDM further observes that novelty consistently decreases as training progresses. Following standard practice, we thus omit novelty reporting on QM9, as do prior works including EquiFM, END, GeoLDM, and UniGEM.
2. On meaningful novelty in broader chemical space: To provide a more meaningful evaluation, we report novelty scores on GEOM-Drugs, a dataset sampling sparsely from a much larger chemical space. As shown below, our method (SLDM) achieves higher novelty than EDM (implemented in UniGEM), while maintaining strong validity and uniqueness: || Atm_stab(%)|Validity(%)|Valid&Unique(%)|Valid&Unique&Novel(%)| |-|-|-|-|-| |EDM (U) (T=1000)|84.93|98.36| 98.31|98.18| |SLDM (T=1000)|88.19|99.96|99.91|99.72| |SLDM (T=50)|89.16|99.52|99.49|99.42|

These results confirm that SLDM can generate novel, valid, and unique molecules in realistic settings. We appreciate the reviewer’s concern regarding the potential risk of memorization in a low-variance setting. However, SLDM does not remove stochasticity, rather, it carefully balances stochasticity over time through the annealed temperature scheme eq. 9. This ensures early-stage exploration while gradually refining the structure in later stages for stability. This design prevents collapse into duplicated training samples while maintaining high fidelity. As shown in Appendix C.4, we tune the annealing rate to maximize the unique & valid score, achieving an effective balance between diversity and chemical stability.

Ablation of constant sigma

Reply: We find our method remain robust across a reasonable range of σ values. We only trained for 1280 epoch due to time limitation of rebuttal.

Random seeds for tab.1

Other issues on clarity:

Line 154: What is “more stable generative process” and why?

Reply: By “more stable,” we refer to the fact that our schedule injects significantly less noise/slower SNR drop compared to EDM in the early steps of the reverse process, resulting in a more gradule, less chaotic trajectory. This stability avoiding large, abrupt shifts that lead to invalid conformations in small t.

Fig2 : backward or forward process?

Reply: It is the forward noising process. Showing the reverse sampling process would not ensure that all models generate the same molecule, making comparison unclear.

Fig2: GeoBFN faster?

Reply: We respectfully note that this is a misinterpretation. The figure shows the diffusion process from t = 1 to t = 0 (left to right). In diffusion models, even when using fewer sampling steps, all methods simulate the full trajectory from t = 1 to t = 0 by discretizing it differently, not by stopping early (e.g., at t = 0.5). GeoBFN’s schedule allocates many steps focusing on fine-tuning geometry, while compressing the challenging structure formation phase into very few early steps, limiting model capacity.

L145-147: why Langevin dynamics make trajectories linear?

Reply: We clarify that Langevin dynamics does not directly make the sampling trajectory more linear, but rather mitigates the impact of deviations from the ideal trajectory at early steps, i.e. improving robustness to initial sampling error. We agree the wording in the original text may be misleading and will revise the manuscript accordingly.

An ablation support this, SLDM maintains high validity and stability (surpassing other baselines in paper) even when the initial distribution deviates from the theoretical optimum(ε: standard Gaussian noise) (Sampling steps T=50)

Eq.7: where does β come from?

Reply: Eq.7 approximates the true reverse transition distribution p(xt−Δt∣xt) using a Gaussian with a mean and variance parameterized by β. While the true reverse distribution is intractable and non-Gaussian, it is possible to derive its true conditional mean analytically under our forward process assumptions, as shown in Eq.41. To ensure our Gaussian approximation is as faithful as possible, we determine β s.t. the mean of Eq.7 matches this true conditional mean.

Alg 2: why inner condition?

Reply: The inner condition (i.e., skipping noise injection at the final step) is a common technique used to improve the final sample quality. Intuitively, this step allows the sampler to directly apply the learned score function to denoise the sample, effectively implementing Tweedie’s formula.This idea has been adopted in previous works such as DDPM [1, Alg.2] and improved score-based models [2, Alg.1], and we follow the same practice here.

[1] Denoising Diffusion Probabilistic Models

[2] Improved Techniques for Training Score-Based Generative Models

We additionally provide ablation results below to show the effect of this choice:

Why longer training time?

Reply: Thanks for the suggestion. We confirm that SLDM already outperforms UniGEM under the same number of epochs and training time:

While longer training leads to slight performance gains, we found that the mol_stable and atm_stable metrics remain stable around 2000 epochs on QM9 and 13 epochs on GEOM-Drugs during the training process. This indicates that extended training time is not a necessary factor for SLDM’s performance advantage.

If there are any further questions or clarifications needed, we would be glad to continue the discussion.

Thank you for your detailed responses. However, I still remain unconvinced about them.

1. Lack a direct comparison with UniGEM

I understood what UniGEM [1] proposed. However, your framework still relies significantly on the UniGEM framework. The results you presented in the rebuttal are unconvincing, as UniGEM's results can be applied to different diffusion models. I would expect to see if SLDM can achieve good performance without the UniGEM approach, meaning predicting atom types and coordinates separately as EDM did. If your model performs well, I would be satisfied with your contributions. If not, I would be afraid that your approach to success relies significantly on the UniGEM's success.

In addition, the claim that UniGEM is not a new generative model is somewhat of a bold statement. I read the UniGEM paper, where the authors claimed that it is the first diffusion-based unified model to successfully integrate molecular generation and property prediction, delivering superior performance in both tasks, as stated in the paper's abstract. Even though they are not a new diffusion model as you said, an ablation study of work on using SLDM alone without UniGEM will strengthen your contributions.

2. Incorporate the novelty metric? SLDM may struggle to explore the whole distribution

Novelty metrics are not meaningful on the QM9 dataset. I have seen some papers using this fact to avoid evaluating the novelty metric on QM9. I think this is not a good excuse and not a good practice. What would make sense to learn and generate your training data again? If it's not a good dataset, it should be excluded from the evaluation.

Moreover, GeoBFN [2] is a framework that evaluates the novelty metric on QM9, achieving reasonable results ranging from 65% to 70%. You compared SLDM with GeoBFN on QM9, but ignored the novel metric.

In contrast, I think GEOM-Drugs is not a good dataset to evaluate the novelty metric, as its search space is very huge due to a wide range of atom sizes. To the best of my knowledge, most diffusion models can easily achieve 100% novelty of this dataset. That's also a reason why several models ignored the novelty metric on GEOM-Drugs.

Again, I am not yet convinced about the linear straight-line diffusion model with a small noise variance in your experimental results. It would be beneficial to measure the novelty metric of this approach on the challenging dataset QM9, as you did for GEOM-Drugs.

3. FIGURE 2

I am still unclear on your answers to Figure 2. The figure shows the diffusion process from t = 1 to t = 0 (left to right)' which is exactly what I read about the figure from left to right. I could not interpret your phrase, the diffusion process from t = 1 to t = 0. When I think about a diffusion process, it is a forward noise process that should go from t=0 to t=1. Could you make your interpretation clearer?

It is the forward noise process. Showing the reverse sampling process would not ensure that all models generate the same molecule, making comparison unclear. I still could not understand your phrase. If, as you said, this is a forward process, I still see the same molecules on the right side of the figure. Why would you say they are not the same molecules?

I began to become skeptical about Figure 2. I thought that you would try to show the generation/ sampling process as GeoBFN did in their Figure 3 [2], in which I read that GeoBFN generated stable molecules faster than EDM. However, through the rebuttal, it seems not to be what I thought. And I think it's a good way to visualize, as GeoBFN did in Figure 3, because SLDM also tried to generate stable molecules faster. So, I am still unsatisfied with your answers on Figure 2. In addition, I think this is the most important and intuitive figure of your manuscript, visualizing the performance of different approaches.

I understand that you have attempted to strengthen your approach through theoretical development and additional experimental results in the rebuttal. I would appreciate this and consider it a good point. There are still some minor things that can be discussed in your manuscript. However, given the limitation on the rebuttal period, it would be sufficient to clarify the above questions. Thank you!

[1] Feng et al. UniGEM: A Unified Approach to Generation and Property Prediction for Molecules. In ICLR 2025.

[2] Song et al. Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks. In ICLR 2024.

Thank you for your thoughtful feedback! I believe your novel results, which importantly prove that SLMD can work independently with GeoBFN. Thus, my primary concern has been solved. However, my belief in Figure 3 of GeoBFN remains unchanged. It should be a generation process, as explicitly stated in Section 3.3 and the GeoBFN author's discussion with the Reviewer mPXU on OpenReview. As the authors have provided substantial novel results during the rebuttal, I will increase the score to 5 and recommend for acceptance.

Best,

Dear reviewer,

We sincerely thank you for your dedication during the review process. Your professional insights and constructive suggestions have significantly improved the overall quality of the manuscript. We have carefully recorded all points raised and will incorporate them into the revised version to ensure better rigor and clarity. Thank you again for your time, expertise, and thoughtful suggestions.

Wishing you all the best!

We sincerely thank the reviewer for the thoughtful comments and for recognizing the motivation and contributions of our work. If deemed useful, we will include the discussion in our paper.

I suspect directly adapting the progress in fast image generation, e.g. consistency model, to molecular generation can already give significant speed-up. More such experimental comparison is favorable.

Reply: Thank you for this insightful suggestion. We fully agree that progress in accelerating diffusion models from other domains is potentially inspiring for molecular generation. We have discussed their connection with SLDM in Appendix D.2.

However, in practice there are nontrivial domain-specific challenges that limit their effectiveness in molecular settings. For example, a recent blog post “Equivariant Diffusion for Molecule Generation in 3D using Consistency Models” attempted to apply the Consistency Models to the 3D molecule generation task using the EDM framework (which shares the same architecture and evaluation protocol as ours). Although they achieved up to 24× speed-up, the best atom stability reported was only 19%, which is significantly lower than both our method and most baselines (>95%). In our view, this reflects a domain gap that requires careful adaptation, beyond simple plug-and-play transfer, underscoring the importance of developing specialized and efficient generative models for molecular data — which is precisely the goal of our work. We believe transferting image domain techniques to molecules is an important future research.

In Appendix A.4, I don't quite understand the logic jumping from the toy example of delta distribution to a general data distribution. Is it just a loose analogy, without rigorous mathematical reasoning or intuition?

Reply: The delta-distribution case in Appendix A.4 serves as an intuitive starting point, not a derivation for general data distributions. Specifically, our goal is to design a schedule that yields as linear sampling trajectories for any data distributions as possible. As a necessary condition, the linearality should holds for any special data distribution. Due to its analytical tractability, we use the delta distribution case to derive a necessary condition for the forward process: a constant noise level (σ = const) and a linearly interpolating mean (μ(t) linear). The necessary condition and boundary condition(line85-86) determine the SLDM schedule.

We then move to the general case and provide a separate, formal analysis (Thm3.1) showing that SLDM produces approximately linear trajectories even for realistic data distributions. So the toy case is used only as motivation, while the general result is proved independently. We believe this structured approach, from intuition to formal validation, provides both interpretability and generality.

If there are any further questions or clarifications needed, we would be glad to continue the discussion. If our responses have addressed your concerns or contributed to a clearer understanding of our work, we would deeply appreciate your consideration for a higher score or confidence, which would greatly support the visibility and potential impact of our work within the community!

Thank you for your recognition and the thoughtful suggestions, many of them inspired us to improve both the clarity and completeness of the paper. We will carefully revise the manuscript to address these points.

Unigem rationale discussion

Reply: We would first like to clarify that the generation framework used in UniGEM is not the contribution of our paper, but we are happy to share our understanding on this design choice. We fully agree with the reviewer that atom identities typically determine the geometry. However in de no vo molecule generation, both atom types and coordinates need to be modeled. Typical way is to jointly iteratively refine atom types and coordinates by diffusion model. However, modeling discrete variables (like atom types) in diffusion models is challenging. UniGEM notes that atom type generation by diffusion can lead to instability and mode oscillation (i.e., frequent switching between atom types during denoising), which in turn negatively impacts coordinate generation. By first generating continuous coordinates alone, UniGEM avoids this problem and achieves more stable and accurate 3D predictions.

This idea has been further supported by [1] (Sec. 4.3 last paragraph), showing that generating positions first and predicting bonds/types later leads to better results.

Regarding:"Is the idea here to learn geometries implicitly contain atom-type information and purely from data, then fill in atom types post-hoc?" Yes, that is indeed the key assumption behind UniGEM.

[1] Piloting Structure-Based Drug Design via Modality-Specific Optimal Schedule

Conditional generation and baseline (classifier).

Reply: The classifier is not part of the generative model, but is used post hoc to evaluate whether generated molecules match the conditioning properties. We apologize for any confusion and will clarify the setup in the revised version. All settings follow baselines. Briefly:

• QM9 is split into two halves: one to train the property predictor 𝜙, the other for training the conditional generative model.

• 𝜙 evaluates generated molecules; conditional accuracy is measured by MAE between predicted and target properties.

• Conditioning machanism: property values are concatenated to atom types as EGNN's invariant input. Though simple, this method performs well and even outperforms recent guidance-based approaches: |Category|Model|Cv|µ|α| ∆ε|HOMO|LUMO| |-|-|-|-|-|-|-|-| |Classifier-guidance|EEGSDE|0.941|0.777|2.50|487|302 |447| |Training-free guidance|TFG-Flow|1.750|0.817|2.32|804|364 |941 | |Conditional generation|Cond-Flow|1.520|0.962 | 3.10|805|435|693 | |"| Cond-EDM|1.065|1.123|2.78|671|371|601| |"| SLDM|0.745|0.797|1.46|440|320|348|

Classifier-guided methods require an auxiliary classifier at each diffusion step, which incurs high computational cost. Since our goal is to evaluate the ability and efficiency of the generative model, we chose not to include such methods, but we’re glad to include them if you find it helpful.

Lack of mol_stable metric in GeomDrugs and compare to VoxMol

Reply: We agree that molecular stability is a key metric for 3D molecule generation. For GEOM-Drugs, the low molecular stability is not specific to our method. As noted in [1], EDM-style bonding heuristics are not well-suited for this dataset due to its diverse structures. Even ground truth molecules score poorly under this evaluation (see Table below). Despite this, our model (SLDM) shows clear improvement, validating the benefit of our low-temperature sampling (lines 189–193).

[1] GEOM-Drugs Revisited: Toward More Chemically Accurate Benchmarks for 3D Molecule Generation

Regarding VoxMol, we appreciate the reference. VoxMol adopts a different evaluation setup, which use via OpenBabel to get bond and allows charged atoms. We re-evaluated SLDM under their protocol and outperforms VoxMol:

Allowing charged atoms may overly relax evaluation and may be unsuitable for models that do not generate charges. Thus, we did not adopt this metric, but we’re happy to include the comparison in the appendix if you find it helpful.

Slow training

Reply: Thanks for point this out. We compared the performance of SLDM with other baselines, such as UniGEM.

We observes SLDM already achieves state-of-the art result under the same number of training epochs and training time. While longer training leads to slight performance gains, we find that the mol_stable and atm_stable metrics remain stable around 2000 epochs on QM9 and 13 epochs on GEOM-Drugs during the training process. This indicates that extended training time is not a necessary factor for SLDM’s performance advantage.

Clarify why the SLDM formulation results in the deviation from linear sampling at early steps, ultimately correcting to linear?

Reply: The deviation from a linear trajectory in SLDM primarily originates from the trajectory crossing phenomenon in the forward diffusion process, as discussed in Rectified Flow. Although the forward noising process is defined as $x_t=(1−t)x_0+σ(t)ϵ$ which suggests a linear mean trajectory $(1 - t) x_0$, the added stochasticity introduces crossing between diffusion paths, i.e. the same $x_t$ can be derived from multiple different $x_0$. This crossing effect becomes especially pronounced at early sampling steps ($t\to1$), where the noise dominates and leads to high uncertainty in the reverse mapping.

Specifically, in SLDM, the reverse ODE is given by Eq.36 $\frac{dx_t}{dt} = -𝐸[x_0|x_t]$. If $𝐸[x_0|x_t]$ were constant, the sampling path would be linear. But due to forward crossing, $𝐸[𝑥_0∣𝑥_t]$ shifts over time: from dataset average at high t to more precise $𝑥_0$ as t→0, resulting in curved trajectories. This behavior is quantitatively characterized in Theorem 3.1 and visualized in Figure 6.

Notably, this issue is common in diffusion models. SLDM controls this deviation within initial generative steps (Theorem 3.1), and our stochastic sampling helps mitigate error accumulation over time. Other approaches often rely on costly training techniques like distillation or solving OT to overcome this issue, as discussed in Appendix D.2.

Clarify Figure 2 setup, memorizing issue

Reply: We apologize for the confusion. Figure 2 shows the forward noising process. We used this for consistent comparison across EDM, GeoBFN, and SLDM on the same target molecule. That said, we also visualized the reverse sampling process (not shown here due to rebuttal limitations on links/media), and found consistent trends:

• EDM produces very noisy intermediates throughout;
• GeoBFN quickly settles, barely changing after early steps, underusing later steps;
• SLDM shows less noisy and more progressive refinement.

Importantly, our sampling method is stochastic, allowing flexible rearrangement of atomic positions. This helps avoid rigid memorization and promotes diversity. To demonstrate this, we show that SLDM trained on GEOM-Drugs can generate novel molecules not seen in the training data:

As for annealed Langevin temperature in Eq.9, it enables fine-grained control over the stochasticity by injecting more randomness early (for exploration), and reducing it later (for precision), which we found helpful for balancing chemical stability and diversity.

If there are any further questions or clarifications needed, we would be glad to continue the discussion. If our responses have addressed your concerns or contributed to a clearer understanding of our work, we would deeply appreciate your consideration for a higher score or confidence, which would greatly support the visibility and potential impact of our work within the community!

Dear reviewer,

We sincerely thank you for your time and effort in evaluating our work. We are delighted that our clarifications have resolved your concerns, and we truly value your engagement in this process.

Wishing you all the best!

We sincerely thank the reviewer for the thoughtful comments. We will carefully revise the manuscript to address these points.

Fast sampling from diffusion models usually relies on ODE solvers instead of directly sampling from the SDE trajectory. Could the authors add ODE results for the experiments in Section 4? I wonder if using ODE solvers can get even faster generative model.

Reply: Indeed, ODE-based sampling is a common strategy to accelerate generation in diffusion models. In our work, this is taken into account when comparing with the baseline EquiFM, which leverages flow matching with multiple efficient ODE solvers. We have already compared with the best result of EquiFM in Figure 1 and Table 1, showing that SLDM outperforms in both efficiency and sample quality under comparable settings. To further address your suggestion, we will include EquiFM results in Table 2 for the conditional generation task on QM9. The following table summarizes the results (lower is better):

It would be better if the authors can provide some qualitative examples of the generated molecules to help readers intuitively understand the generation quality.

Reply: Thank you for your constructive suggestion. We have randomly generated 40 molecules respectively using our trained models on QM9 and GEOM-Drugs in the unconditional setting, and visualized their 3D structures using PyMOL. These examples clearly demonstrate the structural plausibility and diversity of our generated molecules. Unfortunately, the rebuttal phase does not allow us to include external links or supplementary material. However, we will definitely include representative visualizations in the appendix of the revised manuscript to provide readers with a more intuitive understanding of the generation quality.

We sincerely thank the reviewer for recognizing the strong performance and contribution of our work. We've open-sourced our code at submission, as we genuinely hope our method can be useful to the community. We’d greatly appreciate a more favorable score to help increase the visibility and impact of this contribution! If there are any further questions or clarifications needed, we would be glad to continue the discussion.

In light of the rebuttal and the other reviews/rebuttals, please edit this comment and update with reasoning for /why/ you keep the rating unchanged.