Response to Weakness 1, Weakness 3, and Question 2

W1. Comparison with Recent Edge-Modeling Methods and Potential Integration of Edge Features

W3. Difficulty in Claiming State-of-the-Art Performance ...

Q2. Could you perform comparisons with recent models like EQGAT-Diff, JODO ...

We appreciate the reviewer's comment.

First, We would like to highlight that the suggested algorithms are realized with different assumptions in terms of data modeling and experimental settings.

1. Additional data features. MiDi, EQGAT-Diff, JODO, and SemlaFlow all modeled bonds explicitly or implicitly, which required extra features of the molecular data.

2. Filtered large molecules. Particularly, these algorithms, including MiDi, EQGAT-Diff, and SemlaFlow, filtered out large-size molecules (more than 72 atoms) in the GEOM-DRUG dataset. Such assumption may limit their applicability in modeling edges for large-size molecules.

3. Incomplete molecule modeling. Many of them ignore the H atom in modeling the molecule, which also reduced the difficulty of the problem to some extent.

It is unfair to claim that our methods are not SOTA as the data modeling assumptions and our experimental settings differ.

Second, in response to Weakness 1, we implement our method with the integration of edge modeling and compare it to this type of method under the same experimental setting and datasets. We presented the results in the Table below:

Considering topological metrics, 3D metrics, and performance metrics, the proposed method with bond modeling also shows competitive performance compared to a series of methods that also model bonds, including MiDi, JODO, EQGAT-Diff, and SemlaFlow. Specifically, our method achieves the best validity, novelty, significance, sampling steps, and energy. We could not provide some results of JODO, EQGAT-Diff, and SemlaFlow, as some of them are not open-source, and the original paper did not report novelty and energy.

Third, we would like to highlight the difference between our method and SemlaFlow.

1. Equivariance. We proposed a new equivariant optimal transport method for dealing with the multi-modal property of molecules. In contrast, SemlaFlow simply utilized existing equivariant optimal transport [1], both of which do not consider the multi-modality of molecules.

2. Various sizes of molecules. To deal with various sizes of molecules, we proposed the flow refinement and purification mechanism, whereas SemlaFlow utilized several techniques, like layer normalization, to circumvent this challenge.

3. Conditional generation. SemlaFlow did not provide experiments on conditional generation. In contrast, our method achieves a better trade-off between accuracy and efficiency in conditional molecule generation.

[1] Song, Yuxuan, et al. "Equivariant flow matching with hybrid probability transport for 3d molecule generation." Advances in Neural Information Processing Systems 36 (2023).

We appreciate the reviewer’s recognition of the strengths in our work. We are glad that the reviewer found our significant theoretical development in joint optimal transport and performance improvement and comprehensive comparisons. We address the reviewer's concerns below.

Response to the Weakness 2 and Question 1:

W2. Concerns About Atom Stability Metric and 3D Evaluation: ...

Q1. Could you provide a more comprehensive evaluation of the generated 3D molecular structures? ...

We appreciate the reviewer's suggestions about the evaluation metrics. Following your suggestion, we supplemented the energy using GFN-xTB as an extra metric on the GEOM-DRUG dataset. We presented the supplemented results in the Table below:

The experimental results show that the energy obtained by the compared methods achieves competitive performance. Among baselines focusing on fast molecule generation (GeoBFN and EquiFM), the proposed method GOAT achieves the lowest energy with faster sampling speed.

We are grateful that the reviewer acknowledges the strengths of our work. We are glad that the reviewer found that our design is very clever and reasonable, discussion on the computational cost is comprehensive, and our approach achieves a satisfiable performance. We address the reviewer's concerns below.

Response to Weakness 1

While the paper includes some ablation studies ...

In response to your concern, we added experiments evaluating the significance of the equivariant autoencoder (w/o EAE). The new results are added in row (w/o EAE) below:

The findings indicate that without the equivariant autoencoder (w/o EAE), the transport cost is significantly higher, and the generation time nearly quadruples compared to the GOAT framework. This underscores the importance of jointly optimizing transport for atom coordinates and features.

Response to Weakness 2

The training process itself may be time-consuming...

We agree with the reviewer’s comment, which we have acknowledged in our manuscript. We will investigate methods to enhance the training efficiency of our framework.

Response to Question 1

It would be interesting if the paper could provide a comparison...

We thank the reviewer for the insightful suggestions. We presented a comparison of the distance between generated molecules and the initial noise, including compared methods, the proposed GOAT, and its variants (w/o ODT and w/o OMT). We presented the experimental results in the Table:

The experimental results on the distance between molecules and noises validate that the proposed method achieves the minimum transport distance from the noise and thereby also verifies the superiority of our method in generation speed.

Response to Question 2

The use of optimal transport may raise concerns about ...

We appreciate the reviewer’s valuable suggestions and agree that diversity is a crucial evaluation metric. However, we note that the COV-R and MAT-R metrics focus specifically on conformation generation tasks, where the accuracy of generated conformations is assessed against a set of reference conformers. In our molecular generation task, there are no predefined target conformations. Utilizing a training set would result in evaluating the generation of existing molecules, which contradicts the objective of de novo drug design. Consequently, preparing a reference set is challenging.

Nonetheless, our adopted metric of uniqueness addresses this concern. Uniqueness measures the fraction of unique molecules among all valid generated samples. In our evaluation with 10,000 samples, the proposed method achieved an impressive 99.0% uniqueness, indicating that approximately 9,900 distinct molecules were generated, demonstrating our method's capability to produce diverse molecular structures.

We recognize that our explanation may not completely cover your comment. If you could share more details, we would be happy to provide a more focused response.

