Bridging Geometric States via Geometric Diffusion Bridge
Authors: Shengjie Luo
Yixian Xu
Di He
Shuxin Zheng
Tie-Yan Liu
Liwei Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: Bridging Geometric States
Generative Modeling
Geometric Deep Learning
Diffusion Bridge
Doob's h-transform
Equivariance
Primary area: generative_models

Abstract: The accurate prediction of geometric state evolution in complex systems is critical for advancing scientific domains such as quantum chemistry and material modeling. Traditional experimental and computational methods face challenges in terms of environmental constraints and computational demands, while current deep learning approaches still fall short in terms of precision and generality. In this work, we introduce the Geometric Diffusion Bridge (GDB), a novel generative modeling framework that accurately bridges initial and target geometric states. GDB leverages a probabilistic approach to evolve geometric state distributions, employing an equivariant diffusion bridge derived by a modified version of Doob's $h$-transform for connecting geometric states. This tailored diffusion process is anchored by initial and target geometric states as fixed endpoints and governed by equivariant transition kernels. Moreover, trajectory data can be seamlessly leveraged in our GDB framework by using a chain of equivariant diffusion bridges, providing a more detailed and accurate characterization of evolution dynamics. Theoretically, we conduct a thorough examination to confirm our framework's ability to preserve joint distributions of geometric states and capability to completely model the underlying dynamics inducing trajectory distributions with negligible error. Experimental evaluations across various real-world scenarios show that GDB surpasses existing state-of-the-art approaches, opening up a new pathway for accurately bridging geometric states and tackling crucial scientific challenges with improved accuracy and applicability.

Bridging Geometric States via Geometric Diffusion Bridge
Shengjie Luo'; Yixian Xul:4> Di He Shuxin Zheng', Tie-Yan Liu', Liwei Wang 1.3t
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology. Peking University
2 Microsoft Research A14Science *Center for Data Science, Peking University Pazhou Laboratory (Huangpu), Guangzhou, Guangdong 510555, China edu. en, xyx050Retu. pku.edu. edu.cn,
lussj@stu, pku.edu. {shuz, tyliu) croooft.ooooooooomoooo com, {dihe, wanglw)@pku. edu.
Abstract
The accurate prediction of geometric state evolution in complex systems is critical for advancing scientific domains such as quantum chemistry and material modeling. Traditional experimental and computational methods face challenges in terms environmental constraints and computational demands, while current deep learning approaches still fall short in terms of precision and generality. this work, we troduce the Geometric Diffusion Bridge (GDB), novel generative modeling frame- work that accurately bridges initial and target geometric states. GDB leverages a probabilistic approach to evolve geometric state distributions, employing an equiv- ariant di diffusion bridge derived by a modified version of Doob's 's h-trrasffrmm for con- necting geometric states. This tailored diffusion process anchored by initial and target geometric states as fixed endpoints and governed by equivariant transition ker- nels. Moreover, trajectory data can be seamlessly leveraged in our GDB framework by using a chain of equivariant diffusion bridges, providing a more detailed and conduct thor-
accurate characterization of evolution dynamics. Theoretically, ough examination to confirm our framework's ' ability to preserve joint distributions of geometric states and capability to completely model the underlying dynamics inducing trajectory distributions with negligible error. Experimental evaluations across various real-world scenarios show that GDB surpasses existing state-of-the- art approaches, opening new pathway for accurately bridging geometric states and tackling crucial scientific challenges with improved accuracy and applicability.
1 Introduction
Predicting the evolution of the geometric state of system is essential across various scientific domains [46, 88, 55, 17, 20, 101], offering valuable insights into difficult tasks such as drug discov- ery [25, 29], reaction modeling [9, 24), and catalyst analysis [13, 105]. Despite its critical importance, accurately predicting future geometric states of interest is challenging. Experimental approaches often face obstacles due to strict environmental requirements and physical limits of instruments [102, 3, 69]. Computational approaches seek to solve the problem by simulating the dynamics based on underlying equations [81, 88]. Though providing greater flexibility, such calculations are typically driven by first pritcirinciple methods or empirical laws, either requiring extensive computational costs [68] or sacrificing accuracy [40].
"Equal contribution. Cooresponndecce to: He<dihedpku.eds edu, cn>, Liwei Wang <wanglu@pku, edd. edu, en>,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In recent years, deep learning has emerged as a pivotal tool in scientific discovery for many fields [43, 23, 69, 107], offering new avenues for tackling this problem. One line of approach aims to train models to predict target geometric states (e.g., equilibrium states) from initial states directly and develop neu- ral network architectures that respect inherent symmetries of geometric states, such as the equivariance of rotation and translation (104, 31, 8, 87, 89, 103]. However, this paradigm requires encoding the iter- ative evolution into a single-step prediction model, which lacks the ability to fully capture the system's s underlying dynamics and potentially leading to reduced accuracy. Another line of research trains machine learning force fields (MLFFs) to simulate the trajectory of geometric states over time [32, 34, 6, 70, 5,5,, showing better affliiency-accuran balance [15, 13, 105, 84). Nevertheless, MLFFs are typically trained to predict intermediate labels, such as the force of the (local) current state. During inference, states are iteratively updated step by step. Since small local errors can accumulate, reliable predictions over long trajectories highly depend on the quality of intermediate labels, which cannot be guaranteed [7, 106, 30). Therefore, an ideal solution that can precisely bridge initial and target geometric states and effectively leverage trajectory data (if available) as guidance is in great demand.
In this work, we introduce Geometric Diffussin Bridge (GDB), general framework for bridging geometric states through generative modeling. From a probabilistic perspective, predicting target geometric states from initial states requires modeling the joint state distribution across different time steps. The diffusion models [37, 99] are standard choices to achieve this goal. However, these methods ideally generate data by denoising samples drawn from Gaussian prior distribution, which makes it challenging to bridge pre-given geometric states or leverage trajectories in a unified manner. To address the issue, we establish a novel equivariant diffusion bridge by developing a modified version of Doob's 's h--ransform [82, 81, 16]. The proposed stochastic differential equation (SDE) is anchored by initial and target geometric states simultaneously model the joint state distribution and is governed by equivariant transition kernels to satisfy symmetry constraints, Intriguingly, we further demonstrate that this framework can seamlessly leverage trajectory data to improve prediction. With available trajectory data, we can construct chains of equivariant diffusion bridges, each modeling one segment in the trajectory. The segments are interconnected by properly setting the boundary conditions, allowing complete modeling of trajectory data. For model training, we derive scalable and simulation-fis matching objective similar to [59, 61, 77], which requires no computational overhead when trajectory data is leveraged.
Overall, our GDB framework offers a unified solution that precisely bridges geometric states by modeling the joint state distribution and comprehensively leverages available trajectories as fine- grained depiction of dynamics for enhanced performance. Mathematically, we prove that the joint distribution of geometric states across different time steps can be completely preserved by our (chains of) equivariant diffusion bridge technique, confirming its expressiveness in bridging geometric states and underscoring the necessity of design choices in our framework. Furthermore, under mild and practical assumptions. we prove that our framework can approximate the underlying dynamics governing the evolution of geometric state trajectories with negligible error in convergence, remarking on the completeness and usefulness of our framework in different scenarios. These advantages show the superiority of our framework over existing approaches.
Practically, we provide comprehensive guidance for implementing our GDB framework in real- world applications. To verify its effectiveness and generality, we conduct extensive experiments covering diverse data modalities (simple molecules & adsorbate-catalyst complex), scales (small, medium and large scales) and scenarios (with & without trajectory guidance). Numerical results show that our GDB framework consistently outperforms existing state-of-the-art machine learning approaches by large margin. In particular, our method even surpasses strong MLFF baselines that are trained on 10x more data in the challenging structure relaxation task of OC22 [105], and trajectory guidance can further enhance our performance. The significantly superior performance demonstrates the high capacity of our framework capture the complex evolution dynamics of geometric states and determine valuable and crucial geometric states of interest in critical real-world challenges.
2 Background
2.1 Problem Definition
Our task of interest is to capture the evolution of geometric states, i.e., predicting future states from initial states. Formally, let S denote system consisting of set of objects located in the

---

