Accelerating 3D Molecule Generative Models with Trajectory Diagnosis

Weaknesses 1：Relationships to prior works

We thank the reviewer for their careful assessment of our contribution. While our framework follows prior research [1, 2], our core contribution lies in identifying and addressing a fundamental, previously unappreciated characteristic of 3D molecular generation that has significant implications for efficiency.

Driven by this novel insight, we introduce two distinct and highly effective techniques. Our empirical results demonstrate that this dual approach delivers a magnitude-order improvement in sampling speed over the state-of-the-art. Furthermore, our techniques are complementary and can be integrated into various generative frameworks, highlighting their broad utility.

Beyond the demonstrated speedup in 3D molecule generation, we believe our work offers insights into essential elements for modeling 3D molecular structures. Specifically, our findings suggest that a strong inductive bias towards geometric priors can significantly simplify the inherent complexity of 3D geometric generation.

Weaknesses 2. Out-of-distribution generation

Thanks for your insightful comment! Current generative framework, training objectives and even evaluation metrics tend to encourage models to fit the empirical distribution of the training data. As you mentioned, we also recognize the importance of generating novel and valid molecules beyond the training set, which holds significant practical value. However, evaluating generalization ability remains an open challenge. To address this, we conducted additional empirical studies and highlight three key points to demonstrate the model’s generalization capabilities:

Novelty and Validity as proxies for generalization:

As shown in Table 1 in the original manuscript, our method achieves state-of-the-art novelty metric: V&N&U (67.03% on QM9) while maintaining high validity: 96.04%. This suggests that the geometric prior does not collapse to memorizing training structures but instead guides exploration of chemically plausible regions beyond the training distribution.

Generalization across datasets:

As mentioned in the main text, we have explicitly tested the prior’s transferability to external dataset. We constructed the geometric-prior using training data from GEOM-DRUG, and plugged it into the generative model trained on QM9—a dataset with distinct structural patterns (smaller molecules, different atom types). Using 12 NFEs we generated 10000 molecules and evaluated on QM9:

As the results show, using prior constructed from external dataset achieves the same accelerating effect as original MolTD, demonstrating that the prior captures fundamental geometric structures shared across datasets.

Generalization to a number of atoms that do not exist in the training set

For this task, we generate molecules with 28 atoms, as QM9 does not contain a molecule with 28 atoms and most of its molecules have less than 28 atoms (>99%). We construct prior using GEOM-DRUG dataset and plugged it on MolTD trained and evaluated on QM9, and generated 1000 molecules with 28 atoms. The results show competitive performance with SOTA using only 12 steps:

Q1: How the novelty metric is computed?

We appreciate the opportunity to clarify the computation of the novelty metric, which consist of three steps: First, we select the generated molecules that can be successfully transformed to SMILES strings using RDKit. Then, among these valid SMILES strings, we select a subset of the unique SMILES strings. Finally, if a unique and valid SMILES strings did not appears in the training set, then we called it a novel molecules, and the novelty metric is the percentage of novel molecules among the uniqe ones. Furthermore, the metric V&N&U is the percentage of novel molecules among all generated molecules.

Q2: How rotational symmetry evolves throughout the generative process?

Thanks for the insightful question! As detailed in the main text, the geometric-informed prior utilized the EOT algorithm to align the geometric structures of molecules, which eliminates the variations due to rotation and translation. As a result, the geometric-informed prior simplifies the multiple rotational symmetries to just one orientation, and accelerates the stabilizing of rotation. Proposition 4.5 formalizes the intuition as the rotational-equivariant property of generated molecules.

Furthermore, we provide empirical evidence for the acceleration effect. Similar to the main text, we analyzed the generative trajectory with respect to rotation, by measuring the rotation angle between intermediate molecules to the final molecules. The result shows that MolTD requires less rotational transform to align intermediate molecules to the final molecules, demonstrating the acceleration effect of our method in terms of rotation symmetry.

Following the NeurIPS 2025 instruction, we can not upload image in this stage of rebuttal. We will update the results in the revised manuscript.

Reference

[1] Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation. In NeurIPS23.

[2] Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks. In ICLR24.

Thank you for your detailed answers. I think this is a good paper! I am happy to keep my score and raise my confidence score to 5.

Thank the thorough review and insightful suggestions! Your feedback will help us significantly improve our work. We address your concerns as follows:

W1: Conditional generation

We apply our method to structure-based drug design, integrating it with MolCRAFT [1], a state-of-the-art method in this domain. The prior is constructed using Algorithm 1 from the set of target ligands within the training set. Given that MolCRAFT also utilizes BFNs as its generative backbone, we were able to seamlessly apply our consistency parameter objective to its parameter space. We adhered to MolCRAFT's established evaluation settings to report our results:

For Vina-based metrics, lower value indicate better performance. For others, larger value indicate better performance. Our method demonstrates comparable or superior performance with significantly fewer NFEs. This outcome underscores its considerable practical value.

W2: V&U&N metric on DRUG dataset

On DRUG, both our model and all compared works consistently achieve near-perfect scores for uniqueness and novelty. Our model achieves 99.7% Uniqueness and 99.9% Novelty. Given this saturation, we followed the established practice of prior works [2,3] by primarily reporting the validity metric.

Furthermore, we are one of the few works to report molecular stability on the DRUG dataset, as detailed in Table 3. Including this metric provide a more comprehensive evaluation of generation quality.

W3: Actual runtime

The following table presents the wall-clock times required to generate 1000 samples on the QM9 dataset, using a single RTX 3090 GPU with a batch size of 64:

As the results demonstrate, MolTD achieves a sampling speed over 100x faster than diffusion-based models. Furthermore, the practical speed-up is even greater than what is reflected by the reduction in NFEs alone, as MolTD benefit from a more efficient implementation.

W4: Clarity and Terminology

We have revised the paper based on your valuable comments, a provide a point-to-point response in the following:

1. Line 149: We change permutations to transpositions to improve clarity.
2. Line 24: Compared to official results reported by AlphaFold3
3. Line 39: We change the phrase to '..., 3D molecular generative models that generate the structure holistically, requires determining the permutation order ...', which is more rigorous.
4. Line 40: We change the phrase to 'This two-phase segmentation is unique to 3D geometric generation and...'
5. Line 53–57: We have revised the paragraph to more precisely articulate our contribution and provide high-level insights.
6. Line 104: We have corrected it in the revised manuscripts.

Q1: Which metrics in Table 1 is more indicative?

Both Validity and Stability are crucial for assessing a model's capacity to generate chemically sound molecules. However, Stability is arguably a more robust indicator of model performance than Validity, primarily because the Validity is susceptible to being artificially inflated.

We borrow the discussion from [2] : 'Experimentally, it is observed that validity could artificially be increased by reducing the number of bonds. For example, predicting only single bonds was enough to obtain close to 100% of valid molecules on GEOM-DRUGS. On the contrary, the stability metrics directly model hydrogens and cannot be tricked as easily.'

Furthermore, the Novelty metric, which quantifies the percentage of generated molecules absent from the training set, also requires careful interpretation. While high novelty is generally desired, an extremely high score can paradoxically suggest an underfitting to the training data distribution.

Given these considerations, we believe that the V&U&N metric is less indicative of a model's true generation performance. It should be viewed as a signal that ideally falls within a reasonable range. Our method achieves a V&U&N of 67.03%, aligning well with the 50-70% range reported by previous works in Table .

We report two additional metrics to evaluate quality and diversity: Validity & Uniqueness (V&U) and Atom Stability & Validity & Uniqueness (S&V&U):

As the result shows, our method achieves significant speed-up while maintaining superior generation quality.

Q2: Definition of NFEs

One NFE denotes one forward pass of the neural network, which is equivalent to one sampling step in the denoising process. This definition is consistent across all compared models.

Q3 Exclusion of GOAT in Figure 2

We agree that GOAT is a key state-of-the-art model, and we have aimed to compare against it thoroughly.

As shown in Table 1, we conducted a comprehensive comparison with GOAT on all standard metrics, where our method demonstrates a significant speed-up while maintaining comparable generation quality.

The analysis in Figure 2, however, requires analyzing intermediate results, and thus requires model's implementation. Unfortunately, at the time of our experiments, the official implementation of GOAT is still under actively updating. Attempting our own reproduction, especially within a limited timeline, would risk introducing implementation biases.

We are keen to include GOAT in Figure 2 based official implementation or a faithful reproduction.

Q4 More explanation on the third subfigure in Figure 2

