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

Predicting low-energy molecular conformations given a molecular graph is an important but challenging task in computational drug discovery. Existing state-of-the-art approaches either resort to large scale transformer-based models that diffuse over conformer fields, or use computationally expensive methods to generate initial structures and diffuse over torsion angles. In this work, we introduce Equivariant Transformer Flow (ET-Flow). We showcase that a well-designed flow matching approach with equivariance and harmonic prior alleviates the need for complex internal geometry calculations and large architectures, contrary to the prevailing methods in the field. Our approach results in a straightforward and scalable method that directly operates on all-atom coordinates with minimal assumptions. With the advantages of equivariance and flow matching, ET-Flow significantly increases the precision and physical validity of the generated conformers, while being a lighter model and faster at inference. Code is available https://github.com/shenoynikhil/ETFlow.

Paper

References (74)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer eBNM4/10 · confidence 4/52024-07-10

Summary

The paper proposes the Equivariant Transformer Flow (ET-Flow) to generate high-quality molecule conformations. The authors use rotational alignment, stochastic sampling, and chirality correction to improve the flow matching framework for this task. Additionally, the paper modifies the TorchMD-NET equivariant transformer architecture to parameterize the target vector field. Experiments on the GEOM dataset show that ET-Flow can perform well on the conformation prediction task. Furthermore, conformations generated by ET-Flow can also be used to predict ensemble properties with great performance.

Strengths

1. Very clear writing. 2. The method is simple and easy to understand.

Weaknesses

1. 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. 2. 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. [1]. Klein, Leon, Andreas Krämer, and Frank Noé. "Equivariant flow matching." Advances in Neural Information Processing Systems 36 (2024).

Questions

See the sections above.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

Reviewer AAmo7/10 · confidence 4/52024-07-13

Summary

The paper proposes the Equivariant Transformer Flow (ET-Flow) which predicts low-energy molecular conformations given the molecular graphs. Unlike existing methods that rely on large transformer-based models for conformed fields or complex internal geometry calculations, ET-Flow leverages flow matching with equivariance and harmonic prior which directly operates on all-atom coordinates with minimal assumptions. The extensive experimental results illustrate that ET-Flow achieves state-of-the-art performance in molecular conformer generation benchmarks with fewer parameters and faster inference times, outperforming or matching previous methods while maintaining high accuracy and efficiency.

Strengths

1. The proposed approach represents a significant innovation by leveraging flow matching with equivariance and harmonic prior in order to simplify the conformer generation process while maintaining high accuracy. 1. ET-Flow achieves SOTA model performance on molecular conformer generation benchmarks, outperforming or matching existing methods with much fewer parameters and faster inference times without sacrificing accuracy. The high accuracy and efficiency make it a promising candidate for practical applications. 1. The paper provides comprehensive experimental results, comparing ET-Flow with several leading approaches across various datasets. The inclusion of various evaluation metrics strengthens the validity of the findings.

Weaknesses

1. ET-Flow achieves SOTA or near-SOTA performance on the benchmarks. However, my concern is that TorchMD-Net is already a strong model with a well-established architecture that leverages equivariant transformers for molecular modeling. Given this, 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. 1. The need for a post hoc chirality correction step suggests that the model does not inherently handle stereochemistry (baseline models like MCF do not require such an explicit correction). Such a weakness may result in issues in practical applications. 1. ET-Flow primarily combines TorchMD-Net with the flow matching technique. While both components are robust and effective, 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.

Questions

1. The proposed model utilizes an equivariant model architecture. Meanwhile, it is widely argued that equivariance is not a requirement for molecular machine learning (as suggested in the MCF paper). Without equivariance, the model is more easily generalized and scaled, and the symmetry could be implicitly learned from the data. I would like to ask the authors about the thoughts on this. In particular, I'm interested in the scaling performance (with respect to data or model size) of ET-Flow compared to MCF. 1. Have the authors considered conducting data augmentation to improve the diversity of generated conformations and enhance recall metrics?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations of recall performance and additional chirality correction steps. No potential negative societal impact is involved.

Reviewer 281V6/10 · confidence 4/52024-07-16

Summary

The paper describes an equivariant flow matching model for conformer generation. The stated contributions are accurate - the model performance is state-of-the-art and largely due to good engineering.

Strengths

The paper is well written and the evaluations follow other published work. Informative ablation studies are performed. The evaluation of ensemble property averages is particularly appreciated. I appreciate the authors highlight the modifications required for stable training.

Weaknesses