We appreciate that the reviewer acknowledges the strengths of our work. We are glad that the reviewer found that our motivation is well-defined and strong and the performance is strong as well. We address the reviewer's concerns below.

Response to Weakness 1:

In Eq. (5), the authors seem to sample permutation...

Thank you for your insightful question. We somewhat agree with the reviewer that the optimized permutation, rotation, and translation matrices derived from Eq. (5) serve as a means of data augmentation [1]. As discussed in [1], data augmentation can somehow eliminate the equivariance requirement in model design. In addition, our model achieves a strict equivariance with an equivariant graph neural network.

Specifically, we can first establish that the latent data $z_1$ is equivariant with respect to the data point $\mathbf{g}_1$ as implemented in the equivariant autoencoder.

Following the approach in [2], we then show that the training loss in Eq. (8) is invariant to translations and rotations of the noise $z_{0}$ and equivariant with respect to both the sampled noise $z_{0}$ and the latent data point $z_{1}$. Therefore, the optimized permutation, rotation, and translation do not compromise the model's equivariance.

[1] Brehmer, Johann, et al. "Does equivariance matter at scale?." arXiv preprint arXiv:2410.23179 (2024). [2] Song, Yuxuan, et al. "Equivariant flow matching with hybrid probability transport for 3d molecule generation." Advances in Neural Information Processing Systems 36 (2023).

Response to Weakness 2:

I am wondering maybe it is possible to distill the trained model...

We appreciate the reviewer’s insightful comments. We conducted experiments to evaluate the performance of the distilled model using the distillation method introduced in Rectified Flow, and the results are presented below.

The experimental results indicate that distillation improves generation speed; the distilled model achieves the best S-Time. Although the validity is lower, the distilled model produces valid, unique, and novel molecules at the fastest speed. However, the stability and validity of the atoms remain unsatisfactory. We conjecture that the distillation on molecular generators requires additional designs, and we plan to investigate the distillation method in molecular generation further in our future research.

We are glad that the reviewer recognizes the strengths of our work. We appreciate the reviewer found that our problem statement is clear, the entire writing is comprehensive, and our experiments are abundant to illustrate the effectiveness. We address the reviewer's concerns as follows.

Response to Weakness 1

W1. In equation (5), we still have the same $\pi$ for both coordinates and features right? The reasoning for this part is missing.

Thanks for the astute observation. Coordinates and features together depict the atomic-level geometry of a molecule. Therefore, a unified optimal $\pi$ for both coordinates and features is needed for obtaining an optimal transport, leading to fast and high-quality molecule generation.

Response to Weakness 1 and Question 1

W1. I do not see why after equivariance autoencoder mapping to the latent space, we thus have a unified optimal permutation.

Q1. Why using equivariance autoencoder to produce latent codes can yield in unified permutation?

Thanks for the insightful question. We would like to clarify that the equivariance autoencoder is designed to deal with the multi-modal properties of molecules, mapping coordinates and atom features into a unified, continuous latent space. Simply mapping with the equivariance autoencoder indeed could not directly yield a unified optimal permutation. It is critical to achieve an optimal permutation to minimize the transport cost; we realize such an objective using the Hungarian algorithms.

Response to Weakness 2, 4, and Question 3:

W2. I am not sure the distribution coupling, $\Gamma$ is ...

W4. Is Geometric Probability Distribution a formal terminology ...

Q3. Please clarify the terminology issues ...

Thanks for pointing out the terminology issues. We address the reviewer's comments one by one below.

1. In response to W2 & Q3: we specify the notation $\Gamma$ for the pairing points between two molecule distributions. Specifically, we use $(\mathbf{g}_0,\mathbf{g}_1)$ to denote the pairing points, where $\mathbf{g}_0$ is sampled from the noise distribution and $\mathbf{g}_1$ is sampled from the data distribution. As in Eq. (5), we consider permutation, rotation, and translation for geometric optimal transport cost; the purpose is to preserve geometric symmetries while minimizing the transport cost.

2. In response to W4 & Q3: The purpose of using geometric probability distribution is to reflect the geometric property of molecules. We are unaware of introducing such confusion to the readers. For clarity, we have added an explanation of this terminology in the manuscript (a footnote on page 3).

We understand that our explanation might not fully address your comment. You could provide further details or specify the aspects that are unclear. We would be glad to further offer a more targeted explanation.

Response to Weakness 3 and Question 2:

Maybe limited novelty: In EquiFM[1], …

Can I assume your paper is based on …

We thank the reviewer for giving us the chance to re-clarify our novelty. Simply combining EOT and Equivariant Autoencoder does not work. Therefore, our main contributions could not be concluded as EOT + Equivariant Autoencoder. Firstly, EOT (or EquiFM [1]) indeed attempts to accelerate the generation process with the notion of flow matching. However, it didn't consider the multi-modal properties of molecules, leading to sub-optimal transport. Different from EquiFM [1], we consider both modalities by mapping the continuous atom coordinates and categorical atom types with the Equivariant Autoencoder. Secondly, though with the mapping, searching for optimal coupling is still challenging due to the various sizes of 3D molecules. To address the challenges, we propose flow refinement and purification to approximate optimal transport and accelerate molecule generation.

[1] Song, Yuxuan, et al. "Equivariant flow matching with hybrid probability transport for 3d molecule generation." Advances in Neural Information Processing Systems 36 (2023).

Response to Weakness 5:

The abbreviation EAE in line 251…

Yes. EAE represents Equivaraint Autoencoder, which we have updated on page 5 of the manuscript.

Thanks to the authors for their clarifications. I maintain my rating of leaning toward acceptance.