We have provided a detailed explanation on the subfigure in the revised manuscript, and a point-to-point response in the following:

• The boundary of $x_2^1$ refers to one sheet of a two-sheeted hyperboloid with foci at $x_0^i$ and $x_0^j$, passing through the $x_2^1$. This surface defines the decision boundary for whether a permutation is needed between the i-th and j-th rows, supported by Proposition 4.3 and 4.4.
• These circles around x01 and x21 represent the neighborhood $V(x_0^i; r)$, where a specifical radius $r$ is chosen to ensure that these neighborhood are non-overlapping. Within $V(x_0^i; r)$, the optimal matching is the identity permutation.
• Phase I denotes the permutation phase, where permutation is needed. Phase II corresponds to the adjustment phase, where no further permutation is needed during Phase II.

Q5: Geometric-Informed Prior

1. Construction: for size of molecules (number of atoms), prior is created by:

• Alignment: Using Equivariant Optimal Transport to align all geometric structures within the subset.
• Averaging: After alignment, we compute the mean of all atomic features (including spatial coordinates) across all molecules in the subset.

2. Number of Representatives: We construct one representative structure for each molecule size contained in molecules.

3. Inference: Following the previous frameworks [2,3], we first decide the number of atoms based on the distribution of training data, which is done by sampling from the histogram of molecule size. Once the number of atoms is fixed, we retrieve the corresponding prior and plug it in Equation (7) to create $\theta_p$.

Q6: Consistency Parameter

1. Training Phase Separation: Consistency training can suffer from high variance during training, because the consistency objective enforces self-consistent outputs across all time steps, creating complex optimization goals [4] .

By introducing $t_{stable}$, we apply the consistency objective only during the adjustment phase—once a stable atomic structure has been largely determined. This simplifies the learning task and substantially reduces training variance.

We trained our model on QM9 with and without the $t_{stable}$ threshold from Equation (9), repeated each experiment 5 times and reported the mean and standard deviation.

The results demonstrate the effectiveness of acceleration and variance-reduction of our method.

2. Terminology and Proposition 4.6: Proposition 4.6 establishes that the original BFN objective intrinsically trains the network to predict ground truth molecules from intermediate time steps, which is conceptually identical to that of consistency models. Consequently, both the original BFN and the consistency parameter objective are utilized to stabilize the training process, without introducing conflicting optimization goals.

3. How does the method handle cases where the structure does not stabilize? For the permutation phase, the network is trained using only the BFN loss, meaning the acceleration observed in this phase primarily stems from the geometric-informed prior.

Q8: Include ablation result in Figure 2

Thank for the constructive suggestion. We include visualization in Figure 2 in the revised manuscript

Reference

[1] Qu, Yanru, et al. "MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space." ICML 2024

[2] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." ICML 2022

[3] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." ICML 2023

[4] Song, Yang, and Prafulla Dhariwal. "Improved Techniques for Training Consistency Models." ICLR 2024

W1: Wall clock time comparison

We appreciate the opportunity to clarify the computational efficiency of our model. Apart from the fewer sampling steps, the efficiency advantages in wall-clock time are more significant. The following table presents the wall-clock times required to generate 1000 samples on the QM9 dataset using various generative models, all performed on a single RTX 3090 GPU with a batch size of 64:

As the results demonstrate, MolTD achieves a sampling speed over 100x faster than diffusion-based models like EDM and 20x faster than the flow-matching model EquiFM. Furthermore, the practical speed-up is even greater than what is reflected by the reduction in NFEs alone, as MolTD benefit from a more efficient implementation.

W2: How does the distribution of molecule size affect the quality of the generated molecules?

Thanks for the insightful question! Our approach leverages geometric-informed priors tailored to the number of atoms (N) in a molecule, since molecules with different numbers of atoms show distinct structural patterns (Figure 6 in appendix). This approach allows us to extract fine-grained geometric information within each cluster of molecular structures.

Our method was evaluated on the public QM9 and DRUG datasets, following the established settings of previous work [1,2]. The number of molecules available for each size varies within these datasets. Two factors contribute to improved generation quality:

• A Better Prior: A structural prior derived from more extensive data better captures the full spectrum of geometric features.
• A Better Generative Model: More training data yields a more robust generative model. To empirically validate this, we conducted an experiment on the QM9 dataset. For molecules with 12 and 19 atoms respectively, we generated 1000 samples and evaluated the results:

As the results illustrate, molecule sizes corresponding to a higher distribution of training data exhibit demonstrably better generation quality compared to those with lower distribution.

However, we highlight that our method could boost the generation efficiency, even with only 10 samples to construct the prior. We randomly sample 10 molecules for each N to create prior and generated 10000 molecules on QM9:

This result demonstrates the substantial practical value of the geometric-informed prior, even in data-scarce settings.

Q1: How to transform prior to another dataset, the difference of prior constructed from different dataset and its impact on generation quality

We thank the reviewer for this insightful question. Our method utilizes distinct structural priors for each molecule size, derived from the DRUG dataset, and applies them to the QM9-trained model during inference. The effectiveness of this cross-dataset prior depends on the molecule size (N).

For small N, the space of valid molecular geometries is relatively constrained. In this scenario, priors derived from different datasets (QM9 vs. DRUG) are quite similar, and we observe only a marginal performance difference when applying the DRUG prior.

For large N, however, the geometric configuration space becomes vast and complex. The DRUG dataset provides significantly more data for these larger sizes, leading to a much richer and more informative prior. A prior constructed from this much larger and more diverse set captures a broader spectrum of valid geometric features, leading to a notable improvement in generation quality.

To empirically validate this, we conducted an experiment for larger molecules (N = 25 to 29), generating 1000 samples for each size and comparing the results when using priors derived from QM9 versus those from DRUG.

These results demonstrate general applicability of geometric-prior, which could utilize external dataset to improve generation quality.

Q2: Typo in line 44: analysis -> analyze

Thanks for your careful assessment and for catching this typo! We will correct it in the revised manuscript.

Reference:

[1] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." ICML 2022

[2] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." ICML 2023

W1: Comparison with acceleration methods from graph generation

We thank the reviewer for this insightful suggestion regarding comparisons with graph generation acceleration methods.

While related, a key distinction makes direct comparison challenging: 2D graph generation primarily models discrete states (e.g., the adjacency matrix and node types), whereas 3D molecule generation must model continuous states (i.e., the 3D atomic coordinates), which adds a significant layer of complexity. This fundamental difference in data modalities means that acceleration techniques are often not directly transferable.

Nonetheless, we agree that a comparison is valuable. We evaluated our method against two state-of-the-art 2D graph generation models known for their efficiency:

• GraphBFN [1] : Using probabilistic interpolation for acceleration of graph generation
• DruM [2]: Utilizes a mixture of diffusion processes to model graph topology.

The results show that while 2D methods perform well on stability metrics—a natural result of focusing on the simpler, discrete task of graph generation—our method (MolTD) achieves comparable stability with a drastically lower NFE (12 vs. 100-1000), highlighting its high efficiency.

W2: Terminology choice of 'Permutation' stage

We thank the reviewer of this useful comment! We agree that the 'permutation stage' could be ambiguous in the context of our full pipeline, and 'geometry alignment stage' is more consice. We will update it in the next version.

Q1: Application on Structure-based Drug Design

Thanks for the insightful comments! We have applied our method to structure-based drug design, integrating it with MolCRAFT [3], a state-of-the-art method in this domain.

The geometric-informed prior is constructed using Algorithm 1 from the set of target ligands within the training set. Given that MolCRAFT also utilizes BFNs as its generative backbone, we were able to seamlessly apply our consistency parameter objective to its parameter space. We adhered to MolCRAFT's established evaluation settings to report our results:

For Vina-based metrics, lower value indicate better performance. For others, larger value indicate better performance. Our method demonstrates comparable or superior performance with significantly fewer NFEs. This outcome underscores its considerable practical value.

Q2 & W3: Relation to protein structure generation

Thanks your for the insightful question! AlphaFold3, the state-of-the-art model for protein structure prediction, utilized the diffusion models to generate spatial coordinates of proteins. Our ablation study in Figure 3 demonstrates that our method can be successfully applied to diffusion-based models, yielding significant acceleration. Thus, we believe the protein structure generation could be similarly decomposed as permutation phase and adjustment phase, and applying our method with proper modification will boost the generation efficiency. Furthermore, we anticipate the permutation phase would be less computationally intensive for proteins, as the fixed amino acid sequence already provides strong conditional information that constrains the possible ordering of nodes.

