Efficient Molecular Conformer Generation with SO(3)-Averaged Flow Matching and Reflow

We would like to thank the reviewer for reviewing and acknowledging the novelty of the SO(3)-Averaged Flow. Please see the response below:

Theoretical Claims

The major motivation behind the development of AvgFlow is to eliminate the need for data augmentation through rotation by training the model to learn the flow from interpolant $x_t$ to ground truth $x_1$ averaged over the SO(3) group. Sec 3.1 is the major contribution, as it mathematically derives a closed-form solution for the SO(3)-Averaged Flow so that training can be efficient without much of computational overhead. We have also attached the Python implementation of solving the AvgFlow objective in Sec A.3 to accompany the mathematical derivation. We would be happy to take further questions from the reviewer about specific steps in Sec. 3.1.

Weakness

1. We want to emphasize that the major motivation of this paper is to improve the efficiency of diffusion/flow-based models to the level of cheminformatics tools for conformer generation through algorithmic innovations. Given the need of ultra-large scale (~$10^8$-$10^9$ compounds) virtual screening, diffusion/flow conformer generation model would become practical with improved sampling efficiency. Bearing the motivation in mind, we chose to implement a compact 4.7M-parameter equivariant GNN model that achieves significant sampling speedup while maintaining good generation quality. Comparing to larger transformer-based baselines, the model's speedup in few-step sampling compensates for reduced generation quality: the 2-step $\mathrm{AvgFlow_Reflow}$ model outperforms 3-step MCF-B in precision metrics. It also achieves ~58% of the 5-step ETFlow with ~40x speedup in sampling. We believe that the current 4.7M model has fulfilled this motivation.

   That being said, we agree with the reviewer that achieving the SOTA in generation quality is valuable. We have implemented a more scalable diffusion transformer (DiT) with pairwise biased attention, similar to AlphaFold3 [1] and Proteína [2], for conformer generation. Due to the limited time, we trained a 52M-parameter DiT model with AvgFlow objective for only ~124k steps and benchmarked it on the Drugs test set (see results here). We have also compared the performance of the DiT model trained with AvgFlow, Kabsch, and Conditional OT (same experiment as Sec 4.1 of manuscript, see results here). The benchmark results of DiT demonstrates that:

   • AvgFlow can be a better training objective than Conditional OT and Kabsch even with non-equivariant model architecture.

   • A transformer-based architecture with more parameters trained with AvgFlow can achieve performance on-par with other SOTA models, even with a relatively limited number of training steps.

   The performance of the DiT model is expected to improve further as it is trained for more steps. It can also be scaled up to similar size to MCF-L for better performance if resources allow. We are currently working on reflow fine-tuning the DiT model and we will share results once completed.

2. Clarification about Table 3:

   • We agree with the reviewer that baselines should be compared with same number of sampling steps when possible. We took the benchmark values from corresponding papers because their checkpoints were not released before this manuscript was submitted. We have now benchmarked the 2-step and 1-step generation results of ETFlow and MCF-B/L (see updated Table 3 here). Our reflow model outperforms MCF for 2-step generation across all metrics. More importantly, our distill model outperforms both ETFlow and MCF by a large margin for 1-step generation, especially in the recall metrics. Despite the high 2-step generation quality of ETFlow, our reflow and distill models are still 16x and 32x faster in sampling, respectively.

   • We want to respectfully disagree with the reviewer's suggestion that the effect of implementation should be excluded when comparing sampling speed. We would argue that faster implementation should be counted as a technical contribution as it validates a potential direction for optimization and acceleration.

Answers to questions

1. Please see above answers regarding experiments with larger model.

2. Yes, we are working on reflow/distill fine-tuning the DiT model (non-equivariant transformer architecture) mentioned above. These experiments can be time-consuming as they require generating ($X_0'$, $X_1'$) couples. We will report results once completed.

References:

[1] Abramson et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature (2024)

[2] Geffner et al. "Proteina: Scaling Flow-based Protein Structure Generative Models." ICLR (2025).

We want to thank the reviewer for the comprehensive review. Please see below for responses:

Essential references

The references suggested by the reviewer are indeed relevant to this paper. However, we want to point out that we have explicitly compared our model to GeoMol for both the QM9 and Drugs benchmarks. We have also discussed DiSCO as a related work in the field of molecular conformer optimization in Sec 2.1 (line 60-62). From our perspective, DiSCO is a great method to optimize molecular conformer, which is a fundamentally different application compared to generating conformer from scratch (noise). Therefore, we did not benchmark explicitly against DiSCO. As far as we know, Equiformer has not yet been used for conformer generation. It would be nice if the reviewer could point us to related literature.

Weaknesses

Conformer generation in virtual screening

We want to respectfully correct the reviewer's statement about that the conformer generation being an offline task during virtual screening. During ultra-large virtual screening campaign (~$10^8$-$10^9$ compounds), conformer generation is indeed an online task, rendering the acceleration crucial. Our method, which achieves significant speedup, makes diffusion/flow model more practical for the use in virtual screening.

Degradation of performance after reflow

We want to argue that the goal of reflow is to improve the model's performance in few-steps sampling, which is critical for achieving speedup in generation. Comparing to the model before reflow, the model after reflow demonstrates significantly better performance in <5-step sampling. We also want to emphasize that the reflow technique is model architecture-agnostic and can be applied to other flow-based models to accelerate generation. That being said, we are planning on reducing possible performance degradation by excluding lower quality generated conformer ($X_1'$) from the reflow fine-tuning dataset. In that way, we expect to alleviate the problem of sampling error propagating to fine-tuning stage.

