Efficient molecular conformer generation with SO(3) averaged flow-matching and reflow

We would like to thank the reviewer for their comprehensive review and feedback. We are delighted to see that the reviewer highlights the novelty of our Averaged Flow method in conformer generation. Please see below for the responses to the weaknesses and questions raised by the reviewer:

Weakness 1. Details of the Averaged Flow methods.

To address the concern of the reviewer about the clarity regarding the Averaged Flow objective, we rewrote the Sec. 3.1 with thorough derivation and explanation on how it is applied to the conformer generation task. We have also attached a code snippet in the Appendix to directly show how the Averaged Flow is implemented.

Through the revision of Sec. 3.1, we now set a clear definition of $\hat{x}$ as ground truth conformers. We explain how we train the model in practice by sampling conformers to approximate the expectation of flow to conformer ensembles. The potential bias introduced during the sampling process can be alleviated by factoring the whole conformer ensemble in the calculation of partition function $Z$, which we have implemented in the code snippet (Appendix 3.1).

Weakness 2. Presentation not clear enough for Averaged Flow and lack of novelty with direct application of reflow.

To address the concern of the reviewer about the novelty contribution, we have provided more details about the derivation of the Averaged Flow objective as well as the python implementation of it. Furthermore, in Figure 3, we demonstrated that to enable very few step sampling ($N_{step}$<5), reflow is necessary to maintain generation quality. Moreover, the reflow technique is model architecture independent. That being said, with our contribution in this work to confirm the effectiveness of reflow in molecular conformer generation tasks, powerful models such as ET-Flow can also be fine-tuned through reflow to enable faster sampling in future works. We now emphasize this at the end of Sec. 4.4.

Weakness minor.

We thank the reviewer for pointing out some formatting and grammatical issues. We have thoroughly checked the manuscript and fixed the issues.

Question 1. Physical plausibility of generated conformers.

Following the suggestion of the reviewer, we tested the PoseBusters passing rate for $\mathrm{AvgFlow}$ (102 steps), $\mathrm{AvgFlow_{Reflow}}$ (2-steps), $\mathrm{Tor. Diff_{3steps}}$ , and $\mathrm{Tor. Diff_{5steps}}$. Taking the $\mathrm{max}(64, 2 \times N_{\mathrm{conformer}} )$ generated conformers of all molecules in the GEOM-Drugs dataset, the average passing rates are 93.8%, 67.1%, 63.6%, 68.7% for $\mathrm{AvgFlow}$, $\mathrm{AvgFlow_{Reflow}}$, $\mathrm{Tor. Diff_{3steps}}$, $\mathrm{Tor. Diff_{5steps}}$, respectively. We can see that $\mathrm{AvgFlow}$ achieves a very high PoseBusters passing rate on generated conformers. 2-step generation of $\mathrm{AvgFlow_{Reflow}}$ can match $\mathrm{Tor. Diff_{5steps}}$ in passing rate, which is promising considering the $\mathrm{Tor. Diff}$ starts from ~100% valid RDKit generated conformers while $\mathrm{AvgFlow_{Reflow}}$ starts from Gaussian noise.

Question 2. More detailed description of baselines.

We have added more detailed descriptions of baseline to the beginning of Sec. 4. We also emphasized that we demonstrate the effectiveness of both Averaged Flow and reflow through comparing the performance of the same model architecture trained with different objectives and before/after reflow. The following paragraph can be found on page 7 of revised manuscript.

There are three types of baselines in this work, including (a) methods with fast inference speed such as cheminformatics tools (RDKit, OMEGA) and regression model GeoMol, (b) lightweight diffusion model with reduced degree of freedom (Torsional Diffusion), and (c) large transformer-based diffusion or flow model operating on Euclidean atomistic coordinates (MCF and ET-Flow). Moreover, to fairly validate the effectiveness of the AvgFlow objective, we compare the performance of our NequIP-like architecture (Appendix 1.2) trained with different objectives. Similarly, we compare the performance of the same model architecture before and after reflow+distillation to show the necessity of reflow for few-step generation.

Dear reviewer 4ywq:

Since the discussion period will end tomorrow (Dec 2nd), we want to check if our response has addressed your questions and concerns regarding our paper. Please let us know if you have any follow-up comment or question regarding our manuscript. Again, thank you for the time spent on reviewing and discussing the manuscript.

We appreciate for the response from the reviewer and score raising. To address the concerns from the reviewer:

1. More background of Lie groups and Haar measure.

We thank the reviewer for bringing this to attention. We would add more background of concept Lie groups (our propose objective is flow averaged over SO(3), which is a Lie group) and Haar measure (used for defining the integral of data distribution over the SO(3) group). Some citations that will be added to the manuscript if it is accepted include:

[1] Anthony Zee. Group theory in a nutshell for physicists, volume 17. Princeton University Press, 2016.

[2] Leopoldo Nachbin and Lulu Bechtolsheim. The Haar Integral. 1965.

[3] Gregory S Chirikjian and Alexander B Kyatkin. Engineering applications of noncommutative harmonic analysis: with emphasis on rotation and motion groups. CRC press, 2000.

We hope the references above can help to provide more background information of this specific domain to readers.

2. For Averaged Flow, clearly distinguishing our contribution from prior works.

We want to clarify that our contribution on Averaged Flow are three-fold:

1. Conceptualized and derived the novel SO(3)-Averaged Flow objective for flow-matching training.
2. Implemented the closed-form solution of Averaged Flow. Here, the closed-form solution of integral over $R$ as proposed by Mohlin et al. (Eq.5) is a critical component in our implementation.
3. Validated the effectiveness of Averaged Flow in the case of molecular conformer generation through comprehensive experiments.

3. Circumstances that Averaged Flow can be applied on and constraints.

Theoretically, Averaged Flow can be extended to other applications when the generated sample is "correct" no matter how it is rotated. Some examples are protein structure generation and point cloud generation. The constraint here is the Gaussian prior as guessed by the reviewer because the implementation is built upon the closed-form solution in Eq.5. Some variants of standard Gaussian prior can also be accepted including the Harmonic Prior (non-isotropic Gaussian based on covalent bond between atoms). We have included the implementation of Averaged Flow with Harmonic Prior in the Appendix as well. We will elaborate on application circumstances and constraints in the manuscript if it is accepted.

Again, we thank the reviewer for engaging in the discussion. We hope that our response can clarify the concerns of the reviewer and improve the confidence of rating.

We would like to thank the reviewer for their comprehensive review. We appreciate that the reviewer acknowledges the advantage of both Averaged Flow and reflow in conformer generation. We are also glad that the reviewer consider our presentation is clear and paper is easy to follow. Please see below for responses to the weaknesses or questions raised by the reviewer:

1. Scaling and the "model-agnostic" feature of the SO(3)-Averaged Flow

We agree with the reviewer that models with more parameters might yield better results compared to our current 4.7M parameters NequIP-like architecture. One of the major goals of our work is to demonstrate that Averaged Flow as an alternative training objective can help the model to converge faster to better performance. In that sense, we did not compare across different models, but rather compared the same model trained with different objectives (Fig. 2). Since Averaged Flow is a new training objective which does not require specific model architecture design, any modern neural network architecture can theoretically be used to parameterize $v^{\theta}_{t} (x_t)$ in Eq.7. We will surely explore various model architectures as well as investigate the scaling law of models in future work.

2. Necessity of reflow with performance degradation

The reviewer raised a highly valuable question about the performance drop after reflow. In Fig.3 and Sec. 4.4, we compared the performance of our model before and after reflow at different steps. In short, the conclusion is that reflow is necessary to maintain a reasonably good generation quality when $N_{step}<10$. For any model architecture, reducing the number of ODE steps from 10 to 2 means 5x faster inference. Without reflow, such acceleration is not feasible without overly compromising the generation quality. We also want to highlight that AvgFlow_reflow with 2-steps sampling outperforms all cheminformatics tools. It also outperforms or matches Torsional Diffusion (5-steps), MCF-S (3-steps) and MCF-B (3-steps) with >20x faster inference. We believe the points we listed above are strong evidence that reflow is valuable in pushing the boundary of speed-quality trade-off of conformer generation. Moreover, our work serves as a confirmation that the reflow+distillation algorithm is also effective in straightening the flow trajectory when generating molecular conformers with equivariant GNN. With the findings in our work, models like ET-Flow can also be fine-tuned for faster inference.

The coverage metrics used to evaluate the models’ performance are highly sensitive to model errors accumulated during reflow. Specifically, we use generated samples ($X_1’$) for finetuning the model during reflow. If some generated samples are not within coverage threshold with the ground truth conformers, undesired bias can be introduced during reflow and errors will be accumulated, leading to a performance degradation after reflow. One potential solution we will explore in future works is to filter the generated $X_1’$ by including only those with low RMSD to ground truth conformers in the reflow fine-tuning dataset. We will include the improvement discussion in the manuscript (Sec. 4.3).

3. Quality-speed trade-off

We agree with the reviewer that an ideal solution would improve both accuracy and inference speed at the same time. However, in practice, the trade-off between generation quality and speed is inevitable. For molecular conformer generation, both cheminformatics/quantum chemistry tools and deep learning methods are facing the same issue. Through reflow, we are demonstrating a method to push the boundaries of such trader-off by significantly reducing the inference speed while maintaining a reasonably good generation quality.

4. Details of SO(3)-Averaged Flow. Is it augmenting the training set with different rotations? More explanation on the theoretical motivation.

We want to clarify that the AvgFlow proposed in our paper is a new flow-matching objective instead of a data augmentation strategy (i.e. adding rotations to ground truth conformers). Our objective is used to train models to learn the averaged flow to all rotations of conformers, which is analytically solved. Following the suggestion of the reviewer, we have provided more details in Sec. 3.1 and the code implementation in the Appendix 3.1.

We appreciate that the reviewer acknowledge the novelty of Averaged Flow. We hope the response below can address the questions raised by the reviewer:

1. Derivation of $u_t(x)$ from $\left[\partial_\alpha \log Z_t(x_t,\alpha)\right]_{\alpha=0}$

Indeed, $\left[\partial_\alpha \log Z_t(x_t,\alpha)\right]_{\alpha=0}$ is not exactly $u_t(x)$, but the non-trivial part of $u_t(x)$. To be more specific, the $-x$ and $\frac{1}{1-t}$ in the term $\frac{\hat{x}R^{T}-x}{1-t}$ of $u_t(x)$ are trivial when computing the average of this term wrt to $R$.

Eq.8 of the manuscript (shown below) provides detail of how to derive the vector field $u_t(x)$ from $\left[\partial_\alpha \log Z_t(x_t,\alpha)\right]_{\alpha=0}$

Eq.8: $u_t(x_t) = (\left[\partial_\alpha \log Z_t(x_t,\alpha)\right]_{\alpha=0} - x_t)/(1-t).$

2. Math detail of Eq.4 to Eq.6

Generally speaking, Eq.6 is derived through expanding the quadratics in the exponent in Eq.4. Specifically:

$$\begin{aligned} -\frac12 \frac{\| x - t \hat x R^T \|^2_\Sigma}{(1-t)^2} + \alpha \cdot (\hat{x}R^{T}) &= -\frac12 \frac{\| x - t \hat x R^T \|^2_\Sigma}{(1-t)^2} + \mathrm{tr} (\alpha^T \hat x R^T) \\\\ &= -\frac12 \frac{\mathrm{tr}(x^T \Sigma x) + t^2 \mathrm{tr}(\hat x^T \Sigma \hat x R^T R) - 2 t \ \mathrm{tr}(x^T \Sigma \hat x R^T)}{(1-t)^2} + \mathrm{tr} (\alpha^T \hat x R^T) \\\\ &= \text{tr}((\alpha^T + \frac{t}{(1-t)^2} x^T \Sigma) \hat{x} R^T) - \underbrace{ \frac{\mathrm{tr}(x^T \Sigma x) + t^2\mathrm{tr}(\hat x^T \Sigma \hat x)}{2(1-t)^2}}_{\text{Denoted as ``constant in }\alpha\text{" in Eq.6}} \end{aligned}$$

Please note that note that $R^T R$ is the identity.

Since we are not able to make revision to the manuscript, we will include the details steps from Eq.4 to Eq.6 in the appendix of camera-ready version of the manuscript if the paper is accepted. Thank you for the understanding.

We would like to thank the reviewer for the comprehensive review. We appreciate that the reviewer consider our work to be well-motivated, and our comparison between Averaged Flow and other objectives to be sound. We also acknowledge the reviewers asks for more details of the SO(3)-Averaged flow and implementation. Please see below for the responses to the weakness and questions:

Weakness: Implementation detail not clear.

For the reference of the reviewer, we have now updated the manuscript with the Python implementation of Averaged Flow in Appendix 3.1.

1. Definition of $\hat{x}$.

Yes, in the specific problem of molecular conformer generation, $\hat{x}$ is a ground truth conformer from the dataset. To avoid confusion, we have added more explanation in Sec. 3.1 including the definition of $\hat{x}$.

2. Dimentionality of ($A_t$, $b_t$).

To provide more details and avoid confusion, we have changed some notations and rewriten the Averaged Flow derivation in Sec. 3.1. Please refer to the modified manuscript for details.

3. Clarification of norm in Eq.4.

Yes, we can confirm it's matrix norm.

4. Computation time of Averaged Flow objective.

This question is of high importance as we do not want to add too much computational overhead when training with the new objective. We have attached the Python implementation of Averaged Flow in Appendix. We benchmarked the computation time of solving for the averaged flow with batched graphs containing $\\{1, 10, 100, 1000 \\}$ graphs (Table 5 in Appendix 3.2). Each graph in a batch contains 50 nodes, higher than the average number of atoms in GEOM-Drugs molecules. Compared to conditional OT, which computes the flow as linear interpolation between data and noise, our method adds only a very small overhead to the computational time: 0.5 ms with batch size 100 and 1 conformer. It also scales well with the number of graphs in a batch. Therefore, we can conclude that the Averaged Flow objective will not be a speed bottleneck during training.

5. Scaling of model and harmonic prior.

We want to emphasize that we demonstrate the effectiveness of both Averaged Flow through comparing the performance of the same model architecture trained with different objectives. That being said, we acknowledge that Averaged Flow can be used for training other model architectures such as the one used in ET-Flow. We will surely explore the model architecture and model scaling laws in the future work. We also agree with the reviewer that the harmonic prior can be beneficial for the models’ performance. Therefore, we implemented a variant of Averaged Flow that is able to start from the harmonic prior (a non-isotropic Gaussian based on covalent bond). Please refer to the code snippet in Appendix 3.1 for details.

6. Sampling speed agains ET-Flow-SS.

We have now included the inference speed benchmark of ET-Flow (5-steps) in Table 3. We choose ET-Flow instead of ET-Flow-SS because the checkpoint of ET-Flow has not been released so we can only take the available performance results from the paper. ET-Flow is the only one reported for few-step generation. ET-Flow has very high coverage and precision with 5-steps generation, but we are not able to benchmark its performance with even less ODE steps. As we have shown in Fig.3, $N_{step}$<5 can significantly impact the generation quality. Compared to ET-Flow, our $\mathrm{AvgFlow_{Reflow}}$ is ~40x faster in inference. With the confirmation from our work that reflow can further reduce sample steps, we can explore the possibility of improving ET-Flow sampling speed in the future.

Thank you for your detailed response on my questions and concerns. I appreciate your effort in updating the manuscript with the details to better understand what is going on under the hood when performing Averaged Flow in Section 3.1 as well as the appendix by providing the python code.

I also believe that the harmonic prior is a strong inductive bias from which Averaged Flow can benefit, regardless of the used model architecture. As I mentioned before, I appreciate the authors effort to evaluate the proposed Average Flow training objective within the same network architecture against OT as well as OT+Kabsch Alignment. I believe the scaling the network combined with Harmonic Prior is a valuable avenue for future work.

Since other reviewers mentioned the lack of details of the implementation in the first round of rebuttal phase, and the authors have resolved this issue, I am willing to increase my score to 8.

We would like to thank the reviewer for their feedback. We appreciate that the reviewer highlights our contribution in balancing compute cost and performance, which is the core motivation of this work. Please see below for the responses to the raised questions:

1. Experimental support for "model-independency" of the Averaged Flow and reflow+distillation method.

We agree with the reviewer that emphasizing on Averaged Flow and reflow+distillation is “model-independent” requires more experimental support with multiple model architectures being trained using our framework. The focus of the manuscript is to demonstrate the effectiveness of Averaged Flow and reflow in improving training and inference of model for faster conformer generation. Therefore, we removed the third point in the main contribution claim. However, we still believe that both methods can theoretically be extended to other model architectures for several reasons:

1. Averaged Flow is simply an alternative objective to the condition OT (linear interpolant), which does not require extra design of neural network architecture. Model architectures that are powerful enough to be used in flow-matching tasks (such as ET-Flow) can be used in Eq.7 to parameterize $v^{\theta}_{t} (x_t)$.

2. Reflow+distillation has been applied to domains other than image generation, such as point cloud generation. Our work serves as a confirmation that the reflow+distillation algorithm is also effective in straightening the flow trajectory when generating molecular conformers with equivariant GNN. With the findings in our work, models like ET-Flow can also be finetuned for faster inference.

3. That being said, experiments on different model architectures to search for the optimal ones are definitely critical and will be pursued in future works.

2. Performance and computational cost change of ET-Flow with fewer steps.

We thank the reviewer for raising the question about inference speed comparison between ET-Flow and our model. We tested the inference speed of ET-Flow with 5-steps. To be consistent, the inference speed is benchmarked with batchsize=1. The Inference time of ET-Flow is added to Table 3 of the manuscript. We can see that $\mathrm{AvgFlow_{Reflow}}$ is ~40x faster than ET-Flow$\mathrm{_{5step}}$.

ET-Flow maintains high generation quality with Nstep=5. However, our analysis in Fig.3 shows that the generation quality can drop significantly when $N_{step}<5$. Since the checkpoint of ET-Flow has not been released, we are not able to benchmark its generation quality with $<5$ ODE steps. We hope that the reflow and distillation algorithm in our paper can be extended to finetune powerful models like ET-Flow to further reduce sampling steps and accelerate inference.

3. Details of SO(3)-Averaged Flow derivation.

We have added a thorough explanation of how the new objective is calculated with a closed-form solution. The mathematical derivation is also accompanied by code implementation which is added to the Appendix 2.1 of the manuscript.

4. Comparison of different conformer generation methods and different steps.

It is a great suggestion from the reviewer to compare methods with different steps of sampling. We want to emphasize that one goal of this work is to enable few-step conformer generation through reflow. Therefore, we focus more on the generation quality with a very limited number of steps $N_{step}<5$. In Fig.3, we demonstrate that the benchmark performance of $\mathrm{AvgFlow_{Reflow}}$ surpasses the model by a large margin when $N_{step}<5$. Then, in Table 3, we compare $\mathrm{AvgFlow_{Reflow}}$ with other models sampled with 3-5 steps and demonstrate that reflow effectively alleviated the performance degradation in few-steps generation.

Dear reviewer mA9y:

Since the discussion period will end tomorrow (Dec 2nd), we want to check if our response has addressed your questions and concerns regarding our paper. Please let us know if you have any follow-up comment or question regarding our manuscript. Again, thank you for the time spent on reviewing and discussing the manuscript.

We would like to thank the reviewer for the thorough review and acknowledging our technical contribution. Please see below for the point-to-point response to questions raised by the reviewer:

1. SO(3)-Averaged Flow details, implementation, and comparison to FrameAveraging.

We agree with the reviewer that sharing the code of Averaged Flow implementation can help to clarify the sampling and averaging procedure. We have now attached the code in Appendix 2.1.

We have to clarify that the Averaged Flow objective is analytically calculated without sampling rotations. The rewritten Sec. 3.1 of the manuscript provides more detailed mathematical derivation. We hope the detailed derivation along with the code can improve the transparency of our method to readers.

We thank the reviewer for bringing the FrameAveraging paper to attention. Based on our understanding, the goal of FrameAveraging is to adapt backbone networks to be invariant or equivariant to desired symmetries. FA achieves that by averaging a subset of group operators (called Frame) which they show yield the same equivariance/invariance as averaging over the entire group. In contrast, the goal of Averaged Flow is not to change symmetry equivariance/invariance of the backbone model, but rather to variance reduced training objective. Unlike FA that averaged over a subset of group operators, Averaged Flow analytically computes the averaging of SO(3) with closed-form solution.

2. Explanation needed to justify SO(3)-Averaged Flow when backbone is equivariant.

We want to clarify that the AvgFlow proposed in our paper is a new flow-matching objective instead of a data augmentation strategy (i.e. adding rotations to ground truth conformers). Our objective is used to train models to learn the averaged flow to all rotations of conformers, which is analytically solved. To clarify the confusion, we provided more details in Sec. 3.1 and code implementation in Appendix 2.1.

For the conformer generation task, the generated sample can be freely rotated while still being a “correct” conformer. Starting from a rotationally-invariant distribution such as standard Gaussian, the vector field (flow) of atoms at time-step t can point to any rotation of the ground truth conformer and still be “correct”. Therefore, instead of data augmentation by randomly rotating the ground truth conformers during training, we analytically solved for the SO(3)-averaged flow and used that as an alternative variance-reduced training objective. The AvgFlow objective helps the model to converge faster to better performance compared with the linear interpolation or Kabsch-aligning objective to the starting noise.

3. Explanation of Averaged Flow being "model-agnostic" and comparison to other methods such as DiCSO.

Since the AvgFlow objective proposed in our manuscript is a training objective, any modern neural network architecture can theoretically be used to parameterize $v^{\theta}_{t} (x_t)$ in Eq.7. To clarify, the phrase “model agnostic” specifically means “architecture-agnostic” instead of claiming AvgFlow can transport atom coordinates sampled from any arbitrary distribution to ground truth conformational distribution.

We thank the reviewer for bringing up the DiSCO paper. We believe that DiSCO can be very useful to optimize conformers generated by a training model or cheminformatics tools to lower-energy states. From our perspective, DiSCO serves more as an optimization tool like xTB relaxation instead of generating from pure Gaussian noise. We have included the discussion of DiSCO in the related work section.

4. Necessity of reflow+distillation.

The reviewer raised a highly valuable question about the performance drop after reflow. In Fig.3 and Sec. 4.4, we compared the performance of our model before and after reflow at different steps. In short, the conclusion is that reflow is necessary to maintain a reasonably good generation quality when $N_{step}<10$. For any model architecture, reducing the number of ODE steps from 10 to 2 means 5x faster inference. Without reflow, such acceleration is not feasible without overly compromising the generation quality. We also want to highlight that $\mathrm{AvgFlow_{Reflow}}$ with 2-steps sampling outperforms all cheminformatics tools. It also outperforms or matches Torsional Diffusion (5-steps), MCF-S (3-steps) and MCF-B (3-steps) with >20x faster inference. We believe the points we listed above are strong evidence that reflow is valuable in pushing the boundary of speed-quality trade-off of conformer generation.

The coverage metrics used to evaluate the models’ performance are highly sensitive to model errors accumulated during reflow. Specifically, we use generated samples ($X_1’$) for fine-tuning the model during reflow. If some generated samples are not within coverage threshold with the ground truth conformers, undesired bias can be introduced during reflow and errors will be accumulated, leading to a performance degradation after reflow. One potential solution we will explore in future works is to filter the generated ($X_1’$) by including only those with low RMSD to ground truth conformers in the reflow fine-tuning dataset. We will include the improvement discussion in the manuscript (Sec. 4.3).

5. Computation time of SO(3)-Averaged Flow

We thank the reviewer for the suggestion. To clarify, instead of sampling multiple rotations of training data before averaging, we analytically solve for the SO(3)-Averaged Flow (implementation details added to the Appendix) and use that as the training objective. Therefore, the training data size is not increased during each epoch of training.

We benchmarked the computation time of solving for the averaged flow with batched graphs containing $\\{ 1, 10, 100, 1000 \\}$ graphs (Table 5). Each graph in a batch contains 50 nodes, higher than the average number of atoms in GEOM-Drugs molecules. Compared to conditional OT, which computes the flow as linear interpolation between data and noise, our method adds only a very small overhead to the computational time: 0.5 ms with batch size 100 and 1 conformer. It also scales well with the number of graphs in a batch. Therefore, we can conclude that the Averaged Flow objective will not be a speed bottleneck during training.

Dear reviewer GPaB:

Since the discussion period will end tomorrow (Dec 2nd), we want to check if our response has addressed your questions and concerns regarding our paper. Please let us know if you have any follow-up comment or question regarding our manuscript. Again, thank you for the time spent on reviewing and discussing the manuscript.