We are actively exploring this promising direction as part of our ongoing research.

Q3: Discussion on the relation between permutation stage and refinement stage to graph structure generation and graph feature generation

We thank the reviewer for this insightful question, which highlights an interesting analogy between our method and graph generation:

• Our Permutation Phase is analogous to generating the 2D graph topology (e.g., an adjacency matrix), as it establishes the fundamental geometrical structure.
• Our Refinement Phase is then similar to the 2D-to-3D "lifting" step, where 2D topology is mapped to 3D geometric space.

However, we wish to clarify a critical distinction: In 3D molecular generation models, we do not explicitly model edges or an adjacency matrix. Instead, the final molecular topology is implicitly defined by the geometric configuration of the generated atoms. While in graph generation, the topology is directly generated as the adjacency matrix. This distinction may suggest why specialized methods, such as those we introduce for the permutation and adjustment phases, may be more beneficial for advancing 3D molecule generation compared with general graph generation.

Q4: Using permutation invariant loss

We thank the reviewer for this insightful suggestion. The suggested approach, which involves aligning molecular structures before applying a permutation-invariant objective, is a valid strategy similar to methods like EquiFM.

In our preliminary experiments, we explored this exact direction. We used Equivariant Optimal Transport (EOT) to align molecular structures first, with the goal of simplifying the subsequent generation task. However, we found this approach to be less effective than our current method.

The primary issue was that pre-aligning the spatial features made it significantly harder for the model to predict the discrete atom types. As a result, the generation process became less efficient, nullifying the potential acceleration benefits. This performance trade-off is empirically demonstrated in Figure 2 of our manuscript, represented by EquiFM.

Consequently, we opted for our current architecture, which we found addresses the challenges of both discrete and continuous generation more effectively, leading to a more robust and efficient model.

Q5. Comparison with fragment-based molecular generation

This is a great question! Adapting fragment-based methods to 3D generation is a challenge due to the combinatorial complexity of assembling fragments in 3D space. Consequently, few models have demonstrated high performance in this area.

For our comparison, we selected two representative state-of-the-art models:

• Symphony: An auto-regressive model with E(3)-equivariance that iteratively adds molecular fragments using a spherical harmonics framework.[4]
• HierDiff: A hierarchical diffusion model that generates a coarse fragment-based structure before refining it to a full atomic representation.[5]

We recognize that these models were developed for distinct tasks and were originally evaluated with different metrics. Symphony is primirly evaluated for validity and statistical difference between training data and generated molecules, while HierDiff is foused on drug-likeness. To ensure a fair and direct comparison, we report the Validity and Uniqueness, which are standard metrics consistently used across all compared methods.

As the results show, our methods achieve superior generation quality and efficiency.

Q6: Impact of geometric-informed prior on diversity

We thank the reviewer for raising this important point regarding sample diversity.

You are correct that a structural prior can introduce a trade-off between sample quality and diversity. Our method is designed to make this trade-off fully controllable.

As shown in our ablation study in Figure 5, the influence of the prior can be precisely modulated. By applying a prior with lower accuracy level (or equivalently, inject the prior at the first few steps in generation), the model retains significant freedom to explore diverse chemical space while still benefiting from prior information.

Crucially, our primary results demonstrate that this control does not hinder the model's ability to generate a diverse set of molecules. As reported in Table 1, our method achieves Validity, Uniqueness, and Novelty scores (67.03 combined) that are comparable to state-of-the-art models. Furthemore, we compare the Validity & Uniqueness on QM9 in the following:

This confirms that our approach successfully explores the vast molecular space to generate valid and diverse structures.

Reference

[1] Song, Yuxuan, et al. "Smooth Probabilistic Interpolation Benefits Generative Modeling for Discrete Graphs." ICML 2025

[2] Jo, Jaehyeong, Dongki Kim, and Sung Ju Hwang. "Graph generation with destination-predicting diffusion mixture." (2023).

[3] Qu, Yanru, et al. "MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space." ICML 2024

[4] Daigavane, Ameya, et al. "Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation." ICLR 2023

[5] Qiang, Bo, et al. "Coarse-to-fine: a hierarchical diffusion model for molecule generation in 3d." ICML 2023