WGFormer: An SE(3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction

For your comments, below we answer them one by one.

W1: No discussion about efficiency

A1: Firstly, the inefficiency of energy-based/quantum-based methods (e.g., molecular dynamics simulation and density functional theory calculation) has already been a well-established consensus (also has been discussed in lines 12-20 of our paper) in this field [1][2][3], which is why numerous learning-based methods are proposed to approximate their accuracy.

Secondly, as shown in Figure 4 of our paper, we have demonstrated that our WGFormer is significantly faster than existing learning-based methods, reducing the runtime per molecule by over 50%.

W2&W5: Lack interpretability and the relationship with the conformational energy minimization process

A2&W5: As we have mentioned in lines 30-35, our WGFormer is a Wasserstein gradient flow-driven SE(3)-Transformer architecture. It optimizes molecular conformations by minimizing an energy function (i.e., Eq.17) defined on the latent mixture models of atoms. The corresponding proofs, including the relationship between our energy function and conformation optimization (lines 220-242), have been provided in Section 4. To further verify our claim, we have conducted analytic experiments in our response to Reviewer dW7c (https://openreview.net/forum?id=2wUQttiab3&noteId=bSAlOuwM9f).

These experiments demonstrate that:

1. Our WGFormer is indeed minimizing the energy function defined on the latent mixture models of atoms;

2. Minimizing this energy function indeed helps optimize the physically-meaningful energy of molecular conformation, and these two kinds of energy are highly correlated;

3. Minimizing this energy function is indeed closely relevant to improving the final metrics.

W3: Clarification of relevant statements

A3: The interpretability we emphasize here refers to the fact that the feedforward computation of existing model architectures (e.g., those baselines) cannot effectively correspond to the process of molecular conformation optimization. In contrast, WGFormer optimizes molecular conformations by minimizing an energy function defined on the latent mixture models of atoms, whose feedforward computation is the Euler step solving a continuity equation.

The DSMBind you mentioned works to predict protein-ligand binding energy, which is unrelated to our work in either problem setting or technical route.

W4: Unclear motivation

A4: The motivation of our model design is clear --- as shown in lines 66-85, we aim to build a new SE(3)-equivalent model with interpretable architecture for obtaining molecular ground-state conformations from their low-quality conformations. To achieve this aim, we formulate the task as a conformation energy optimization problem, so that the feedforward computation of the model needs to be the Wasserstein gradient flow of the contunity equation in Eq.13. All our architectural improvements, including modifying QKV mechanism and applying Sinkhorn-based attention (as shown in lines 206-219), serve for this aim.

W6: A lot of same work and lack contribution

A6: We respectfully disagree with this viewpoint. In fact, all other reviewers have widely recognized the contribution of our work, including interesting and important relations with energy-based methods (Reviewer dW7c), the first attempt to predict molecular ground-state conformation through the lens of Wasserstein gradient flow (Reviewer NScR), and the fact that performance improvements are sound and lean in a promising direction in terms of architectural exploration (Reviewer FN8S).

In the aspect of model architecture, we propose a first Wasserstein gradient flow-driven SE(3)-Transformer for molecular modeling. It is the first SE(3)-equivariant model that corresponds to a Wasserstein gradient flow in the latent space of atoms, which not only enhances the model interpretability but also improves computational efficiency.

In the aspect of theory, we build the connection between the model architecture and the Wasserstein gradient flow and provide the conditions for the connection. In addition, as shown in Eq.17, we analyze the latent energy function in depth and explain its rationality from the perspective of entropic OT.

In the aspect of application, currently, few attempts are made to predict molecular ground-state conformations. Our work provides a strong and interpretable solution to this challenging and important problem.

In summary, from any perspective, we believe our work deserves a higher score. We respectfully ask you for re-evaluating our work in light of the responses above and the comments of the other reviewers.

[1] Learning gradient fields for molecular conformation generation, ICML 2021.

[2] Learning neural generative dynamics for molecular conformation generation. arXiv preprint 2021.

[3] Energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 2022.

Thanks for your positive feedback and constructive comments. Below, we resolve your concerns one by one.

W1: Test on larger datasets.

A1: As shown in Table 1 of our paper, in addition to QM9, our WGFormer also achieves SOTA performance on the Molecule3D dataset, which contains about four million molecules and is larger than DRUGS. In addition, each molecule in DRUGS has multiple conformations, and we don't know which one is the ground-state conformation (even don't know whether the ground-state conformation is among them). So, this dataset is not very suitable for our task.

Nevertheless, we follow your suggestion, adding an experiment on DRUGS. We select the conformation with the highest Boltzmann weights as its ground-state conformation, then train and test different models. The results below demonstrate WGFormer's superiority in accuracy and efficiency on DRUGS.

W2: The novelty and significance of proposed model architecture

A2: Although attention-based models have been widely used in molecule and protein generation, our WGFormer is still valuable because of the following reasons:

Firstly, as shown in Figure 2 and lines 206-214, WGFormer introduces new QKV architecture, leading to a new cross-head interaction mechanism with fewer model parameters. Secondly, as shown in Appendix A, WGFormer applies the Sinkhorn-based attention module while maintaining the SE(3)-equivariance property, which enhances the interpretability of the model and makes it suitable for 3D molecular modeling. Combining these two improvements jointly, WGFormer can be interpreted as Wasserstein gradient flow for the latent mixture model of atoms.

No matter whether the original SE(3)-Transformer works well or not, our architectural improvements have led to better performance with fewer trainable parameters and better interpretability, demonstrating the rationality of our design.

W3: 1) Can finite Sinkhorn iterations be interpreted as Wasserstein gradient flow? 2) The energy functional depends on $k$ and changes for every layer.