paree-dimensiona Euclidean space. We use H E Raxd to denote the objects with features, where the number of objects, and is the feature dimension. For object let r, E R denote its Cartesian coordinate. We define the system as S = (H, R), where R :...... F.). This data structure ubiquitously corresponds to various real-world systems such as molecules and proteins [17, 20, 101]. In practice, the geometric state is governed by physical laws and evolves over time, and we denote the geometric state at a given time t as R² = (ri...., }. Given a system S = (H, 1(²)) ) at time to. our goal is to predict 5th = (H, Rt) at future time i As example, in molecular system, R can be the equilibrium state of interest evolved from the initial state Rt
In this problem, inherent symmetries in geometric states should be considered. For example, rotation that is applied to the coordinate system at time to should also be applied subsequent time steps. These symmetries are related to the concept of equivariance in group theory [19, 18, 91]. Formally, let @ : X -> y' denote a function mapping between two spaces. Given group G. let px and x² denote its group representations, which describe how the group elements act on these spaces. A function @ x [g)] o (p* [g)[ Vg G, E X. When x² = T (identity transformation), it is also known as invariance. SE(3) group, which pertains to translations (T(3)) and rotations (3)(3) in 3D Euclidean space, is one of the most widely used groups and is employed in our framework.
>> y is said to be equivariant if it satisfies the following condition:
2.2 Diffusion Models
Diffusion models [95, 37, 99] have emerged as the state-of-the-art generative modeling approaches across various domains [83, 85, 47, 115, 113, 117]. The main idea of this method is to construct a diffusion process that maps data to noise, and train models to reverse such process by using a tractable objective.
Formally, to model the data distribution Idata (X), where X Rd we construct a diffusion process (Xt) which is represented as a sequence of random variables indexed by time steps. We set Xo ~ Idama (X) and X Pptior (X), where Porior (X) has a tractable form to generate samples efficiently, c.g. standard Gaussian distribution. Mathematically, we model (X))eefo.T) as the solution to the following stochastic differential equation (SDE):
dx, f(X,, + (t)
(1)
where f(.,)): IR [0, R is a vector-valued function called the drift coefficient, (.): [0, 7] +
R is scalar function known the diffusion coefficient, and (Be) is the standard Wiener process (a.ka.,, Brownian motion) [26]. We hereafter denote by Pa (X) the marginal distribution of X, Let p(x', ft(xx, t) denote the transition density function such that P(X,' E A = x) = ((p(', f'',x, t)dx' for any Borel set A. By simulating this diffusion process forward in time, the distribution of X, will become Ppesee (X) at the final time T. In the literature, there exist various design choices of the SDE formulation in Eqn. (1) such that it transports the data distribution into the fixed prior distribution [98, 37, 99, 72, 97, 47].
In order to sample X,- ~ (X) = Idata ()), an intriguing fact can be leveraged: the reverse of a diffusion process is also a diffusion process [2]. This reverse process runs backward in time and can be formulated by the following time-reversal SDE:
dx,- [f(x.,1) - x (e.P)(.)) di + (I)dB,
(2)
where Vx log (X) denote the score of the marginal distribution at time the score known for all time, then we can derive the reverse diffusion process from Eqn. (2), sample from Potior PPt) and simulate this process to generate samples from the data distribution Gdata (X). In particular, the score Vx log Pt (X) can be estimated by training a parameterized model so (X, t) with denoising score matching objective [98, 97]. theory, the minimizer of this objective approximates the ground-truth score [99] and this objective is tractable.
Geometric Diffusion Bridge
As discussed in the introduction, effectively capturing the evolution of geometric states crucial, for which three desiderata should be carefully considered:
3

---

Coupling Preservation: From probabilistic perspective, the evolution of geometric states transports their distribution from ( to (dma((*)) and we are interested in modeling the distribution of target geometric states given the initial states, i.e., Qdata ( (St))) ==
Gdata (R' H ) which can be achieved by preserving the coupling of geometric states, i.e., Idata (R R [H). For brevity, we hereafter omit the condition of H because it keeps the same along the evolution and can be easily incorporated into the models.
Symmetry Constraints: Since the law governing the evolution unchanged regardless of how the system is rotated translated, the distribution of the geometric states should satisfy symmetry constraints, i.e,, quam (p= (g) R (g)
Idma ( and
qdata (p (g) [ 1. (g)"))))("")) (R))) Qdata RR RR1) for all g SE(3), R R.
Trajectory Guidance: Trajectories of geometric states are sometimes accessible and provide fine meeegrained descriptions of the evolution dynamics. For completeness, it is crucial to develop unified framework that can characterize and leverage trajectory data guidance for better bridging geometric states and capturing the evolution.
However, existing approaches typically have their limitations for this task, which we thoroughly discuss in Sec. and summarize into Table 1. In this section, we introduce Geometric Diffusion Bridge (GDB), general framework for bridging geometric states through generative modeling. We will elaborate on key techniques for completely preserving couping under symmetry constraints (Sec. 3.1), and demonstrate how framework can be seamlessly extended to leverage trajectory data (Sec. 3.2). Theoretically, we conduct a thorough analysis on the capability of our unified framework, showing its completeness and superiority. All proofs of theorems are presented in Appendix B. A detailed guidance of practical implementing our framework is further provided (Sec. 3.3).
Table I: Comparisons of different candidates for bridging geometric states
Methods
Symmetry Constraints Coupling Preservation Trajectory guidance
Direct Prediction [104, 31, 87, 89,
MI.FFs [90, 33, 6, 34, 58]
x
Geometric Diffusion Model (115, 38, 114]
Geometric Diffusion Bridge (ours)
3.1 Equivariant Diffusion Bridge
Our key design lies in the construction of equivariant diffusion bridge, a tailored diffusion process (RE)EEI for bridging initial states Rd [R*) and target states R RT "qdaal (R*1 R² completely preserving coupling of geometric states and satisfying symmetry constraints. Firstly, we investigate necessary conditions for a diffusion process on geometric states to meet the symmetric constraints:
Proposition 3.1. Let R denote the space of geometric states and fπ (-..) R [0, T + R denote the drift coefficient on R. Let (W") "eeoo.r denote the Wiener process on R. Given an SDE on geometric
states dR' fR(R! t)dt + (t) R q(R"), its transition density Pie (z', f/|2, t), 2, E R 1, t). Vg E SE(3). < i.e., pr (R("", ('[ .t) = PR (9) (x(1((((())))))) RR],
t, S T, if these conditions satisfied: (1) (RR)) is equivariant and T(3)- -iv-aiiiriinn;; (3) the transition density of (W" )<<(o.r] is -333i-eqquvariann.
(2) fR(:,t) is SO(3)-
Using Proposition 3.1, we can obtain a diffusion process that respect symmetry constraints by properly considering conditions for key components. Next, we modify a useful tool in probability theory called Doob's 's h-transform [82, 81, 16], which plays an essential role in the construction of our equivariant diffusion bridge for preserving coupling of geometric states:
Proposition 3.2. Let (2/,,'122z, be the transition density of the SDE in Proposition 3.1. Let hR( : R x [o, T] > R50 be a smooth function satisfying: (1) h t) (i33)-inaaiaa,; (2) lRR (2, PR ('|z, 1)/r (z', l')dz'. Then we can derive the following lRR -transffrrmedd SDE
on geometric states:
dR' = [R .0) R log ho (R' dt (l)dW'
(3)
12.t) 1) equals to pr(2', 1112, ha(='d)) with -E333)-quuvariaa transition density (2', (/

---

Proposition 3.2 provides an equivariant version of Doob's h--rransform, which can be used to guide free SDE on geometric states hit an event almost surely. For example, if we set (:.t) (2(2, T| 1) E RR i.e., the transition density of the original SDE evaluated at R z, then the -rg-transformed SDE in Eqn. (3) arrives at the specific geometric state almost surely at the final time (see Proposition B.7 in the appendix for more details). Therefore, if we derive a proper (::)) function under the symmetry constraints, our target process (R)) )e([0,2] can be constructed:
Theorem 3.3 (Equivariant Diffusion Bridge). Let dR fR(RR' (R', (Rt,RR)))t + (t) be an SDE on geometric states with transition density PR (2', ,,'2,,,,,, 2, E R. satisfying the conditions in P Proposition 3. Let luR (z, 1; z0) = f pR can derive the following hR bn-trannformed SDE:
(= dz'. By using Proposition 3.2, we (00(===T//) T(,, (R("""T
dR' = [RR(R', + o²( qx) (RR .) R logpR (R", ,T(RR,))R" .O)R", (IR,,RY ,R']] (l)
which corresponds to a process (R²) (R))ee0..r) R ~ Qdma (Poo satisfying the the following properties:
let ((-.-) R x + R20 denote the joint distribution induced by (R' )ee(,.T) then ((R,RR) equals to Idaro R;
its transition density qR (R', t' (''R ,R (p)))(((((()))) (o) (2)) qdoooo ~qgdo ()
VOSt,t' geSE(e) "
We call the tailored diffusion process (R") (ejo, an equivariant diffusion bridge.
According to Theorem 3.3, given an initial geometric state R we can predict target geometric states R by simulating the equivariant diffusion bridge (R' (R from R R which arrives at RT Qdau (R (Rt)) However, the score Eg (R" ( ..RR' .0) |V log Po T R) R in Eqn. (4) is not tractable in general. Inspired by the score matching objective in diffusion mod- els [99], we use parameterized model (R', t; Rd) to estimate the score by using the following
training objective:
[(0)= (20.21) (aam))))TanRRRo ).Rt (RS ((t))lvo(R,, 20) VR log PR (21, T'R t)]
(5)
where 1 (00 T) (the uniform distribution [0, T), and A(-) [0, R2o is positive weighting function. Theoretically, we prove that the minimizer of Eqn. (5) approximates the ground- truth score (see Appendix B.5 for more details). Moreover, this objective is tractable because the transition density pR and qR can be designed to have simple and explicit forms such as Gaussian, which we will elaborate on in Sec. 3.3.
3.2 Chain of Equivariant Diffusion Bridges for Leveraging Trajectory Guidance
In this subsection, we elaborate on how to leverage trajectories of geometric states as a fine-grained guidance in our framework. Let (R') (RE))) denote trajectory of N + geometric states and
(aa ..., ..., RR) denote the joint probability density function of geometric states in a trajectory.
In practice, the markov property of trajectories typically holds [109, 78]. Under this assumption, ( ..., RiN) can be equivalently reformulated into (f(") I. ( by the chain rule of probability. If Gagg [R') R can be well modeled, we capture the distribution of trajectories of geometric states completely.
According to Theorem 3.3, given R ~ (f.), an equivariant diffusion bridge (R')te )eee), j0,7j can be constructed to model the joint distribution Graj ( Ae-) and hence /ing (ft (ft is preserved. Therefore, if we construct a series of interconnected equivariant diffusion bridges, the distribution of trajectories can be modeled:
Theorem 3.4 (Chain of Equivariant Diffusion Bridges). Let
denote a series of N equivaraint diffuusion bridges defined in Theorem 3.3. For the i-th bridge (Rt) if we set
g 1 (2'[=o)
(I) hige (z, t; 20) I P (2', T t) then the joint distribution qR (R., R Rf ..... ,Rk_) induced by (RR) tejo, hejN- equals to
pr (2', dz'; (2) Rd ~ Ling( (RO), R. RT A.L... VO < < N,
quay (Ro ..... RN). We call this process a chain of equivariant diffusion bridges.
5

---

In this way, a chain of equivariant diffusion bridge can be used to model prior trajectory data, and simulating this chain not only bridges initial and target geometric states but also yields intermediate evolving states. Similarly, we can also use a parameterized model to estimate the scores of bridges in this chain. Instead of having only one objective in all time steps, we now have N bridges in total, which categorize the time span into N groups with different time-dependent objectives. Therefore, by properly specifying time steps and initial conditions, the objective in Eqn. (5) can be seamlessly extended (see Appendix B.7 for more details its provable guarantee):
(0) EE (B00..)EN))) ( Mad R
... RN) R. (()) 1; 1;2z)) VR log l (# TIR EII,; (6)
where L((, x - -x R afe (R) 2' 24, 0).
Lastly, we provide the following theoretical result, which further characterizes our framework's expressiveness to completely model the underlying dynamics that induce the trajectory distributions:
Theorem 3.5. Assume (ft' DEEN) is sampled by simulating a prior SDE on geometric states di''
~ (R')dt + +dW.. Let μo denote the path measure of this prior SDE when t E [iT, (i + 1)77.
Building upon (R")ee(n let {uR hi denote the path measure of our chain of equivariant diffusion bridges. Under mild assumptions, we have lim max (L( 0. N-+00
It is noteworthy that the assumption of the prior SDE existence holds in various real-world applications. For example, in geometry optimization, we can formulate the iterative updating process of a molecular system as dR = aa n- V (RR")dd dt + 3dW', where V (RR') denotes the potential energy at R' and a; B are step sizes [88]. From Theorem 3.5, such prior SDE serves as the underlying law governing the evolution dynamic, and our chain of equivariant diffusion bridges constructed from empirical trajectory data can well approximate it, showing the completeness of our framework.
3.3 Practical Implementation
In this subsection, we elaborate on how to practically implement our framework. According to Eqn. (5), it is necessary to carefully design (1) tractable distribution @R (R', f/2,,, T; z0, 0) for sam- pling Rt; (2) closed-form matching objective VR log T f).
Matching objective. Inspired by diffusion models that use Gaussian transition kernels for tractable computation, we design the SDE on geometric states in Proposition 3.1 to be:
dR' adW with transition density pR ar<(z't'( (2)22,)) (20, o²(t' I))
(7)
The explicit form of the objective can be directly calculated, i.e., R
T ()
- R!
-T
Sampling distribution. According to Theorem 3.3, the transition density QR RR. 1/2,, T; 20. 0) can be calculated by using the Doob's 's h-transform in Proposition 3.2, i.e., QR (R' ///2,, T; z0. 0) PR R t|22, T) lox (R' d:s0).
Jo (x1, Moreover, he is determined by Adata and pR. which is already specified
in Eqn. (7). Therefore, we can also calculate qR (R', f| T; 20, 0) N( I 20, R² -(T-t) I).
T
Symmetry constraints. In proposition 3.1, we have several conditions that should be satisfied to meet the symmetry constraints. Firstly, since a parameterized model vo vo(RR, ,t; RD) used to estimate the score of our equivariant diffusion bridge, it should be -3quiqaqui and T(3)-invariant. Besides, we follow [50, 115] to consider CoM-free systems: given R = fri.....Inh we define F = Σn F2 and the CoM-free version of R = {ri F...., In - F}. To sample from (220, I' with 20 E R consisting of n objects, we (1) sample e = {4)] by i..dd drawing e NN Is): (2) calculate the CoM-free e' of e; (3) obtain 2p oe'.
Trajectory guidance. According to Eqn. (6), both pkz, and of for all iE[N-1] should be determined. Similarly, we set (1 +')=N(R' (t)(T))' which further induces (R² =111=111, 21 0) = N( Tut zi, "))
Combining all the above design choices, we have the following algorithms for training our Geometric Diffusion Bridge (Alg. 3) and leveraging trajectory guidance if available (Alg. 4). After the model

---

-1) is well trained, we leverage ODE numerical solvers [12] to simulate the bridge process by using its equivalent probability flow ODE [99]. In this way, we can effectively and deterministically predict future geometric states of interest from initial states in an efficient iterative process. Lastly, it is also noteworthy that our framework general to be implemented by using other advanced design strategies [99, 47, 48], which we leave as future work.
Algorithm 1 Training
Algorithm 2 Training with trajectory guidance
repeat
1: repeat
2: (zo, 21) ~ Qines (R'o, R
2: (zu...., EN) quml R R
T), t~ ( --
3: 1~ [0
c0.1
4:
c~N((,11)
R'= f.2 + T-t 20 +
T are
5:
R ="-$#20
+ + 7 2: +
(P(T-e)
T Ove
6: Take gradient descent step on
6:
Take gradient descent step on
Vod(t) (((--)) vo(R tu
VeX(t)
vR t; of(T-19
until converged
until converged
4 Experiments
In this section, we empirically study the effectiveness our Geometric Diffusion Bridge on crucial real-world challenges requiring bridging geometric states. In particular, we carefully design several experiments covering different types of data, scales and scenarios, as shown in Table 2. Due to space limits, present more details in Appendix D.
Table 2: Summary of experimental setup.
Dataset
Task Description
Data Type
Trajectory data Training set size
QM9 [79]
Equilibrium State Prediction
Simple molecule
110,000
Molecule3D [116]
Equilibrium State Prediction
Simple molecule
X
2.339.788
OC22, IS2RS [13]
Structure Relaxation
Adsorbate-Catalyst complex
45,890
4.1 Equilibrium State Prediction
Task. Equilibrium states typically represent local minima on the Born- Nomm-Oppehheime potential energy surface molecular system [54], which correspond to its most stable geometric state and play essential role in determining its properties in various aspects [4, 21]. In this task, our goal is to accurately predict the equilibrium state from the initial geometric state molecular system.
Dataset. Two popular datasets are used: (1) QM9 [79] is a medium-scale dataset that has been widely used for molecular modeling, consisting of 330,000 organic molecules. In convention, 110k, 10k, and 11k molecules are used for train/valid/test sets respectively; (2) Molecule3D [116] a large- scale dataset curated from the PubChemQC project [67, 71], consisting of 1,899,647 molecules in total and its train/valid/test splitting ratio is 2 In particular, both random and scaffold splitting methods are adopted to thoroughly evaluate the in-distribution and out-of-distribution performance. For each molecule, an initial geometric state generated by using fast and coarse force field [73, 52] and geometry optimization is conducted to obtain DFT-calculated equilibrium geometric structure.
Setting. In this task, we parameterize vo 1; R by extending Graph-Transforen based equivariant network [92, 63] to encode both time steps and initial geometric states as conditions. For inference, we use 10 time steps with the Fuler solver [12]. Following [111], we choose several strong baselines for comprehensive comparison, and use three metrics for measuring the error between predicted target states and ground-truth states: C-RMSD, D-MAE and D-RMSE, The detailed descriptions of the baselines, evaluation metrics and training settings are presented in Appendix D.I.
Results. Results on QM9 and Molecule3D are shown in Table and 4 respectively. can be easily seen that our GDB framework consistently surpasses all basclines by significantly large margin
7

---

Table 3: Results on the QM9 dataset (A). We report the official results of baselines from [111]
Validation D-RMSEJ
Test
D-MAEJ
C-RMSD,
D-MAEL
D-RMSEL
C-RMSD
RDKit DG
0.358
0.616
0.722
0.358
0.615
0.722
RDKit ETKDG
0.355
0.621
0.691
0.355
0.621
0.689
GINE [39]
0.357
0.673
0.685
0.357
0.669
0.693
GATv2 [10]
0.339
0.663
0.661
0.339
0.659
0.666
GPS [80]
0.326
0.644
0.662
0.326
0.640
0.666
GTMGC [III|
0.262
0,468
0.362
0,264
0,470
0,367
GDB (ours)
0.092
0.218
0.143
0.096
0.223
0.148
Table 4: Results on the Molecule3D dataset (A). We report the official results of baselines from
Validation D-RMSEL
Test
D-MAEL
C-RMSDL
D-MAEL
D-RMSEL
C-RMSD]
(a) Random Split
RDKit DG
0.581
0.930
1.054
0.582
0.932
1,055
RDKit ETKDG
0.575
0.941
0.998
0.576
0,942
0.999
DeeprrGCIN-DGGNN [116]
0.509
0.849
0.571
0.961
GINE [39]
0.590
1.014
1.116
0.592
1.018
1.116
GATv2 [10]
0,563
1.082
0,564
0,986
1,083
GPS [80]
0,528
0.909
1.036
0.529
0.911
1.038
GTMGC [I1I]
0.432
0.719
0,712
0.433
0,721
0,713
GDB (ours)
0.374
0.631
0,622
0,376
0.626
0,619
(b) Scaffold Split
RDKit DG
0.542
0.872
1.001
0.524
0.857
0.973
RDKit ETKDG
0,531
0.874
0.928
0.511
0.859
0.898
DeeperGCN-DAGNN [116]
0.617
0,763
1.176
GINE [39]
0.883
1.517
1.407
1.400
2,224
1.960
GATv2 [10]
0.778
1,385
1.254
1.238
2.069
1.752
GPS [80]
0.538
0.885
1.031
0.657
1,091
1.136
GTMGC [11I]
0.406
0.675
0,678
0,400
0.679
0,693
GDB (ours)
0.335
0.587
0.592
0.341
0.608
0.603
QM9, e.g., 60.5% relative C-RMSD reduction on valid/test sets respectively, establishing a new ofte--of-theart performance. Similar trends also can be observed in Molecule3D, i.e., 22.6%//3.222 relative C-RMSD reduction for valid/test sets the random split and 12.7% 3.0% reduction for the scaffold split, largely outperforming the best baseline. These significant error reduction results show the superiority our GDB framework for bridging geometric states, and its generality on both medium and auree-ccle challenges. Moreover, framework performs consistently across valid and tests of both random and scaffold splits, further verifying its robustness in challenging scenarios.
4.2 Structure Relaxation
Task.
Catalyst discovery is crucial for various applications. Adsorbate candidates are placed on catalyst surfaces and evolve through structure relaxation to adsorption states, in which the adsorption structures can be determined for measuring catalyst activity and selectivity. Our goal is thus to accurately predict adsorption states from initial states of adsorbate-catalyst complexes.
Dataset. We adopt Open Catalyst 2022 (OC22) dataset [105], which has great significance for the development of Oxygen Evolution Reaction (OER) catalysts. Each data is in the form of the adsorbate- catalyst complex. Both initial and adsorption states with trajectories connecting them are provided. The training set consists of 45,890 catalys--adsorbtte complexes. To better evaluate the model's perfor- mance, the validation and test sets consider the in-distribution (ID) and out-i--distribuution (OOD) set- tings which unseen catalysts, containing approximately 2,624 and 2,780 complexes respectively.
Setting, Following [105], we use the Average Distance within Threshold (ADwT) the evaluation metric, which reflects the percentage of structures with an atom position MAE below thresholds. We parameterize vo (R'R,, R") by using GemNet-OC [34], which also serves verification that

---

Table Results on the OC22 IS2RS Validation set. "OC20+0022" denotes using both OC20 [13] and OC22 data; "OC20=+0C22" means pre-training on OC20 data then fine-tuning on OC22 data; "OC22-only" means only using OC22 data. We report the official results of baselines from [105]
Model
ADwT [%] (ID)
ADwT [%] 1 (OOD)
Avg [%] 1
OC20+0C22
SpinConv [94]
55.79
47.31
51.55
GemNet-OC [34]
60.99
53.85
57.42
0C20-00C22
SpinConv [94]
56.69
45.78
51.23
GemNet-OC: [34]
58.03
48.33
53.18
GemNet-OC--arge [34]
59.69
51.66
55.67
OC22-only
baseline
44.77
42.59
43.68
SpinConv [94]
54.53
40.45
47.49
GemNet-dT [32]
59.68
51.25
55.46
GemNet-OC [34]
60.69
52.90
56.79
GDB (ours)
63.01
55.78
59.39
trajectory guidance
62.14
54.94
58.54
Rd condition
60.17
49.26
54.71
our framework is compatible with different backbone models. For inference, we also use 10 time steps with the Euler solver. Following [105], we choose strong MLFF baselines trained on force field data for a challenging comparison. The detailed descriptions of baselines and settings are presented in Appendix D.2.
Results. In Table 5, our GDB significantly outperforms the best baseline, e.g., 3.3%/3.6%% 999//.66//344% rela- tive improvement on the ADwT metric of ID, OOD and Avg respectively. It noteworthy that the best baseline the GemNet-OC force field trained on both OC20 and OC22 data, which is 10 times more than OC22 data only. Nevertheless, our framework still achieves better performance on predicting the adsorption geometric states. Moreover, our framework without using any trajectory data still can achieve better performance compared to the best baseline, e.g., 58.54 v.s. 57 57.42 Avg[%) All the re- sults on this challenging task further demonstrate the superiority and completeness of our framework.
Ablation study. Furthermore, we conduct ablation studies to examine key designs of our framework in Table 5. Firstly, we can see that using trajectory guidance indeed improves the performance of our framework, e.g., 1. 1.4% relative improvement on Avg ADwT. Moreover, we also investigate the impact of R condition in vo(R f; R2), which plays an essential role in preserving the joint distribution of geometric states. Without this condition, we can see a significant drop, e.g., 6.5%/10.39 relative ADwT drop on Avg/OOD respectively. Overall, these ablation studies serve strong supports on the necessity of developing a unified framework that can precisely bridge geometric states by preserving their joint distributions and effectively leverage trajectory data as guidance for enhanced performance.
Related Works
Direct Prediction. One line of approach for bridging geometric states direct prediction, i.e., train- ing model to directly predict target geometric states given initial states as input. Models that carefully respect symmetry constraints such as the equivariance to 3D rotations and translations are typically used, which are called Geometric Equivariant Networks [11, 36, 120, 27]. Different techniques have been explored to encode such priors, which mainly include vector operations such as scalar and vector product [35, 87, 89, 41, 103, 14], e.g., the scalar-vector product used in EGNN [87]. and tensor product based operations [104, 31, 8, 57, 64]. Despite its simplicity and efficiency, direct prediction requires encoding the iterative evolution of geometric states into single-step prediction model, which lacks the ability to capture the underlying dynamics and cannot leverage trajectories of geometric states.
Machine I Learning Force Field. Another line of approach called machine learning force field (MLFF) [106. 5, 6, 70, 75, 58), which trained to predict intermediate labels, such the potential

---

energy or force of the (local) current geometric state instead. After training, MLFFs can be used to simulate the trajectory of geometric states over time based on underlying equations. Using Geometric Equivariant Networks the backbone, MLFFs typically satisfy the symmetry constraints. Besides, trajectory data with additional energy or force labels can directly be used for training ML.Fs. However, this paradigm highly depends on the existence and quality of intermediate labels since small local errors in energy or force prediction can accumulate along the simulation process [7, 106, 30]. Moreover, there exists no guarantee that MLFFs can completely model joint state distributions, which is another limitation for bridging geometric states.
Geometric Diffusion Models. recent years, diffusion models [37, 99] have emerged with state-
of-the-art generative modeling performance across various domains [85, 108, 51, 56]. In geometric domain, diffusion models are typically used for molecule conformation generation [115, 114, 38) and protein design [108, 117]. By properly design the noising process and model architectures, symmetry constraints on the transition kernel and prior distribution can be satisfied, which guarantees the generated data is sampled from noto-translational invariant distributions [115, 38). In addition the score-based formulation, recent advances further extend new techniques such as flow matching [59, 61, 1] to satisfy symmetry constraints for these generation tasks [49, 100]. Nevertheless, there exists no guaratee that these approaches can model the joint distribution of geometric states [61, 96]. And how to leverage trajectory data as guidance for bridging geometric states is also challenging.
Other techniques. MoreRed [45] trains diffusion model on equilibrium molecule conformations with a time step predictor, and directly use it for bridging any conformations to their equilibrium states. GTMGC [111] instead develop a Graph Transformer to directly predict equilibrium conformations from their 2D graph forms. Both of them are limited to the equilibrium conformation prediction task, cannot preserve the joint state distribution and leverage trajectory data. EGNO [112] is a concurrent work that develops a neural operator based approach to model dynamics of trajectories. By carefully designing temporal convolution fourier spaces, EGNO can learn from trajectory data. However, this tailored approach cannot be directly used without trajectory guidance. To preserve joint data distributions, [22, 121] coincide with us to leverage Doob's ss lu-transform to repurposing standard diffusion processes, but they do not respect symmetry constraints and cannot leverage trajectories. There also exist recent works that study the diffusion bridge framework [76, 93] and apply it to various domains such as images and graphs [110, 62, 42]. Compared to all above approaches, our GDB framework stands out as a unique and ideal solution that can precisely bridge geometric states and effectively leverage trajectory data (if available) in unified manner.
Conclusion
In this work, we introduce Geometric Diffusion Bridge (GDB), a general framework for bridging geometric states through generative modeling. We leverage modified version of Doob's h-transform to constructe an equivariant diffusion bridge for bridging initial and target geometric states. Trajectory data can further be seamlessly leveraged guidance by using chain of equivariant diffusion bridges, allowing complete modeling of trajectory data. Mathematically, we conduct a comprehensive theoretical analysis showing our framework's 's ability to preserve joint distributions of geometric states and capability to completely model the evolution dynamics. Empirical comparisons on different settings show that our GDB significantly surpasses existing state-of-the-art approaches and ablation studies further underscore the necessity of several key designs in our framework. In the future, it is worth exploring better implementation strategies of our framework for enhanced performance, and applying our GDB to other critical challenges involving bringing geometric states.
Broader Impacts and Limitations
This work newly proposes a general framework bridge geometric states, which has great signifi- cance in various scientific domains. Our experimental results have also demonstrated considerable pos- itive potential for various applications, such as catalyst discovery and molecule optimization, which can significantly contribute to the advancement of renewable energy processes and chemistry discov- ery. However, is essential to acknowledge the potential negative impacts including the development of toxic drugs and materials. Thus, stringent measures should be implemented to mitigate these risks.
10

---

There also exist some limitations to our work. For the sake of generality, we do not experiment with advanced implementation strategies of training objectives and sampling algorithms, which leave room for further improvement. Besides, the employment of Transformer-based architectures may also limit the efficiency of our framework. This has also become a common issue transformer-based diffusion models, which we have earmarked for future research.
Acknowledgements
We thank all the anonymous reviewers for the very careful and detailed reviews as well as the valuable suggestions. Their help has further enhanced our work. Liwei Wang is supported by National Science and Technology Major Project (2022ZD01 14902) and National Science Foundation of China (NSFC62276005). Di He is supported by National Science Foundation of China (NSFC62376007).
References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden- Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 2023.
[2] Brian DO Anderson, Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-322 1982.
[3] Muratahan Aykol, Joseph H. Montoya, and Jens Hummelshoj. Rational solid-state synthesis routes for inorganic materials. Journal of the American Chemical Sociery, 143(24):9244-2255 9259, 2021.
[4] Keld L Bak, Jürgen Gauss, Poul Jørgensen, Jeppe Olsen, Trygve Helgaker, and John F Stanton. The accurate determination of molecular equilibrium structures. The Journal of Chemical Physics, 114(15):6548 6556, 2001.
[5] Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1423-1121111 11436, 2022.
[6] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
[7] Jorg Behler. Perspective: Machine learning potentials for atomistic simulations. The Journal of chemical physics, 145(17), 2016.
[8] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022.
[9] Linda Broadbelt, Scott M Stark, and Michael Klein. Computer generated pyrolysis modeling: on-the-fly generation of species, reactions, and rates. Industrial & Engineering Chemistry Research, 33(())790-799 1994.
[10] Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? arXiv preprint arXiv:2105.. 14491, 2021.
[11] Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478. 2021.
[12] John Charles Butcher. Numerical methods for ordinary differential equations. John Wiley & Sons, 2016.
[13] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Doming Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11(10))6059-600) 2021.

---

[14] Tianlang Chen, Shengjie Luo, Di He, Shuxin Zheng, Tie- Yan Liu, and Liwei Wang. GeoM- Former: general architecture for geometric molecular representation learning, Forty-first International Conference on Machine Learning, 2024.
[15] Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Robert Müller. Machine learning of accurate energy- conrgyccnnsrivngg molecular force fields. Science advances,
2017.
[16] Kai Lai Chung and John B Walsh. Markov processes, Brownian motion, and time symmetry, volume 249. Springer Science & Business Media, 2006.
[17] Jonathan Clayden, Nick Greeves, and Stuart Warren, Organic chemistry. Oxford University Press, USA, 2012.
[18] John F Cornwell. Group theory in physics: An introduction. Academic press, 1997.
[19] F Albert Cotton. Chemical applications of group theory. John Wiley & Sons, 1991.
[20] F Albert Cotton, Geoffrey Wilkinson, Carlos A Murillo, and Manfred Bochmann. Advanced inorganic chemistry. John Wiley & Sons, 1999.
[21] Attila G Császár, Gábor Czakó, Tibor Furtenbacher, Jonathan Tennyson, Viktor Szalay, Sergei V Shirin, Nikolai F Zobov, and Oleg L Polyansky. On equilibrium structures of the water molecule. The Journal of chemical physics, 122(21). 2005.
[22] Valentin De Bortoli, Guan-Horgg Liu, Tianrong Chen, Evangelos Theodorou, and Weilie Nie. Augmented bridge matching. arXiv preprint arXiv:23/1.06974 2023.
[23] Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmalcki, Diego de Las Casas, et al. Mag- netic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414-
419, 2022.
[24] Amanda Dewyer, Alonso Argüelles, and Paul M Zimmerman. Methods for exploring reac- tion space in molecular systems. Wiley Intendisciplinary Reviews: Computational Molecular Science, 8(2) 3344, 2018
[25] Jacob D Durrant and J Andrew McCammon. Molecular dynamics simulations and drug discovery. BMC biology, 9:1-9, 2011.
[26] Rick Durrett. Probability: theory and examples, volume 49. Cambridge university press, 2019.
[27] Alexandre Duval, Simon V Mathis, Chaitanya K Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D) Malliaros, Taco Cohen, Pietro Liò, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arXiv:2312.075 2023.
[28] Andreas Eberle. Stochastic analysis.
[29] Ferran Feixas, Steffen Lindert, William Sinko, and J Andrew McCammon. Exploring the role of receptor flexibility in structure-basee drug discovery. Biophysical chemistry, 186:31-45,
2014.
[30] Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Komez-Bombarelli and Tommi S. Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. Transactions on Machine Learning Research, 2023. Survey Certification.
[31] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances neural information processing systems, 33:1970-1981, 2020.
[32] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.
12

---

[33] Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123. 2020.
[34] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das. Gemnet-OC: Developing graph neural networks for large and diverse molecular simulation datasets. Transactions on Machine Learning Research, 2022.
[35] Mojtaba Haghighatlari, Jie Li, Xingyi Guan, Oufan Zhang, Akshaya Das, Christopher J Stein, Farnaz Heidar-Zzddeh, Meili Liu, Martin Head- Gordon, Luke Bertels, et Newtonnet: newtonian message passing network for deep learning of interatomic potentials and forces. Digital Discovery, 13))3333343, 2022.
[36] Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey. arXiv preprint arXiv:2202.07236, 2022.
[37] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances neural information processing systems, 33:6840-6883 2020.
[38] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation 3d. In International conference on machine learning. pages 8867-8887. PMLR, 2022.
[39] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Perey Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv: 12265, 2019.
[40] Frank Jensen. Introduction to computational chemistry. John wiley & sons, 2017.
[41] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2021.
[42] Jaehyeong Jo, Dongki Kim, and Sung Ju Hwang. Graph generation with cictinatiion-redditii diffusion mixture, 2024.
[43] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 589, 2021.
[44] Wolfgang Kabsch. A discussion of the solution for the best rotation to relate two sets of vectors. Acta Crystallographiic Section Crystal Physics, Diffraction, Theoretical and General Crystallography, 34(5):827-828, 1978.
[45] Khaled Kahouli, Stefaan Simon Pierre Hessmann, Robert Müller, Shinichi Nakajima, Stefan Gugler, and Niklas Wolf Andreas Gebauer. Molecular relaxation by reverse diffusion with time step prediction. arXiv preprint arXiv:2404. 10935, 2024.
[46] Martin Karplus and J Andrew McCammon. Molecular dynamics simulations of biomolecules. Nature structural biology, 9()):44-655) 2002.
[47] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of ceffusiio--aaer generative models. Advances Neurul Information Processing Systems, 35:26565-26577, 2022.
[48] Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances Neural Information Processing Systems, volume 36, pages 65484-65516 Curran Associates, Inc., 2023.
[49] Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances Neural Information Processing Systems, 36, 2024.
13

---

[50] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361-5370. PMLR, 2020.
[51] Zhifeng Kong, Wei Ping. Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatiee diffusion model for audio synthesis. In International Conference on Learning Representations, 2021.
[52] Greg Landrum. Rdkit: Open-source cheminformatics software. Github, 2016.
[53] Christian Léonard. Girsanov theory under a finite entropy condition. In Séminaire de Proba- bilités XL/V, pages 429-465. Springer, 2012.
[54] Ira N Levine, Daryle H Busch, and Harrison Shull. Quantum chemistry, volume 6. Pearson Prentice Hall Upper Saddle River, NJ, 2009.
[55] Raphael Levine. Molecular reaction dynamics. Cambridge University Press, 2009.
[56] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy $ Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances Neural Information Process- ing Systems, 35:4328-4343, 2022.
[57] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. 11990, 2022.
[58) Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In The Twelfth Interna- tional Conference Learning Representations, 2024.
[59] Yaron Lipman, Ricky T. Q. Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023.
[60] Meng Liu, Cong Fu, Xuan Zhang, Limei Wang. Yaochen Xie, Hao Yuan, Youzhi Luo, Zhao Xu, Shenglong Xu, and Shuiwang Ji. Fast quantum property prediction via deeper 2d and 3d graph networks. arXiv preprint arXiv: 106.0851, 2021.
[61] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
[62] Xingchao Liu, Lemeng Wu, Mao Ye, and qiang liu. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2023.
[63] Shuqi Lu, Zhifeng Gao, Di He, Linfeng Zhang, and Guolin Ke. Highly accurate quantum chemical property prediction with uni-mol+. arXiv preprint arXiv:2303. 16982, 2023.
[64] Shengjie Luo, Tianlang Chen, and Aditi S. Krishnapriyan. Enabling efficient equivariant oper- ations in the fourier basis via gaunt tensor products. In The Twelfth International Conference on Learning Representations, 2024.
[65] Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. In The Eleventh International Conference on Learning Representations, 2023.
[66] Shengjie Luo, Shanda Li, Shuxin Zheng, Tie-Yan Liu, Liwei Wang, and Di He. Your trans- former may not be as powerful as you expect. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[67] Nakata Maho. The pubchemqe project: A large chemical database from the first principle calculations. In AIP conference proceedings, volume 1702, page 090058. AIP Publishing LLC., 2015.
14

---

[68] Richard M Martin. Electronic structure: basic theory and practical methods. Cambridge university press, 2020.
[69] Amil Merchant, Simon Batzner, Samuel S Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nature, 624(7990):80-B8 2023.
[70] Albert Musaclian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
[71] Maho Nakata and Tomomi Shimazaki. Pubchemqe project: a large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57(6):1300-1309 2017.
[72] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162-8171. PMLR, 2021.
[73] Noel M O' Boyle, Michael Banck, Craig A James, Chris Morley, Tim Vandermeersch, and Geoffrey R Hutchison. Open babel: An open chemical toolbox. Journal of cheminformatics, 3:1-14, 2011.
[74] Bernt Øksendal and Bernt Øksendal. Stochastic differential equations. Springer, 2003.
[75] Saro Passaro and Lawrence Zitnick. Reducing so (3) convolutions to so (2) for efficient equivariant gnns. In International Conference Machine Leurning, pages 27420-27438. PMLR, 2023.
[76] Stefano Peluchetti.
Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023.
[77] Stefano Peluchetti. Non-denoising forward-time diffusions. arXiv preprint arXiv:2312.14589 2023.
[78] Jan-Hendrik Prinz. Hao Wu, Marco Sarich, Bettina Keller, Martin Senne, Martin Held, John D Chodera, Christof Schütte, and Frank Noé. Markov models of molecular kinetics: Generation and validation. The Journal of chemical physics, 134(17), 2011.
[79] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[80] Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for a general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35:14501-14515, 2022.
[81] Dennis C Rapaport. The art of molecular dynamics simulation. Cambridge university press, 2004.
[82] L Chris G Rogers and David Williams. Diffussions, Markov processes and martingales: Volume
2, Itô calculus, volume 2. Cambridge university press, 2000.
[83] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjorn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10684-10695, June 2022.
[84] David Rosenberger, Justin Smith, and Angel Garcia. Modeling peptides with classical and novel machine learning force fields: A comparison. The Journal Physical Chemistry B, 125(14):3598-35 -612,, 2021.
[85] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479-36499 2022.
15

---

[86] Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
[87] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[88] H Bernhard Schlegel. Geometry optimization. Wiley Interdisciplinary Reviews: Computational Molecular Science, 1(5):790-800 809, 2011.
[89] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Muchine Learning, pages 9377-9388, PMLR, 2021.
[90] Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
[91] William Raymond Scott. Group theory. Courier Corporation, 2012.
[92] Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tan Liu. Benchmarking graphormer on large-scale molecular modeling datasets, arXiv preprint arXiv:2203. 04810, 2022.
[93] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
[94] Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C Lawrence Zitnick. Rotation invariant graph neural networks using spin convolu- tions. arXiv preprint arXiv:2106. 00.009575 2021.
[95] Jascha Sohl- Dickkiics,,,,, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsu- pervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2265, Lille, France, 07-09 Jul 2015. PMLR.
[96] Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion schrödinger bridges. Uncertainty in Artificial Intelligence, pages 1985-1995, PMLR, 2023.
[97] Jiaming Song. Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[98] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[99] Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[100] Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei-Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024.
[101] Howard Stephen Stoker and G Lynn Carlson. General, organic, and biological chemistry. Houghton Mifilin, 2004.
[102] Challapalli Suryanarayana. Experimental techniques in materials and mechanics. Crc Press, 2011.
[103] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference Learning Representations, 2022.
16

---

[104] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-an and cunslationeqqivarriinn neural networks for 3d point 2018.
clouds. arXiv preprint arXiv: arXiv:
[105] Richard Tran, Janice Lan, Muhammed Shuaibi, Brandon M Wood, Siddharth Goyal, Ab- hishek Das, Javier Kieras-Doming Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysts. ACS Catalysis, 13(5):3066 -3044, 2023.
[106] Oliver Unke, Stefan Chmiela, Huziel Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schutt, Alexandre Tkatchenko, and Hobu-R Muller. Machine learning force fields. Chemical Reviews, 121(16):10142-0011 2021.
[107] Hanchen Wang. Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang, Ziming Liu, Payal Chandak, Shengchao Liu, Peter Van Katwyk, Andreca Deac, et al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972):47-66, 2023.
[108] Joseph L Watson, David Juergens, Nathaniel Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620079766:1089-11100 2023.
[109] E Weinan and Eric Vanden-Eijnden, Transition-pall theory and path-finding algorithms for the study of rare events. Annual review of physical chemistry, 61 (2010):391-422)) 2010.
[110] Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and qiang liu. Diffusion-basse molecule generation with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Newral Information Processing Systems, 2022.
[111] Guikun Xu, Yongquan Jiang. PengChuan Lei, Yan Yang, and Jim Chen. Gtmge: Using graph transformer predict molecule's ground-state conformation. In The Twelfth International Conference on Learning Representations, 2023.
[112] Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzade- nesheli, Jure Leskovec, Stefano Ermon, and Anima Anandkumar. Equivariant graph neural operator for modeling 3d dynamics. arXiv preprint arXiv:2401.1103/ 2024.
[113] Minkai Xu, Alexander S Powers, Ron O. Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3D molecule generation. In Andreas Krause. Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceed- ings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 38592-38610, PMLR, 23-29 Jul 2023.
[114] Minkai Xu, Alexander Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference on Machine Learning, pages 38592-3888 10010. PMLR, 2023.
[115] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[116] Zhao Xu, Youzhi Luo, Xuan Zhang, Xinyi Xu, Yaochen Xie, Meng Liu, Kaleb Dickerson, Cheng Deng, Maho Nakata, and Shuiwang Ji. Molecule3d: A benchmark for predicting 3d geometries from molecular graphs. arXiv preprint 2rX.v22110.01717, 2021.
[117] Jason Yim, Brian Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 40001-40039. PMLR, 23-29 Jul 2023.
17

---

[118] Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke. Di He. Yanming Shen, and Tan Liu. Do transformers really perform badly for graph representation? In A. Beygelzimer, Y. Dauphin, P. Liang. and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[119] Bohang Zhang, Shengjie Luo, Liwei Wang, and Di He. Rethinking the expressive power of GNNs via graph biconnectivity. In The Eleventh International Conference on Learning Representations, 2023.
[120] Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.088. 08423, 2023.
[121] Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models. The Twelfth International Conference Learning Representations, 2024.
18

---

A Organization of the Appendix
The supplementary material is organized as follows. In Appendix B, we first recall definitions and tools from stochastic calculus and then give the proofs of all theorems. In Appendix C, we give the derivation of our practical objective function and our sampling algorithms. In Appendix D, we give some details of our experiments, including comprehensive introduction to the datasets, baselines, metrics and settings.
B Proof of Theorems
B.1 Review of Stochastic Calculus
Let (Xt))ejo, be stochastic process. We use ('ll, (1:11111000 22, i (n) to denote its condi- tional density functiin satisfying
, E A|x,; 21, X1, .2...., Xt, = Ixn)= = / p(a', f/ t1; 22, f2:...:: : In )dr
ln- (Xx) a Markov process, p(a', f'1 $1; 22, tgj...................... nn tn) = p(x', o(x't'' ,t tx), which is also called transition density function.
for any Borel set A, where 4 12 <
One of the most important results of stochastic calculus is the Ito's formula. The precise statements follows.
Theorem B.1 (Ito's formula for Brownian Motion). Let B, be the d- dimensiiinna Brownian Motion. Assume I is a bounded real valued function with continuous second-order partial derivatives, i.e. E (TR.). Then the Itos 's formula is given by
[(B,)= I(Bn) + !'"
VS(B.) -B. 1 (B).B..
We follow [86] for the proof of Doob's h-transform. The infinitesimal generator of the Markov process plays an important role in the proof of the Doob's b-transform. The precise definitions are as follows.
Definition B.2. (Generator of Process) The infinitesimal generator A₁ of a stochastic process (X,) for a function 6(x) is
X 6(x)
A(o(z) lim s-0+
(9)
where à is a suitably regular function. For an Itô process defined as the solution to the SDE dX, ()d, + (I)dB,, the generator is
d
0
4 = Σ
F (z,
Σ 2 1=1
0
(t)
(10)
Ox;
Or.
The Fokker-Planck's Equation an useful tool to track the evolution of the transition density function associated with SDE. The precise statements are as follows.
Proposition B.3. (Fokker- Piatee---lacc'' Equation) Let p(x', ,'a,,,
t) be the transition density function of the SDE dx, f(X,, t)dl + (I)dB,, Then p(æ', 1//,,,, satisfies the Bokker-Plande 's Equution
Öp(r, t/ 0) - O(f t)p(a, t/ 0)) + Σ d2))) d
O p2p(x, t/ 0)
(11)
Øt
Ox;
Oz?
with the initial condition p(x, O,z0. 0) = 5(x - no). The Fokker-Planck Equation can also be written in compact form using the generator A.:
ar p(e, 120. 0) = Aūp(a, 1/20, 0),
(12)
where Aī is the adjoint operator of A:
((f'
d
(2(-) σ (1)
A;
(13)
Ox:
2 1=1
Ox
19

---

When the terminal is fixed, the evolution of the transition density function can also given by PDE, which is called the Backward Kolmogorov Equation. We give the precise statement as follows. Proposition B.4. (Backward Kolmogorov Equation) Let pla', (/,, 1) be the transition density func- tion of the SDE dX.- = D)dL + a (L)dB, Then p(æ' (1 1) satisfies the Backward Kolmogorov Equation
0p( 8) Os
d Σ (xr,))
0p(2t,112,8 8) Ox,
θ (o( 0,
Σ (s)
(14)
Ox?
with the initial condition p(a,, t| ò(x 16---11 21). written in compuct form using the generator As:
The Backward Kolmogorov Equation can also be
_/200 + As
plac, 8),,8) O.
(15)
B.2 Proof of Proposition 3.1
Proposition B.5. Let R. denote the space of geometric states and (.(;))) R [0, T] - R denote the drift coefficient on R. Let (W")eejorr) denote the Wiener process on R. Given an SDE on geometric
states dR' = fR(R² t)dt + (t)dW'' R ~ q(R"), its transition density PR pR(z', f/,2,, t), 2, 2' E R is SEEi)--qquvvarian, i.e., PR( I' f) pR(0R (g)[R( 1,8'|p (R) f), Vg E SE(3), vO <
t < T, if the following conditions are satisfied: (1) q(R)) is -invariinti (2) (.
is -equieequi and T(3) ((3)-inaai;nn;; (3) the transition density of (Weee1 is SE(3)- equivariant.
Proof. In this section, we view R = (rI...., rn E R r1 @ rr 3 which is the concatenation of r,. So from this perspective, the space R is isomorphic to the Euclidean space 3m Then [e is the Wiener process with dimension d = 3n.
For any g E SE(3), p (g) can be characterized by an orthogonal matrix O(g) E R3x3 satisfying det (O(g) = 1, and a translation vector t IR3 Then the representation of SE(3) on R given by
[R) Ox(g) + tR.
(16)
where (R diag Og)). (g)) tR = t o e ... R3 It's obvious that O (g)
is also an orthogonal matrix in 3 satisfying O (g) = (g))
According to Proposition B.3, the evolution of the transition density function is given by the Fokker-Planck's Equation
@xR t/ 0) d (f' (z, t))R(() f/ 0))
1 d
82 (pR(2, to,t 0))
=())
(17)
Øt.
Oa:
2
Ox?
with the initial condition pr (z, 0/ 0) = 8(x zo).
Let (g(g)z z tR. yo O 20 + tR. then have
PR (p R (g)(z) (g) [zo] 0) pr (g)r + tR, (g) + tR, 0) PR (v, Auo, 0). (18)
The evolution of the transition density function (2 tlyo, 0) is also given by the Fokker-Planck's Equation:
8p (y, 1/30, 0)
a (F' (1l,) (v. f/50, 0))
d O² (px (v, f/20, 0))
Σ
21)
(19)
Øt
Oy:
2 i=1
dyf
with the boundary condition
Px (y. 0/8o: 0) = 5(y vo) = 6(x 2o). Since y = 0g + tR. we
have x = (g)( tx... Then by the chain rule, we have
0
dz θ
Σ o (2))
a daz
Σ (R dxy
i
(20)
Dy
Oy: dz
20

---

Since fx(;,t) is
and function, we have
(y. (Ox()))))+++)+++++++++++))>))= tz, (2,t) t) (OR = ((O(()))
(21)
Then the Fokker-Planck's equation becomes
Opr (v. flyo. 0)
Öt
Σ
d
a (f,, O) (l00, 01)
By:
+
02)(1)) Σ
2
1
82 (px(u,tluo)))
Oyf
(22)
d d
(R(()) (((OR(9))a as 1 (I, DRR (y. 1100: 0))
(23)
dz;
ΣΣΣ OR (R (g))as a
b (p. (Ox(g))s 0))
+
(24)
Ork
dr,
d
O( (I, (v t 1130. 0))
OR(O))
(25)
dx,
+ 2
(t)
ΣΣ>
d
d
a (px (v. to 0))
(26)
Dzz
Or,
Since (R(9) is an orthogonal matrix, the columns of (g) are orthogonal to each other, i.e.
Σ (0 = ررر
d
0
j * k,
(27)
1
j = k.
So the Bokker-Planckk equation can be simplified to
Opx (,, t/yo, 0)
d(f (I, t) (v, t| 0))
(28)
Ø1
Or,
k=1
1
á a (px (v. t/ 0))
+
(1)
8gk Dok
(29)
2
Oaz
J=1 km
0(Eg (a, (), fl 0))
d 82 (px(v.t/y000) 0))
dx,
2
01) 1=1 A(z,)²
(30)
which is same as Eqn.(17). Since the boundary condition pr (y, Olyn, 4n) = 5ly - yo) = 8(x
anp) PR (2 0 Jrn 0), then pr (v, t flyo, to) PR (x, faoro fo), Vt [0, T). Thus have proved that pR R R' t) = Pre (p (g) R 1.p (g) [*). t). Vg E SE(3). 4A S <tt xt's' ST.
B.3 Proof of Proposition 3.2
Proposition B.6 (Doob's s h--transform.. Let PR P (2,))) be the transition density of the SDE in Proposition 3.1. Let he le((; (()-) R x 00, T Ryo be smooth function satisfying: (I) hg(,,)) f) is (in3))innrriant; (2) (x(2z t) = PT (z', 2211 (z', ,')dz'. We can derive the following hv-
transformed SDE on geometric states:
dR
(R' t) o² (t) V R log ho (R' t)] dt + (tt)dW
(31)
with transition density p (2',
(=( 1'| (|, = p2(22, fl hx(2,t)
preserving the symmetry constraints.
Proof. We use the definition of the infinitesimal generator to prove the proposition. The infinitesimal generator of p (x', t' l/xxt) t) for a function 6(x) is given by
E R
R
= r (xx)
())())) = lim 8-20
(32)
21

---

Since (2',,t = pn pr f22,,t) kml='s)) Ao
so we have
Ea ( R = x] (
11,,,1400 8) R'' = x]
(33)
h(x,t)
Then of)(2) can be simplified as
( R² = ((xx)
A[)()) = lim
(34)
*-+0+
8
B|o(R +*) )R tt+ 8) RR' = x| ((x))((t))
= lim
(35)
8-00-
sh(z,t)
1
Oh(z, t)
p(x)+ + Σ @h(zz t)
0x(z)
Ōz: o(x) + h(xx t)
tz)
(36)
h(x,t)
0t
1 d
Oh t)
(() + Σ (t) ah(z. t) dx(z))
0
(t)
(37)
2
dz
dx. Ozi
1
d
8 (((z) h(z,t)]
+
(t)
(38)
2
dx²
1
Oh(z, t)
0ō(z)
((z)) + (A,h(z, ((x)
Σ
h(z, t)
Ox. (I,
(39)
h(x,t)
Öt
i=1
(t) Oh(z, t) 0o(z)
I Φo(z) Ou?
Σ
(1)
h(z, t)]
(40)
dz: Ox:
Since h(x,,) 1 pR( (lr,,)//l t')dx, we have
a @ A.) M(z,t)=)) = / (Pulee (o(al,t1z,, (lx,t) + Aipla', (op(a,,1x.)) tlp,,t ) Ma'' dddr. (
(41)
Øt
According to the Backward Kolmogorov Equation (Proposition B.4), we get
Öp(a',t' t((x,t)
+ Aop(z', f/1,,)t) =0. O.
(42)
Øl
So we get
(
a + Ar) h(z, ()= =0
(43)
Then Att (xx) can be simplified as
0o(x)
d
ah(x, t) a6(z)
A² h(x, t)
h(x, 1)
p (xx + Σ (1)
(44)
dx:
Ox: Ox.
1=1
d
8 PO(L))
+
(t)
t]
(45)
dx²
1=1
d 05(x) pl (x, + Σ (t) h(xxt) Ør. dxi
1 th(a, 06(x))
+
1
Σ
(t)
Po(x)
(46)
Dzg
dx?
i=1
((x,)) + 02))
2 log h(x, 1) )
00(x))
+ 2
1 Σ
22)))
P
(47)
Øz;
Ox:
dxk
i=1
So show that
Glog h(x,)) (xx, + +1) Ox;
0 +
2 (1) 82
A =
Σ
Σ
(48)
dx.
1=1
i=1
22

---

According to the correspondence between SDE and its generator (Definition B.2), the equation above implies that the t-transformed SDE is given by
dR' [R (R', t) o VR) log (R', t)] dt + (t)dd.
(49)
Additionally, we need to show that the h-transformed transition density function satisfies the symmet- ric constraints. First, we show that if lis, (0) is -in33-invariant, then l(;; () is also SEE33-invariant
vt E j0, T). For any g E SE(3), assume p" (g)[z) (g)g)z + tR, where O (g) is an orthogonal matrix and det(Ope (g)) 1. Since loR (2,t) satisfies
1R(2,() = f pR(('s 100,, tol, 1o)d=''
(50)
then we have
/ (202),0)1=()))))))
( (g)(=),t)=
(51)
= f I pr
(52)
(53)
By Proposition 3.1, Pr (0%10000((0011111000)000 pr (12 let
21 = (g)) [2/],, then
/ PR (10
hx (g)la)))
(54)
/ pR t) (g) (O)
(55)
/ pR t)
(56)
(2(2,t).
(57)
So h(,,t) is SE(3)- -ivvariatt Vt [0, T). h(,,t) is well-defined under these symmetric constraints. Then show (2', preserves the symmetric constraints:
hR (p (o)231))
( (1 (1]==),
(58)
h (x<()()))=))))))))))))))))))))))
hR ( =pR('/,t|12,
(59)
PPR(2 (2,, (x2(2't'))
(60)
(2
= =Pk('',t(12,))
(61)
Thus have proved that
p ( (o)(=') 1' (g)[=),)) (0) = ((),))
(62)
which implies that / (=', (1),)) preserves the symmetric constraints for any g SE(3). So the proof is completed.
Next, we show how construct a SDE with a fixed terminal point as an simple application of the Doob's 's b-transform. The result of this example is very useful to construct diffusion bridge.
Proposition B.7. Assume the original SDE is given by dx' = f(X', (o Let hig l t) = Pr (y. T/. f) which is the transition density function of the original SDE evaluated X, y. Then the h-transformed SDE
dR' [f(RR t) o² (t) R t)] dt (f)dW,,
(63)
arrive y almost surely at the final time.
23

---

Proof. The original SDE is given by
dx' = f(x', t)dt + at)dW,,
(64)
First, we need to verify that hi (z, t) satisfies the condition
( 1) - / pre(a' tola, (1( to)dx'.
(65)
Since hR (x, = pR (y, Tax t), have
/ pR po (1x,, ,(,,,, ( 10=1 / pr (12,, Max (')d''.
(66)
Then by the Kolmooo-Kllmogooo''s equation
f pR(2', ,1,,, t)pR Ti'' t')d'' = (R.
(67)
we get
/ pre('', 1/1,,, t) (2', t') = pr T(x,t) (2,
(68)
So the condition is satisfied. Then use the result of the Proposition 3.2. The h-transformed SDE is given by
dR [f(R", ,)) (t) PogpR( T(R, .)] dt + (f)dW.
(69)
And the h-transformed transition density function satisfies
f p (x 1/, t)dx'
/ pr pr (1x,, u 1') dx
(70)
hR (x, t) / pR(a', (t),,tt) PR (y. T1x t') pR (y, T x t) P(X E A|X, = x, XT y),
dr
(71)
(72)
where we use the Bayes' theorem to deduce the last equality and A is an arbitrary Borel set. Since R, a process conditioning on X = y, then RT = y almost surly.
B.4 Proof of Theorem 3.3
Theorem B.8 (Equivariant Diffusion Bridge). Given SDE geometric states dR' (R (t)dt + (t)d with transition density (2(2', l'|z, 2, E R satisfying the conditions in Proposition 3.1. Let hR,2, (z) pR(2', Tiz,t) t)
gsads (200) px(= (T dz'. By using Proposition 3.2,
can derive the following loR lur-trunsformed SDE:
dR' = [[x 1) + 0² (1) .TRR (IRR,RRO) R log (R", TTRR,O)RR "R]]" dt + (t)dw*
(73)
which corresponds to process (R.) )eEoo,2) R ~ qaaaa (R'o) satisfying the following properties:
let ((::) R x R + R220 denote the joint distribution induced by (R' )eejo.TT then g(R", ) equals to Idara RPto RRt);;
transition density qR (R' ( FRR,,:R ()) (gESE()) qdaua (Iet))
(p)
RR)).
We call the tailored diffussio process (R') an equivariant diffusion bridge.
Proof. Let l (2; =
T;
pR. 0j. then we define
(=
(x(2,t;20)= = /
Qdata 20) dz' Vt [0,7).
p (12,,
(74)
(x, T 0)
24

---

So we can easily show that he (e, zu) satisfies the condition
bx(2,t;20) / prelk' 1 T12, 2o)de' . Vt E [0, V2, 20 E R.
(75)
Then we can use the result of Theorem 3.2 get the h-transformed SDE. By Theorem 3.2, the h-transformed SDE is
dR' = [R + (2) log li
dt (D)d
(76)
Next, we need to find the explicit form V log (R') )").
V log (2, = V (z
(77)
h (z, ti 20)
/ V=Preez', T.pr((',Tiz,, (2 (2', T 20, 0)
Adata (ona(2'|20)) dz'.
(78)
(2, 20)
The h-transformed density function is
hR (T',TTTTi)))) 20) hR (z, 20)
(il',TT|,tt 20, 0) p2( T12,t) (
(79)
Idatal (2' (20)
= PP(22,TTZ
(80)
(2', Ti 20, 0) (2, 20)
Then we have
lRR(2, 1; z0) / V.pR(E',TT) 12,,) PR (=', ; T 0)
Adata (2a 20)20)
V log hR (R(2 (2)
dz'
(81)
V=PR(2',T|es T22,t) 2R (5', T|z E; 20, 0) de
(82)
(2 T22,)
v. log T22 t) (z', T| 20, O)dz'.
(83)
So get explicit form of VR log hv (R', t:: R"):
VRi log hoR (R', t; R = E (R. (TR) VR log (R". , Then the h-transformed SDE becomes
T 20. R*).
(84)
dR' = (R', t) + (t) yn (R3 ( V log PR (R² ,'' t) R']] df + (t)dd
(85)
Since 1 (z, 0; 20) = p TZz, 0)
(R 2o.0) dz' qaana (2' (zo)dd = 1, then
(on(='
Adata (2 (2'120)
( 0) = (2(22' 22, 0)
(2 T; zo. 0)Le (20, 0; 20) = Adiaa (z')z))
(86)
so which means 9R RR. R 0) qamaa (R T RR).. Since the initial distribution R Odma R
QR R = Idata RR)) So we can deduce that
(R.. RT) qu (R. (RR))RRR (R., T Idaal (R" qdda (R. R Idata (R., R²).
(87)
Finally, we need to show that the transition density function satisfies the corresponding symmetric constrains. Since (222"', T; 20)
9tm (z'l2zzz)
is SE(3) i.e.
he R (g)kk)T;; R (=)) = hx T; 20), Vig E SE(3),
(88)
we can show that h(.,t;) ( is also (en(3)-ivvarint Vt E [0, using the following property
ho (2 t; 20) = / |pR(s (1e,t, (2 2o)ds'.
/
(89)
25

---

For any E SE(3), assume p² [ = OR(9)z + tR, where (g is an orthogonal matrix satisfying det(Ox (g) then we have
1 (0)l=), ( [
(90)
= = 10
(91)
(92)
By Proposition 3.1, pr (19 PR (2) let
21 -o) then
Po (p ( (g)l=)) (g)[=))
(93)
/ PR |,
(94)
I pr (21, h 11
(95)
/ pR tul,, t) 12,, tu: 22
(96)
hi (z, t; zu).
(97)
So h(:, t; is SF((3) innvariant A1 [0, T). Then we show (R(2', 11| 20, 0) preserves the symmet- ric constraints:
(0 (g)
(98)
(9)( (11 h
(g) 1 (g)[zo)
(99)
pp( hR x (00
(100)
h (g)
hR (
(2' (2', 6/'tleo (2,t)
(101)
1 t;
(2' f' t0, 20,
(102)
which completes our proof.
B.5
Objective Function of the Equivariant Diffusion Bridge
Lemma B.9. Let XI, X,, Y, Z be random variables. Then the optimal approximation of Y based on (X) 2 is ("",, X.) arg min E f'X.,.... ,x,))1 Y X,,.
Proof. Denote
= (X,,... Xn). We show the following decomposition first:
E| - f(X)1l Y - B| + [[|[|] - -
(103)
We can compute E| Y - f(X) 12 directly by
E Y /(x)11 = EB| Y E| X + B f(x)
(104)
= E Y E X 2 + E
E| 1²
(105)
+ EEYY Y E B X f(X))
(106)
Since
E(Y - E
- - E -
Y Y - f(X))
O,
(107)
have
E Y 1(x)12 = Y - EY E [ f(X)
(108)
+ E(Y - E|YX E f(X))
(109)
= EE| Y -E 2 [X|] - [2
(110)
2 E Y - 11.1I.......
x.11
(111)
26

---

The inequality becomes equality if and only if f ... X,,) = X,)- So the the optimal approximation of based on {X] is E| , i.c.
(X1, X.) mig f(X1, x X,).
(112)
Proposition B.10. The training objective function of Equivariant Diffusion Bridge is:
((0) [x0.0 )-on ,R',,R') .R"1) R.R' (R', R, T ((t)]lvo B log (21, R t)]?
(113)
where T). Then the optimal parameter 0 arg min ((0) satisfies
Va- (4;4)= = (RE, Vr log (R.", ,T)R' ())R
(114)
Proof. Let L(0) = ((-2 where
(0) ( (Rto R R.R (la (R (R))) V logpR (21, T D)
(115)
Then by Lemma B.9, Vo RR, L; z0) Egg (RT, R R RR log (R. T () R R*] mini-
mize C' (0), [0, T.. Since X(t) 2 0, the optimal parameter 0 arg min [(0) satisfies
Vø- (R', 20) = Egg RRAR)))AR)) R logpR (R. T R R* At [0,
(116)
B.6 Proof of Theorem 3.4
Theorem B.11 (Chain Equivariant Diffusion Bridges). Let (R())c(o.T) Ne(N-1] denote a series
of N equivaraint diffusion bridges defined in Theorem 3.3. For the i-th bridge (R) [, if we set
4+1 (20)
(I) (2, 20) f pR (=', Ti 1) then the joint distribution qR (Rg, R Rt ....... R induced by {(R' )eeo, M equals to
Px (='2""2 (z..) dz'; (2) Ry Lrng (fe),, R R 4O < N.
qray R ...., RN). We call this process a chain of equivariant diffusion bridges.
Proof. By Theorem 3.3, the transition density function of (R)rEjo,T) satisfies o (R. R?)
( [R.), vO < i < N 1. The ground truth probability density function has the decompo-
sition Lag (Ro) I 9 [ Ri-1) Then we use the boundary condition. R ~ eiag (F()) R
R V0 < N, have
N
QR R R
R
9%
(R.)
II
OR
(R
R
(117)
0% (R!)
II
o
[
R
(118)
I=
R
R
R
(119)
R
RT
R.T.T. R
(120)
So the joint distribution UR (R., RT, RT",
R induced by (RR Jejo)rhhee equals to
Utaj R .... RNN.
27

---

B.7 Objective of the Chain of Equivariant Diffusion Bridge
Proposition B.12. The training objective function of the Chain of Equivariant Diffusion Bridge is:
(() B (za..... Aaag R RN),,, ((t)]lvo(R) ,t; (;) Ray log log ( T I)???
(121)
where ~ U(O,N x xT,
1i,, - -t-l=ttixXip t (R" (' T',s 2i, 0). Then the optimal
parameter 8* = arg min (t) satisfies
Vg- Vg. (R. 2;2 = E v/ (RF TiR (R)) Vv; log . R) R.
(122)
Proof. Let ('( EuuuONTT ((t) where
(0 = (2o..... J /2 AN ,.
Vo Vo (R. 1 zi) V Ry log log 1 T 1']1, (123)
where ~ 21(10, N x T), -=
R de (R! = T; 0). Then by Lemma
B.9, Vø- (R, 11 2 gf (RT. ..RR) VR: log pie (R T RR. R.] minimize (,))), VI
[0, NT.]. Since X(t) 2 0, so the optimal parameter 0" arg min [()) satisfies
Vg. RR 1;2 = E vie (R. (R.,, minn an)) VR log
PR R l R) R.).
T
(124)
B.8 Proof of Theorem 3.5
In this paper, we choose the Brownian bridge as our matching target. Let's first recall the definition and properties of the Brownian bridge. A Brownian bridge (Xx)ee(orr) with the initial position Xo and the terminal position X7 is given by the following SDE
XT X,
dx,
dt + adW.,
(125)
-
where W, is the standard wiener process. The solution of the Brownian bridge is given by
X ~ N ((1 - x) + + o cet(1-t))
(126)
Next, we recall the definition of the KL Divergence:
Definition B.13 (KL Divergence). The relative entropy (or Seclbback-Leiblee Divergence) Kf between two probability density functions fxx and g(x) is defined by:
KL(fllg) -f f(x:) log
f(x) dr. g(x)
(127)
In general, let P and Q be two probability measures space X. Assume TP is absolutely continuous with respect to Q then the Kullback-eiihlee Divergence between and Q is defined as follows
dP
KL(P
(2)
log
dP,
(128)
aQ
where dQ
dF is the Radon-Nikodym derivative P with respect to Q
When we need to compute the KI. divergence between the path measures associated with two SDEs, the Girsanov's ss theorem [53] is useful tool to get the Radon-Nikodym derivative between the two measure. The precise statements are as follows.
Theorem B.14 (Girsanov's Theorem). Let W, be a d-dimensional Wiener process defined on (R,F, IP)., Let H, be a d-dimensional F, Fa-adapted process such that
H 2dt < x0, a.s.
(129)
28

---

Define
= exp (L. HH.,d,, .dW, 2
L [H,]] [
(130)
Assume Ze is a martingale. Define the probability measure Q on Fi by
dQ = Zr.P.
(131)
M, = W, So H,ds, then M, is d- dimensional Wiener process with respect to Q
In practice, the condition that Ze is martingale is hard to vertify. So the condition is often replaced by the Novikov's condition
exp /'"
]]
|H ||a.|/dd
< <00.
(132)
For more discussions and applications of the Girsanov's theorem, please see [86, 74, 28). Now we can give the precise assumptions and proof of Theorem 3.5 using properties of Brownian Bridge and Theorem B.14.
Theorem B.15. Assume (R') (iE(N) is sampled by simulating prior SDE on geometric states dR
~HH dt + odWi. Let μú denote the path measure this prior SDE when t E [iT,, (i + 1Y7]. Building upon (EEN) lel (UR]NE(N-1) denote the path measure of our chain of equivariant diffussion bridges. Assume {ukz (uR-eENN-1) is composed of chain of the Brownian Bridge. Assume the total time is NT = 1. Under the following assumptions:
+)) Rd TR is scalar function with continuous second-order partial derivative; The drift function is Lipschitz: there exist constant L such that
IV (x) V Hi (v)ll = L. | yll. Va, TRd.
H*)() satisfies (=)) < K(1 |x||), Vz
E|| < M, Vt NT]}
ltt)=
E R is continuous function on E [0, NT]:
The Novikov's 's condition:
NT V (W))
]
exp
< <00;
2
The function H2 satisfies the following regulaity condition: there exist a constant C such [VH (2) < C, Vo Rd,
that H
( (x)
then we have lim max (L N
Proof. Let p= be the probability density function associated with the ground truth SDE di² f (R', (od R R Let (R))
denote series of N equivaraint diffusion bridges defined in Theorem 3.4. Then by theorem 3.4, dre (R., T .T...... R induced by (RR)) ( equals to pk (R. R ....... R Additionally, the conditional probability density function qR (R. R R.) for S < (i + 1)T, associated with the Brownian bridge
R R --
di' adW, where t' = -t- -- iT. Then by the chain rule of KL divergence
JR:
(133)
KT((411(4) KL(p; (R (+11T R R llok ( (+1)T R
(134)
E
(R R
[KL( (- (++I) R R )llk (- (++)) R.T))
(135)
29

---

Since p (R 11+ R dz (R (++)) )
we have
KL(Millu p (ft R
( (1+1)T R P (++1)T R ))
(136)
Since the prior SDE is time homogeneous, we can only consider the case i 0 without loss of generality, Let be the path measure of the Brownian motion a W' on space R. Since the condition of Theorem B.14 14 is satisfied, then we can use Theorem 1. and get
d/6( R = exp (- (-" (H) dw
-
[VHHR(OW")/Paa dut(1R99..
20
(137)
Then we can use the Ito's formula (Theorem B. to simplify the expression
d (R)) exp (# ()()) + ( (V" + (1 do( (R)).
(138)
To simplify our notation, we denote
1 ZT = exp ( ( -H + / (V" / 1 ( )| )t
(139)
Let F, g be measurable functions on C(O, 7), Rd, respectively. Then by the disintegration of Wiener measure into pinned Wiener measures (path measure of the Brownian Bridge), we have
1R [Fg(o [F = /Bur |Bu Re (r (Zrlg))pp
E MG [20)
Rr,
(140)
where pT (T is the transition density function of a W Let F = 1, we get
f B Pec- (r=)) (rllg 1zo
dr / (f
(141)
So we have Ey fo nt -x)=))))) [ZT] = p5 (IT Let g= = 1, then we get
f F [ P1 [F
(142)
So we can conclude that
d (-IRR.RR
Pr( R Ph (R Rd)
exp( + (Hip (FR))
exp (1.50 (v
V (-)l))a)
du( [RT,RT)
exp(
./2 (H (R*)
(143)
Note that ( R., RT)
vt- R RT). Now we calculate the KL divergence by
KL(HÔllp(R)
Egg (R² Re)
(L RT RR 1 (-R" R$)]
(144)
+/ (H (Ro) exp
Pr ( R
CT
S MA (R)
log
+
(145)
Ph (R Ro)
(HH (RT)
2
exp
(
)
CT
= E pg (
log
E 0²
H (R
E
H
(R )
+
PG (
(146)
M() - (T))
CT
E (R) (R. Ip R R +
+
(147)
0²
2
< h(0) h(T) CT
(148)
0²
2
When N - 00, T = 1
- 0. Since h(t)) continuous by our assumption, then KT.
N
0.
30

---

C Derivation of Practical Objective Function
In this subsection, we show the implementation details of our framework. We set T in all the experiments.
Matching objective. We design the SDE on geometric states in Proposition 3.1 to be:
dR' = odW with trunsition density The explicit form of the objective is
p=( 1) N o² (1' - 1)1)
(149)
21 R
R log pR( 1 1) Va ol
(150)
02 - 1)
Then the h-transformed SDE becomes
R R dt odW.
dR'
(151)
1
which is known as the Brownian bridge. The corresponding h-transformed density is
@R (R', f22,, 20, 0) Ni(zi + t) oo((11 t)1).
(152)
In practice, we do not use OR (R., 21 1; 2g. 0) 8R 20) our initial distribution. We use 2R. (R", 0| 1; zo, 0) (= I) instead. Since the solution of the Brownian bridge is given by
R (1 t)R tR +o/t(1 - t)
(153)
where (0, I), then the marginal distribution of R becomes (V - =0 + (21, - (1-1)o11)
We use this distribution to sample geometric state R in the training stage.
Trajectory guidance. Similarly, we set = N Pi (2141. T ,t') NRRR (T-'t') when we use the trajectory guidance. So the h-transformed SDE becomes
R R - dt addW'
dR
(154)
which is a Brownian bridge with T N. Then associated density function is
T-t
2 (T t)
dh (R" 2141, T; 2,, 0) = N T
21 o.
(155)
T
Additionally, we set o; decays linearly with respect to f. i.e. o, = =Nyio, where o is a hyperpa- rameter. Again, in training stage, we set dx (R), 0/21, 10 0) = (222,,01)) as initial distribution,
N
and the terminal distribution is d/R (R. 0 31, 1; 20, 0) = N221,, of I), which is same as the initial distribution of the next bridge.
Sampling Algorithm We use the ODE-based method to generate samples at inference time. After the training process, the neural network is trained as described in Algorithm 3 and Algorithm 4. When the network trained without trajectory guidance, we simulate the following ODE to generate samples:
dR'
= vo R²), Rd ~ qdara (1 E [0,7].
(156)
di
When the network is trained with trajectory guidance, we solve the following ODE to generate samples:
d = vo (R RL+JT,, RD qamaa (1*)), E [0 x T].
(157)
di
Denote a black box ODE solver by Solver(v, t). Solver(v, t) takes a vector field v and a time point as inputs, then returns the solution of the ODE
dx dt
v(X (10) X =00
(158)
at the specific time I, i.e. Solver(v, I) X Combining all the above design choices, we have the following algorithms for sampling our Geometrie Diffusion Bridge (Algorithm 5) and leveraging trajectory guidance if available (Algorithm 6).
31

---

Algorithm 3 Training
Algorithm 4 Training with trajectory guidance
repeat
I: repeat
2: (2o, 21) ~ qanes ( (1")
2: (Zu...., 2N) ~ qeml R
R
3: t=
3:
t~u N x) 11,
= -
4: (0,
c~N((,
5: R' t.s T-I 20 +
(1(
de
= = I 2: +
(P(T-1))
7 ove
6: Take gradient descent step on
6:
Take gradient descent step on
VoA(1)
D- (T--))
vo (o(RR44200) 1:20
Vet(t)
$-11RR -7(1-1)
Vol R 1;24
until converged
7: until converged
Algorithm 5 Sampling
Algorithm 6 &Sampling with trajectory guidance
Initial geometric state 20 ~ Gaazs (R" )) a Require: Initial geometric state 20 qlam (R')),
Require:
trained neural network Vs. a numerical ODE solver Solver(v,
trained neural network Ve. numerical ODE solver Solver(v,()
I: R
I: R
2: NT Solver (,, (R:++! NT)
R' (RI,
Ensure: R.
Ensure: R
D Experiments
D.1 Equilibrium State Prediction
Dataset. QM9 [79] is a quantum chemistry benchmark consisting of 134k stable small organic molecules, which has been widely used for molecular modeling. These molecules correspond to the subset of all 133,885 species out of the GDB-17 chemical universe of 166 billion organic molecules. In convention, 110k, 10k, and 11k molecules The geometric conformations that are minimal in energy are provided in the QM9 dataset. The equilibrium conformation and its relative properties are all calculated the 33LYP/6-31G(2df)) level of quantum chemistry.
used for train/valid/test sets respectively.
Molecule3D [116] is large-scale dataset curated from the PubChemQC project [67, 71], consisting of 3,899,647 molecules in total, 2,339,788 788 molecules in training set, 779,929 molecules in the validation set, 779,930 molecules in the test set, and its train/valid/test splitting ratio is 6 2 2. For each molecule, the 2D atom graph, the 3D equilibrium geometric conformation, and four extra properties provided. In particular, both random and scaffold splitting methods adopted to thoroughly evaluate the in-distribution and geometric state is generated by using fast and coarse force field [73, 52] and geometry optimization is conducted to obtain B3LYP/6-31G level DFF-calculated equilibrium geometric structure.
performance. For each molecule, an initial
Baselines. We comprehensively compare our GDB framework with previous equilibrium confor- mation prediction methods. Following [111], we use DG and ETKDG algorithms implemented by RDkit our fundamental baselines. The benchmark [116] used the DeeperGCN-DAGNE framework [60] which proposed a deep graph neural network architecture to predict 3D geometric conformation of the molecule based on its 2D graph structure, and got impressive performance the Molecule3D dataset. GINE [39] proposed a method for pretraining GNN to improve the performance and capacity of GNN. GATv2 [10] proposed dynamic graph attention mechanism and improved the performance of the graph attention network on several tasks. GPS [80] proposed a general framework that supported multiple types of encodings with efficiency and scalability guarantees in both small and large graph prediction tasks. GTMGC [111] proposed a novel neural network based
(GT) [118 (118, 66, 119, 65) to predict the equilibrium conformation of the molecule in 3D based on its 2D graph structure.
Metric. Following [116], three metrics are adopted to evaluate predictions of equilibrium states: (1) C-RMSD: given prediction R [f. J.W, which is rigidly aligned to the ground-truth R {rī 1 by the Kabsch algorithm [44], Root Mean Square Deviation between their atoms is calculated,
ie.. C-RMSD R*)
E |l;;
r/ 11: (2) D-RMSE: based on Å and R = {r;}
interatomic distances can be calculated, i.c., {d. N and {d N. Root Mean Square Error be-
32

---

tween these distances is calculated, i.e., D-RMSE((4) (41))))) V Ny E (d; d));: (3) D-MAE((d)). (d; (d きे Ei=1d: d;
Settings. In this task, we parameterize vo tt; R") by extending a Graph- Traph--Taanfommer based equivariant network [92, 63] to encode both time steps and initial geometric states as conditions. For training, we use Adam W as the optimizer, and set the hyper-parameter e to 1e-8 and ( to (0.9,0.999). The gradient clip norm is set to 5.0. The peak learning rate is set le-4. The batch size is set to 512. The weight decay is set to 0.0. The model is trained for 500k steps with 30k-step warm-up stage, After the warm-up stage, the learning rate decays linearly to zero. The noise scale a set to 0.5. For inference, we use 10 time steps with the Euler solver [12]. All models are trained on 16 NVIDIA V100 GPU.
D.2 Structure Relaxation
Dataset. Open Catalyst 2022 (OC22) dataset [105] is widely used dasaset, which has great significance for the development of Oxygen Evolution Reaction (OER) catalysts. Each data in the dataset is in the form of the adsorbate-catalyst complex. Both initial and adsorption states with trajectories connecting them are provided. The dataset consists 62,331 Density Functional Theory (DFT) relaxations trajectories, and about 9,854,504 single-point DFT calculations across range of oxide materials, coverages, and adsorbates. The training set consists of 45,890 catalyst-adsorbate complexes. To better evaluate the model's performance, the validation and test sets consider the in-distribution (ID) and distiriiitttionn (OOD) settings which use unseen catalysts, containing approximately 2,624 and 2,780 complexes respectively.
Baselines. Following [105], we choose strong MLFF baselines trained on force field data for a challenging comparison. Spinconv [94] introduced a novel approach called spin convolution to model angular information between sets of neighboring atoms in graph neural network and got impressive performance in molecular simulation tasks. Gemnet [32] proposed multiple structural improvements for geometric GNN with theoretical insights, which significantly improved the experimental perfor- mance as well. Based on Gemnet's framework, Gemnet-OC [34] modified the architecture of the network and improved the experimental performance on more diverse tasks.
In [105], there are still other baseline setting. [105] introduce a large-scale dataset Open Catalyst 2020 (OC20), which consists of 1,281, 1,281,040 Density Functional Theory (DFT) relaxations and $64,890,000 000 single point evaluations to help training the baseline model. [105] presented baselines using both OC20 and OC22 data in training stage and baselines using only OC20/OC22 for comparison.
Metric. Following [105], we use the Average Distance within Threshold (ADwT) as the evaluation metric, which reflects the percentage of structures with an atom position MAE below thresholds.
To be more precise, the ADWT metric across thresholds ranging from 0.01 to 3 0.5. in increments of Á. The computation of ADwT metric count the percentage of structures with an atom position MAE below the threshold.
Settings. In this task, We parameterize vo (R 4; RD) by using GemNet-OC [34], which also serves as a verification that our framework is compatible with different backbone models. For training, we use AdamW the optimizer, and set the hyper-parameter gradient clip norm is set to 10.0. The peak learning rate is set to 5e-4. The batch size is set to 64. The weight decay is set 0.0. The model is trained for 200k steps. After the warm-up stage, the learning rate decays linearly to zero. The noise scale o is set to 0.5. The trajectory length is set to
to 1e-8 and (81, (2) (0.9,0.999). The
N 10. For inference, we also use 10 time steps with the Euler solver [12]. All models are trained on 8 NVIDIA A100 GPU.
33

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section 3, 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We have discussed several future directions in Section 3 and 6
34

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions and complete proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 4 and Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
35

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: The code and model checkpoints will be publicly available after the submission is acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //wpss.cc/publcc//eeceee////p/// healliic/guieeeeeeee /CodeSubmiii httpa//gggidee//eeuuutiiiiii ssi ssionPolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
36

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section and Appendix D. Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: There exists little randomness in all the experiments of this submission, which means that results of using different random seeds are almost the same.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section and Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
37

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //tpp//////urp..../// cc cc/public//eeeeccccccccccccccict wttpeii//thiiiuieellleee
Answer: [Yes]
Justification: The research in this work conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
38

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
39

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
40


Reviews:
[{"review_id": "PxQHToTFl4", "reviewer": "Reviewer_Sjdg", "summary": "In this paper, the authors introduce a Geometric Diffusion Bridge (GDB) framework, which aims to predict the evolution of geometric states in complex systems accurately, crucial for fields such as quantum chemistry and material modeling. Traditional methods face computational challenges, while deep learning approaches lack precision and generality. The authors use Doob’s h-transform to construct an equivariant diffusion bridge. By applying Doob’s h-transform, the authors adjust the SDE to ensure that the process starts from an initial geometric state and is conditioned to reach a target geometric state. This ensures that the transformed process respects the symmetry constraints of the geometric states, leading to more accurate and physically meaningful predictions.", "questions": "See above", "limitations": "No limitations are addressed in the paper by the authors", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The framework utilizes an equivariant diffusion bridge derived from a modified Doob’s h-transform. This ensures that the diffusion process respects symmetry constraints, making the predictions more robust and reliable.\n+ The paper provides a theoretical framework analysis about preserving symmetries and accurately modeling evolution dynamics.\n+ Experimental evaluations show that GDB is better than state-of-the-art approaches in various real-world scenarios, including equilibrium state prediction and structure relaxation tasks.\n+ The framework achieves significant error reduction compared to strong baseline models, particularly in challenging tasks such as structure relaxation in the Open Catalyst 2022 dataset", "weaknesses": "- The framework, especially when leveraging trajectory data, might introduce significant computational overhead. The simulation-free matching objective is designed to be efficient, but the overall framework’s computational demands might still be high\n- Some mathematical notations and definitions in the paper could be made clearer. For instance, explicitly defining all variables and functions used in the modified Doob’s h-transform and constructing equivariant diffusion bridges would improve readability and understanding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing both the theoretical analysis and practical effectiveness of our GDB framework. We also appreciate your suggestions which can improve our work further. Here are our responses to your questions.\n\n>**Regarding the computational cost of our GDB framework**\n\nThanks for the question. Indeed, our GDB framework is efficient. Here, we discuss the reasons in the following two settings: \n\n- Setting 1: Bridging geometric states **without trajectory guidance**. In this setting, our training objective (Equ.(5) and Algorithm 1) and inference process are similar to typical diffusion models, which do not bring additional computational overhead. \n\n- Setting 2: Bridging geometric states **with trajectory guidance**. In this setting, trajectory data is used during training. From Equ.(6) and Algorithm 2 in our paper, it can be seen that the loss calculation at each iteration is still efficient (compared to Algorithm 1) with no computational overhead. For inference, we simulate the ODE in time interval $[0, NT]$. In practice, we set $T=1/N$ so the time interval is still $[0,1]$, which is kept the same as setting 1. The intuition is that we split the total time interval $[0,1]$ into $N$ parts for our Chain of Equivariant Diffusion Bridges with $N$ chains. Since the total length of the time interval is the same, predicting the future geometric state from the initial geometric state by using our model trained with trajectory guidance does not bring significant computational overhead.\n\nWe find your question highlights the need for a detailed table to introduce the inference algorithm. We will incorporate this in the next version. We hope this addition clarifies the efficiency of our framework and highlights the improvement in model performance.\n\n>**Regarding the mathematical notations and definitions**\n\nThanks for your suggestion of explicitly defining all variables and functions. We will revise our paper and add a notation paragraph to elaborate on all mentioned notations and concepts in Sections 3.1 & 3.2 for improved clarity and readability. \n\nWe thank you again for your efforts in reviewing our paper, and we have replied to each of your concerns. We sincerely look forward to your re-evaluation of our submission based on our responses.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "Cw5ZHsPvpj", "reviewer": "Reviewer_znC7", "summary": "This paper proposes a type of diffusion model that captures the evolution of geometric states. The model is characterized by a diffusion SDE that couples the initial state with the target state, in the middle of which trajectory guidance is enabled when such data present. The framework is designed to yield equivariant density similar to other geometric diffusion models. Experiments on equilibrium state prediction with or without trajectory data have been performed to verify the applicability of the proposed approach.", "questions": "Q1. The evaluation protocol on QM9 and Molecule3D, especially to compare with direct prediction approaches, is not a common practice. A more convincing benchmark protocol would be to compare with methods such as GeoDiff [1] on molecule generation tasks since they are also generative models. Since the paper is positioned to tackle generative modeling, the experiments should also be designed to align with the goal.\n\nQ2. Could the authors provided detailed sampling algorithm this approach adopts? If the model uses sampling approach similar to other diffusion models, there should be related discussions on sampling steps/sampling time the method consumes.\n\nQ3. A more reasonable baseline would be to directly apply existing bridge models (e.g., [2]) to the current task by switching the backbone to the one this paper adopts. This would help the audience understand the unique contribution of this work since both bridge models and equivariant (geometric) diffusion models have been proposed in literature.\n\nQ4. Ablation studies such as investigating the importance of preserving equivariance of the density modeled should be included. This would help justify the necessity of the proposed components.\n\n\n[1] Xu et al. GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation. In ICLR'22.\n\n[2] Zhou et al. Denoising Diffusion Bridge Models. In ICLR'24.", "limitations": "There seem to be no discussions on limitations in the paper. It would be better to discuss potential limitations from perspectives such as scalability and sampling time.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The distinction between existing works has been elaborated in Table 1, which is clear.\n\n2. The method is designed with an option to leverage additional trajectory data, which is quite interesting.", "weaknesses": "1. The experimental setup and comparison with baselines on equilibrium state prediction is a bit troublesome which requires more clarification or additional comparisons. Please refer to Q1.\n\n2. The presentation is a bit unclear. Please refer to Q2.\n\n3. Additional baselines may be considered. The baselines selected in the paper are not closely connected to the proposed approach. See Q3.\n\n4. Missing ablation studies. In the current shape it is unclear where the performance gain comes from. See Q4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for spending time reviewing our paper. We would like to first address your misunderstanding by **clarifying the task of our interest**: to capture/predict the evolution of geometric states, i.e., *predicting future states from initial states*. This goal has been carefully stated at the very beginning of our paper (please refer to (1) line 1 in Abstract; (2) line 22 in Sec 1 Introduction; (3) lines 83-84 in Sec 2.1 Problem Definition). We also provide its formal definition. See lines 84-91. \n\nBased on these statements, we would like to argue that **our scope is different from conventional molecular generation problems**, such as the one introduced in GeoDiff mentioned in the review comments. In Sec 3.1 of GeoDiff, the definition of molecule generation task is \"Given multiple **graphs $G$**, and for each $G$ given **a set of** conformations $C$, ..., the goal is learning a generative model $p_\\theta(C|G)$\". From the definition, we can see: \n\n- The input in GeoDiff is the molecular graph, while in ours, the input is the geometric state of a system at a given time $t_0$.\n\n- The output of the molecule generation task is a set of sampled conformations $C$ from Boltzmann distribution, while the output of our task is the geometric state at a given time $t_1$ evolved from time $t_0$.\n\n- The goal of GeoDiff is to learn a generative model for the underlying Boltzmann distribution of molecular conformations, while our task is to capture the evolution of geometric states over time. \n\nDue to these differences, our task has **several unique difficulties, including how to leverage trajectory data and preserve coupling between states over time while satisfying SE(3) at the same time**. Our work tackled these challenges and provided promising results. \n\nWe guess that the misunderstanding may partly be due to the lack of detailed task descriptions for the related works. We hope the above clarifications can help the reviewer better understand our task definition and our goal. We will update these discussions in the next version of our paper to avoid ambiguity and misunderstanding.\n\n>**Regarding the experimental setup and evaluation protocol (Weakness 1 & Q1)** \n\n**First, we believe our evaluation protocol aligns well with our task**. To measure how accurate a model is in predicting a system's future geometric state from its initial geometric state, we follow existing works in this direction [1, 2, 3] to use C-RMSD, D-MAE, and D-RMSE on QM9/Molecule3D and ADwT on OC22 which comprehensively measure the position/distance error between predicted geometric state and ground-truth geometric state. \n\nThe reviewer suggested that we should follow GeoDiff and use their proposed metric. We are afraid that it is not a proper choice. GeoDiff (In Sec.5.2) calculates metrics between two **sets** of conformations, which aims at measuring the difference between the learned distribution and ground truth distribution rather than how accurate a target geometric state is predicted when specifying the initial geometric state.\n\n**Second, we believe strong and latest baselines in this direction have been included**. We followed the most recent works (e.g., GTMGC in ICLR24) and benchmarks (e.g., OC22) to compare our approach. The reviewer suggested that we should use GeoDiff as the baseline. But again, it is not a proper choice as GeoDiff cannot leverage trajectory guidance if provided and cannot take the initial geometric state as input to predict the target geometric state, which has been clearly discussed at the beginning of our response. \n\n>**Regarding the details of our sampling algorithm (Weakness 2 & Q2)** \n\nWe have already discussed the sampling/inference method (ODE solver) of our GDB framework in lines 261-265 and presented details of the ODE solver (here we use the basic Euler solver) and sampling steps (10 steps in total) in Appendix C. Following your suggestion, we will provide an additional pseudo-code of our inference process in the next version of our paper.\n\n>**Regarding additional baseline comparisons and the unique contribution of our work (Weakness 3 & Q3)** \n\nFirst, we would like to clarify the reviewer's concern: \"The baselines selected in the paper are not closely connected to the proposed approach.\" As stated in our responses to Weakness 1 & Q1, our task of interest is to predict the future geometric state of a system from its initial geometric state. In both Sec 1 Introduction and Sec 5 Related Works, we carefully review existing approaches for this task, including experimental approaches, traditional computational methods, direct prediction, and machine learning force fields. Although \"both bridge models and equivariant (geometric) diffusion models have been proposed in literature\" as the reviewer mentioned, there exist few works using generative modeling techniques to comprehensively investigate our task of interest due to the difficulty of unified satisfying the three desiderata (Coupling Preservation, Symmetry Constraints and Trajectory Guidance, lines 132-145). Following existing works, we use strong baselines of this task for comparisons, which is common practice.\n\nFollowing your suggestion, we compare our approach with Denoising Diffusion Bridge Models by switching the backbones on OC22. Our GDB framework significantly outperforms this new baseline by 6.37% on average ADwT. Moreover, our approach can further leverage trajectory guidance to achieve better performance, i.e., 7.93% improvement compared to this baseline. We will update these results in the next version of our paper, which we believe can help the audience better capture our GDB framework's effectiveness.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "2TQfyxw16c", "reviewer": "Reviewer_U4NU", "summary": "This paper proposes a generative model for bridging initial and target geometric states using diffusion bridge. This work introduces an equivariant diffusion bridge based on equivariant transition kernels for symmetry constraints. The proposed method was validated on diverse settings including simple molecules and adsorbate-catalyst complex, outperforming previous MLFF baselines.", "questions": "- What is the reason for using deterministic process (i.e., probability flow ODE) instead of the original stochastic process? Does ODE results in better performance?\n- Is GDB scalable to geometric states of high dimensions? While analysis on this may not be necessary, it could strengthen the work.", "limitations": "While the paper discusses future direction for the proposed method, specific limitations of the work is not specified. One potential issue might the scalability of GDB as the model has transformer architecture, and other issue could be long inference time which is a typical problem of diffusion models.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The motivation of using diffusion bridge to bridge initial and target geometrical states is reasonable. \n- Using diffusion bridge model for equilibrium state prediction and structure relaxation is novel to the best of my knowledge, and the paper shows that GDB significantly outperforms previous methods with diverse datasets.\n- Equivariant design of bridge process is based on solid theory.\n- The paper is well written except for some missing relevant works on diffusion bridge.", "weaknesses": "- Related works on diffusion bridges or diffusion mixtures were not discussed. Diffusion bridges has been studied in [1,2,3,4] with applications to molecules, graphs, point clouds, and images, and more recent works have studied general framework for diffusion bridges [5, 6] which is worth discussing. While GDB has a contribution for using diffusion bridges in new tasks, discussing related works and clarifying the novel contributions is necessary in particular for strengthening the contribution of this work.\n- Contribution seems limited as using diffusion bridge as generative modeling was already studied [1,2,3,4], in particular deriving diffusion bridges using Doob's h-transform. Designing an equivariant diffusion process (not necessarily bridge) specifically in SE(3) group has been covered in [7,8, 9]. What is the difference of designing equivariant diffusion bridges compared to equivariant diffusion processes?\n\n[1] Peluchetti, Diffusion Bridge Mixture Transports, Schrodinger Bridge Problems and Generative Modeling, JMLR 2023\n[2] Liu et al., Learning Diffusion Bridges on Constrained Domains, ICLR 2023\n[3] Wu et al., Diffusion-based Molecule Generation with Informative Prior Bridges, NeurIPS 2022\n[4] Jo et al., Graph Generation with Destination-Predicting Diffusion Mixture, arXiv 2023 \n[5] Albergo et al., Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, arXiv 2023\n[6] Shi et al., Diffusion Schrodinger Bridge Matching, NeurIPS 2023\n[7] Xu et al., GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, ICLR 2022\n[8] Xu et al., Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023\n[9] Yim et al., SE(3) diffusion model with application to protein backbone generation, ICML 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the motivation, contributions, and theoretical analysis of our GDB framework. We also appreciate your suggestions which can improve our work further. Here are our responses to your questions.\n\n>**Regarding discussions of related works.**\n\nThank you for listing these related works [1-6] on diffusion bridges. We agree that adding more discussions on these related works would help the audience to understand our contributions better. We will carefully cite all these works, compare them with our work regarding their contributions, and update these discussions in the next version of our paper. \n\n>**Regarding our contributions.**\n\nOur work cannot be considered a simple extension or follow-up work of constructing diffusion bridges with SE(3) symmetry constraint ([1-4,7-9]). As described throughout the whole manuscript, the task we tackle is to capture the evolution of geometric states. This task has unique difficulties compared to conventional generative models, including **how to leverage trajectory data and preserve coupling between geometric states over time**. Simply combining existing diffusion bridge approaches with equivariant diffusion processes cannot meet the requirement.\n\nTo the best of our knowledge, our GDB framework is the first approach to leverage the characteristics of diffusion bridges for flexibly incorporating trajectory guidance. We conduct theoretical analysis to show that our approach can preserve the joint distribution of the trajectory (Theorem 3.4) with strong expressiveness guarantees (Theorem 3.5). Empirical results on the large-scale real-world benchmark OC22 also verify the superiority of trajectory guidance enabled by our framework.\n\nMoreover, our GDB framework preserves the coupling of geometric states, which is also crucial and necessary in modeling the evolution of geometric states (please refer to line 132 in section 3 and Theorem 3.3 for more details). SE(3) diffusion models [7,8,9] are restricted to transport from the standard Gaussian distribution to the target distribution of molecules or proteins, which are obviously unable to preserve the coupling of geometric states. Previous works on diffusion bridges [1,2,3,4,5,6], which the reviewer has mentioned, do not provide guarantees on coupling preservation, especially for distributions of geometric states. Our paper provides detailed derivations and theoretical analysis for the coupling preservation of geometric states, which we believe add value to the communities of both geometric deep learning and diffusion bridge-based approaches.\n\nWe will carefully add the above discussions to the next version of our paper to let the general audience better understand our contributions.\n\n>**Regarding the usage of the ODE sampler.**\n\nAs stated in lines 259-265 in our paper, we leverage ODE solvers for inference due to efficiency considerations. This is motivated by [10], which has also become a common choice in the literature of diffusion models for efficiency sampling. Our framework can also be implemented by using other advanced sampling strategies, which we leave as future works.\n\n[10]. Song, Yang, et al. \"Score-based generative modeling through stochastic differential equations.\" ICLR 2021.\n\n>**Regarding the scalability of GDB to geometric states of high dimensions.**\n\nThank you for the suggestion. In Table 2 of our paper, we briefly introduce the types and scales of datasets we used, covering both medium-scale and large-scale real-world benchmarks. Here we present more details on the data dimensions of these benchmarks:\n\n- QM9 is a medium-scale dataset that consists of \\~130,000 organic molecules (see line 277 in section 4.1). Molecules in this dataset typically consist of 9 heavy atoms C,O,N,F (not counting hydrogens), so the dimension of geometric states is $9\\times 3 = 27$.\n\n- Molecule3D is a large-scale dataset curated from the PubChemQC project, consisting of 3,899,647 molecules in total (see line 279 in section 4.1). In the Molecule3D dataset, each molecule includes 29.11 atoms on average. The dimension of geometric states is around 87 on average. \n\n- The Open Catalyst 2022 (OC22) dataset consists of 62,331 Density Functional Theory (DFT) relaxations, which have great significance for the development of Oxygen Evolution Reaction (OER) catalysts. Each data is in the form of the adsorbate-catalyst complex (see line 305 in section 4.2 for more details), consisting of hundred-scale atoms with periodic boundary conditions being considered. The dimension of geometric states in this dataset is thus in hundred to thousand scales.\n    \nFrom the above statistics, we can see that our GDB framework is general to be applied to accurately predict the evolution of both low-dimensional and high-dimensional geometric states. Bridging geometric states of higher dimensions, such as protein conformation states, indeed has great significance in various scientific problems like transition state discovery. In our ongoing work, we successfully extend our GDB framework to this higher-dimensional problem with satisfactory preliminary results, which demonstrate the promising scalability of GDB for bridging geometric states of higher dimensions.\n\n>**Regarding the discussion of limitations**\n\nThank you for the suggestion. For the sake of generality, we do not experiment with advanced implementation strategies of training objectives and sampling algorithms, which leave room for further improvement. Besides, the employment of Transformer-based architectures may also limit the efficiency of our framework. This has also become a common issue in transformer-based diffusion models. We will organize these discussions into a single Limitation section and update it in the next version of our paper.\n\nWe thank you again for your efforts in reviewing our paper, and we have replied to each of your comments. We look forward to your re-evaluation of our submission based on our responses and updated results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "YhMppRX4vs", "reviewer": "Reviewer_dV6w", "summary": "The paper introduces the Geometric Diffusion Bridge (GDB), a novel framework designed to generate the evolution of geometric states in geometric (coordinate) systems. GDB uses a diffusion bridge connecting initial and target geometric states with equivariant transition kernels, preserving symmetry and joint state distributions. Furthermore, GDB can use a chain of equivariant diffusion bridges to leverage trajectory data for more accurate dynamic modeling.", "questions": "See weaknesses.", "limitations": "The authors note the need of exploring better implementation strategies for their framework to enhance performance.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The presentation of theorems in Section 3.1 is clear and straightforward, establishing a solid theoretical foundation for GDB. The authors effectively derive theorems and integrate them with point cloud states.\n- GDB demonstrates strong performance across various tasks, including QM9, Molecule3D, and OpenCatalyst IS2RS.", "weaknesses": "I have no complaints regarding the technical and experimental sections, as they are well-written. However, I wonder existing works, such as [1] and [2], also use diffusion bridges over molecular data. What advantages does your approach have over theirs?\n\n[1] Diffusion-based Molecule Generation with Informative Prior Bridges. Lemeng Wu, et al. NeurIPS 2022.\n\n[2] DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization. Danyeong Lee, et al. AAAI 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing both the theoretical analysis and practical effectiveness of our GDB framework. We also appreciate your suggestions which can improve our work further. \n\nOur proposed method has the following advantages compared to your list of works [1, 2]. \n\n- First, our proposed method can make good use of trajectory data during training (trajectory guidance), which [1, 2] cannot. Trajectory data provide valuable insights into the evolution of geometric states. We conducted a theoretical analysis, which showed that our approach can preserve the joint distribution of the trajectory (Theorem 3.4) with strong expressiveness guarantees (Theorem 3.5). Empirical results on the large-scale real-world benchmark OC22 also verify the superiority of trajectory guidance enabled by our framework.\n\n- Second, our GDB framework can preserve the coupling of geometric states, which is also crucial and necessary in modeling the evolution of geometric states (please refer to line 132 in section 3 and Theorem 3.3 for more details).  Diffusion bridge works the reviewer mentioned [1, 2] do not guarantee coupling preservation. Our paper provides detailed derivations and theoretical analysis for the coupling preservation of geometric states, which we believe add value to the communities of both geometric deep learning and diffusion bridge-based approaches.\n\nBesides these contributions, we also provide detailed derivations and theoretical guarantees for satisfying the symmetry constraints of the equivariant diffusion bridge. Although both the equivariant diffusion process and diffusion bridge have been studied, few works exist to thoroughly introduce methodologies for combining the power from both sides, while our work makes it complete.\n\nWe will carefully cite the listed works and add the above discussions to the next version of our paper to let the general audience better understand our contributions.\n\n[1] Diffusion-based Molecule Generation with Informative Prior Bridges. Lemeng Wu, et al. NeurIPS 2022.\n\n[2] DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization. Danyeong Lee, et al. AAAI 2024.\n\nWe thank you again for your efforts in reviewing our paper, and we have replied to each of your comments. We look forward to your re-evaluation of our submission based on our responses and updated results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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 introduces the Geometric Diffusion Bridge (GDB), a novel framework for predicting the evolution of geometric states in complex systems using an equivariant diffusion bridge approach. The reviewers generally acknowledge the technical soundness and the innovative aspects of the proposed method, particularly its ability to preserve symmetry constraints and leverage trajectory data for enhanced dynamic modeling. Reviewer Sjdg highlights that “GDB surpasses existing state-of-the-art approaches in various real-world scenarios,” noting its significant error reduction in challenging tasks. Reviewer dV6w appreciates the “solid theoretical foundation for GDB” and the strong empirical performance across multiple datasets. Despite concerns from Reviewer znC7 about the experimental setup and missing baselines, the authors provided extensive clarifications, including comparisons to additional relevant works and ablation studies to demonstrate the framework’s efficacy. Reviewer U4NU acknowledges the novelty of applying diffusion bridges to bridge geometric states, albeit suggesting that related works could be better discussed to strengthen the contributions. On balance, the paper presents a significant advancement in the field with robust theoretical backing and promising empirical results, meriting acceptance.
Decision: Accept (poster)