Responses to comments

With the increasing demand of billion-level virtual screening, the efficiency gain would be critically impactful to the application of flow-based conformer generation model. We would add more discussion about the impact in addition to line 42-43 of right column. Please see the previous answer for performance degradation of reflow and proposed future solution. We will perform a thorough grammar check for the future version of manuscript. We will release the model upon publication.

Answers to questions

1. Many molecules in the Drugs test set are with >100 ground truth conformers. Therefore the method's ability to generalize to flexible molecules is reflected well by the benchmark results in Table 2.

2. Theoretically, AvgFlow as an training objective should not be affected by molecular size and bonds because those features are not used in the close-form solution.

3. We want to clarify that the major contributions of this work, including AvgFlow and reflow/distillation techniques, are training schemes rather than new model architectures. We believe the proper way of showing the advantage of AvgFlow is by comparing same architecture trained with different flow-matching objective (as shown in Fig.1). Similarly, we show in Fig.3 the necessity of reflow/distillation. To strengthen this point, we have recently trained a diffusion transformer with pairwise biased attention with the AvgFlow which achieves on par performance MCF and ETFlow (please see details in the response to reviewer iSgP). This demonstrates that the improvements brought by AvgFlow are architecture-independent.

4. We believe that RL-based fine-tuning can help models generate better conformers (more consistently at lower energy state) but may not improve sampling efficiency, which is the primary motivation of this work.

5. Yes the model can be extended to larger-scale drug discovery pipeline as it achieves significant speedup for flow-based model in conformer generation. The generation quality of our 1-step distilled model has surpassed cheminformatic tools.

We thank the reviewer for acknowledging the novelty of integrating reflow and distillation method for accelerating the sampling of conformer generation model. Please see our response below to other questions and comments:

Weakness

1. We want to emphasize that the major motivation of this paper is to improve the efficiency of diffusion/flow-based models to the level of cheminformatics tools for conformer generation through algorithmic innovations. Given the need of ultra-large scale (~$10^8$-$10^9$ compounds) virtual screening, diffusion/flow conformer generation model would become practical with improved sampling efficiency. Bearing the motivation in mind, we chose to implement a compact 4.7M-parameter equivariant GNN model that achieves significant sampling speedup while maintaining good generation quality. Comparing to larger transformer-based baselines, the model's speedup in few-step sampling compensates for reduced generation quality: the 2-step $\mathrm{AvgFlow_Reflow}$ model outperforms 3-step MCF-B in precision metrics. It also achieves ~58% of the 5-step ETFlow with ~40x speedup in sampling. We believe that the current 4.7M model has fulfilled this motivation.

   That being said, we agree with the reviewer that achieving the SOTA in generation quality is valuable. We have implemented a more scalable diffusion transformer (DiT) with pairwise biased attention, similar to AlphaFold3 [1] and Proteína [2], for conformer generation. Due to the limited time, we trained a 52M-parameter DiT model with AvgFlow objective for only ~124k steps and benchmarked it on the Drugs test set (see results here). We have also compared the performance of the DiT model trained with AvgFlow, Kabsch, and Conditional OT (same experiment as Sec 4.1 of manuscript, see results here). The benchmark results of DiT demonstrates that:

   • AvgFlow can be a better training objective than Conditional OT and Kabsch even with non-equivariant model architecture.

   • A transformer-based architecture with more parameters trained with AvgFlow can achieve performance on-par with other SOTA models, even with a relatively limited number of training steps.

   The performance of the DiT model is expected to improve further as it is trained for more steps. It can also be scaled up to similar size to MCF-L for better performance if resources allow. We are currently working on reflow fine-tuning the DiT model and we will share results once completed.

2. We agree with the reviewer that an OOD evaluation on the GEOM-XL dataset can examine the model's generalizability to large molecules. The benchmark results of the 4.7M $\mathrm{AvgFlow}$ model on GEOM-XL is attached here. In general, our model has slightly higher mean AMR but on-par median AMR for both recall and precision compared with SOTA model MCF and ETFlow. We also want to mention that the DiT model we recently trained with AvgFlow objective is expected to generalize better to larger molecules thanks to its scalable architecture. We will benchmark its performance on GEOM-XL as well.

Answer to the question

We understand the confusion of the reviewer about the derivation in Sec 3.1. To clarify:

• For the conformer generation problem, we define orbits $\hat{x}$ as low-energy conformers of a given molecule. Therefore the integral $\int d\hat{x}\ \hat{q}(\hat{x})$ in Eq. 2, representing the entire conformer ensemble, can be written as $\sum_{\hat x \in \mathcal{X}} \hat q(\hat x)$, where $\mathcal{X}$ is the set of conformers and $\hat q(\hat x)$ is the weight associated with each conformer (also elaborated by line 144-152, right column).
• Your understanding is correct that for a given conformer $\hat{x}$, the rotated molecule $g \cdot \hat{x}$ has the same energy state.
• The AvgFlow method proposed is capable of integrating over all conformers (orbits) and the SO(3) group. In practice, we only integrate over the SO(3) group during training and sample one conformer in each epoch to approximate the expectation of the conformer ensemble. (line 168-170, right column). Hence, the flow-matching objective proposed in this work is SO(3)-Averaged Flow.

References:

[1] Abramson et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature (2024)

[2] Geffner et al. "Proteina: Scaling Flow-based Protein Structure Generative Models." ICLR (2025).