A3: Firstly, using finite Sinkhorn iterations is not strong evidence that the model cannot be interpreted as Wasserstein gradient flow. The gap between theoretical analysis and practical implementation is natural --- for convex optimization, we stop an algorithm in finite iterations, but it does not mean its analysis based on infinite series is meaningless. In practice, the Sinkhorn algorithm converges very fast to the optimum. The Sinkformer in [2] merely applies 3-5 Sinkhorn iterations. Our WGFormer follows the same setting.

Secondly, $k$ changes do not mean the energy functional changes for every layer. In our paper, the energy functional is shown in Eq.(17), which is formulated as an entropic OT problem. Each WGFormer layer optimizes the same energy functional, leading to the Wasserstein gradient flow (lines 220-251, left column).

W4&W5: Interpretations of architecture improvement and Figure 5

A4&A5: Our WGFormer can optimize molecular conformations by minimizing an energy function defined on the latent mixture models of atoms. To further verify our claim, we have conducted additional experiments in our response to Reviewer dW7c (https://openreview.net/forum?id=2wUQttiab3&noteId=bSAlOuwM9f). These experiments demonstrate that:

1. WGFormer indeed minimizes a latent energy function for atoms' probability measure;

2. Minimizing this latent energy helps optimize the physical energy of conformation, highly correlated with the final metrics.

For Figure 5, as shown in lines 426-439, we first train a WGFormer with $L=30$ layers and test it using different layer numbers ($L=1,...,30$). The performance is improved as the number of layers increases, indicating that each layer is an Euler step to minimize the energy function in Eq.17. When applying a non-WGFormer architecture with $L$ layers, the performance may not be improved consistently when increasing the number of layers, as shown in this anonymous link (https://anonymous.4open.science/r/WGFormer-comparison/Non-WGFormer.pdf).

W6: Chemical property metrics

A6: Following [1], we compare predicted and real ground-state conformations on their energy (in kcal/mol), dipole moment (in debye), and HOMO-LUMO gap (in kcal/mol). The MAEs of WGFormer and typical baselines are shown below, further demonstrating WGFormer's superiority.

Hope the above responses help enhance your confidence to further support our work.

[1] Torsional diffusion for molecular conformer generation, NeurIPS 2022.

[2] Sinkformer: Transformers with doubly stochastic attention, AISTATS 2022.

Thanks for your comments. Below, we resolve your concerns one by one.

Q1: Comparisons with other generative models

A1: Existing generative models generate multiple conformations by sampling. To make such models (e.g., GeoDiff [2] and TorsionDiff [3]) applicable for generating molecular ground-state conformations, we follow the commonly-used protocol in [1], generating multiple conformations and only preserving the one with the lowest energy.

Using the above strategy, we train and test GeoDiff and TorsionDiff on QM9. The table below shows that WGFormer outperforms these two models, achieving lower prediction errors and much higher inference speed.

Q2: The differences with other methods

A2: The differences between our WGFormer and the methods you mentioned can be summarized as follows:

1) Interpretable architecture:

As shown in lines 36-49, existing conformation optimization methods (e.g., ConfOpt-Two/Three Atoms) merely apply neural networks to approximate the gradients of atoms' motions when optimizing a molecule's conformation. The feedforward computations of their models do not correspond to minimizing a meaningful energy function. In contrast, our WGFormer is a Wasserstein gradient flow-driven SE(3)-Transformer. Its feedforward computation exactly corresponds to the Wasserstein gradient flow of the latent mixture model of atoms (i.e., the evolution of atoms' probability measure in the latent space). Each WGFormer layer is an Euler step minimizing an energy function of the molecular conformation, thereby significantly improving performance and interpretability.

2) Focus on ground-state conformation:

High-quality conformation $\neq$ Ground-state conformation. The AlphaFold series predicts protein structures but cannot ensure the structures are in the ground state. Although AlphaFold3 can predict biological complexes, it does not ensure ground-state conformations, either. WGFormer predicts the ground-state conformation, which determines basic molecular properties (shown in lines 43-48). Therefore, the AlphaFold series is not correlated with our work.

Q3: Does Wasserstein gradient need to compute gradient? The comparison on computational consumption is required.

A3: Firstly, we would like to explain the key concepts clearly. As shown in lines 238-258, Wasserstein gradient flow captures a probability measure's evolution in the Wasserstein space (e.g., the change of the latent mixture model of atoms over time) rather than the gradient of each atom. The detailed differences between Wasserstein gradient flow and gradient flow can be found in [4], which has been cited in the paper.

Secondly, Wasserstein gradient flow is achieved by the model architecture itself rather than additional computation. As shown in Section 4.1 and A2, the feedforward computation of WGFormer corresponds to the Euler step solving a continuity equation (i.e., Eq.10) and minimizes an energy function (i.e., Proposition 4.1).

Thirdly, WGFormer improves SE(3)-Transformer's architecture without increasing the complexity. We can train it on a single 3090 GPU, and we have demonstrated its superiority on inference speed in Figure 4 of our paper and the Table in A1.

Q4: Comparison with flow-matching based methods?

A4: Our work, i.e., developing a Wasserstein gradient flow-based model, is different from the flow-matching learning strategy:

Theory: Wasserstein gradient flow optimizes an energy function of probability measures in the Wasserstein space, while flow-matching aims to model the velocity field of particles in the sample space, which does not model the energy of particles or their distribution evolution explicitly.

Tech route: Our work focuses on improving model architecture and making it interpretable as Wasserstein gradient flow. We do not change the model's learning paradigm. Flow-matching is a model-agnostic learning strategy for fitting the velocity field of data.

Implementation: WGFormer is learned to predict the ground-state conformation, and its architecture ensures the feedforward computation leads to lower energy. Flow-matching methods focus on generating "valid" rather than energy-minimized conformations. In general, they cannot ensure the energy is minimized in the inference phase.

Hope the above responses can resolve your concerns and help re-evaluate our work.

[1] REBIND: Enhancing Ground-state Molecular Conformation Prediction via Force-Based Graph Rewiring, ICLR 2025.

[2] Geodiff: A geometric diffusion model for molecular conformation generation, ICLR 2022.

[3] Torsional diffusion for molecular conformer generation, NeurIPS 2022.

[4] {Euclidean, metric, and Wasserstein} gradient flows: an overview. Bulletin of Mathematical Sciences, 2017.

Thanks for your appreciation of our work. Below, we resolve your concerns one by one.

Q1: Visualize the latent energy function values achieved through different numbers of layers.

A1: As we have mentioned in lines 30-35 of our paper, our WGFormer is a Wasserstein gradient flow-driven SE(3)-Transformer architecture. In particular, it can optimize molecular conformations by minimizing an energy function defined on the latent mixture models of atoms (i.e., Eq.17 in the paper, denoted as latent energy for short), and the corresponding proofs have been provided in Section 4 of our paper.

To further verify our claim, given a trained WGFormer with 30 layers, we have conducted a series of analytic experiments in the inference phase.

Firstly, following your insightful suggestion, we randomly sample some RDKit-based molecular conformations and pass them through the trained WGFormers. Given the output (i.e., $\mathbf{X}$ and $\mathbf{R}$) of each layer, we can calculate the latent energy per layer by solving Eq.17 (using the Sinkhorn-scaling algorithm). Then, we plot the curve of the latent energy varying with the number of layers. As demonstrated by the figure in this anonymous link (https://anonymous.4open.science/r/WGFormer-energy/Latent_Energy.pdf), the latent energy decreases gradually as the number of layers increases, which strongly validates that our WGFormer is indeed minimizing the energy function defined on the latent mixture models of atoms.

Q2: Verify that the latent energy is closely related to the physical energy and final metrics.

A2: Furthermore, we employ the widely used xTB tool [1] to calculate the physical energy values of the molecular conformations obtained through different layers (i.e., given the $\mathbf{X}$ and $\mathbf{R}$ obtained by each layer, we can pass them through the trained decoder and obtain the corresponding molecular conformation) and analyze its correlation with the latent energy values obtained by Eq.17. In particular, taking the physical and latent energy obtained in the first layer as the references, we record the relative changes of the two kinds of energy w.r.t. number of layers and their correlations in the table below:

Here, a strong linear correlation is indicated by the Pearson correlation coefficient (0.885±0.033), while the slightly higher distance correlation (0.906±0.018) suggests additional nonlinear dependencies. These results further validate the interpretability of our WGFormer --- minimizing the latent energy defined on the atoms' mixture model helps optimize the physically-meaningful energy of molecular conformation.

Finally, as shown in Figure 5 of our paper, all evaluation metrics (D-MAE, D-RMSE, and C-RMSD) are improved steadily as the number of layers increases, having validated that minimizing the latent energy value can effectively improve the final metrics.

In general, the above results effectively verify the rationality and interpretability of our WGFormer, further supporting our theoretical claims proposed in Section 4 from the experimental point of view.

Thank you once again for your valuable and insightful review, which has made our work more complete and convincing. We will add the above analytic experiments to the revised paper. We hope our responses resolve your concerns and make you more confident in supporting our work.

Reference:

[1] Bannwarth C, Ehlert S, Grimme S. GFN2-xTB—An accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-dependent dispersion contributions. Journal of chemical theory and computation, 2019, 15(3): 1652-1671.