ET-Flow: Equivariant Flow-Matching for Molecular Conformer Generation

Lack of experiments. The experiments on the GEOM dataset are not enough to show the empirical performance of the framework. Some further experiments on large-scale datasets with more data and larger system size are necessary. I suggest the authors do experiments on OC20/OC22(Open Catalyst 2020/2022) datasets to test the performance of the ET-Flow framework.

Thank you for your valuable feedback. Our primary focus in this work is on sampling molecular conformers, and therefore, the dataset and benchmark metrics we used align with the standards established by previous works that addressed the same objective from different perspectives[1,2,3]. While we appreciate the suggestion to test our framework on the OC20/OC22 datasets, this would be a nontrivial task due to the fundamentally different nature of the catalyst datasets. Consequently, existing methods for molecular conformer generation do not benchmark performance on OC20/22 and therefore is beyond the scope of our work. However, we agree that evaluating our framework on a diverse set of datasets is an important and promising direction for future research.

The novelty of this work is not enough. I think the idea and method of this paper is very similar to [1], but in this work, the authors use a transformer-based model. Additionally, this paper proposes several tricks such as rotational alignment, stochastic sampling, and chirality correction. So the authors should clarify the novelty of this work.

We argue that the correct integration and modification of existing methodologies to achieve significant improvements should not be overlooked. One of the strengths of our work lies in identifying the core problem with existing methods and devising a straightforward approach with precise engineering, which is often crucial for empirically important tasks, especially in the application of machine learning to chemistry. We include ablation studies in table 2 of the global rebuttal demonstrating the effect of each design choice on the result.

[1] Ganea, O., Pattanaik, L., Coley, C., Barzilay, R., Jensen, K., Green, W. and Jaakkola, T., 2021. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances in Neural Information Processing Systems, 34, pp.13757-13769.

[2] Jing, B., Corso, G., Chang, J., Barzilay, R. and Jaakkola, T., 2022. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35, pp.24240-24253.

[3] Wang, Y., Elhag, A. A., Jaitly, N., Susskind, J. M., & Bautista, M. Á. Swallowing the Bitter Pill: Simplified Scalable Conformer Generation. In Forty-first International Conference on Machine Learning.

Dear Reviewer,

We first want to thank the reviewer for taking the time to review our work and ask thought-provoking questions. We hope that we are able to address said questions and concerns below.

It can be challenging to attribute the performance improvements of ET-Flow to the flow matching approach. Ablation studies focusing on the flow matching component or results from simpler architectures might be useful.

We first want to acknowledge that the performance of our model is not only attributed to the use of flow matching but is a result of a combination of different components such as choice of prior, rotational alignment, sampling method. We additionally included ablation studies demonstrating the effect of these design choices in table 2 in global rebuttal.

We agree that leveraging the expressive equivariant architecture combined with apt modification and engineering is indeed one of the key reasons for improved performance. However, the strengths of our approach, such as faster inference and flexible choice of prior are enabled by adopting the flow matching framework. Therefore, we view our work as addressing different challenges to create a harmonious integration of these components.

The need for a post hoc chirality correction step. Such a weakness may result in issues in practical applications.

TorchMDNet is O(3) equivariant, thus necessitating a post-hoc chirality correction (CC) step to break the reflection symmetry. We conduct an additional experiment with SO(3) equivariance, alleviating the need for an additional correction step. This is achieved by modifying the vector output of TorchMDNet using the cross product in Equation(18). The empirical results are shown in table 3 for GEOM-DRUGs and table 4 for GEOM-QM9 in global rebuttal.

For GEOM-QM9, we show that the O(3) with CC and SO(3) TorchMDNET achieve roughly similar results, however in the case of GEOM-DRUGS the advantage of O(3) with CC is more pronounced. An important point to highlight is the negligible cost of the CC step as highlighted in figure 1 in the global rebuttal pdf. In practice, one may choose to use ET-Flow combined with CC or directly use the SO(3) version of the model.

While we acknowledge the reviewer's concern, we respectfully disagree about the practical weaknesses. MCF uses data augmentation and more model parameters to learn physical symmetries. As shown in figure 1 of the attached PDF in the global rebuttal, our method’s inference time remains significantly faster, even with the CC step. Furthermore, ET-Flow’s inference time primarily depends on the number of parameters.

The novelty of the approach is limited as it largely builds on existing methodologies. This integration, although leading to performance improvements, does not significantly advance the state-of-the-art in terms of methodological innovation.

We argue that the correct integration and modification of existing methodologies to achieve significant improvements should not be overlooked. One of the strengths of our work lies in identifying the core problem with existing methods and devising a straightforward approach with precise engineering, which is often crucial for empirically important tasks, especially in the application of machine learning to chemistry. We also conduct design choice ablations as shown in Table 2 in the global rebuttal.

It is widely argued that equivariance is not a requirement for molecular machine learning. Without equivariance, the model is more easily generalized and scaled, and the symmetry could be implicitly learned from the data. I'm interested in the scaling performance (with respect to data or model size) of ET-Flow compared to MCF.

The MCF paper advocates for a more general approach to molecular generative tasks by not embedding geometric inductive biases into the model, which is a valid and valuable perspective. However, our stance is that in the application of AI to scientific domains, leveraging known inductive biases is crucial for effective performance (scale vs performance tradeoff). Explicitly incorporating inductive biases ensures that the model inherently respects physical symmetries, leading to more reliable and physically consistent predictions. This approach is particularly important in domains where data can be scarce or expensive to obtain, as the inductive biases help the model generalize better from limited information.