Very little evaluation of out-of-distribution performance is considered (a small evaluation is in the appendix). 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). The paper would be stronger if generalization performance was more comprehensively evaluated. Although efficiency is one of the main claims, there's no evaluation of inference time. I think "DRUGS" is used as a short-hand for GEOM-DRUGS - this is confusing, use consistent nomenclature throughout.

Questions

What is the out-of-distribution performance? How well does the model generalize to different chemotypes (not just larger) and how does this compare to conventional and other generative models? 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. Can you perform evaluations that demonstrate these conformers are better for downstream tasks than those generated through conventional approaches?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

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.

Authorsrebuttal2024-08-13

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, | Method | Recall Coverage (mean) | Recall Coverage (median) | Recall AMR (mean) | Recall AMR (median) | Precision Coverage (mean) | Precision Coverage (median) | Precision AMR (mean) | Precision AMR (median) | |--------|------------------------|--------------------------|-------------------|---------------------|---------------------------|-----------------------------|-----------------------|------------------------| | ET-Flow (Random Split) | 94.99 | 100.00 | 0.083 | 0.035 | 91.00 | 100.00 | 0.116 | 0.047 | | ET-Flow (Scaffold Split) | 95.00 | 100.00 | 0.083 | 0.029 | 90.25 | 100.00 | 0.124 | 0.053 | 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.

Reviewer 281V2024-08-07

The additional evaluations (out of distribution, inference time, and ablation studies) definitely strengthen the paper (assuming they get included). A way to evaluate generalization to novel chemotypes is to perform a scaffold split of the training data. Downstream tasks from conformer generation include molecular docking, shape similarity, and pharmacophore search. While I would not expect a comprehensive evaluation of ET-Flow's performance in these tasks in this paper, if the goal of the authors is to develop methods that advance drug discovery and not just to get a bold number in an ML conference paper, they should be evaluated and compared to standard approaches (e.g. RDKit ensembles - see https://pubs.acs.org/doi/full/10.1021/acs.jcim.3c01245).

Authorsrebuttal2024-08-12

We would like to thank the reviewer for the valuable feedback. We will indeed incorporate the additional evaluations into the manuscript. Currently, we are conducting an experiment on GEOM-QM9 using a scaffold split, and we expect to upload the results within a day. We also plan to run a scaffold split experiment on GEOM-DRUGS, but given time and compute constraints, we will not be able to provide those results immediately. We will include both scaffold-split experiments in the manuscript. Are there any additional experiments or discussions the reviewer would recommend to further enhance the quality of the paper?

Reviewer eBNM2024-08-11

Thanks for your reply. As you have mentioned in the paper, your goal is to construct a "scalable equivariant model that generates energy-minimized conformers given a molecular graph". I think you can test your framework in the OC20/22 datasets because the IS2RS task (predict the relaxation structure based on the initial structure) is equivalent to "generate an energy-minimized conformer". Indeed, it is a nontrivial task but it's necessary to test your method on large-scale datasets with more data and larger system size. I thank you again for clarifying the novelty of your work.

Authorsrebuttal2024-08-12

Thanks for the nice suggestion! In OC20, one aims to find the relaxed structure of a molecule conditioned on a slab, which differs fundamentally from our work. The catalyst periodicity requires specialized approaches, such as AdsorbDiff[1]. To compare the OC20 dataset size and diversity with our datasets: - The GEOM-DRUGS training dataset contains ~240,000 molecules and up to 30 conformations each, resulting in 5.7 million conformations. OC20 contains 82 molecules. - GEOM-DRUGS molecules (\~44 atoms) are larger than the average size of the adsorbates (structure to be relaxed) in OC20 (\~5 atoms) - To evaluate on larger molecules, we also evaluate on GEOM-XL (>100 atoms) with larger molecules than the training distribution. The results are shown in Appendix C.1 table 7. Do the experiments on Geom-XL address your concerns sufficiently under these considerations? We appreciate your insights and look forward to your feedback on this approach. [1] Kolluru, A., & Kitchin, J. R., 2024. AdsorbDiff: Adsorbate Placement via Conditional Denoising Diffusion. arXiv preprint arXiv:2405.03962.

Reviewer AAmo2024-08-12

Thank the authors for the response. The additional experimental results and thorough analysis provided in the rebuttal effectively address most of my concerns and strongly support the paper's ideas. With that, I'd like to raise my score.

Authorsrebuttal2024-08-12

We sincerely thank the reviewer for their thoughtful feedback and insightful discussion. This has greatly contributed to enhancing the quality of our work.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC