WGFormer: An SE(3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction
Authors: Fanmeng Wang
Minjie Cheng
Hongteng Xu
Venue: ICML 2025 poster (2025)
Keywords: Molecular ground-state conformation
Molecular modeling
Wasserstein gradient flow
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: This work introduces a Wasserstein gradient flow-driven SE(3)-Transformer to achieve both accurate and efficient molecular ground-state conformation prediction.

Abstract: Predicting molecular ground-state conformation (i.e., energy-minimized conformation) is crucial for many chemical applications such as molecular docking and property prediction. 
Classic energy-based simulation is time-consuming when solving this problem, while existing learning-based methods have advantages in computational efficiency but sacrifice accuracy and interpretability.
In this work, we propose a novel and effective method to bridge the energy-based simulation and the learning-based strategy, which designs and learns a Wasserstein gradient flow-driven SE(3)-Transformer, called WGFormer, for ground-state conformation prediction. 
Specifically, our method tackles this task within an auto-encoding framework, which encodes low-quality conformations by the proposed WGFormer and decodes corresponding ground-state conformations by an MLP.
The architecture of WGFormer corresponds to Wasserstein gradient flows --- it optimizes conformations by minimizing an energy function defined on the latent mixture models of atoms, thereby significantly improving performance and interpretability.
Extensive experiments demonstrate that our method consistently outperforms state-of-the-art competitors, providing a new and insightful paradigm to predict ground-state conformation.
The code is available at https://github.com/FanmengWang/WGFormer.

WGFormer: An Srans-Tranee Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction
Fanmeng Wang'
Minjie Cheng'
Hongteng Xu 123
Abstract
minimized state the potential energy surface. It deter- mines many important molecular properties (Xu et al., 2023) and thus plays crucial role in various downstream applica- tions like molecular docking (Chatterjee et al., 2023; Paggi et al., 2024) and molecular property prediction (Moon et al., 2023; Liu et al., 2024). Traditionally, molecular ground- state conformation can be obtained through various energy- based simulation methods, e.g., molecular dynamics (MD) simulation (Hollingsworth & Dror, 2018) and density func- tional theory (DFT) calculation (Pracht et al., 2020). How- ever, these methods are too expensive and computationally slow to meet the growing enormous demands (Wang et al., 2023a; Cen et al., 2024; Li et al., 2025a). Although some cheminformatic tools like RDKit (Landrum al., 2013) have been developed to generate molecular conformations efficiently, they suffer from undesired precision when ob- taining ground-state conformations.
Predicting molecular ground-state conformation (i.e., energy-minimized conformation) is crucial for many chemical applications such as molecular docking and property prediction. Classic energy- based simulation is consumnnn when solv- ing this problem, while existing learning-based methods have advantages in computational effi- ciency but sacrifice accuracy and interpretabillty. In this work, propose novel and effective method to bridge the energy-based simulation and the learning-basee based strategy, which designs and learns a Wasserstein gradient flow-driven SE(3)- Transformer, called WGFormer, for ground-state conformation prediction. Specifically, our method tackles this task within auto-encoding frame- work, which encodes low-quality conformations by the proposed WGFormer and decodes corre- sponding ground-state conformations by an MLP. The architecture of WGFormer corresponds to Wasserstein gradient flows - optimizes con- formations by minimizing an energy function defined the latent mixture models of atoms, thereby significantly improving performance and interpretability. Extensive experiments demon- strate that our method consistently outperforms state-of-the-art competitors, providing new and insightful paradigm to predict ground-state con- formation. The code available https:/ :// github. com/ annengWang /WGF ormer.
Recently, with the significant advancement of artificial intel- ligence in the scientific field (Wang et al., 2023b; Abramson et al., 2024), learning-based methods have emerged as a promising solution to tackle this task (Kim et al., 2025). Early methods like ConfVAE (Xu al., 2021b) employ ious generative models to generate multiple potential low- energy conformations and then still require screening the energy-minimized conformation among them, which lack the guarantee of obtaining ground-state conformations due to the uncertainty of the two-stage process (Xu et al., 2023). In this context, subsequent efforts have shifted to design specialized methods to predict molecular ground-state con- formation directly (Xu et al., 2021d). Among them, some state-of-thes -ftt--r-- methods, e.g., ConfOpt (Guan et al., 2022), formulate ground-state state conformation prediction as opti- mization problem, which takes low-quality conformations as input and optimizes these conformations accordingly. Es- sentially, these methods aim to apply neural networks to predict the gradient field of the conformational energy land- scape and thus imitate gradient descent update the input conformations. However, their neural networks are designed empirically, making the feedforward computations gener- ally not correspond to minimizing a physically meaningful energy function conformations. Therefore, such empiri- cal neural network architectures not only harm the model interpretability but also lead to sub-optimal performance.
1. Introduction
Molecular ground-state conformation represents the most stable 3D molecular structure corresponding to the energy-
Gaoling School of Artificial Intelligence, Renmin University of China "Beijing Key Laboratory of Research Large Mod- els and Intelligent Governance 'Engineering Research Center of next-Generation Intelligent Search and Recommendation. MOE, Correspondence to: Hongteng Xu <hongtengxu @ ruc edu.en>.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

WGFormer: An SE(3)- for Molecular Ground-Stale Conformation Prediction
WGFormer
WGFormer
WGFomer
Decoder
Encoder
Initial RDKit 3D Conformation
3D Ground-State Conformation
A Latent Mixture Model of Atoms
Minimization of Energy Function
Updated Mixture Model
Figure 1. An illustration of the proposed model architecture, in which WGFormer corresponds Wasserstein gradient flows and minimizes a physically reasonable energy function defined on the latent mixture models atoms.
In this work, considering that low-quality conformations can be readily available (Zhou et al., 2022; Feng et al., 2024), and specialized model architectures designed for 3D molecules have made great progress (Fuchs al., 2020: Satorras al., 2021), we design and learn Wasserstein gradient flow-driven Transffraa,,,,m,r,, called WGFormer, to obtain molecular ground-state conformation through op- timizing the corresponding low-quality conformation. As illustrated in Figure 1, our method tackles this task within an auto-encoding framework, where we first encode low- quality conformations (i.e., those generated by RDKit) through our WGFormer, and then decode the corresponding ground-state conformations through a simple multi-layer layer perceptron (MLP). Inspired by the Sinkformer in (Sander et al., 2022), WGFormer adjusts the query. key, and value matrices and applies the Sinkhorn-scaling algorithm (Cuturi, 2013) to compute attention maps that involve relational in- formation. Such an architecture makes our WGFormer work as Wasserstein gradient flows (Santambrogio. 2017) under specific conditions, which minimizes energy function defined on the latent mixture models of atoms, significantly improving performance and interpretability.
2014) to generate conformations. CGCF (Xu et al., 2021a) and ConfGF (Shi et al., 2021) are further proposed to en- hance performance through normalizing flows (Kobyzev et al., 2020) and score-based generative models (Song et al., 2020), respectively. Recently, some methods like GeoD- iff (Xu al., 2021c) and TorsionDiff (Jing al., 2022) leverage popular diffusion models (Ho et al., 2020) to gener- ate conformations. However, these methods are essentially designed to generate multiple "reasonable" conformations rather than directly obtaining ground-state conformations.
Focusing on molecular ground-state conformation predic- tion, the work in (Xu et al., 2021d) first proposes the Molecule3D benchmark with ground-state conformations determined through density functional theory (DFT) calcu- lations and develops a graph neural network-based solution. GTMGC (Xu et al., 2023) formally defines this task and pre- diets molecular ground-state conformation based on Graph Transformer (Ying et al., 2021). REBIND (Kim et al., 2025) improves GTMGC by adding edges to molecular graphs based on the Lennard-Jones potential, which captures non- bonded interactions for low-degree atoms. However, these methods predict ground-state conformations merely based on 2D molecular graphs, whose input information may be insufficient to obtain ground-state conformations.
To the best of our knowledge, our method makes the first attempt to predict molecular ground-state conformation through the lens of Wasserstein gradient flows, which leads to new and interpretable model architecture for molecular modeling. Extensive experiments consistently demonstrate that our method significantly outperforms state-of-the-art methods across various datasets, achieving promising molec- ular ground-state conformation prediction results. Mean- while, comprehensive ablation studies and analyses are also conducted to offer valuable insights into our method, thus validating its rationality and superiority.
Recently, some methods like ConfOpt (Guan et al., 2022) formulate molecular ground-state conformation prediction from the perspective of conformation optimization. They take low-quality conformations as input and apply various Transformer-basee models to predict the gradient field of the conformational energy landscape (Tsypin et al., 2024). Accordingly, the feedforward steps of their models corre- spond to optimizing the input conformation by gradient descent. However, the empirical architectures of the models prevent them from minimizing a physically meaningful en- ergy function, thus resulting in limited model interpretabil- ity and sub-optimal performance. Although some modified Transformer- -based models have been proposed from the per- spective of differential equations (Dutta et al., 2021; Sander et al., 2022), they are seldom applied to model structured data like molecular conformations and often ignore the re- quirement of eqqiivarrae (Li et al., 2025b).
2. Related Work
2. Molecular Conformation Generation
In the past few years, many molecular conformation genera- tion methods have been developed (Han et al., 2024; Wang et al., 2025). Early methods like CVGAE (Mansimov et al., 2019) apply variational autoencoders (Kingma Welling,
2

---

WGFormer: An Trans-Transtormer for Molecular Ground-Stale Conformation Prediction
2.2. Applications of Ground-State Conformation
vector (7(0) d. l|ē; - Gtyle is the Fuclidean distance between the i-th and j-th atom within the initial low-quality conformation, and (:: pr, or) is Gaussian kernel func- tion (Scholkopf et al., 1997) with learnable mean value I E "R." and standard deviation E RR. while they, E RR and Dayy, E ER" represent the learnable weight and bias for each atom type pair (0,, U2). Note that, different from Uni- Mol, we ensure that tu,ny = Mayn: and Dayny Meyms leading to symmetric relational representation.
Since molecular ground-state conformation represents the menggy-miimiize state on the potential energy surface, where the inter-atomic forces are balanced equilib- rium (Kim et al., 2025), it plays a crucial role in various downstream applications (Xu et al., 2023). For example, considering many physical, chemical, and biological prop- crties of a molecule are determined by its ground-state con- formation, these conformations are extensively utilized in molecular property prediction (Moon et al., 2023; Liu et al., 2024) and molecular docking (Chatterjee et al., 2023; Paggi et al., 2024). Furthermore, while 3D molecular pretraining has greatly enhanced the performance of various molecule- related tasks (Zhou et al., 2022; Wang et al., 2024), it has also dramatically increased the demand for high-quality 3D molecular structure data, underscoring the critical need to develop accurate and efficient methods for obtaining molec- ular ground-state conformations.
Then, we derive the final atom-level and relational represen- tations (X(4) R( by passing the initial atom-level and relational representations { (x((), R through L proposed WGFormer layers, i.e.,
x(1) x (L) R(4) = WGFormerz (X (x) R.),,
(2)
where WGFormer, refers to the architecture formed by stacking L WGFormer layers together.
Decoder. Given the initial relational representation R [r]]], J. the final relational representation R(()) [r([) and the initial low-quality conformation Č = [2,], we design an MLP-based decoder to obtain the predicted ground-state conformation C = [ci]. For i = I,... N, we have
3. Proposed Model
3.1. An Auto-Encoding Framework for Ground-State Conformation Prediction
We denote the molecular graph with atoms as G = (V, E), where V = {v:} } is the set of atoms and E {es N is the set of edges representing interatomic bonds. For the corresponding ground-state conformation, each atom in V is embedded by 3D coordinate vector e E RR and the ground- state conformation can be represented as C' = [c,] E IR N In addition, the interatomic distances are further denoted
MIP (L)
(0) 1j (a - 22))
2 Σ 1=1
(3)
N
where é E RR' is the 3D coordinate of the i-th atom in the predicted ground-state conformation and ē E R3 is the 3D coordinate of the i-th atom in the initial low-quality confor- mation. The proposed decoder predicts the residual between é and è; by the weighted average of the displacement vec- tors from è; to the other 3D coordinates (i.e., {č, 2, } The weights of the displacement vectors are determined by a simple multi-layer perception (MLP), which takes the residual between Tij bined with Eq. (1), we can easily prove that our whole auto- encoding framework is SEuii--qui,a and the detailed proof is provided in Appendix A.
as D = [d.y]] E RNN, where dij = ||c: eyll2 is the Euclidean distance between the i-th and j-th atom.
Following previous works (Yang et al., 2024; Feng al., 2024), we apply the cheminformatic tool RDKit (Landrum et al., 2013) to obtain initial low-quality conformation for the molecular graph G. denoted as C = [2:] E N Given the graph structure G and the initial conformation C, we aim to design and learn model to predict the ground- state conformation, i.e., C = (G, C), where Č = [cs] E Nx3 is the predicted ground-state conformation and go is the target model with learnable parameter 0. In this work, we tackle this task from the perspective of conformation optimization, designing go in an auto-encoding framework.
(L) and rij
(0) input. In addition, com-
The proposed WGFormer plays a central role in our model. Therefore, in the following content, we will introduce its architecture in Section 3.2 and further interpret it from the perspective of Wasserstein gradient flows in Section 4.
Encoder. Following the pipeline of Uni-Mol (Zhou et al., 2022), we first extract the initial level representation (0) (o E NxD and interatomic relational rep- resentation R0)) = [r] E NNNN from G and C, respectively. For i = 1.... ...,,N. N, we have
3.2. The Architecture of WGFormer
Inspired by the Sinkformer in (Sander et al., 2022), the proposed WGFormer implements attention maps based on the Sinkhor-sccaling algorithm (Cuturi, 2013). Consider the WGFormer in the (-th layer, which takes {x (1-),, R(-1)) as input and outputs ( (x(1) R('} accordingly. Specifically, it contains H attention heads, each of which corresponds to a channel of the tensor R(I-1), i.e., R(A,,A) e R NNN for
a (0) f(ve), rij (0) N Vide + Magiiit, o),
(1)
where f embeds each atom type o; to a D-dimensional
3

---

WGFormer: An for Molecular Ground-Star Conformation Prediction
Sinkhorn-based Attention Head
4. Revisiting WGFormer through the Lens of Wasserstein Gradient Flows
Q
((.h)
R
K(hhh)
4.1. Continuous Counterpart of Attention Head
((-1))
W
Sinkhorn
Let's consider WGFormer with single attention head (i.e., = omitted), whose parameter matrix is W E DDDx, Given a conformation with N atom-level representations, we intro- duce a mixture model of the atoms, i.c.,
= 1 and the superseripts I and h in Eq. (4) can be
v(AA)
W0T y(u
T
Figure 2. An illustration of the h-th head in the 1-th WGFormer.
J E M(R and μ = Σn1 M4.
N
h L..... H. As illustrated in Figure 2, the attention head is implemented by the following five steps:
where M(RD) is the space of the measures defined in RP. and 14 M(R is the measure associated with the i-th atom. For an arbitrary measure pair (p:, we describe the relation between μ4; and byy quantitatively as
1) ((t.k) = (((AA) = ((-1))
2) (thh) = WOTT))
3) v(ah) -X A (a,))
(4)
Quam Tayy
Ti = a WW 3 = +. Aæ
88)
4) R(M)) R(-IN,
VD.
where ã; = f æd (z) for 1, .... N. The collection of the relations leads to a symmetric matrix R [ra]]
5) ((lhh) = R. (R V((
where ( = [W Will E RDD is a learn- able matrix and the submatrix (((A)) E R DD Da corre- sponds to the 1-th attention head, with Da = D/H. Mean- while, (W (())(())) E TR. Dx means that we first com- pute ( ( denoted as (1) = (A.1.)) A (I. where (L.A) E Dx Du for 1, ..., H, and then preserv- ing the h-th submatrix A (1.A). Such operation integrates the functionality of feature-level fusion into the attention heads, where the value matrix of each head involves the model parameters shared across different heads. In addi- tion, H. (-) denotes the Sinkhorn-scaling algorithm (Cuturi, 2013), passing the attention matrix that involves relational information through an elementwise exponential operation and then normalizing its rows and columns alternatively in M steps, i.e., YR E NNN.
When each μ4 is a Dirac measure, i.e., /4; 8n,, we can
equivalently represent the relations by function:
r(z,, x²) = -E -
2ij==1 - re (2, x') V Vx, x' ER. (9)
In such a situation, the update of the atom-level represen- tations in Eq. (6) is the ResNet equation (He et al., 2016), i.e., X + X + T, which works as a discrete Euler scheme of the ordinary differential equation (ODE) (Weinan, 2017; Satorras al., 2021; Sander et al., 2022), i...
dz, T.M (a. Vi N.,.,N dl
(10)
Here, 2t) denotes the i-th atom's 's representation time t, and TyM (M()) is an operator for the measure = - 14
x M R No N, o ......... N. N₂ (exp(
(5)
M steps
ΣN &n,, which is defined as
where N,(R) R (RINx)) and (R) R s
N1NxN R) denote row and column normalization operations, and @ denotes element-wise division of matrix. Obviously, RD means the exponential function, i.e., ko (R) = exp( RR).
kM An''d()) T (2:) = JRD
(11)
where M (xx x') is the continuous counterpart of the Sinkhorn module R.M in Eq. (4), i.e.,
Finally, given (T((M), R( 1 WGFormer further con- catenates them together to derive (1) and R() as
(1) (1-1) + Coneat( (oneat((TTLii)) (1.))) 1)))
(x x²) = exp(x Ax' + +r(x, x²));
(6)
R ( = Concat( (( |
m even, (12)
km+1 km+1(x, 2')
I km )do(w")
k" (a.
where Concat(-) denotes the concatenation operation. Stacking L WGFormer modules together and passing (x(0) (0) ROO} through them, we derive the final representa- tions (L) R( as expressed in Eq. (2).
T T )dp(m))
m is odd.
Obviously, the collection of (TM (x:)) is equivalent to the matrix T derived by Eq. (4).

---

WGFormer: An SE(3)-Transformee for Molecular Ground-State Conformation Prediction
4.2. Wasserstein Gradient Flows When M 0 or + oo
determined by (XAXT + R). As shown in Eq. (17), this
term is equivalent to maximize the expected atomic distance, where the matrix Dx = 1 W x; W æ rúl E R g X captures the atomic distances in the latent space. In other words, minimizing E can prevent atoms from being concentrated together in the latent space. On the other hand, the second term in Eq. (17) regularizes the entropy of P. which helps build dense and coherent relations among the atoms in the latent space. In our opinion, minimizing such an energy function involving the above two terms is reason- able for conformation optimization for the following two reasons. Firstly, the atoms of the ground-state conformation should have reasonable interatomic distances, avoiding the high energy caused by too close distances. Secondly, the energy of conformation calculated based on dense atomic pairs rather than sparse chemical bond-based connections.
The ODE in Eq. (10) can be equivalently written as conti- nuity equation (Renardy & Rogers, 2006), i.e.,
0.p + div(u(Ta)) p = O.
(13)
has been known that TPd can be the Wasserstein gradient of an energy function E' :M(PP) at ps i.c.,
/ = -W- == =( ( 6 8p
(w)).
(14)
where 6E (1) denotes the first variation of E at μ4, which is 8μ
a function satisfying dE(u+ee)) c=- = S 6E sp for every öy
perturbation p (Santambrogio, 2017). In this situation, the continuity equation in Eq. (13) captures the evolution of the measure y that minimizes the energy function E, and each (e) follows dz - VwF(()) (Jordan et al., 1998).
In summary, by stacking I, WGFormer layers and setting a large value for the hyperparameter M, our model achieves the Wasserstein gradient flows defined on the measure of atoms. Given the initial measure y(()) corresponding to (0) and the initial relational information R! each layer updates the measure, i.e., p((+2) arg min, E (u) + W2(n, (((+)), where :2(+:)) denotes the 2-order Wasser- stein distance, and update the relational information R [may] based on the new measure (by Eq. (8)).
dt
Based on Proposition in (Sander al., 2022), we demon- strate that for μ Σ de: WGFormer corresponds to Wasserstein gradient flows when M = 0 or M -
Proposition 4.1. For μ p= mu=Ei Ei= 8z,, the T:M in Eq. (11) the Wasserstein gradient flow, i.e,, TM = - EM (n), when M- = or M - 00. The corresponding energy functions are
When each μ.; is a non-Dirac measure, our WGFormer can still achieve Wasserstein gradient flows when M 0. For each non-Dirac μ4, we can define an operator as
1 / 10 (4) - //2
(15)
T (x) Σ exp(x Ax' + rry) (xy Ar'd (18)
[log [logh" (o x' (a, (TT)
(y) = min 2 xEll,,
(16)
1 ʃ 0: (xx, x') log (0 x²) du d1(2))d)((2))
Then, we have the following theorem: Proposition 4.2. Let E : MMRD)N + R be an energy function for {us)e which is defined as
= 2
xo (x))
where 11 = f E M(R" x RD) Sa dι( x') = u(x'), Se. is the entropy of n.
u(x)}, and (T)) = log (T)dd
N Σ exp(x Ax' + (r())u((()dpp((a)) (19) 1-j=1
E= = 1
Note that, E (μ) is an interpretable energy function for conformation optimization. As shown Eq. (16), it cor- responds to an entropic optimal transport problem (Cuturi, 2013; Peyré et al., 2019). Because μ >i=1 Smc, we can rewrite E (μ) in the following discrete format:
For i = 1,..., N. To is the Wasserstein gradient flow that min,y E, i.e, To
pi Elys (Ughigs).
However, for non-Dirac 24's. WGFormer does not cor- respond to Wasserstein gradient flows in general when M + do because in this case the relational information Tij in Eq. (8) cannot be rewritten as the function in Eq. (9). The proofs of the above propositions are shown in Appendix B.
E. (n) min (XXX + R. (P, log P) Pell,
+ max Pell,
Dxx, P) - log log P)
(17)
Atomic distanse
Entropy
5. Learning Algorithm
where P E 111 means =ly=
= 1N. and P is the coupling controlling the pairwise relations among the atoms in the latent space. In particular, the first term in Eq. (17) penalizes the cost associated with the interatomic distances
Given the predicted ground-state conformation Ò, we can derive the corresponding interatomic distances D. Exist- ing learning- methods merely consider penalizing the
5

---

WGFormer: An Trans-rranffrmer for Molecular Ground-Stale Conformation Prediction
GIMGC
ContOpk Two/tto
Coulopt ThreAton
WGFOnI
Referro
Figure 3. The visual comparison for the ground-state conformations predicted by our WGFormer and the top-3 baselines.
mean-absolutteenror (MAE) between D and the ground truth D, which may not provide sufficient supervision and thus limits the model performance (Guan et al., 2022; Xu et al., 2023; Kim et al., 2025). In this work, we design a multi-task learning strategy to train the proposed model, making the model fit the ground-state conformation C and the corresponding interatomic distances D jointly. In par- ticular, given C and C, firstly align C to i.e.
loss supervising the predicted ground-state conformation. À 2 0 controls the significance of the auxiliary loss.
6. Experiments
demonstrate the effectiveness of our method, we com- pare with sate-of-the-and methods. Meanwhile, also conduct comprehensive ablation studies and analyses to val- idate its rationality and superiority. Experimental details and results are provided in this section and Appendix C.
C. C 1
N En C. Cu =
1 N
Σn
Ei (20)
1=
= -CoRootoCo, Co) = Co VO
6.1. Experimental Setup
where Co and C zero-mean conformations, and C* is the aligned ground truth. Rot(Co, Co) is the rotation matrix derived by Procrustes analysis (Wang & Mahadevan, 2008), in which U and V are obtained from the singular value decomposition (SVD) of Co Co.
previous work (Xu et al., 2023), employ the widely used Molecule3D and QM9 datasets to validate our method. In particular, Molecule3D (Xu et al., d) is a large-scale dataset, comprising approximately four million molccules, and provides two dataset-splitting re- sults (i.e., random random-splitting and scaffold -sifftdd-llllitigg) to ensure comprehensive evaluation. QM9 (Ramakrishnan al., 2014) is small-scale quantum chemistry dataset, compris- ing approximately 130,000 molecules with up to 9 heavy atoms. and adopts the identical dataset-splitting strategy de- scribed (Liao & Smidt, 2023). Fach molecule these datasets corresponds to the ground-state conformation deter- mined through density functional theory (DFT) calculations.
Datasets. Following
The final loss function is defined follows:
DIII +A + IC CI
C = N2 D
(21)
MAF(D.D)
RMSD(C..CC)
where MAE(D, D) denotes the mean-absolute-error be- tween D and the ground truth D, which is the main loss supervising the corresponding interatomic distances. RMSD((C, C*) denotes the root-mean mneean-square-deviiie between Co and the ground truth C*, which the auxiliary
Evaluation. To guarantee comprehensive and fair evalua- tion, employ the same metrics used in (Xu et al., 2023), including hean-absolute-error of distances (D-MAE), root-
6

---

WGFormer: An Trans--rr for Molecular Ground-Stale Conformation Prediction
Table The performance comparison for various methods. For each dataset, bold the hest result and underline the second-best one.
Validation
Test
Dataset
Method
Model
D-MAE
D-RMSE &
C-RMSD
D-MAE 4
D-RMSE
C-RMSD +
GINE
0.590
1.014
1.116
0.592
1.018
1.116
GATv2
0.563
0.983
1.082
0,564
0,986
1.083
2D Method
GPS
0.528
0.909
1.036
0.529
0.911
1.038
GTMGC
0.432
0.719
0.712
0.433
0.721
0.713
Molecule3D
SE(3)- Transffrrffrr
0,466
0.712
0.800
0,467
0.774
0.802
(random)
EGNN
0.461
0.704
0.798
0,462
0.766
0.799
3D Method
ContOpt-TwoAAtom
0.438
0.668
0.748
0.438
0.670
0.749
ConfOpt-ThreeAAtom
0.429
0.659
0.734
0.430
0.661
0.736
WGFormer (ours)
0.391
0.649
0.662
0.392
0.652
0.664
GINE
0.883
1.517
1.407
1.400
2.224
1.960
GATv2
0.778
1.385
1.254
1.238
2.069
1.752
2D Method
GPS
0.538
0.885
1.031
0.657
1.091
1.136
GTMGC
0.406
0.675
0.678
0.400
0.679
0.693
Molecule3D
0.460
0.676
0.775
0,456
0.678
0.747
(scaffold)
EGNN
0.448
0.666
0.758
0.442
0.670
0.741
Method
TwoAtom
0.408
0.626
0.708
0.402
0.628
0.698
ConfOpt-ThreeeAtom
0.401
0.619
0.697
0.395
0.622
0.691
WGFormer (ours)
0.363
0.599
0.618
0.360
0.610
0.627
GINE
0.357
0.673
0.685
0.357
0.669
0.693
GATv2
0.339
0.663
0.661
0.339
0.659
0.666
Method
GPS
0.326
0.644
0.662
0.326
0.640
0.666
GTMGC
0,262
0.468
0.362
0,264
0,470
0.367
Transffrraf
0.254
0.451
0.296
0.256
0.455
0.303
QM9
EGNN
0.248
0.442
0,257
0,251
0,449
0.265
3D Method
ConlOpt-TwoAtom
0.245
0.439
0.245
0,248
0.444
0.254
ConSOpt-ThreeeAtom
0.241
0.433
0.237
0,244
0.438
0.246
WGFormer (ours)
0.223
0.416
0.198
0.227
0.422
0.206
mean suan-squared-error of distances (D-RMSE), and root-mean- square-deviation of coordinates (C-RMSD), to evaluate the performance our method. The specific definitions of these evaluation metrics are expressed as follows:
2023) and its various variants that use GINE (Hu et al., 2020), GATv2 (Brody et al., 2021), and GPS (Rampášek et al., 2022) as backbones. The second category is the 3D methods predicting ground-state conformations via confor- mation optimization, including ConfOpt (Guan al., 2022) (TwoAtom and ThreeAtom versions) and its various vari- ants that use (E(3)-Transformer (Fuchs et al., 2020) and EGNN (Satorras et al., 2021) as backbones. All these base- lines are implemented based on their default settings.
D-MAE(D, DDD')
de digl.
(22)
N
D-RMSE(D,D")
N Σ (dig -dty,
(23)
Model architecture. Our model comprises 30 WGFormer layers, each equipped with 64 attention heads. The atom- level and relational representation dimensions set 512 and 64, respectively. We train the whole model on single RTX 3090 GPU and select the optimal checkpoint based on performance evaluated on the corresponding validation set.
C-RMSD(Č,C')
Σ cc c/11.
(24)
where Č Č (c) E R Nx3 is predicted ground-state confor- the ground-truth
mation men) C* [cğ] R x3 N3 aligned through Eq. (20), N' is the number of atoms in the molecule, and d,3 |é; م d' || </llz
6.2. Comparisons
Table presents the performance of various methods, show- ing that our WGFormer consistently outperforms all base- lines across all evaluation metrics on all datasets. Specifi- cally, compared to the best baseline, our WGFormer reduces C-RMSD (the most critical metric in this task) by nearly 10% on Molecule3D and 15% on QM9, demonstrating its
Baselines. To demonstrate the effectiveness of our method, state-of-the-ar methods in two typical categories are em- ployed as our baselines. The first category the 2D meth- ods predicting ground-state conformations merely based on 2D molecular graphs, including GTMGC (Xu et al.,

---

WGFormer: An for Molecular Ground-Stale Conformation Prediction
1.42
GTMGC
Table 2. The ablation study of multi-task learning strategy.
1.4
ConfOgt-TwoAtom
TonnOpt-ThreeAtom WGFormer (ours)
[MAE
LRMSD
D-MAE
D-RMSE + C-RMSD 4
Average Time (s/mal)
1.2
x
0.652
0.921
0.716
Validation
x
0.384
0.657
0.688
0.86
0.391
0.649
0.662
0.0
x
7
0.653
0.924
0.718
Test
v
0.386
0.660
0,689
0.6
0.392
0.652
0.664
0.4
0.35
0.2
0.17
Table 3.
The ablation study of our attention head.
0.0
Molecule3D (random)
Mo (scaffold)
QM9
Sinkhorn QKV
Module Adjust
D-MAE D-RMSE 4 C-RMSD +
x 0.379
0.628
0,661
x 0.378
0.631
0,656
Figure 4. The efficiency comparison for our WGFormer and the top-3 baselines on each dataset.
Validation
r
0.379
0.625
0.658
0.363
0.599
0.618
0.375
0.636
0.662
v
x
0.376
0.644
0.662
exceptional capability. Mcanwhile, our WGFormer also achieves consistent dominance across both random and scaf- fold splits of Molecule3D, demonstrating its effectiveness under different data-splittigg strategies. In addition, we have visualized some ground-state conformations predicted by our WGFormer and the top-3 baselines (GTMGC, ConfOpt- in Figure As shown
Test
x
v
0.376
0.635
0.665
v
0.360
0.610
0.627
over atomic 3D coordinates, proving the importance of em- phasizing optimizing interatomic distances during training. Moreover, compared with merely minimizing a single loss, minimizing them jointly achieves better performance, espe- cially in C-RMSD (the most critical metric), thus supporting the rationality of our multi-task learning strategy.
TwoAtom, and ConfOpt-ThhreeAle in this figure, the ground-stare state conformations predicted by our WGFormer align more closely with the ground truth, further demonstrating its superiority.
Figure 4 compares the efficiency of our WGFormer and the top-3 baselines by calculating the average time taken per molecule on each dataset. Even if our model applies 30 WGFormer layers, it is significantly faster than base- lines, reducing the runtime per molecule by over 50%. One important reason for this superior performance is the sim- plicity of our WGFomer architecture, which can update the atom-level and relational representations simultaneously. On the contrary, each layer ConfOpt applies two modules to update the atom-level and relational representations se- further considers the
Rationality of WGFormer architecture. In addition to our thcoretical analysis in Section 4, we also validate the ratio- nality of our model architecture through experiments. Here, we conduct ablation study for the proposed Sinkhorn- based attention head on scaffold-spli Molecule3D. As de- scribed in Section 3.2, our WGFormer applies the Sinkhorn- scaling algorithm with the proposed "QKV" adjustment. The results in Table show that i) each single modification does not harm performance, and éi) applying the two modi- fications jointly leads to significant performance improve- ments, especially in C-RMSDD (the most critical metric). These empirical findings align closely with theoreti- cal analysis in Section 4, further validating the rationality behind the proposed architecture design.
quentially, and ConfOpt-Threed relations within the triplets of atoms, significantly increasing the model complexity and computational cost. In summary, the above quantitative and qualitative comparisons high- light the superiority of our WGFormer in both accuracy and efficiency for ground-state state conformation prediction.
Correlative analysis of latent energy minimization and conformation optimization. As analyzed theoretically in Section 4, our WGFormer optimizes conformations by min- imizing an energy function (i.e., Eq. (17)) defined on the latent mixture models of atoms. To validate that minimizing this latent energy correlates with conformation optimization, we conduct a series of analytic experiments on QM9. In particular, given a trained WGFormer with 30 layers, we first randomly sample some molecules from the test set and pass them through the trained WGFormer.
6.3. Ablation Studies and Analyses
Rationality of multi-task learning, We conduct an abla- tion study of our multi-task learning strategy on random- split Molecule3D. As shown in Table 2, merely minimizing RM, C*) performs significantly worse than mini- mizing MAE(D,D) This phenomenon aligns with pre- vious works prioritizing optimizing interatomic distances
8

---

WGFormer: An (B(3))Traasformee for Molecular Ground-Stan State Conformation Prediction
Latent Energy
equiasea
Latent Energy
13900
Wifomme Ju-foeer---00 1-10)
WGF (1-301
TeYn0 [1-18
Energy
Potential Energy
-14500
14570
Latent Energy
Potential Energy
-1391022
D-MAE
0:40
D-RMSE
0.60
CRRMSD
20
14590
13020
0.30
0.50
0.30
244000
1.0
0.5 Time
1.0
0.5
-14620
Time
Time
25
10
20
Layers
Layers
Figure 6. The comparison for WGFormerio WGFormer and Trans-Trann on their optimization processes from = 0 to = 1. For each model and = I,... L, the point at f = corresponds to the model's performance achieved by I layers.
Figure 5. Some cases about the change of latent and potential en- ergy as the number of layers increases, where latent energy is obtained by solving Eq. (17) and potential energy is obtained by xTB (Bannwarth et al., 2019). More cases are provided in Figure
interval [0, 1], and accordingly, the step size is At = L.
Table 4. The mean energy change across layers.
We quantitatively analyze the impact of layer numbers and thus further verify the optimization nature of our model. In particular, we train two models with 10 and 30 WGFormer layers on QM9, respectively. Given each trained model, we evaluate its performance using different layer numbers on the test set, so that we can observe the optimization process achieved by each model to analyze the convergence rate and optimization efficiency. As illustrated in Figure 6, the performance of our models improves steadily as the number of layers increases. Meanwhile, the 30-layer WGFormer model performs better than the 10-layer WGFormer model, which matches the common sense of optimization: Learn- ing large number of layers means applying fine-grained optimization with more steps, which leads to lower energy and improved model performance while increasing runtime. Besides, these results further validate that our model indeed achieves conformation optimization.
Layers
10
15
20
25
Latent energy -3.629 -7.729
-8.512
-8.932
99.195 -00,885
Potential energy -9.135 118.1918.999111111
19.9555
-34.814
45.204--522 -52.378
Then, based on the outputs (i.e., X and R) of each layer, we can calculate the latent energy per layer by solving Eq. (17). As illustrated in Figure 5, the latent energy decreases gradu- ally as the number of layers increases, which validates that our WGFormer is indeed minimizing this latent energy.
Furthermore, considering that conformation optimization corresponds to minimizing the potential energy as the ground-state conformation is the mergy-minimiizd confor- mation, we employ the widely used xTB tool (Bannwarth et al., 2019) to calculate the corresponding potential energy of conformations obtained by different layers (i.c., passing the outputs of each layer through the trained decoder). As shown in Figure 5, the potential energy indeed decreases gradually the number of layers increases. Taking the latent and potential energy of the first layer as references, we calculate their relative change across increasing layers (the mean values are reported in Table 4) to analyze the cor- relation between the latent and potential energy. Here, the high Pearson correlation coefficient (0.885 + 0.033) indi- cates strong linear relationship, and the marginally higher distance correlation (0.906 + 0.018) suggests the presence of additional nonlinear dependencies. These results validate that minimizing this latent energy indeed drives potential energy minimization, thereby being highly correlated with conformation optimization.
In addition, as shown in Figure 6, when applying a non- WGFormer architecture (e.g.. an original removing our modifications) and keeping other settings un- changed, its performance may not be improved consistently when increasing the number of layers. This phenomenon further demonstrates the lack of interpretability in existing methods, thereby highlighting the value of our WGFormer.
7. Conclusion
In this work, we propose WGFormer, a Wasserstein gradient flow-driven SE(3)- Transformer that can effectively tackle the molecular ground-state conformation prediction task. In particular, our WGFormer can be interpreted as Wasserstein gradient flows, which optimizes molecular conformation by minimizing a physically reasonable energy function defined the latent mixture models of atoms, thereby significantly improving performance and interpretability. In the future, we will continue to explore the applications of our method in various conformation-relan study the mathematics of our WGFormer to design more powerful and interpretable model architectures.
Impact analysis of layer numbers on conformation opti- mization. As discussed earlier, our model corresponds to minimizing the latent energy function defined in Eq. (17), and each layer works as an update step in the Fuler scheme of ODE. Specifically, without the loss of generality, we as- sume that the optimization progress starts at t = O and ends tt = 1. In this case, learning our model with L WGFormer layers results in applying optimization steps in the time
tasks. Besides, we will further

---

WGFormer: An Trans-rrrne for Molecular Ground-State Conformation Prediction
Acknowledgements
national Conference on Neural Information Processing
Systems- Volume 2, pp. 2292-2300, 2013.
This work was supported by the National Natural Science Foundation of China (92270110), the Fundamental Research Funds for the Central Universities, the Research Funds of Renmin University of China, and the Public Computing Cloud, Renmin University of China. We also acknowledge the support provided by the fund for building world-class universities (disciplines) of Renmin University of China and by the funds from Beijing Key Laboratory of Research on Large Models and Intelligent Governance, Engineering Research Center of Next-Generation Intelligent Search and Recommendation, Ministry of Education, and from Intelli- gent Social Governance Interdisciplinary Platform, Major Innovation & Planning Interdisciplinary Platform for the "Double-First Class" Initiative, Renmin University of China.
Dutta, S., Gautam, T., Chakrabarti, S., and Chakraborty, T. Redesigning the transformer architecture with insights from multi multi-particle dynamical systems. Advances in Neural Information Processing Systems, 34:5531-5544, 2021.
Feng, S., Ni, Y., Li, M., Huang, Y., Ma, Z.-M., Ma, W.-Y., and Lan, Y. Unicorn: A unified contrastive learning ap- proach for multi-view molecular representation learning. In International Conference on Machine Learning, pp. 13256-13277. PMLR, 2024.
Fuchs, F., Worrall, D., Fischer, V., and Welling, M. Se (3)- tranffssfore:::: 3d roto-transkan trosssaanllaion equivariant attention networks. Advances in neural information processing systems, 33:1970-1981, 2020.
Impact Statement
Guan, J., Qian, W. W., Liu, Q.. Ma, W.-Y., Ma, J., and Peng, J. Energy-inspired molecular conformation optimization. In International Conference Learning Representations, 2022.
This paper presents work that aims advance the field of Machine Learning and AI for Science, especially for molec- ular modeling and ground-state conformation prediction. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Han, J., Cen, J., Wu, L., Li, Z., Kong, X., Jiao, R., Yu, Z., Xu, T., Wu, F., Wang, Z., et al. A survey of geomet- ric graph neural networks: Data structures, models and applications. arXiv preprint arXiv:2403.00488 2024.
References
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-777777777 2016.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6851, 2020.
Bannwarth, C., Ehlert, S., and Grimme, S. Gfn2-xtb-an accurate and broadly parametrized self-consistent tight- binding quantum chemical method with multipole elec- trostatics and -reeed-ddeenndent dispersion contributions. Journal of chemical theory and computation, 1663):1652- 1671, 2019.
Hollingsworth, A. and Dror, R. O. Molecular dynamics simulation for all. Neuron, 99(6): 121-1199,1 1143, 2018.
Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovee, J. Strategies for pre-training graph neu- ral networks. In International Conference on Learning Represenaations, 2020.
Brody, S., Alon, U., and Yahav, E. How attentive are graph attention networks? In International Conference on Learning Representations, 2021.
Cen, J., Li, A., Lin, N., Ren, Y., Wang, Z., and Huang, W. Are high-degree representations really unnecessary in cquivariant graph neural networks? Advances in Neural Information Processing Systems,
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 24240-24253, 2022.
2024.
Chatterjee, A., Walters, R., Shafi, Z., Ahmed, O. S., Sebek, M., Gysi, D., Yu, R., Eliassi-Rad, T., Barabási, A.L and Menichetti, G. Improving the generalizability of protein- ligand binding predictions with ai-bind. Nature communications, 14(1):1989, 2023.
Jordan, R., Kinderlehrer, D., and Otto, F. The variational formulation of the fokker-planck equation. SIAM journal on mathematical analysis, 29(1):1-11) 1998.
Kim, T., Sco, H., Ahn, S., and Yang, E. Rebind: Enhancing ground-state molecular conformation prediction via force- based graph rewiring. In The Thirteenth International Conference on Learning Representations, 2025.
Cuturi, M. Sinkhorn distances: lightspeed computation of optimal transport. In Proceedings of the 27th Inter-
10

---

WGFormer: An Trans-Trr for Molecular Ground-Stai Conformation Prediction
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference Learning Repre- sentations, 2014.
Renardy, M. and Rogers, R. C. An introduction partial differential equations, volume 13. Springer Science & Business Media, 2006.
Kobyzev, I., Prince, S. J., and Brubaker, M. A. Normalizing flows: An introduction and review of current methods. IEEE Iransactions on pattern analysis and machine intel- 2020.
Sander, M. E., Ablin, P., Blondel, M., and Peyré, G. Sink- formers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pp. 3515-3530 PMLR, 2022.
ligence, 43(11))3964-33)
Landrum, G. et al. Rdkit: A software suite for cheminformat- ics, computational chemistry, and predictive modeling. Greg Landrum, 8(31.10):528) 2013.
Santambrogio, {Euclidean, metric, and Wasserstein] gradient flows: an overview. Bulletin of Mathematical Sciences, 7:87-154, 2017.
Li, Z., Cen, J., Huang, W., Wang, T., and Song, L. Size- generalizable RNA structure evaluation by exploring hi- erarchical geometries. In The Thirteenth International Conference on Learning Representations, 2025a.
Satorras, V. G., Hoogeboom, E., and Welling, M. E(n) equiv- ariant graph neural networks. In International conference on machine learning, pp. 9323 9332.. PMI.,, 2021.
Zi, Cen, J., Su, B., Xu, T., Rong, Y., Zhao, D., and Huang. W. Large Language-geoment model: When LLM meets cquivariance. In Ferty- International Conference on Machine Learning, 2025b.
Scholkopf, B., Sung, K.-K., Burges, C. J., Girosi, F. Niyogi, P., Poggio, T., and Vapnik, V. Comparing support vector machines with gaussian kernels to radial basis function classifiers. IEEE transactions on Signal Processing, 45 11:2758-2765 1997.
I Liao, Y.-l.. and Smidt, T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In Interna- tional Conference on Learning Representations, 2023.
Shi, C., Luo, S., Xu, M., and Tang, J. Learning gradient fields for molecular conformation generation. In Interna- tional conference on machine learning, pp. 9558-9568. PMLR, 2021.
Liu, M., Li, C., Chen, R., Cao, D., and Zeng, X. Geometric deep learning for drug discovery. Expert Systems with Applications, 240:1 1 12498,, 2024.
Song, Y. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Mansimov, E., Mahmood, O., Kang. S., and Cho, K. Molec- ular geometry prediction using deep generative graph neural network. Scientific reports, 91):20381, 2019.
Moon, K., Im, H.-J., and Kwon, S. 3d graph contrastive learning for molecular property prediction. Bioinformat- ics, btaa33,,,37,, 2023.
Tsypin, A., Ugadiarov, L. A., Khrabrov, K., Telepov, A.. Rumiantsev, E., Skrynnik, A., Panov, A., Vetrov, D. P., Tutubalina, E., and Kadurin, A. Gradual optimization learning for conformational energy minimization. The Twelfth International Conference on Learning Represen- tations, 2024.
Paggi, J. M.. Pandit, A., and Dror, R. O. The art and science of molecular docking. Annual Review of Biochemistry, 93, 2024.
Peyré, G., Cuturi, M., et al. Computational optimal trans- port: With applications to data science. Foundations and Trends@ in Machine Learning, 1 116566):355-077 2019.
Wang, C. and Mahadevan, S. Manifold alignment using pro- crustes analysis. In Proceedings of the 25th international conference on Machine learning, pp. 1120-1127, 2008
Pracht, P., Bohle, F., and Grimme, S. Automated explo- ration of the low-energy chemical space with fast quan- tum chemical methods. Physical Chemistry Chemical Physics, (114)):719-71111 2020.
Wang, F., Xu, H., Chen, X., Lu, S., Deng, Y., and Huang, W. Mperformer: An se (3) transformer- molecular per- ceptron. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pp. 2512-2522, 2023a.
Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Wang, F,, Guo, W., Cheng, M., Yuan, S., Xu, H., and Gao, Z. Mmpolymer: A multimodal multitask pretraining frame- work for polymer property prediction. In Proceedings of the 33rd ACM International Conference on Information and Knowledge Management, pp. 2336-2346, 2024.
Rampäšck, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D. Recipe for a general, powerful, seal- able graph transformer. Advances in Neural Information Processing Systems, 35:14501 -14511111555,, 2022.

---

WGFormer: An Trans-rrrne for Molccular Ground-Stafe Conformation Prediction
Wang, F., Guo, W., Ou, Q. Wang, H., Lin, H., Xu, H., and Gao, Z. Polyconf: Unlocking polymer conforma- tion generation through hierarchical generative models. In Forty-second International Conference on Machine Learning, 2025.
Wang, H., Fu, T., Du, Y., Gao, W. Huang, K., Liu, Z Chandak, P., Liu, S., Van Katwyk, P., Deac, A., et al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972)447-60, 2023b.
Weinan, W. proposal machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Xu, G.. Jiang. Y., Lei, P.. Yang. Y., and Chen, Gtmge: Us- ing graph transformer predict molecule's s ground-state conformation. In The Twelfth International Conference on Learning Representations, 2023.
Xu, M., Luo, S., Bengio, Y., Peng, J., and Tang. Learning neural generative dynamics for molecular conformation generation. In International Conference Learning Representations, 2021a.
Xu, M., Wang, W., Luo, S., Shi, C., Bengio, Y., Gomez- Bombarelli, R., and Tang, An end-to-end framework for molecular conformation generation via bilevel program- ming. In International conference machine learning, pp. 11537-11547. PMLR, 2021b.
Xu, M., Yu, L.. Song. Y., Shi, C., Ermon, S., and Tang, J. Geodiff: geometric diffusion model for molecular conformation generation. In International Conference Learning Representationss 2021c
Xu, Z., Luo, Y., Zhang, X., Xu, X., Xie, Y., Liu, M., Dicker- son, K., Deng, C., Nakata, M., and Ji, S. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint arXiv:21100017117 2021d.
Yang. J., Zheng, K., Long, S., Nie, Z., Zhang, M., Dai, X., Ma, W.Y. and Zhou, H. Mol-ae: Auto-encoder based molecular representation learning with 3d cloze test objective. In International Conference on Machine Learning, pp. 56793-56811. PMLR, 2024.
Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877-28888 2021.
Zhou, G., Gao, Z., Ding, Q. Zheng, H., Xu, H., Wei. Z., Zhang, L.,, and Ke, G. Uni-mol: A universal 3d molecu- lar representation learning framework. In The Eleventh International Conference on Learning Representations, 2022.
12

---

WGFormer: An Trans-rrrne for Molecular Ground-Star Statee Conformation Prediction
A. Squivvaque
To prove our whole auto-encoding framework -eqiivquivariint,, we will verify that the input low-quie low-quality conformation Č [c] E Rx* undergo an transorrantoonnttioo (Q. t), where Q I is rotation matrix and IR is translation,
the predicted ground-state conformation é = [&]] E R x3 will also transform in the same manner. Specifically, let input
[ Q²: + t.. we aim to demonstrate that the corresponding output satisfies ع [c.] [Qa + t..
In the encoding phase, the initial atom-level representation x(0)
[ [ RND can be extracted as follows:
x (0)
f(::);
(25)
where /' is an embedding function dependent only on the atom type Since atom type is independent of the geometric coordinates, the initial atom-level representation (0)
[=:]]] remains unchanged under the transforoanoo
Meanwhile, the initial interatomic relational representation R [rip]] RNNNH can be extracted as follows:
N(d4 Uvay, + Venny μl M,
(26)
where dij |2: 2gg l| the interatomic distance, which remains unchanged under the SE(3)-rrantooomaati (Q, t), i.e.,
d ج ć 2l IIQQ; + +t) -Q2 + t)ll2 Q[ QBIl1 l|ē: og 12 = d.j.
(27)
[0)] The initial interatomic relational representation (2(0) [rlg also remains unchanged under the SE(3))transformaiio
Furthermore, since the initial atom-level representation X (0) [ (z) and interatomic relational representation R! [r(]]
remains unchanged under the (553)))tansformmtion the final atom-level representation X [x= and interatomic
( will also remain unchanged under the SE(3)-transformmttin. In other words, X (0) [x:0] () = [r(]) [0) ), x (L) = [x] and R = =r are all SE(3)-invariant.
relational representation R L)
[r
In the decoding phase, the predicted ground-state conformation غ = [2:] can be computed as follows:
(0) é ē Σn NN MLP(r (L) - (c 22)
(28)
N
Under the dransfoattooooottoonn (Q, t), the input coordinates o and č are transformed as:
C Qč: +t. č Qč, +t,
(29)
The relative displacements in the decoding equation are then given by:
c d = (Qa: + t) (Qé, t) Q(2: 22);
(30)
As the MLP operates on SE(3)-invariant inputs (r-() - rij (0)). ), the updated predicted coordinates become:
& = Σ, ML (L) r ( 2))
(31)
N
Expanding é and ě in terms of Q and t, we have:
& = (Qè; +t)+ Σn N MIP((.
(0) Q(( -
rij
(32)
Since the rotation matrix Q is linear transformation, factors out from the weighted sum:
& = - Σ
Mr (1)
rig (0) (ā= - -.))
+ t,
(33)
N
13

---

WGFormer: An Trans-Tran for Molecular Ground-Stafe State Conformation Prediction
The term in parentheses is simply the original é; i.e.,
C Qè: + t,,
(34)
Therefore, our whole auto-encoding framework is -quivaaau.n
B. Proofs of Key Theoretical Results
In principle, the proofs of our theoretical results are based on the techniques used in (Sander et al., 2022). However, because of considering the interatomic relations and their updates, our WGFormer corresponds Wasserstein gradient flows under stricter condition, and the energy function it minimizes considers the relational information of molecular conformation.
B.1. Proof of Proposition 4.1
Proof. Recall that μ Eim Pi, Mi E MR for i 1,,., N. In addition, impose the standard sufficient regularity assumption on each [4:- When each component measure a Dirac measure, i.e., μt Oz;, we have = Ei=1 6n.- such a situation, we rewrite the energy function in Eq. (15) as
1 (+) - - /JSoxxe exp(# Aæ' r(æ, x r')) dp(z))d((z)) 1 N Σ 15=1
exp(ac] Ax, + re)).
(35)
Accordingly, its first variation at j is
8E (1)( = f exp(x Ax' + r((a,a'))du(()) = Σ
N
exp(x Ar, + r(x,
(36)
8p
j=1
In this case, the operator T (z) becomes
(z) = / exp(x Ax' + +r) Ar'du(z')) = Σ
exp(x Ax, + 2xxx,2)))Az, j=1
(37)
_80
Then, we have
N V exp(x Ax, + +rx
N
Vw V (p)(2) =
Σ
exp(x T Ax, + +r(x, = (x).
(38)
j=1
In general, given μ = Ei=l pis we define an energy function { with M + as
N
E ((u)) = inf (E
Σ
((xxx))(e Ax' + +n)) + +( x') log ( x (z)dd(())))
if=1
(39)
inf =EIT(L.))
Σ RR XRD
((x, Ax' (n) H(),,
ij=1
where H(T) - -/goxx x <(2, x' (...=ddd(()))) )du(z)dp(z')) is the entropy of the coupling æ. Essentially. Eq, (39) corresponds to an entropic optimal transport (EOT) problem defined for the mixture model μ.
Furthermore, when μ = - 8R,, we can equivalently rewrite as a function of μ and in a classic EOT format, i.c.,
1 min rEll(U)))
E (n)
(x:
Aw' TAx'
r(x, x))) (a of (d(())))(((a)) H H())
cost: log (ax,rr))
1
1: (x, x') log (=0 da(z)dp(x²)
k0 (xxxx')
(40)
2
1
max f(x) (z))((((z) SEC(R")
2
14

---

WGFormer: An Sea3)-Transformer for Molecular Ground Shouud Conformation Prediction
Here, the second equality in Eq. (40) means applying the Sinkhorn-scaling algorithm solve the EOT problem (with M 00), where & is the optimal coupling. The third equality in Eq. (40) based on the duality of the EOT problem (Peyré et al., 2019) and f is the soft c-transform defined as
(a log / exp(log k2 (xr, + (()))))(a)))
(41)
In such a situation, we can follow the proof in (Sander et al., 2022): It has been known that when f is optimized, we have fr and (x, x²) = exp(log 1. (xx + f(x:) f(x)). As a result, we have
Vw (VE Vxf(x)
exp(f(x:) f(x') + log (xx a²) log (, a aa)dd((a")
RD
JRD
he (x. V log ko (x, x' e')dp(z))
(42)
1:00 (x, (x) du(x')
RD
T (x).
B.2. Proof of Proposition 4.2
Proof. Given non-Dirac measures {14);-) we can fix each component measure except μ4, and minimize the energy function in Eq. (19), i.e., ming, Flu;; (1z)szs). According to the work (Santambrogio, 2017), 7), the first variation of E(ps; (15 (z4))))) at μ4 is
8E (2))
exp(x: Ax' + ra) + exp(x" Ax + rx)dd((()))
8p,
x
exp(æ Ax' + +ri)) exp(x" Ax rri)dp (x')
(43)
N
exp(x T Ax' + (r))d(((a')
where the second equality in Eq. (43) because both A and R [rig] are symmetric matrices. By differentiation under the integral, we have
(
SE (1:) (xr)
V E(u (4gbss)(a))
5p4s
/ V exp(x Ax' + rx))d
Σ
1=1
(44)
/
Σ
exp(x= T Ax' + ris)) Ax' Am)
du(('))
J=1
-7 Tu (z).
15

---

WGFormer: An Trans--ran for Molecular Ground-Stale Conformation Prediction
Table The scale overview of Molccule3D and QM9 datasets.
Dataset
Splitring
Training
Validation
Test
randon-splitinn
2,339,745
779,918
779,915
Molecule3D
scaffold-splitingg
2,339,724
779,929
779,928
QM9
default
108,831
9,900
10,697
B.3. The Details of Eq. (17)
Note that, X E RN W E DDD, and A = WWT E RDDD Denote Y = XW @ XW, where @ represents the Hadamard product of matrix. The details of Eq. (17) is shown below:
P
arg max (Dx. P) (P, log P) PElli
arg min Pelli
Dx P) + (P, log P)
(R 2 (YIDxN + INNDY 2XWW X xT,, + (P, log P)
1
arg min
Pelli
2
(45)
1 - (XAX) + R. R. log P) (YIDNN, P)
= arg min Pelli
constant due the doubly stochastic constraint
arg min (XAX + R. P) + (P, P) Pelli
where (YIDXN, P) = tr DxN PT) = (YYDDxx))))) because P E 1T
C. More Experimental Details and Results
C.1. Datasets
As mentioned in the main body, also employ the widely used Molecule3D and QM9 datasets in work remain consistent with previous works, thus evaluating the performance of our method fairly, The introductions of the above two datasets are provided below:
Molecule3D. The first benchmark introduced by (Xu et al., 2021d) for the molecular ground-state conformation prediction task. It comprises approximately 4 million molecules, each with the corresponding ground-state conformation determined through density functional theory (DFT) calculations. Additionally, it provides two dataset-splitinng results (i.e., -andttis and ecaffold-splitinng to ensure a comprehensive evaluation.
QM9. A small-scale quantum chemistry dataset proposed in (Ramakrishnan et al., 2014), and the work in (Xu et al., 2023) further applies it to the molecular ground-state state conformation prediction task. It comprises approximately 130,000 molecules with up to 9 heavy atoms, each with the corresponding ground-state conformation determined through density functional theory (DFT) calculations. Besides, adopts the dataset-splitting strategy in (Liao Smidt, 2023).
Table 5 provides a scale overview of Molecule3D and QM9 datasets, outlining their division into training, validation, and test sets. For the Molecule3D dataset, both random-splitinn and scaffold-splitiing strategies are applied, resulting in about 2.3 million molecules for training and approximately 0.78 million for both validation and test sets. Meanwhile, the QM9 dataset utilizes the default splitting strategy described (Liao & Smidt, 2023), consisting of 108,831 molecules in the training set, 9,900 in the validation set, and 10,697 in the test set.
In addition, Figure 7 further visualizes the distribution of atom count (including the hydrogen atom) across these datasets. It demonstrates that molecules in the Molecule3D dataset generally contain more atoms than those in the QM9 dataset, thereby increasing the difficulty predicting their corresponding ground-state molecular conformations. This observation aligns with the actual performance presented Table 1, thus explaining the decreased performance all methods on the Molecule3I dataset relative to the QM9 dataset.
16

---

WGFormer: An Trans--ran for Molecular Ground-State Conformation Prediction
Training Validation
Training Validation Test
0.15
Training
Validation Test
0.14
0.12
Density
Density
Density
0.10
0.02
0.00
0.01
0.01
0.02
0.00
20
40
60
00
100
120
140
0.00
20
40
50
90
200
120
140
0.00
10
15
20
25
20
Atoms per Molecule
Atoms per Molecule
Atoms per Molecule
(a) Molecule3D )plittiin)
(b) Molecule3D (scaffold-splin -spitting))
(c) QM9
Figure 7. The atom count distribution of Molecule3D and QM9 datasets.
C.2. Evaluation
As mentioned in the main body, to guarantee a comprehensive and fair evaluation, we choose the same metrics introduced by (Xu et al., 2023), including mean-absolltee-rroo of distances (D-MAE), waammnn--uuaee----eee of distances (D-RMSE), and coot-mean-squale cere--suurreeeeeee -enuaredeviation of coordinates (C-RMSD), to evaluate the performance of various models. Here, the specific definitions of these evaluation metrics can be expressed as follows:
N
N
D-MAR(D,D") D*)
(dg digl=
(46)
N
N
N (dij - -od-d dzy)
D-RMSE(D,D")
(47)
N2
N
C-RMSD(Č.C*)
1
Σ 1c: - ci113
(48)
N
where Ō = [&:] RNxx is predicted ground-state conformation (zero-mean), C² = [c]] Nx is the ground-truth aligned
through Eq. (20), N is the number of atoms in the molecule, and dig = llè, - ب dij = ||| - cg E1l22
C.3. Baselines
As mentioned in the main body, we compare our method with 2D methods that predict ground-state conformations merely based on 2D molecular graphs and 3D methods that predict ground state conformations from the conformation optimization perspective to validate the effectiveness of our method. The introductions of typical baselines are provided below:
GTMGC:
The typical 2D method proposed in (Xu et al., 2023), which introduces novel network based on Graph- Transformer (GT) to predict the molecular ground-state conformation from the corresponding 2D molecular graph in an end-to-end manner. In addition, it further designs various 2D basclines by replacing the backbone module with other models like GINE (Hu et al., 2020), GATy2 (Brody et al., 2021), and GPS (Rampášek et al., 2022).
The typical 3D method proposed in (Guan et al., 2022), which introduces variants of -euiiaequva neural networks to predict the gradient field of the conformational energy landscape, and then optimize the input conformation by gradient descent. It can be divided into two variants: ConfOpt-TwoAttom which considers only the relationships between atom pairs, and ConfOpt-Three Atoe,,,,,, which additionally accounts for the relationships among atom triplets. In addition, it further designs various 3D baselines by replacing the backbone module with other models like Transforrae (Fuchs et al., 2020) and EGNN (Satorras et al., 2021).
ConfOpt:
17

---

WGFormer: An
for Molecular Ground-State Conformation Prediction
0.40
0.50
D-MAE
0.25
D-RMSP
0.45
C-RMSD
0.20
0.20
0.35
0.22
3.54
3.06
583
502
5.04
0.06
0.30
0.22
0.54
0.66
Figure 8. The robustness analysis WGFormer data noise on QM9.
Table 6. The specific hyperparameters Molecule3D and QM9 datasets.
Molecule3D
Hyperparameter
QM9
random-splitting
scaffold-splitinng
Batchsize
24
32
32
Learning Rate
5e-6
le-5
2e-5
Auxiliary Loss Weight (A)
0.125
0.125
0.125
Epoch
500
500
500
Adam Eps
le-6
1e-6
le-6
Adam Betas
(0.9, 0.99)
(0.9, 0.99)
(0.9, 0.99)
Num Encoder Layer (L)
30
30
30
Relational Representation Dimension (H)
64
64
64
Atom-level Representation Dimension (D)
512
512
512
Attention Head Dimension (Da)
8
Table 7, The performance of WGFormer with repeated layers on QM9.
Model
D-MAE
D-RMSE +
C-RMSD
GTMGC
0.264
0.470
0.367
ConfOpt
0.244
0.438
0.246
WGFormer (30 different layers)
0.227
0.422
0.206
WGFormer (6 repeated layers)
0.231
0.422
0.219
C.4. Hyperparameters
As mentioned in the main body, model comprises 30 encoder layers (i.e., the number of WGFormer modules), each equipped with 64 attention heads. Meanwhile, the atom-level and relational representation dimensions are set to 512 and 64, respectively. In addition, we train the whole model on a single NVIDIA RTX 3090 GPU and select the optimal checkpoint based on performance evaluated on the corresponding validation set. Here, the specific hyperparameters used across all datasets are summarized in Table 6. Please note that the batch size on the Molecule3D dataset (i.e., 24 for random-splitting and for caffood--plittingg purely decided by GPU memory limitation our work.
C.5. More Experimental Results
Robustness analysis. Given the trained model on QM9, we add Gaussian noise to the input conformations during the inference phase to evaluate our model's robustness. As shown in Figure 8, our model achieves stable performance across various noise levels, demonstrating its robustness data noise. The robustness of our model ensures reliable molecular ground-state conformation prediction even in the presence of perturbations.
Performance with repeated layers. In addition to training our model with 30 different WGFormer layers (i.e., the one in Table 1), we also train another model on QM9, which contains only 6 different WGFormer layers and repeats each layer 5 times. In this case, the repetition of each layer can be interpreted as optimizing the same energy functional with Euler steps. As shown Table 7, this model using repeated layers leads to comparable performance, thus further validating and supporting our theoretical analysis in Section
18

---

WGFormer: An Tran)-Trantf for Molecular Ground-Stale Conformation Prediction
Table 8. The median absolute errors of various property metrics QM9.
Model
E (kcal/mol)
(debye)
Ae (kcal/mol)
GTMGC
0.008
0.014
0.068
ConfOpt
0.006
0.012
0.069
WGFormer (ours)
0.004
0.009
0.053
Table The performance and efficiency comparisons with generative models on QM9.
Model
D-MAE
D-RMSE +
C-RMSD +
Inference Time (s/mol)
GeoDiff
0.278
0.563
0.473
18.058
TorsionDiff
0.378
0.685
0.437
14.359
WGFormer (ours)
0.227
0.422
0.206
0.150
Table 10. The performance and efficiency comparisons on DRUGS.
Model
D-MAE
D-RMSE
C-RMSD +
Inference Time (s/mol)
GTMGC
0.914
1.420
1.657
1.878
ConfOpt
0.825
1.272
1.650
48.135
WGFormer (ours)
0.816
1.281
1.602
0.778
Chemical property metrics. Following the pipeline established in (Jing et al., 2022), we further compare the chemical properties of the predicted and reference ground-state conformations to evaluate the performance of our WGFormer comprehensively. In particular, given each trained model on QM9, we employ the xTB tool (Bannwarth et al., 2019) to compute three popular chemical properties (energy E, dipole moment μ., HOMO-LUMO gap Ae) and calculate the median absolute errors on the test set. As shown in Table 8, our WGFormer significantly outperforms GTMGC (the best 2D baseline) and ConfOpt (the best 3D baseline) across all property metrics, further validating its superiority.
Comparisons with generative models. While existing generative models originally designed to generate multiple conformations, we preserve the one with the lowest energy as the predicted ground-state conformation for comparisons. Here, we train two typical generative models (i.e., GeoDiff (Xu et al., 2021c) and TorsionDiff (Jing et al., 2022)) on QM9 and compare them with our WGFormer on the test set. As shown in Table 9, our WGFormer significantly outperforms the above two generative models across all evaluation metrics, achieving lower prediction errors and higher inference speed. This remarkable result highlights the practicality and superiority of our WGFormer for the molecular ground-state conformation prediction task, where both accuracy and efficiency are critical.
Comparisons the DRUGS dataset. Since each molecule in the original DRUGS dataset has multiple conformations, and we don't t know which one is the ground-state conformation (or even whether the ground-state conformation is among them), this dataset is not well suited for the molecular ground-state conformation prediction task. To make this dataset applicable and further train various models, we select the conformation with the highest Boltzmann weights as the ground- state conformation and adopt the identical dataset-splittnng strategy described (Xu al., 2021c). As shown in Table 10, compared with GTMGC (the best 2D baseline) and ConfOpt (the best 3D baseline), our WGFormer significantly outperforms them, achieving state-of-the-ar performance with the highest inference speed. This demonstrates the generalizability of WGFormer across various datasets, thus enabling it to meet the requirements of real-worke applications effectively.
C.6. More Cases about Energy Change
In Figure 9, we have provided more cases about the change of latent and potential energy as the number of layers increases, where latent energy is obtained by solving Eq. (17) and potential energy is obtained by the xTB tool (Bannwarth et al., 2019). As shown in this figure, both latent and potential energy decrease gradually as the number of layers increases, and their change exhibit remarkably high correlation. These results validate that minimizing this latent energy can drive potential energy minimization, thereby being highly correlated with conformation optimization.
19

---

WGFormer: An Trans-rrnnerrmmer for Molecular Ground-Stafe State Conformation Prediction
-5000
22-
Latent Energy
Potential Energy (kcal/mol)
Latent Energy
-15630
Potential Energy (kcal/mol)
20
Potential Energy
15010
Potential Energy
15640
21
19
Latent Energy
15020
Latent Energy
20
18
15650
15030
19
17
15660
15040
16
18
15050
115770
15
15060
15680
14
16
10
15
10
15 Layers
20
25
30
Layers
15910
Latent Energy
Potential Energy /kcal/mol
Latent Energy Potential Energy
Potential Energy (kcal/mol)
21
Potential Energy
-17340
22
15920
Latent Energy
17350
Latent Energy
20-
15930
19-
17360
159900
18-
-17370
17
15950
16
17380
16
15960
15-
17390
14
5
10
15 Layers
20
25
30
5
10
15 Layers
20
25
30
22
22
Latent Energy
Potential Energy MKcaamm (/lllmoo
Latent Energy
15340
Potential Energy (kcal/mol
Potential Energy
15770
21
Potential Energy
20-
Latent Energy
15775
Latent Energy
20
15350
19
18-
15780
15360
18
16
11788
115370
17
15790
16
15380
10795
15795
15
15390
12
10
15
30
10
15
20
25
Layers
Layers
Figure 9, More cases about the change of latent and potential energy as the number of layers increases.
20


Reviews:
[{"review_id": "rUFvb3oabQ", "reviewer": "Reviewer_dW7c", "summary": "this work proposes \"Wasserstein gradient flow-driven\" transformer, to gradually refine a initial 3D conformation to its ground state.\nthis refinement process is associated with minimizing an energy function, thus enhancing its interpretability, and probably explains its performance improvement.\nthe method is validated with extensive experiments with several ablation studies.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your appreciation of our work. Below, we resolve your concerns one by one.\n\n**Q1: Visualize the latent energy function values achieved through different numbers of layers.**\n\n**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.  \n\nTo further verify our claim, **given a trained WGFormer with 30 layers**, we have conducted a series of analytic experiments in the inference phase.\n\nFirstly, 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.\n\n**Q2: Verify that the latent energy is closely related to the physical energy and final metrics.**\n\n**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:\n\n| The number of layers                  | 5       | 10      | 15      | 20      | 25      | 30      |  \n|--------------------------------------|---------|---------|---------|---------|---------|---------|  \n| Relative Energy Value Change (kcal/mol) | -9.135  | -18.199 | -19.955 | -34.814 | -45.204 | -52.378 |  \n| Relative Latent Energy Value Change     | -3.629  | -7.729  | -8.512  | -8.932  | -9.195  | -10.385 |  \n| Pearson Correlation Coefficient (Energy VS Latent Energy) | 0.885 ± 0.033 |  \n| Distance Correlation (Energy VS Latent Energy)           | 0.906 ± 0.018 |\n\nHere, 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.** \n\nFinally, **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.**\n\nIn 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. \n\nThank 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. \n\n\nReference:\n\n[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.", "claims_and_evidence": "yes for most claims.\n\nstill wondering about one core claim, i.e., \"minimization of the energy function\". several comments:\n* could the author please visualize this process, e.g., plot the energy function value vs the layers, from input to output for illustration\n* any way to verify this \"energy function defined on the latent mixture models of atoms\" is closely relevant to the final metrics?", "methods_and_evaluation_criteria": "looks good to me", "theoretical_claims": "I don't have enough expertise", "experimental_designs_or_analyses": "looks good to me", "relation_to_broader_scientific_literature": "the relation with energy-based methods seem to be interesting and important to me.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "see above", "other_comments_or_suggestions": "see above", "questions_for_authors": "see above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ixuGEkZOBf", "reviewer": "Reviewer_NScR", "summary": "The paper introduces WGFormer, a novel model that combines the strengths of energy-based simulation and learning-based methods for predicting molecular ground-state conformations. WGFormer is built upon an SE(3)-Transformer framework and is driven by Wasserstein gradient flows. In an auto-encoding setup, the model encodes low-quality 3D conformations (e.g., generated via RDKit) and decodes them into ground-state conformations using a lightweight MLP. A key innovation is the customized attention mechanism based on the Sinkhorn-scaling algorithm with adjusted QKV matrices and the omission of the feed-forward network, which together reduce computational cost and promote cross-head feature fusion. The theoretical framework establishes that, under certain conditions, WGFormer operates as a discretized version of Wasserstein gradient flows that minimize a physically meaningful energy function defined on a latent mixture model of atoms. Extensive experiments on datasets such as Molecule3D and QM9 show that WGFormer not only outperforms current state-of-the-art baselines in both accuracy (e.g., lower C-RMSD values) and efficiency but also demonstrates robustness through comprehensive ablation studies.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Below, we resolve your concerns one by one.\n\n**Q1: Comparisons with other generative models**\n\n**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.\n\nUsing 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. \n\n|Method|D-MAE|D-RMSE|C-RMSD|Inference Time (s/mol)|  \n|-|-|-|-|-|  \n|GeoDiff|0.278|0.563|0.473|18.06|  \n|TorsionDiff|0.378|0.685|0.437|14.36|  \n|WGFormer|**0.227**|**0.422**|**0.206**|**0.15**|\n\n**Q2: The differences with other methods**\n\n**A2:** The differences between our WGFormer and the methods you mentioned can be summarized as follows:\n\n**1) Interpretable architecture:**\n\nAs 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.\n\n**2) Focus on ground-state conformation:**\n\nHigh-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.\n\n**Q3: Does Wasserstein gradient need to compute gradient? The comparison on computational consumption is required.**\n\n**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.\n\nSecondly, **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). \n\nThirdly, 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**. \n\n**Q4: Comparison with flow-matching based methods?**\n\n**A4:** Our work, i.e., developing a Wasserstein gradient flow-based model, is different from the flow-matching learning strategy:\n\n**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. \n\n**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.\n\n**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.\n\nHope the above responses can resolve your concerns and help re-evaluate our work.\n\n\n[1] REBIND: Enhancing Ground-state Molecular Conformation Prediction via Force-Based Graph Rewiring, ICLR 2025.\n\n[2] Geodiff: A geometric diffusion model for molecular conformation generation, ICLR 2022.\n\n[3] Torsional diffusion for molecular conformer generation, NeurIPS 2022.\n\n[4] {Euclidean, metric, and Wasserstein} gradient flows: an overview. Bulletin of Mathematical Sciences, 2017.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "- The comparision between WGFormer and other generative model (e.g., GeoDiff) is needed.\n\n- There are many works optimize the low-quality molecule comformation to generate higher-quality molecules, such as AlphaFold3. It utilize the RDKit generated / CCD structures as reference features to predict the ligand structures. What are the difference between WGFormer and these methods?\n\n- It seems that Wasserstein gradient needs the computation of gradient. What is the consumption of WGFormer? Adding the comparison with other generative model is preferred.\n\n- What is the advantage of Wasserstein gradient flow compared with flow-matching based molecule generative model?", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This method makes the first attempt to predict molecular ground-state conformation through the lens of Wasserstein gradient flow.", "essential_references_not_discussed": "[1] Abramson, Josh, et al. \"Accurate structure prediction of biomolecular interactions with AlphaFold 3.\" Nature 630.8016 (2024): 493-500.", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8F5XypW6Lm", "reviewer": "Reviewer_FN8S", "summary": "This paper proposes a transformer architecture for optimizing molecular geometries. The network first processes node features and edge features (encoding the input geometry) via residual update blocks with a bespoke attention mechanism called WGFormer. An interpretation of the WGFormer is provided as a Wasserstein gradient flow in latent space. These features are then used to update the input 3D geometry using relative position vectors. The model is shown to outperform existing methods on this task in terms of MAE and RMSD metrics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive feedback and constructive comments. Below, we resolve your concerns one by one.\n\n**W1: Test on larger datasets.**\n\n**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.\n\nNevertheless, 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.\n\n|Method|D-MAE|D-RMSE|C-RMSD|Inference Time (s/mol)|  \n|-|-|-|-|-|  \n|GTMGC|0.914|1.420|1.657|1.88|  \n|ConfOpt|0.825|**1.272**|1.650|48.14|  \n|WGFormer|**0.816**|1.281|**1.602**|**0.78**|\n\n**W2: The novelty and significance of proposed model architecture**\n\n**A2:** Although attention-based models have been widely used in molecule and protein generation, our WGFormer is still valuable because of the following reasons:\n\nFirstly, 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. \n\nNo 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. \n\n**W3:  1) Can finite Sinkhorn iterations be interpreted as Wasserstein gradient flow? 2) The energy functional depends on $k$ and changes for every layer.**\n\n**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. \n\n**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). \n\n**W4&W5:  Interpretations of architecture improvement and Figure 5**\n\n**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: \n\n1) WGFormer indeed minimizes a latent energy function for atoms' probability measure; \n\n2) Minimizing this latent energy helps optimize the physical energy of conformation, highly correlated with the final metrics. \n\nFor 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).\n\n**W6: Chemical property metrics**\n\n**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. \n\n|Method|Energy|Dipole Moment|HOMO-LUMO Gap|\n|-|-|-|-|\n|GTMGC|0.008| 0.014|0.068|\n|ConfOpt|0.006|0.012|0.069|\n|WGFormer|**0.004**|**0.009**|**0.053**|\n\nHope the above responses help enhance your confidence to further support our work.\n\n[1] Torsional diffusion for molecular conformer generation, NeurIPS 2022.\n\n[2] Sinkformer: Transformers with doubly stochastic attention, AISTATS 2022.", "claims_and_evidence": "The main claim of improved model performance is supported. The results on QM9 are to my knowledge SOTA, but it would be nice to see results on larger and more diverse datasets, e.g., DRUGS.\n\nThe claim that the WGFormer module, and its interpretation as Wasserstein gradient flow, is responsible for this improved performance is significantly less convincing.\n*    The attention based module with explicitly updated dense pair representations bears some similarity to protein structure prediction networks; it is perhaps not so surprising that such an architecture would by itself already outperform existing methods for predicting small molecule structures. Indeed, the paper's ablation results show that even without the Sinkhorn module, the model is already SOTA. \n*    Although the Sinkhorn module helps, the interpretation as Wasserstein gradient flow is unconvincing. First of all, the interpretation only holds in the limit of infinitesimal updates, when residual blocks are reduced to a neural ODE. Second of all, the energy functional depends on $k$ and is therefore changing for every layer.\n*    The analysis in Figure 5 merely shows that adding layers to the network helps performance, which is independent of the Wasserstein gradient flow interpretation. It would be more convincing if the authors could show that the behavior of non-WGFormers are different.\n*    I should stress that an interpretation of the architecture improvement is not necessary to be a strong paper, but if the authors choose to feature it prominently, then it should be strongly supported.", "methods_and_evaluation_criteria": "Yes, the benchmarks are standard. However, larger molecules (DRUGS) have been used in previous works, but are missing here. Also, it would be helpful to show chemical property metrics that are of interest to those computing properties from minimized conformers.", "theoretical_claims": "I did not carefully check the proofs for theoretical claims.", "experimental_designs_or_analyses": "Please see comments in \"Claims and Evidence.\"", "relation_to_broader_scientific_literature": "This paper contributes to the ML literature on learning small molecule geometry optimizers. Although a coherent body of work at ML venues, the broader impact in actual scientific software is not yet clear. In this respect, this submission probably does move the needle much as the improvements are on the same metrics and are incremental (i.e., it is not clear where the threshold for wider applicability is, so it is not clear if it has been reached).\n\nThe paper seems to draw heavy inspiration from Sinkformer, with which I am less familiar. From the point of view of other areas of AI4Science though, the proposed architecture, with its dense pair features used to bias attention, resembles architectures commonly used for proteins. It could be quite interesting to explore further architectural developments along these lines, even if they do not admit clean theoretical interpretations.", "essential_references_not_discussed": "I am not aware of any essential references not discussed.", "other_strengths_and_weaknesses": "No additional comments", "other_comments_or_suggestions": "Despite the concerns about the interpretation of WGFormer, I lean towards accept because the performance improvements are sound and lean in the right direction in terms of architectural exploration.", "questions_for_authors": "No additional questions", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5vYP89Y5Gj", "reviewer": "Reviewer_vjqn", "summary": "This paper proposes a transformer-based neural network architecture to predict the ground state conformation of molecules from some low quality conformation (equivalent to optimization). The encoder part of the model takes as input per-atom embeddings (depending only on atom types) and atom-pair embeddings (distances after Gaussian kernel expansion with some learned weight), and applies the attention mechanism for multiple layers to generate the final representation for each atom. Finally, the decoder predicts deviation in pair-wise distances to update the coordinates of the input conformation.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "For your comments, below we answer them one by one.\n\n**W1: No discussion about efficiency** \n\n**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. \n\nSecondly, 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%. \n\n**W2&W5: Lack interpretability and the relationship with the conformational energy minimization process**\n\n**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). \n\nThese experiments demonstrate that: \n\n1) Our WGFormer is indeed minimizing the energy function defined on the latent mixture models of atoms; \n\n2) Minimizing this energy function indeed helps optimize the physically-meaningful energy of molecular conformation, and these two kinds of energy are highly correlated; \n\n3) Minimizing this energy function is indeed closely relevant to improving the final metrics. \n\n**W3: Clarification of relevant statements**\n\n**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. \n\nThe DSMBind you mentioned works to predict protein-ligand binding energy, which is unrelated to our work in either problem setting or technical route. \n\n**W4: Unclear motivation**\n \n**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.\n\n**W6: A lot of same work and lack contribution**\n\n**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).\n\n**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. \n\n**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.\n\n**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.\n\nIn 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. \n\n[1] Learning gradient fields for molecular conformation generation, ICML 2021.\n\n[2] Learning neural generative dynamics for molecular conformation generation. arXiv preprint 2021.\n\n[3] Energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 2022.", "claims_and_evidence": "- The first claim of this work is efficiency. \"*Classic energy-based simulation is time-consuming when solving this problem while existing learning-based methods have advantages in computational efficiency but sacrifice accuracy and interpretability. In this work, we propose a novel and effective method to bridge the energy-based simulation and the learning-based strategy.*\" This paper identifies WGFormer as a method that is faster than energy-based methods and more accurate than other learning-based methods. However, only accuracy is demonstrated and no discussion about efficiency and especially comparison to quantum-based methods has been provided.\n- Another claim of this work is interpretability. This paper describes the forward pass of the encoder as a process of minimization of some energy function, and suddenly it comes to a conclusion that the latent energy function of this process \"*is an interpretable energy function for conformation optimization*\". This makes no sense as this paper didn't reveal at all how exactly the gradient flow process relates to the physical process of conformation optimization. It seems to me that the conclusion of interpretability is drawn only because both the W.G.flow and the conformation optimization are some kind of energy minimization processes.\n- Further, in the Related Work section, the paper states that \"*the empirical architectures of the models prevent them from minimizing a physically meaningful energy function, resulting in limited model interpretability and sub-optimal performance*\" which is not true. For example, DSMBind [1], an application study of a conformation optimization model, shows that the model's output correlates with physical quantities (binding affinity), which gives interpretability.\n\n[1] Jin, Wengong, et al. \"Unsupervised protein-ligand binding energy prediction via neural euler's rotation equation.\" Advances in Neural Information Processing Systems 36 (2023): 33514-33528.", "methods_and_evaluation_criteria": "There is no major issue with the design of the network. The motivation of getting inspiration from Sinkformer and the three improvement techniques is unclear, though supported by ablation studies.", "theoretical_claims": "The main theoretical claim is the relationship between the WG flow and the conformational energy minimization process, which is not grounded.", "experimental_designs_or_analyses": "There is no major issue with experimental designs.", "relation_to_broader_scientific_literature": "The work focuses on small molecule conformation optimization. There have been a lot of papers in this area. This paper is not very special and does not make unique contribution.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel variant of SE(3) Transformer for the molecule 3D conformation transformation task (from a low-quality conformation to the ground-state conformation). The major design insight is a neural network architecture where each layer can be interpreted as one simulation step of the Wasserstein gradient flow (WGF) of a learnable functional (called "energy") of probability measure on an atom embedding space. The authors make an analogy to the physical energy function minimization process to find the ground-state conformation, in hope to enhance interpretability. The resulting model surpasses existing methods on common conformation prediction benchmarks.

All reviewers and I myself appreciate the significant novelty of the idea, which would inspire a new direction for designing architectures for molecular science. The results also seem promising.

The major concern is on the connection of the WGF of a learned energy functional to a physical energy minimization process. During the rebuttal, the authors visualized the correlation of the learned energy functional with the physical energy function (requested by Reviewer NScR) and additional results that have shared parameter (representing the same energy functional) over several layers (requested by Reviewer FN8S). These results provide more evidence on the WGF interpretation.

Nevertheless, I still sense a gap between the WGF interpretation and physical energy minimization. Conceptually, WGF minimizes a functional of a probability measure, while the physical energy is a function of the molecular conformation. The presented connection works by taking a conformation as a mixture of delta measures of atom embeddings. As at a delta measure, the optimal transport map may not exist, I would wonder if singularities would arise, e.g., if $\\frac{\\delta E}{\\delta \\mu}$ in Eq. (14) is not differentiable. What seems more natural to me following the energy function minimization perspective would be a ResNet whose residue update in each layer is the gradient of an energy function. I hope the authors could discuss the connection to this alternative and perform experiment comparison if necessary.
Decision: Accept (poster)