Bridging Geometric States via Geometric Diffusion Bridge

Authors: shengjie  luo
yixian  xu
di  he
shuxin  zheng
tie-yan  liu
liwei  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:01:38

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.

Main paper: 
 introduction: 
 Predicting the evolution of the geometric state of a system is essential across various scientific domains [46, 88, 55, 17, 20, 101] , offering valuable insights into difficult tasks such as drug discovery [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-principle methods or empirical laws, either requiring extensive computational costs [68] or sacrificing accuracy [40] .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 neural 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 iterative evolution into a single-step prediction model, which lacks the ability to fully capture the system'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, 58] , showing a better efficiency-accuracy 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 Diffusion Bridge (GDB), a 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 a 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 h-transform [82, 81, 16] . The proposed stochastic differential equation (SDE) is anchored by initial and target geometric states to 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 a scalable and simulation-free 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 finegrained 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 a comprehensive guidance for implementing our GDB framework in realworld 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 a large margin. In particular, our method even surpasses strong MLFF baselines that are trained on 10× 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 to capture the complex evolution dynamics of geometric states and determine valuable and crucial geometric states of interest in critical real-world challenges. 
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 a system consisting of a set of objects located in the three-dimensional Euclidean space. We use H ∈ R n×d to denote the objects with features, where n is the number of objects, and d is the feature dimension. For object i, let r i ∈ R 3 denote its Cartesian coordinate. We define the system as S = (H, R), where R = {r 1 , ..., r n }. 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 t = {r t 1 , ..., r t n }. Given a system S t0 = (H, R t0 ) at time t 0 , our goal is to predict S t1 = (H, R t1 ) at a future time t 1 . As an example, in a molecular system, R t1 can be the equilibrium state of interest evolved from the initial state R t0 .In this problem, inherent symmetries in geometric states should be considered. For example, a rotation that is applied to the coordinate system at time t 0 should also be applied to 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 a group G, let ρ X and ρ Y denote its group representations, which describe how the group elements act on these spaces. A function ϕ : X → Y is said to be equivariant if it satisfies the following condition:ρ Y (g)[ϕ(x)] = ϕ ρ X (g)[x] , ∀g ∈ G, x ∈ X . When ρ Y = I Y (identity transformation), it is also known as invariance. SE(3) group, which pertains to translations (T(3)) and rotations (SO(3)) in 3D Euclidean space, is one of the most widely used groups and is employed in our framework.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 q data (X), where X ∈ R d , we construct a diffusion process (X t ) t∈[0,T ] , which is represented as a sequence of random variables indexed by time steps. We set X 0 ∼ q data (X) and X T ∼ p prior (X), where p prior (X) has a tractable form to generate samples efficiently, e.g. standard Gaussian distribution. Mathematically, we model (X t ) t∈[0,T ] as the solution to the following stochastic differential equation (SDE):EQUATIONwheref (•, •) : R d ×[0, T ] → R d is a vector-valued function called the drift coefficient, σ(•) : [0, T ] →R is a scalar function known as the diffusion coefficient, and (B t ) t∈[0,T ] is the standard Wiener process (a.k.a., Brownian motion) [26] . We hereafter denote by p t (X) the marginal distribution of X t . Let p(x ′ , t ′ |x, t) denote the transition density function such that P (X t ′ ∈ A|X t = x) =A p(x ′ , t ′ |x, t)dx ′ for any Borel set A. By simulating this diffusion process forward in time, the distribution of X t will become p prior (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 0 ∼ p 0 (X) := q data (X), 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:EQUATIONwhere ∇ X log p t (X) denote the score of the marginal distribution at time t. If the score is known for all time, then we can derive the reverse diffusion process from Eqn. (2), sample from p prior (X), and simulate this process to generate samples from the data distribution q data (X). In particular, the score ∇ X log p t (X) can be estimated by training a parameterized model s θ (X, t) with a denoising score matching objective [98, 97] . In 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 is crucial, for which three desiderata should be carefully considered:• Coupling Preservation: From a probabilistic perspective, the evolution of geometric states transports their distribution from q data (S t0 ) to q data (S t1 ), and we are interested in modeling the distribution of target geometric states given the initial states, i.e., q data (S t1 |S t0 ) := q data (R t1 |H, R t0 ), which can be achieved by preserving the coupling of geometric states, i.e., q data (R t0 , R t1 |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 is unchanged regardless of how the system is rotated or translated, the distribution of the geometric states should satisfy symmetry constraints, i.e., q data (ρ R (g)[R t1 ]|ρ R (g)[R t0 ]) = q data (R t1 |R t0 ) and q data (ρ R (g)[R t0 ], ρ R (g)[R t1 ]) = q data (R t0 , R t1 ) for all g ∈ SE(3), R t ∈ R.• Trajectory Guidance: Trajectories of geometric states are sometimes accessible and provide fine-grained descriptions of the evolution dynamics. For completeness, it is crucial to develop a unified framework that can characterize and leverage trajectory data as 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. 5 and summarize into Table 1 . In this section, we introduce Geometric Diffusion Bridge (GDB), a general framework for bridging geometric states through generative modeling. We will elaborate on key techniques for completely preserving couping under symmetry constraints (Sec. [104, 31, 87, 89, 8] ✓ ✓ ✗ MLFFs [90, 33, 6, 34, 58] ✓ ✗ ✓ Geometric Diffusion Model [115, 38, 114] ✓ ✗ ✗ Geometric Diffusion Bridge (ours) ✓ ✓ ✓Our key design lies in the construction of equivariant diffusion bridge, a tailored diffusion process (R t ) t∈[0,T ] for bridging initial states R 0 ∼q data (R t0 ) and target states R T ∼q data (R t1 |R t0 ), 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 (•, •) : R×[0, T ] → R denote the drift coefficient on R. Let (W t ) t∈[0,T ] denote the Wiener process on R. Given an SDE on geometric statesdR t = f R (R t , t)dt + σ(t)dW t , R 0 ∼ q(R 0 ), its transition density p R (z ′ , t ′ |z, t), z, z ′ ∈ R is SE(3)-equivariant, i.e., p R (R t ′ , t ′ |R t , t) = p R (ρ R (g)[R t ′ ], t ′ |ρ R (g)[R t ], t), ∀g ∈ SE(3), 0 ≤ t, t ′ ≤ T, if these conditions are satisfied: (1) q(R 0 ) is SE(3)-invariant; (2) f R (•, t) is SO(3)- equivariant and T(3)-invariant; (3) the transition density of (W t ) t∈[0,T ] is SE(3)-equivariant.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 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 p R (z ′ , t ′ |z, t) be the transition density of the SDE in Proposition 3.1. Let h R (•, •) : R × [0, T ] → R >0 be a smooth function satisfying: (1) h R (•, t) is SE(3)-invariant; (2) h R (z, t) = p R (z ′ , t ′ |z, t)h R (z ′ , t ′ )dz ′ .Then we can derive the following h R -transformed SDE on geometric states:EQUATIONwith SE(3)-equivariant transition density p h R (z ′ , t ′ |z, t) equals to p R (z ′ , t ′ |z, t) h R (z ′ ,t ′ ) h R (z,t) = p R (z ′ , T |z, t) qdata(z ′ |z0)p R (z ′ ,T |z0,0) dz ′ . By using Proposition 3.2, we can derive the following h R -transformed SDE:EQUATION) which corresponds to a process (R t ) t∈[0,T ] , R 0 ∼ q data (R t0 ) satisfying the following properties:• let q(•, •) : R × R → R ≥0 denote the joint distribution induced by (R t ) t∈[0,T ] , then q(R 0 , R T ) equals to q data (R t0 , R t1 ); • its transition density q R (R t ′ , t ′ |R t , t; R 0 , 0)=q R (ρ R (g)[R t ′ ], t ′ |ρ R (g)[R t ], t; ρ R (g)[R 0 ], 0), ∀0≤t,t ′ ≤T, g∈SE(3),R 0 ∼q data (R t0 ).We call the tailored diffusion process (R t ) t∈[0,T ] an equivariant diffusion bridge.According to Theorem 3.3, given an initial geometric state R t0 , we can predict target geometric states R t1 by simulating the equivariant diffusion bridge (R t ) t∈[0,T ] from R 0 = R t0 , which arrives at R T ∼ q data (R t1 |R t0 ). However, the score 4) is not tractable in general. Inspired by the score matching objective in diffusion models [99] , we use a parameterized model v θ (R t , t; R 0 ) to estimate the score by using the following training objective:E q R (R T ,T |R t ,t;R 0 ,0) [∇ R t log p R (R T , T |R t , t)|R 0 , R t ] in Eqn. (EQUATIONwhere t ∼ U(0, T ) (the uniform distribution on [0, T ]), and λ(•) : [0, T ] → R ≥0 is a positive weighting function. Theoretically, we prove that the minimizer of Eqn. (5) approximates the groundtruth score (see Appendix B.5 for more details). Moreover, this objective is tractable because the transition density p R and q R can be designed to have simple and explicit forms such as Gaussian, which we will elaborate on in Sec. 3.3.In this subsection, we elaborate on how to leverage trajectories of geometric states as a fine-grained guidance in our framework. Let ( Ri ) i∈[N ] denote a trajectory of N + 1 geometric states and q traj ( R0 , ..., RN ) 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, q traj ( R0 , ..., RN ) can be equivalently reformulated into q 0 traj ( R0 )N i=1 q i traj ( Ri | Ri-1) by the chain rule of probability. If q i traj ( Ri | Ri-1 ) can be well modeled, we can capture the distribution of trajectories of geometric states completely.According to Theorem 3.3, given R 0 ∼ q 0 traj ( R0 ), an equivariant diffusion bridge (R t ) t∈[0,T ] can be constructed to model the joint distribution q traj ( R0 , R1 ) and hence q 1 traj ( R1 | R0 ) 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 {(R t i ) t∈[0,T ] } i∈[N -1] denote a series of N equivaraint diffusion bridges defined in Theorem 3.3. For the i-th bridge (R t i ) t∈[0,T ] , if we set (1) h i R (z, t; z 0 ) = p R (z ′ , T |z, t) q i+1 traj (z ′ |z0) p R (z ′ ,T |z0,0) dz ′ ; (2) R 0 0 ∼ q 0 traj ( R0 ), R 0 i = R T i-1 , ∀0 < i < N , then the joint distribution q R (R 0 0 , R T 0 , R T 1 , • • • , R T N -1 ) induced by {(R t i ) t∈[0,T ] } i∈[N -1]equals to q traj ( R0 , ..., RN ). We call this process a chain of equivariant diffusion bridges.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 on its provable guarantee):EQUATION)wheret ∼ U(0, N × T ), i = ⌊ t T ⌋, t ′ = t -i × T, R t ′ i ∼ q i R (R t ′ i , t ′ |z i+1 , T ; z i , 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 ( Ri ) i∈[N ] is sampled by simulating a prior SDE on geometric states d Rt = -∇H * R ( Rt )dt + σd Wt . Let µ * i denote the path measure of this prior SDE when t ∈ [iT, (i + 1)T ]. Building upon ( Ri ) i∈[N ] , let {µ i R } i∈[N -1]denote the path measure of our chain of equivariant diffusion bridges. Under mild assumptions, we have limN →∞ max i KL(µ * i ||µ i R ) = 0.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 t = -α∇ R t V (R t )dt + βdW t , where V (R t ) denotes the potential energy at R t and α, β 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.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 q R (R t , t|z 1 , T ; z 0 , 0) for samplingR t ; (2) closed-form matching objective ∇ R t log p R (z 1 , T |R t , t).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:EQUATIONThe explicit form of the objective can be directly calculated, i.e., ∇R t log p R (z 1 , T |R t , t) = z1-R t σ 2 (T -t) .Sampling distribution. According to Theorem 3.3, the transition density q R (R t , t|z 1 , T ; z 0 , 0) can be calculated by using the Doob's h-transform in Proposition 3.2, i.e., q R (R t , t|z 1 , T ;z 0 , 0) = p R (R t , t|z 1 , T ) h R (R t ,t;z0) h R (z1,T ;z0). Moreover, h R is determined by q data and p R , which is already specified in Eqn. (7) . Therefore, we can also calculate qR (R t , t|z 1 , T ; z 0 , 0) = N ( t T z 1 + T -t T z 0 , σ 2 t(T -t) T 2 I).Symmetry constraints. In proposition 3.1, we have several conditions that should be satisfied to meet the symmetry constraints. Firstly, since a parameterized model v θ (R t , t; R 0 ) is used to estimate the score of our equivariant diffusion bridge, it should be SO(3)-equivariant and T(3)-invariant.Besides, we follow [50, 115] to consider CoM-free systems: given R = {r 1 , ..., r n }, we define r = 1 n n i=1 r i and the CoM-free version of R = {r 1 -r, ..., r n -r}. To sample from N (z 0 , σ 2 I)with z 0 ∈ R consisting of n objects, we (1) sample ϵ = {ϵ i } n i=1 by i.i.d. drawing ϵ i ∼ N (0, I 3 ); (2) calculate the CoM-free ϵ ′ of ϵ; (3) obtain z 0 + σϵ ′ .Trajectory guidance. According to Eqn. (6), both p i R and q i R for all i∈[N -1] should be determined. Similarly, we setp i R (z i+1 , T |R t ′ , t ′ )=N (R t ′ , σ 2 i (T -t ′ )I), which further induces q i R (R t ′ , t ′ |z i+1 , T ; z i , 0) = N ( t ′ T z i+1 + T -t ′ T z i , σ 2 i t ′ (T -t ′ ) T 2 I).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 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 is general to be implemented by using other advanced design strategies [99, 47, 48] , which we leave as future work. Algorithm 1 Training 1: repeat 2: (z0, z1) ∼ qdata(R t 0 , R t 1 ) 3: t ∼ U [0, T ] 4: ϵ ∼ N (0, I) 5: R t = t T z1 + T -t T z0 + √ t(T -t) T σϵ 6: Take gradient descent step on ∇ θ λ(t) z 1 -R t σ 2 (T -t) -v θ (R t , t; z0)3: t ∼ U (0, N × T ), i = ⌊ t T ⌋, t ′ = t -i × T 4: ϵ ∼ N (0, I) 5: R t ′ i = t ′ T zi+1 + T -t ′ T zi + √ t ′ (T -t ′ ) T σiϵ 6: Take gradient descent step on ∇ θ λ(t) z i+1 -R t ′ i σ 2 i (T -t ′ ) -v θ (R t ′ i , t; zi) 2 7: until converged 4 ExperimentsIn this section, we empirically study the effectiveness of 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, we present more details in Appendix D.  
equilibrium state prediction: 
 Task. Equilibrium states typically represent local minima on the Born-Oppenheimer potential energy surface of a molecular system [54] , which correspond to its most stable geometric state and play an 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 of a 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 130,000 organic molecules. In convention, 110k, 10k, and 11k molecules are used for train/valid/test sets respectively; (2) Molecule3D [116] is a largescale dataset curated from the PubChemQC project [67, 71] , consisting of 3,899,647 molecules in total and its train/valid/test splitting ratio is 6 : 2 : 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 is 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 v θ (R t , t; R 0 ) by extending a Graph-Transformer 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 Euler solver [12] . Following [111] , we choose several strong baselines for a 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.1.Results. Results on QM9 and Molecule3D are shown in Table 3 and 4 respectively. It can be easily seen that our GDB framework consistently surpasses all baselines by a significantly large margin on show the superiority of our GDB framework for bridging geometric states, and its generality on both medium and large-scale challenges. Moreover, our framework performs consistently across valid and tests of both random and scaffold splits, further verifying its robustness in challenging scenarios.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 adsorbatecatalyst complex. Both initial and adsorption states with trajectories connecting them are provided.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 out-of-distribution (OOD) settings which use unseen catalysts, containing approximately 2,624 and 2,780 complexes respectively.Setting. 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. We parameterize v θ (R t , t; R 0 ) by using GemNet-OC [34] , which also serves as a verification that Table 5 : Results on the OC22 IS2RS Validation set. "OC20+OC22" denotes using both OC20 [13] and OC22 data; "OC20→OC22" 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 Results. In Table 5 , our GDB significantly outperforms the best baseline, e.g., 3.3%/3.6%/3.4% relative improvement on the ADwT metric of ID, OOD and Avg respectively. It is noteworthy that the best baseline is 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.42 Avg[%]. All the results 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.4% relative improvement on Avg ADwT. Moreover, we also investigate the impact of R 0 condition in v θ (R t , t; R 0 ), 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.3% relative ADwT drop on Avg/OOD respectively. Overall, these ablation studies serve as 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 is direct prediction, i.e., training a 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 a single-step prediction model, which lacks the ability to capture the underlying dynamics and cannot leverage trajectories of geometric states.Machine Learning Force Field. Another line of approach is called machine learning force field (MLFF) [106, 5, 6, 70, 75, 58] , which are trained to predict intermediate labels, such as 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 as the backbone, MLFFs typically satisfy the symmetry constraints. Besides, trajectory data with additional energy or force labels can directly be used for training MLFFs. 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. In recent years, diffusion models [37, 99] have emerged with stateof-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 roto-translational invariant distributions [115, 38] . In addition to 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 a 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 in 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 h-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 a unified manner.Question: Does the paper provide open access to the data and code, with sufficient instructions 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 NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 In this work, we introduce Geometric Diffusion Bridge (GDB), a general framework for bridging geometric states through generative modeling. We leverage a 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 as guidance by using a chain of equivariant diffusion bridges, allowing complete modeling of trajectory data. Mathematically, we conduct a comprehensive theoretical analysis showing our framework'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.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes] Justification: Section 4 and Appendix D. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we 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 as the institution conducting the review. 
 
 Appendix: This work newly proposes a general framework to bridge geometric states, which has great significance in various scientific domains. Our experimental results have also demonstrated considerable positive potential for various applications, such as catalyst discovery and molecule optimization, which can significantly contribute to the advancement of renewable energy processes and chemistry discovery. However, it 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.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 in transformer-based diffusion models, which we have earmarked for future research.The supplementary material is organized as follows. In Appendix B, we first recall some 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 a comprehensive introduction to the datasets, baselines, metrics and settings.Let (X t ) t∈[0,T ] be a stochastic process. We use p(x ′ , t ′ |x 1 , t 1 ; x 2 , t 2 ; . . . ; x n , t n ) to denote its conditional density function satisfying P (X t ′ ∈ A|X t1 = x 1 , X t2 = x 2 , . . . , X tn = x n ) = A p(x ′ , t ′ |x 1 ,EQUATIONWe 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 h-transform. The precise definitions are as follows.Definition B.2. (Generator of a Process) The infinitesimal generator A t of a stochastic process (X t ) for a function ϕ(x) isEQUATION)where ϕ is a suitably regular function. For an Itô process defined as the solution to the SDE dX t = f (X t , t)dt + σ(t)dB t , the generator isEQUATIONThe Fokker-Planck's Equation is an useful tool to track the evolution of the transition density function associated with an SDE. EQUATIONwith the initial condition p(x, 0|x 0 , 0) = δ(x -x 0 ). The Fokker-Planck's Equation can also be written in a compact form using the generator A t :EQUATIONwhere A * t is the adjoint operator of A:EQUATION)When the terminal is fixed, the evolution of the transition density function can also given by a PDE, which is called the Backward Kolmogorov Equation. We give the precise statement as follows.Proposition EQUATIONwith the initial condition p(x t , t|x, t) = δ(x -x t ). The Backward Kolmogorov Equation can also be written in a compact form using the generator A s : EQUATIONt = f R (R t , t)dt + σ(t)dW t , R 0 ∼ q(R 0 ), its transition density p R (z ′ , t ′ |z, t), z, z ′ ∈ R is SE(3)-equivariant, i.e., p R (R t ′ , t ′ |R t , t) = p R (ρ R (g)[R t ′ ], t ′ |ρ R (g)[R t ], t), ∀g ∈ SE(3), ∀0 ≤ t < t ′ ≤ T, if the following conditions are satisfied: (1) q(R 0 ) is SE(3)-invariant; (2) f R (•, t) is SO(3)-equivariant and T(3)-invariant; (3) the transition density of (W t ) t∈[0,T ] is SE(3)- equivariant.Proof. In this section, we view R = {r 1 , ...,r n } ∈ R as r 1 ⊕ r 2 ⊕ • • • r n ∈ R 3n, which is the concatenation of r i . So from this perspective, the space R is isomorphic to the Euclidean space R 3n . Then (W t ) t∈[0,T ] is the Wiener process with dimension d = 3n.For any g ∈ SE(3), ρ R (g) can be characterized by an orthogonal matrix O(g) ∈ R 3×3 , satisfying det(O(g)) = 1, and a translation vector t ∈ R 3 . Then the representation of SE(3) on R 3n is given by ρ R (gEQUATION) where O R (g) = diag{O(g), O(g), . . . , O(g)}, t R = t ⊕ t ⊕ • • • t ∈ R 3n . It's obvious that O R (g) is also an orthogonal matrix in R 3n×3n , satisfying O -1 R (g) = O T R (g).According to Proposition B.3, the evolution of the transition density function is given by the Fokker-Planck's EquationEQUATIONwith the initial condition p R (x, 0|x 0 , 0) = δ(x -x 0 ).EQUATION)The evolution of the transition density function p R (y, t|y 0 , 0) is also given by the Fokker-Planck's Equation:EQUATIONwith the boundary condition p R (y, 0|y 0 , 0) = δ(y-y 0 ) = δ(x -x 0 ). Since y = O R (g)x + t R , we have x = O -1 R (g)(y -t R ).Then by the chain rule, we haveEQUATION)Since f R (•, t) is a SO(3)-equivariant and T(3)-invariant function, we haveEQUATION)Then the Fokker-Planck's equation becomesEQUATION)EQUATION)EQUATION)Since O R (g) is an orthogonal matrix, the columns of O R (g) are orthogonal to each other, i.e.EQUATIONSo the Fokker-Planck's equation can be simplified to ∂p R (y, t|y 0 , 0) ∂t = -EQUATION)EQUATION)which is same as Eqn. (17) . Since the boundary condition p R (y, 0|y 0 , t 0 ) = δ(y-y 0 ) = δ(x - x 0 ) = p R (x, 0|x 0 , 0), then p R (y, t|y 0 , t 0 ) = p R (x, t|x 0 , t 0 ), ∀t ∈ [0, T ]. Thus we have proved that p R (R t ′ , t ′ |R t , t) = p R (ρ R (g)[R t ′ ], t ′ |ρ R (g)[R t ], t), ∀g ∈ SE(3), ∀0 ≤ t < t ′ ≤ T . B.3 Proof of Proposition 3.2 Proposition B.6 (Doob's h-transform). Let p R (z ′ , t ′ |z, t) be the transition density of the SDE in Proposition 3.1. Let h R (•, •) : R × [0, T ] → R >0 be a smooth function satisfying: (1) h R (•, t) is SE(3)-invariant; (2) h R (z, t) = p R (z ′ , t ′ |z, t)h R (z ′ , t ′ )dz ′ .We can derive the following h Rtransformed SDE on geometric states:EQUATIONwith transition density t) preserving the symmetry constraints.p h R (z ′ , t ′ |z, t) = p R (z ′ , t ′ |z, t) h R (z ′ ,t ′ ) h R (z,Proof. We use the definition of the infinitesimal generator to prove the proposition. The infinitesimal generator of p h R (x ′ , t ′ |x, t) for a function ϕ(x) is given byEQUATION) Since p h R (z ′ , t ′ |z, t) = p R (z ′ , t ′ |z, t) h R (z ′ ,t ′ ) h R (z,t) , so we haveEQUATION)Then A h t ϕ(x) can be simplified asEQUATION= limEQUATIONEQUATION)EQUATIONEQUATION)According to the Backward Kolmogorov Equation (Proposition B.4), we getEQUATION)So we get ∂ ∂tEQUATION)Then A h t ϕ(x) can be simplified asEQUATIONEQUATION)So we show thatEQUATION)So we can easily show that h R (z, t; z 0 ) satisfies the conditionEQUATION)Then we can use the result of Theorem 3.2 to get the h-transformed SDE. By Theorem 3.2, the h-transformed SDE isEQUATIONNext, we need to find the explicit form of ∇ R t log h R (R t , t; R 0 ),EQUATION)The h-transformed density function isEQUATION)Then we haveEQUATION)So we get a explicit form of ∇ R t log h R (R t , t; R 0 ):EQUATION)Then the h-transformed SDE becomesEQUATIONEQUATION)EQUATIONFinally, we need to show that the transition density function satisfies the corresponding symmetric constrains. Since h R (z ′ , T ;EQUATIONwe can show that h(•, t; •) is also SE(3)-invariant ∀t ∈ [0, T ] using the following propertyEQUATIONFor any g ∈ SE(3), assume ρ R (gEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATION= q R (z ′ , t ′ |z, t; z 0 , 0), (102) which completes our proof.Bridge Lemma B.9. Let X 1 , • • • , X n , Y, Z be random variables. Then the optimal approximation of Y based on {X} n i=1 is f * (X 1 , • • • , X n ) = arg min f E∥Y -f (X 1 , • • • , X n )∥ 2 = E[Y|X 1 , • • • , X n ]. Proof. Denote X = (X 1 , • • • , X n ).We show the following decomposition first:EQUATION)We can compute E∥Y -f (X)∥ 2 directly byEQUATIONEQUATIONEQUATIONEQUATIONThe inequality becomes equality if and only if fEQUATIONProposition B.10. The training objective function of Equivariant Diffusion Bridge is:EQUATIONwhere t ∼ U(0, T ). Then the optimal parameter θ * = arg min θ L(θ) satisfiesEQUATIONProof. Let L(θ) = E t∼U (0,T ) λ(t)L t (θ), where EQUATIONEQUATION(R t i ) t∈[0,T ] , if we set (1) h i R (z, t; z 0 ) = p R (z ′ , T |z, t) q i+1 traj (z ′ |z0) p R (z ′ ,T |z0,0) dz ′ ; (2) R 0 0 ∼ q 0 traj ( R0 ), R 0 i = R T i-1 , ∀0 < i < N , then the joint distribution q R (R 0 0 , R T 0 , R T 1 , • • • , R T N -1 ) induced by {(R t ) t∈[0,T ] } i∈[N -1]equals to q traj ( R0 , ..., RN ). We call this process a chain of equivariant diffusion bridges. Proof. By Theorem 3.3, the transition density function of (Rt i ) t∈[0,T ] satisfies q i R (R T i |R 0 i ) = q i traj (R T i |R 0 i ), ∀0 ≤ i ≤ N -1.The ground truth probability density function has the decomposition q 0 traj ( R0 )EQUATIONEQUATIONEQUATIONEQUATION) So the joint distribution q R (R 0 0 , R T 0 , R T 1 , • • • , R T N -1 ) induced by {(R t ) t∈[0,T ] } i∈[N -1]equals to q traj ( R0 , ..., RN ).Proposition B.12. The training objective function of the Chain of Equivariant Diffusion Bridge is:EQUATIONEQUATIONProof.Let L ′ (θ) = E t∼U (0,N T ) λ(t)L ′ t (θ),whereEQUATION)whereEQUATIONB.8 Proof of Theorem 3.5In 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 (X t ) t∈[0,T ] with the initial position X 0 and the terminal position X T is given by the following SDEEQUATIONwhere W t is the standard wiener process. The solution of the Brownian bridge is given byEQUATIONNext, we recall the definition of the KL Divergence: Definition B.13 (KL Divergence). The relative entropy (or Kullback-Leibler Divergence) KL(f ||g) between two probability density functions f (x) and g(x) is defined by:EQUATIONIn general, let P and Q be two probability measures on space X . Assume P is absolutely continuous with respect to Q then the Kullback-Leibler Divergence between P and Q is defined as followsEQUATIONwhere dP dQ is the Radon-Nikodym derivative of P with respect to Q.When we need to compute the KL divergence between the path measures associated with two SDEs, the Girsanov's theorem [53] Let M t = W t - t 0 H s ds, then M t is a d-dimensionalWiener process with respect to Q. In practice, the condition that Z t is a martingale is hard to vertify. So the condition is often replaced by the Novikov's conditionEQUATIONFor 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. ∥∇H * R (x) -∇H * R (y)∥ ≤ L∥x -y∥, ∀x, y ∈ R d ; • H * R (•) satisfies ∥∇H * R (x)∥ ≤ K(1 + ∥x∥), ∀x ∈ R d ; • E∥ Rt ∥ 2 < M, ∀t ∈ [0, N T ]; • h(t) = E[H * R ( Rt )] is a continuous function on t ∈ [0, N T ]; • The Novikov's condition:E exp 1 2 N T 0 ∥∇H * R ( Wt )∥ 2 dt < ∞;• The function H * R satisfies the following regulaity condition: there exist a constantC such that ∇ 2 H * R (x) -∥∇H * R (x)∥ 2 /σ 2 < C, ∀x ∈ R d ;then we have limN →∞ max i KL(µ * i ||µ i R ) = 0.Proof. Let p * be the probability density function associated with the ground truthSDE d Rt = f * R ( Rt , t)dt+σd Wt , R0 = R 0 . Let {(R t i ) t∈[0,T ] } i∈[N -1]denote a series of N equivaraint diffusion bridges defined in Theorem 3.4. Then by theorem 3.4, q R (R 00 , R T 0 , R T 1 , • • • , R T N -1 ) induced by {(R t ) t∈[0,T ] } i∈[N -1] equals to p * R (R 0 0 , R T 0 , R T 1 , • • • , R T N -1). Additionally, the conditional probability density function q R (R t i |R T i , R 0 i ) , for iT ≤ t < (i + 1)T , is associated with the Brownian bridgeEQUATIONwhere t ′ = t -iT . Then by the chain rule of KL divergenceEQUATIONSince p * i ( R(i+1)T , RiT ) = q i R ( R(i+1)T , RiT ), we haveEQUATIONSince 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 σ Wt on space R. Since the condition of Theorem B.14 is satisfied, then we can use Theorem B.14 and getdµ * 0 (•| R0 ) = exp - 1 σ T 0 ∇H * R (σ Wt ) • d Wt - 1 2σ 2 T 0 ∥∇H * R (σ Wt )∥ 2 dt dυ(•| R0 ).(137) Then we can use the Ito's formula (Theorem B.1) to simplify the expressiondµ * 0 (•| R0 ) = exp 1 σ 2 (H * R (σ W0 ) -H * R (σ WT )) + 1 2 T 0 (∇ 2 H * R (σ Wt ) -1 σ 2 ∥∇H * R (σ Wt )∥ 2 )dt dυ(•| R0 ).(138) To simplify our notation, we denoteZ T = exp 1 σ 2 (H * R (σ W0 ) -H * R (σ WT )) + 1 2 T 0 (∇ 2 H * R (σ Wt ) - 1 σ 2 ∥∇H * R (σ Wt )∥ 2 )dt .(139) Let F, g be measurable functions on C[0, T ], R d , respectively. Then by the disintegration of Wiener measure into pinned Wiener measures (path measure of the Brownian Bridge), we have EQUATIONEQUATIONSo we can conclude thatdµ * 0 (•| R0 , RT ) dυ(•| R0 , RT ) = p T ( RT | R0 ) p * 0 ( RT | R0 ) • exp( 1 σ 2 (H * R ( R0 )) exp( 1 σ 2 (H * R ( RT )) exp 1 2 T 0 (∇ 2 H * R (•) -1 σ 2 ∥∇H * R (•)∥ 2 )dt . (143) Note that µ i R (•| R0 , RT ) = υ(•| R0 , RT). Now we can calculate the KL divergence by [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 are used for train/valid/test sets respectively. The geometric conformations that are minimal in energy are provided in the QM9 dataset. The equilibrium conformation and its relative properties are all calculated at the B3LYP/6-31G(2df,p) level of quantum chemistry.EQUATIONEQUATION) When N → ∞, T = 1 N → 0. Since h(t) is continuous by our assumption, then KL(µ * 0 ||µ 0 R ) → 0. Algorithm 3 Training 1: repeat 2: (z0, z1) ∼ qdata(R t 0 , R t 1 ) 3: t ∼ U [0, T ] 4: ϵ ∼ N (0, I) 5: R t = t T z1 + T -t T z0 + √ t(T -t) T σϵ 6: Take gradient descent step on ∇ θ λ(t) z 1 -R t σ 2 (T -t) -v θ (R: t ∼ U (0, N × T ), i = ⌊ t T ⌋, t ′ = t -i × T 4: ϵ ∼ N (0, I) 5: R t ′ i = t ′ T zi+1 + T -t ′ T zi + √ t ′ (T -t ′ ) T σiϵ 6: Take gradient descent step on ∇ θ λ(t) z i+1 -R t ′ i σ 2 i (T -t ′ ) -v θ (R= z0 2: R N T = Solver(v θ (R t , t; R ⌊ t T ⌋T ), t = N T ) Ensure: R N T D Experiments D.1 Equilibrium State Prediction Dataset. QM9Molecule3D [116] is a large-scale dataset curated from the PubChemQC project [67, 71] , consisting of 3,899,647 molecules in total, 2,339,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 are provided. 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 is generated by using fast and coarse force field [73, 52] and geometry optimization is conducted to obtain B3LYP/6-31G* level DFT-calculated equilibrium geometric structure.Baselines. We comprehensively compare our GDB framework with previous equilibrium conformation prediction methods. Following [111] , we use DG and ETKDG algorithms implemented by RDkit as our fundamental baselines. The benchmark [116] used the DeeperGCN-DAGNN 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 on the Molecule3D dataset. GINE [39] proposed a method for pretraining GNN to improve the performance and capacity of GNN. GATv2 [10] proposed a 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 on Graph-Transformer (GT) [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 = {r i } N i=1 which is rigidly aligned to the ground-truth R * = {r * i } N i=1 by the Kabsch algorithm [44] , Root Mean Square Deviation between their atoms is calculated, i.e., C-RMSD( R, R * ) = , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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
- 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.
questions: questions: See above
limitations: limitations: No limitations are addressed in the paper by the authors
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. The presentation is a bit unclear. Please refer to Q2.

3. Additional baselines may be considered. The baselines selected in the paper are not closely connected to the proposed approach. See Q3.

4. Missing ablation studies. In the current shape it is unclear where the performance gain comes from. See Q4.
questions: 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.

Q2. 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.

Q3. 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.

Q4. 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.


[1] Xu et al. GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation. In ICLR'22.

[2] Zhou et al. Denoising Diffusion Bridge Models. In ICLR'24.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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?

[1] Peluchetti, Diffusion Bridge Mixture Transports, Schrodinger Bridge Problems and Generative Modeling, JMLR 2023
[2] Liu et al., Learning Diffusion Bridges on Constrained Domains, ICLR 2023
[3] Wu et al., Diffusion-based Molecule Generation with Informative Prior Bridges, NeurIPS 2022
[4] Jo et al., Graph Generation with Destination-Predicting Diffusion Mixture, arXiv 2023 
[5] Albergo et al., Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, arXiv 2023
[6] Shi et al., Diffusion Schrodinger Bridge Matching, NeurIPS 2023
[7] Xu et al., GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, ICLR 2022
[8] Xu et al., Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023
[9] Yim et al., SE(3) diffusion model with application to protein backbone generation, ICML 2023
questions: 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?
- Is GDB scalable to geometric states of high dimensions? While analysis on this may not be necessary, it could strengthen the work.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

[1] Diffusion-based Molecule Generation with Informative Prior Bridges. Lemeng Wu, et al. NeurIPS 2022.

[2] DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization. Danyeong Lee, et al. AAAI 2024.
questions: questions: See weaknesses.
limitations: limitations: The authors note the need of exploring better implementation strategies for their framework to enhance performance.

Comments: title: Kindly Reminder of the Close of the Discussion Period
Kindly Reminder of the Close of the Discussion Period
Kindly request for feedback and reevaluation
Kindly request for feedback and reevaluation
Official Comment by Reviewer dV6w
Looking forward to your re-evaluation
Looking forward to your re-evaluation
Looking forward to your re-evaluation
Looking forward to your re-evaluation
Rebuttal by Authors (Part 2)
comment: Dear Reviewer znC7,

We would like to express our gratitude for your valuable comments and feedback. As the author-reviewer discussion period is coming to close on Aug 13th, we would greatly appreciate it if you could provide more feedback and reevaluate our work based on our responses and updated results. Thank you very much for your attention to this matter.

Best regards,

Authors
Dear Reviewer U4NU,

We would like to express our gratitude for your valuable comments and feedback. As the author-reviewer discussion period is coming to close on Aug 13th, we would greatly appreciate it if you could provide more feedback and reevaluate our work based on our responses and updated results. Thank you very much for your attention to this matter.

Best regards,

Authors
Dear Reviewer znC7,

Thank you once again for taking the time to review our paper! As the deadline for Reviewer-Author discussions is fast approaching, we would greatly appreciate it if you could kindly reevaluate our work with updated scores. Overall, we have carefully put every effort to address the reviewer's mentioned concerns. Besides, we will definitely follow your suggestion to include updated baseline results and thorough discussions of your mentioned related works in our next version of the paper, which we believe can improve the quality of our submission and address your remaining concerns.

We always believe that the feedback between reviewers and authors would indeed improve the paper's quality, and we will definitely include the updated results and conceptual discussions from the reviewer's suggestions in the revised paper. It would be really nice to see both of us reach a consensus. We sincerely look forward to your reevaluation and feedback!

Best Regards,

Paper 15795 Authors
Dear Reviewer U4NU,

Thank you once again for taking the time to review our paper! As the Reviewer-Author discussion deadline is quickly approaching, we would sincerely appreciate it if you could provide us with further feedback on our responses and kindly reevaluate our work based on our clarification and additional results in our rebuttal.

Following your insightful suggestions, we have thoroughly discussed the related works in the rebuttal and clarified our novel contributions. We also comprehensively illustrate GDB's scalability to geometric states of high dimensions. Additionally, we provide additional details and our motivation for the ODE sampler. Based on these additional results and clarifications, we sincerely hope your concerns can be addressed.

We always believe that the feedback between reviewers and authors would indeed improve the paper's quality, and we will definitely include the related works discussions from the reviewer's suggestions in the revised paper. It would be really nice to see both of us reach a consensus. We sincerely look forward to your reevaluation and feedback!

Best Regards,

Paper 15795 Authors
Thanks for your reply. We are happy to hear that all your concerns have been properly addressed. We thank you again for recognizing our contributions and staying positive about our work.
Thank you for your reply and I am satisfied with the responses. I will keep my positive score.
Thanks for your quick reply and constructive suggestions. We agree with your point that including more generative baselines for comparisons could improve our work. In previous responses, we followed your advice to include Denoising Diffusion Bridge Models with our equivariant backbone, which combines advanced bridge models and equivariant models as a strong generative baseline. We will also follow your suggestion to include DiffMD in our work. However, we have tried our best but found that DiffMD is not open-sourced. We have to reimplement the method and conduct the experiment, which cannot be finished before the deadline. We will definitely update the results in the next version of our paper. We will also carefully cite these generative baselines and provide thorough discussions in our Related Work section.

We use 10-step sampling for the sake of simplicity. In our preliminary experiments, we found that using ten steps already yields good performance. Increasing inference steps still improves model performance slightly but with more computational cost. We believe investigating advanced sampling strategies for our method is important, which we leave as future work.

Through intensive and meaningful discussions with you, we have realized that some terminology could be more precise. For instance, the title '...via generative modeling' is too broad and might give the impression that we are applying conventional generative modeling methods on a standard generative task. **As discussed in the thread, this is not our intention**. We will clarify these parts to reflect our actual focus better.

Thank you again for your valuable suggestions. We sincerely hope the reviewer re-evaluates our work based on these responses and updated results for the next version of our paper.
Thanks for your reply. We are happy to hear that all your concerns have been properly addressed. We thank you again for recognizing our contributions and staying positive about our work.
I thank the authors for the response. However, some concerns are still not addressed.

While the original setting in GeoDiff is molecule generation, from my perspective, it could be easily extended to your setting with the core difference lying in tackling how to condition on an input structure, which can be tackled in a fairly easy way. A related approach that considers such extension is DiffMD [1]. A strong reason why these approaches are compelling is that your approach is claimed to be generative, in which case some generative baselines (or even benchmarks) should be included.

[1] Wu et al. DiffMD: A Geometric Diffusion Model for Molecular Dynamics Simulations. In AAAI 2023.

Regarding Q2, to me it is indeed quite surprising that only 10 steps are needed to obtain high quality samples, as opposed to 100-1000 steps adopted in previous works. Thus it would be interesting to see how the performance changes with the number of sampling steps.
I have read the comments and am satisfied with the responses. However, I would still stick to my original score.
Dear Reviewer dV6w,

Thank you for your time and efforts in reviewing our paper. We have carefully responded to each of your questions. Given that the author-reviewer discussion deadline is approaching, we would greatly appreciate it if you could kindly take a look at our responses and provide your valuable feedback. We are more than happy to discuss more if you still have any concerns.

Thank you once again and we are eagerly looking forward to your re-evaluation of our work.

Paper 15795 Authors
Dear Reviewer U4NU,

Thank you for your time and efforts in reviewing our paper. We have carefully responded to each of your questions. Given that the author-reviewer discussion deadline is approaching, we would greatly appreciate it if you could kindly take a look at our responses and provide your valuable feedback. We are more than happy to discuss more if you still have any concerns.

Thank you once again and we are eagerly looking forward to your re-evaluation of our work.

Paper 15795 Authors
Dear Reviewer znC7,

Thank you for your time and efforts in reviewing our paper. We have carefully responded to each of your questions. Given that the author-reviewer discussion deadline is approaching, we would greatly appreciate it if you could kindly take a look at our responses and provide your valuable feedback. We are more than happy to discuss more if you still have any concerns.

Thank you once again and we are eagerly looking forward to your re-evaluation of our work.

Paper 15795 Authors
Dear Reviewer Sjdg,

Thank you for your time and efforts in reviewing our paper. We have carefully responded to each of your questions. Given that the author-reviewer discussion deadline is approaching, we would greatly appreciate it if you could kindly take a look at our responses and provide your valuable feedback. We are more than happy to discuss more if you still have any concerns.

Thank you once again and we are eagerly looking forward to your re-evaluation of our work.

Paper 15795 Authors
>**Regarding ablation studies (Weakness 4 & Q4)**

We have already provided ablation studies. Please refer to Sec 4.2 (lines 335-342 and Table 5) for detailed descriptions and results. Our ablation studies help us better understand the importance of key designs in our GDB framework, including trajectory guidance and coupling preservation, without which we observe significant performance drops. Following your suggestion, we further conduct an ablation study by investigating the importance of satisfying equivariant constraints, which shows that the performance drop is 14.45% if we remove the equivariant constraints. These ablation studies serve as supporting evidence for the necessity of proposed components in our GDB framework, and we will update these results in the next version of our paper. 

>**Regarding the discussion of limitations**

Thank 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.

We 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 and updated results.