In our work, we used a dataset of comparable size to MCF but kept our model parameters(8.3M) significantly smaller. Although our computational limitations prevented us from scaling up to a similar number of parameters as used in MCF-L(242M), this is an avenue we plan to explore in the near future. Despite these limitations, our approach demonstrates the effective use of equivariant models in achieving competitive results, suggesting that incorporating inductive biases can lead to more efficient and scalable solutions in molecular machine learning.

Have the authors considered conducting data augmentation to enhance diversity and recall metrics?

We believe this suggestion may be coming from the context of its use in MCF and AlphaFold3. In their cases, data augmentation was primarily employed to compensate for the lack of inductive bias, enabling the models to be more expressive while learning equivariance in a soft manner. In contrast, our model inherently incorporates symmetries by design, eliminating the need for such techniques. Additionally, while our recall metrics are slightly lower than those of MCF-M and MCF-L, our precision metrics show a significant improvement over all versions of MCF. In practice, this improvement in precision is crucial because this produces more physically-correct molecules as evidenced in our ensemble properties table having the lowest errors amongst the existing methods.

Dear Reviewer,

We first want to thank the reviewer for taking the time to review our work and ask thought-provoking questions. We hope that we are able to address said questions and concerns below.

Very little evaluation of out-of-distribution performance is considered (a small evaluation is in the appendix).

The paper would be stronger if generalization performance was more comprehensively evaluated.

What is the out-of-distribution performance?

We first want to acknowledge that we followed the evaluation protocols as done with prior work [1, 2, 3]. In the prior works, out-of-distribution (OOD) evaluation is done on GEOM-XL (Appendix C.1 table 7), which consists of molecules >100 atoms, using a model trained on GEOM-DRUGS (average 44 atoms per molecule). We demonstrate that our method performs slightly better than the MCF-S while having 5 million fewer parameters. However, we agree that more experiments that evaluate OOD performance is needed. Accordingly, we also evaluate our model on molecules significantly smaller than the training distribution via GEOM-QM9 (average 18 atoms per molecule) [4]. We provide these results in the table 1 of the global rebuttal statement. We are open to more suggestions and feedback regarding evaluation strategies.

How well does the model generalize to different chemotypes (not just larger) and how does this compare to conventional and other generative models?

Thank you for bringing this to our attention. The idea of exploring different chemotypes is an interesting suggestion, but it is unclear what chemotypes are being referenced. More details about this would be appreciated so we could discuss further. Moreover, we would like to make clear that our domain of application is in drug discovery and therefore we focus on drug like molecules from the GEOM dataset [4].

The ensemble property averages evaluation seems to be done on molecules drawn from the training data (if this is not the case, it needs to be more clear).

We would like to make clear that the ensemble averages experiment was done on the 100 randomly sampled molecules from the test set. We added a statement to section 4.4 to make that more clear.

Although efficiency is one of the main claims, there's no evaluation of inference time.

Thank you for bringing this up. We included a plot of inference time with respect to the number of steps in figure 1 in the attached pdf for global rebuttal. We also include a plot on performance (precision) with respect to inference time in figure 2 in the attached pdf for global rebuttal.

I think "DRUGS" is used as a short-hand for GEOM-DRUGS - this is confusing, use consistent nomenclature throughout.

We made the adjustment to the instances where we use “DRUGS” and just use “GEOM-DRUGS” for consistency.

It would be interesting to see MCF trained using the same frame work as ETflow - e.g. how much of the performance is due to the choices in Table 6 versus the model architecture.

MCF follows a different framework than ETFlow. We would like to ask if the reviewer is implying to evaluate the architecture of MCF (Perceiver IO) with the flow matching framework? Any clarifications on this statement would be much appreciated, thank you.

Some recent work has pointed to the fact that improved recall/precision on these GEOM benchmarks does not necessarily result in better performance on downstream tasks that use conformers. At least some discussion of this limitation would be appreciated.

We would like to point out to the reviewer that the ensemble properties table in Section 4.4 - table 3 in the manuscript indicates that our method has the lowest errors for the chemical properties evaluated against GFN-xTB oracle, demonstrating our model produces more physically-correct molecules compared to other methods. We would be open to benchmarking our model on other metrics if the reviewer could refer us to said metrics and alternative downstream evaluations.

[1] Ganea, O., Pattanaik, L., Coley, C., Barzilay, R., Jensen, K., Green, W. and Jaakkola, T., 2021. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances in Neural Information Processing Systems, 34, pp.13757-13769.

[2] Jing, B., Corso, G., Chang, J., Barzilay, R. and Jaakkola, T., 2022. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35, pp.24240-24253.

[3] Wang, Y., Elhag, A. A., Jaitly, N., Susskind, J. M., & Bautista, M. Á. Swallowing the Bitter Pill: Simplified Scalable Conformer Generation. In Forty-first International Conference on Machine Learning.

[4] Axelrod, S. and Gomez-Bombarelli, R., 2022. GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1), p.185.

For the scaffold split experiment, we divide the GEOM-QM9 smiles based on molecular scaffolds into an 80:10:10 ratio for train, validation, and test sets. Our results based on 1000 randomly sampled molecules from the test set are as follows,

It seems that generalization across scaffolds is possible. We will include the same experiments on GEOM-DRUGS, which are still running.

Thanks for the suggestion - this experiment is a nice addition! Please let us know if they sufficiently address your concern.