Geometric Trajectory Diffusion Models

Authors: jiaqi  han
minkai  xu
aaron  lou
haotian  ye
stefano  ermon
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 22:19:42

Abstract: Generative models have shown great promise in generating 3D geometric systems, which is a fundamental problem in many natural science domains such as molecule and protein design. However, existing approaches only operate on static structures, neglecting the fact that physical systems are always dynamic in nature. In this work, we propose geometric trajectory diffusion models (GeoTDM), the first diffusion model for modeling the temporal distribution of 3D geometric trajectories. Modeling such distribution is challenging as it requires capturing both the complex spatial interactions with physical symmetries and temporal correspondence encapsulated in the dynamics. We theoretically justify that diffusion models with equivariant temporal kernels can lead to density with desired symmetry, and  develop a novel transition kernel leveraging SE(3)-equivariant spatial convolution and temporal attention. Furthermore, to induce an expressive trajectory distribution for conditional generation, we introduce a generalized learnable geometric prior into the forward diffusion process to enhance temporal conditioning. We conduct extensive experiments on both unconditional and conditional generation in various scenarios, including physical simulation, molecular dynamics, and pedestrian motion. Empirical results on a wide suite of metrics demonstrate that GeoTDM can generate realistic geometric trajectories with significantly higher quality.

Main paper: 
 introduction: 
 Machine learning for geometric structures is a fundamental task in many natural science problems ranging from particle systems driven by physical laws [1, 26, 43, 2, 15] to molecular dynamics in biochemistry [22, 16, 45, 10] . Modeling such geometric data is challenging due to the physical symmetry constraint [56, 43] , making it fundamentally different from common scalar non-geometric data such as images and text. With the recent progress of generative models, many works have been proposed in generating 3D geometric structures like small molecules [66, 42, 21, 64] and proteins [59, 23] , showing great promise in solving the equilibrium states of complex systems. Despite this success, these existing methods are limited to synthesizing static structures and neglect the fact that important real-world processes evolve through time. For example, molecules and proteins are not static but always varying with molecular dynamics, which plays a vital role in analyzing possible binding activities [8, 20] . In this paper, we aim to study the generative modeling of geometric trajectories with the additional temporal dimension. While this problem is more practical and important, it is highly non-trivial with several significant challenges. First, geometric dynamics in 3D ubiquitously preserve physical symmetry. With global translation or rotation applied to a trajectory of molecular dynamics, the entire trajectory still describes the same dynamics and the generative model should estimate the same likelihood. Second, trajectories inherently contain the correspondence between frames in different timesteps, requiring generative models to hold a high capacity for capturing the temporal correlations. Last, moving from a single structure to a trajectory c , if available, is leveraged to construct the equivariant prior and as a conditioning signal in EGTN. composed of multiple ones, the distribution we are interested in becomes much higher-dimensional and more diverse, considering both the initial conditions as well as potential uncertainties injected along the evolution of dynamics.To this end, we propose geometric trajectory diffusion models (GeoTDM), a principled method for modeling the temporal distribution of geometric trajectories through diffusion models [49, 52, 53, 18] , the state-of-the-art generative model on various domains such as images [6] , videos [19] , and molecules [64] . Our key innovation lies in designing an equivariant temporal diffusion over geometric trajectories, with the reverse process parameterized by equivariant transition kernels, ensuring the desired physical symmetry of the generated trajectory. To better excavate the complex spatial interactions and temporal correlations, we develop a novel temporal denoising network, where we stack equivariant spatial convolution and temporal attention. Our developments not only guarantee the desirable physical symmetry of the trajectories, but also capture the complex spatial and temporal correspondence encapsulated in the dynamics of geometric systems. Moreover, by leveraging generative modeling, GeoTDM enjoys high versatility in generating diverse yet high-quality geometric trajectories from scratch, performing interpolation and extrapolation, and optimizing noisy trajectory, all under the proposed diffusion framework.In summary, we make the following contributions: 1. We present GeoTDM, a novel temporal diffusion model for generating geometric trajectories. We design the diffusion process to meet the critical equivariance in modeling both unconditional and conditional distributions over geometric trajectories. Notably, we also propose a conditional learnable equivariant prior for enhanced flexibility in temporal conditioning. 2. To fulfill the equivariance of the denoising network, we introduce EGTN, a graph neural network that operates on geometric trajectories, which also permits conditioning upon a given trajectory using equivariant cross-attention, making it suitable to serve as the backbone for GeoTDM. 3. We evaluate our GeoTDM on both unconditional and conditional trajectory generation tasks including particle simulation, molecular dynamics, and pedestrian trajectory prediction. GeoTDM can consistently outperform existing approaches on various metrics, with up to 56.7% lower prediction score for unconditional generation and 16.8% lower forecasting error for conditional generation on molecular dynamics simulation. We also show GeoTDM successfully performs several additional applications, such as temporal interpolation and trajectory optimization. 
related work: 
 Trajectory modeling for geometric systems. Modeling the dynamics of geometric data is challenging since one must capture the interactions between multiple objects. Graph neural networks [11] have emerged as a natural tool to tackle this complexity [26, 41] . Subsequent works [43, 7, 2, 63] discovered equivariance as a critical factor for promoting model generalization. Among these efforts, Radial Fields [27] and EGNN [43] work with equivariant operations between scalars and vectors, while TFN [56] and SE(3)-Transformer [9] generalize to high-order spherical tensors. While considerable progress has been made, they only conduct (time) frame-to-frame prediction, which is subject to error accumulation when performing roll-out inference. Recently, EqMotion [62] approached the problem by learning to predict trajectories. By comparison, our GeoTDM leverages a generative modeling framework, which enables a wider range of tasks such as generation and interpolation.Generative models in geometric domain. There is growing interest in developing generative models for geometric data, e.g.molecule generation [42, 66, 21, 64] , protein generation [59, 24, 69] , and antibody design [30] . Recently, diffusion-based models [21, 64] have been shown to yield superior performance compared to flow-based [42] and VAE-based [65] approaches in many of these tasks. Despite these fruitful achievements, most existing works only produce a snapshot of the geometric system, e.g., a molecule in 3D space, whereas our GeoTDM generalizes to generating a trajectory with multiple frames, e.g., an MD trajectory in 3D. DiffMD [60] specifically tackles MD modeling using Markovian assumption, while GeoTDM directly captures the joint distribution of all frames along the entire trajectory.Temporal diffusion models. Diffusion models have been recently adapted to handle the natural temporality of data in tasks such as video generation [19, 58, 17] , time series forecasting [37, 54] , PDE simulation [38] , human motion synthesis [55, 71] and pedestrain trajectory forecasting [12] . Distinct from these works, GeoTDM models the temporal evolution of geometric data represented as a geometric graph and maintains the aforementioned vital equivariance constraint. 
preliminaries: 
 Diffusion models. Diffusion models [49, 18, 52, 53] are a type of latent variable generative model that feature a Markovian forward diffusion process and reverse denoising process. The forward process progressively perturbs the input x 0 (e.g., image pixels or molecule coordinates) over T steps using a Gaussian transition kernel q(x τ |x τ -1 ) = N (x τ ; √ 1 -β τ x τ -1 , β τ I).Here, {x τ } T τ =1 are latent variables with the same dimension as the input and β τ are predefined using the noise schedule such that x T is close to being distributed as N (0, I). The reverse process maps back from the prior distribution with p(x T ) = N (0, I) using the kernelp θ (x τ -1 |x τ ) = N (x τ -1 ; µ θ (x τ , τ ), σ 2 τ I), where the variances σ2 τ are usually fixed and the mean µ θ is parameterized by a neural network with parameters θ. The model is trained by optimizing the variational lower bound, defined asL vlb = -log p θ (x 0 |x 1 ) + D KL (q(x T |x 0 )∥p(x T )) + T -1 τ =2 D KL (q(x τ -1 |x τ , x 0 )∥p θ (x τ -1 |x τ )).For training stability, [52, 18] suggest the noise-prediction objective:EQUATIONwhere x 0 ∼ p data , τ ∼ Unif(1, T ), the weighting factors λ(τ ) are typically set to 1 to promote sample quality,x τ = √ ᾱτ x 0 + √ 1 -ᾱτ ϵ with ᾱτ := τ s=1 α s = τ s=1 (1 -β s ), and ϵ θ is a specific parameterization of the mean satisfyingµ θ (x τ , τ ) = 1 √ ατ (x τ -βτ √ 1-ᾱτ ϵ θ (x τ , τ )). Equivariance. Functions. A function f is equivariant w.r.t a group G if f (g • x) = g • f (x), ∀g ∈ G. Furthermore, f is invariant if f (g • x) = f (x), ∀g ∈ G [46]. Here we focus on the group SE(3) consisting of all 3D rotations and translations 2 . Each group element g ∈ SE(3) can be represented by a rotation matrix R and a translation r ∈ R3 . For geometrc graph with node features h and coordinates Geometric trajectories and the distributions. We represent a geometric trajectory as (x [T ] , h, E), wherex, if h ′ , x ′ = f (h, x), we expect h ′ , Rx ′ + r = f (h, Rx + r)x [T ] := x (0) , x (1) , • • • , x (T -1) ∈ R T ×N ×Dx is the sequence of temporal geometric coordi- nates, h ∈ R N ×D his the node feature, and E is the set of edges representing the connectivity of the geometric graph. T is the number of time steps and D x , D h refers to the dimension of the coordinate and node feature respectively, with D x normally being 2 or 3 depending on the input data. In this work, we are interested in modeling the distribution of geometric trajectories given the configuration of the geometric graph, i.e., p(x [T ] |h, E).Conditioning. Some applications like trajectory forecasting can be viewed as conditional generative tasks, where we seek to model the distribution of trajectories conditioning on certain observed timesteps, i.e., p(x [T ] |x[Tc] c , h, E) where x [Tc] c∈ R Tc×N ×Dx is the provided trajectory in length T c .Equivariance for geometric trajectories. Since the dynamics must be invariant to rotation or translation, the distribution of the geometric trajectories should also preserve such symmetry. This is formalized by the following invariance constraint:EQUATIONwhere g• x [T ] := [Rx (0) + r, • • • , Rx (T -1) + r].The conditional case should instead preserve:EQUATIONfor all g ∈ SE(3)4 . Intuitively, if the given trajectory is rotated and/or translated, the distribution of the future trajectory should also rotate and/or translate by exactly the same amount. For simplicity, we omit writing the conditions h and E henceforth when describing the distributions of trajectories. 
geometric trajectory diffusion models: 
 In this section, we introduce the machinery of GeoTDM. We first present Equivariant Geometric Trajectory Network (EGTN) in § 4.1, a general purpose backbone operating on geometric trajectories while ensuring equivariance. We then present GeoTDM in § 4.2 for both unconditional and conditional generation using EGTN as the denoising network.Our proposed Equivariant Geometric Trajectory Network (EGTN) is constructed by stacking equivariant spatial aggregation layers and temporal attention layers in an alternated manner, drawing inspirations from spatio-temporal GNNs [70, 61] . In particular, spatial layers characterize the structural interactions within the system and temporal layers model the temporal dependencies along the trajectory. For spatial aggregation, we employ the Equivariant Graph Convolution Layer (EGCL) [43] ,xEQUATIONThe equivariant message passing is conducted independently for each frame t ∈ [T ] := {0, 1, • • • , T -1}, with the goal of passing and fusing the geometric information based on the structure of the graph for each time step. Following such layer, we further develop a temporal layer equipped by self-attention, which has exhibited great promise for sequence modeling [57] , to capture the temporal correlations encapsulated in the dynamics. We first compute Eqs. 5-6,a (t,s) = exp q (t)⊤ k (t,s) u∈[T ] exp q (t)⊤ k (t,u) , (5) h ′(t) = h (t) + s∈[T ]a (t,s) v (t,s) , (6) where q (t) , k (t,s) , v (t,s) are the query, key, and value, respectively. In detail, q (t) = φ q (h (t) ), k (t,s) = φ k (h (s) ) + ψ(t -s), and v (t,s) = φ v (h (s) ) + ψ(t -s), with ψ(t -s) being the sinusodial encoding [57] of the temporal displacement t -s, akin to the relative positional encoding [47] . Incorporating such information is crucial since the model is supposed to distinguish different time spans between two frames on the trajectory. Moreover, compared with directly encoding the absolute time step, our design is beneficial in that it ensures the temporal shift invariance of physical processes. The update of coordinates reuse the attention coefficients a (t,s) and the values v (t,s) ,EQUATIONwhere φ x is an MLP that outputs a scalar to preserve rotation equivariance. The entire network f EGTN , with schematic depicted in Fig. 4 , is constructed by alternating spatial and temporal layers, enjoying equivariance as desired (proof in Appendix A.4):Theorem 4.1 (SE(3)-equivariance). Let x ′[T ] , h ′[T ] = f EGTN x [T ] , h [T ] , E . Then we have g • x ′[T ] , h ′[T ] = f EGTN g • x [T ] , h [T ] , E , ∀g ∈ SE(3).Geometric conditioning. In certain tasks like trajectory forecasting, we are additionally provided with some partially observed trajectories as side input. In order to leverage their geometric information, we augment the unconditional EGTN with equivariant cross-attention, a conditioning technique tailored for geometric trajectories, and more importantly, guaranteeing the crucial equivariance in Theorem 4.1. In principle, our equivariant cross-attention resembles Eqs. 5-7, but instead computes the attention between the conditioning trajectory x[Tc] cand the target x [T ] . In detail, the attention coefficients are recomputed as a (t,s) = exp(q (t)⊤ k (t,s) ) u∈[Tc ]∪[T ] exp(q (t)⊤ k (t,u) ). The updated node feature h ′(t) and coordinate x ′(t) in Eqs. 6-7 are further renewed by the cross-attention terms, yielding h ′′(t) = h ′(t) + s∈[Tc] a (t,s) v (t,s) and x ′′(t) = x ′(t) + s∈[Tc] a (t,s) φ x (v (t,s) )(x (t) -x (s) c ).For unconditional generation, we seek to model the trajectory distribution subject to the SE(3)invariance (Eq. 2). To design a diffusion with the reverse marginal conforming to the invariance, we impose certain constraints to the prior and transition kernel, as depicted in the following theorem. T ) is SE(3)-invariant, the transition kernels p τ -1 (x[T ] τ -1 | x [T ] τ ), ∀τ ∈ {1, • • • , T } are SE(3)-equivariant, then the marginal p τ (x [T ] τ ) at any step τ ∈ {0, • • • , T } is also SE(3)-invariant.Prior in the translation-invariant subspace. Unfortunately, there is no properly normalized distribution w.r.t. Lebesgue measure on the ambient space X := R T ×N ×D that permits translationinvariance [43] . We instead build the prior on a translation-invariant subspace X P ⊂ X induced by a linear transformation P ∈ X ×X with rank(P) = (T N -1)D [36] . Specifically, we choose the prior to be the projection of the Gaussian N (0, I) in X to X P byP = I D ⊗ I T N -1 T N 1 T N 1 ⊤ T N , which corresponds to the function P (x [T ] ) = x [T ] -1 T T -1 t=0 CoM(x (t) ), with CoM(x (t) ) = 1 N N i=1 x (t) ibeing the center-of-mass (CoM) of the system at time t. We denote x := P (x). Then the resulting distribution is a restricted Gaussian (denoted Ñ (0, I)) with the variables supported only on the subspace (see App. A.1), and more importantly, is still isotropic and thus SO(3)-invariant. To sample from the prior, one can alternatively sample x [T ] ∼ N (0, I) ∈ X and then project it to the subspace to obtain the final sample x[T ] = P (x [T ] ) ∈ X P .Transition kernel. To be consistent with the prior, we also parameterize the transition kernel in the subspace X P , given by p θ (x[T ] τ -1 | x[T ] τ ) = Ñ ( μθ (x [T ]τ , τ ), σ 2 τ I). In this way, if the mean function μθ (•) is SO(3)-equivariant, then the transition kernel is also guaranteed SO(3)-equivariance.As suggested by [18] , we re-parameterize μθ (x[T ] τ , τ ) = 1 √ ατ x[T ] τ -βτ √ 1-ᾱτ εθ (x [T ]τ , τ ) , where εθ = P • f θ , with f θ being an SO(3)-equivariant adaptation of our proposed EGTN, fulfilled by subtracting the input coordinates from the output for translation invariance. The diffusion step τ is transformed via time embedding and concatenated to the invariant node features h [T ] in the input.Training and inference. We optimize the VLB for training, which, interestingly, still has a surrogate in the noise-prediction form when specifying the factors λ(τ ) as 1 (proof in App. A.1):EQUATION)The inference process is similar to [18] but with additional applications of P in intermediate steps to keep all samples in the subspace X P . Details are in Alg. 1 and 2.Distinct from the unconditional generation, in the conditional scenario the target distribution should instead be SE(3)-equivariant w.r.t. the given frames, as elucidated in Eq. 3. The following theorem describes the constraints to consider when designing the prior and transition kernel.Theorem 4.3. If the prior p T (x [T ] T |x [Tc] c ) is SE(3)-equivariant, the transition kernels p τ -1 (x [T ] τ -1 |x [T ] τ , x [Tc] c ), ∀τ ∈ {1, • • • , T } are SE(3)-equivariant, the marginal 5 p τ (x [T ] τ |x [Tc] c ), ∀τ ∈ {0, • • • , T } is SE(3)-equivariant.5 Marginal refers to marginalizing the intermediate states in reverse process while still conditioning on x[Tc] c . Flexible equivariant prior. There are in general many valid choices for the prior while satisfying SE(3)-equivariance. We provide a guidance on distinguishing feasible designs when using Gaussianbased prior in the proposition below.Proposition 4.4. N (µ(x [Tc] c ), I) is SE(3)-equivariant w.r.t. x [Tc] c if µ(x [Tc] c ) is SE(3)-equivariant.Proof is in App. A.2. Notably, the mean of the prior x[T ] r := µ(x [Tc]c ) naturally serves as an anchor to transit the geometric information in the provided trajectory to the target distribution we seek to model. For instance, one can choose it as a linear combination of the CoMs of the given frames, i.e., x[T ] r = 1 T ×N ⊗ s∈[Tc] w (s) x(s)c , where s∈[Tc] w (s) = 1 are fixed parameters determined a priori [21, 13] . However, this choice does not leverage temporal consistency of the trajectory and incurs extra effort in optimization, since the model needs to learn to reconstruct the complex structures from points all located at the CoM. In contrast, we propose the following instantiation:EQUATIONfor all t ∈ [T ], where each x(t)r is a point-wise linear combination of x(s) c , an SE(3)-equivariant transformation of the conditioning frames, with w (t,s) ∈ R N being the weights. We first obtainx[Tc] c , ĥ[Tc] c = f η (x [Tc] c , h [Tc]c ) where f η is a lightweight two layer EGTN that aims to synthesize the conditional information. The w (t,s) is then derived as,EQUATIONEQUATIONHere γ ∈ R T are learnable parameters, and w (t,s) is parameterized such that it has a sum of 1 N when s goes through [T c ] to satisfy the constraint in Eq. 9 for translation equivariance. Interestingly, as we formally illustrated in Theorem A.4, our parameterization of the prior theoretically subsumes the CoM-based priors [21, 13] and the fixed point-wise priors when γ, ĥ[Tc] c , andx[Tc] creduce to specific values. Such theoretical result underscores the benefit of our design since it permits the model to dynamically update the prior, leading to better optimization. The parameters η and γ are updated during training with gradients coming from optimizing the variational lower bound.Transition kernel. We need to modify the forward and reverse process such that they both match the proposed prior. The forward process is modified as q(x[T ] τ |x [T ] τ -1 , x [Tc] c ) := N (x [T ] τ ; x r + √ 1 -β τ (x [T ] τ -1 -x r ), β τ I), which ensures q(x [T ]T |x[Tc] c ) matches the equivariant prior x r (proof in App. A.2). The reverse transition kernel is given by p τ -1 (x[T ] τ -1 |x [T ] τ , x [Tc] c ) = N (µ θ (x [T ] τ , τ, x [Tc] c ), σ 2 τ I).Similar to the unconditional case, we also adopt the noise prediction objective by rewriting µ θ (x[T ] τ , x [Tc] c , τ ) = x [T ] r + 1 √ ατ x [T ] τ -x [T ] r -βτ √ 1-ᾱτ ϵ θ (x [T ] τ , x [Tc]c , τ ) . The denoising network ϵ θ is implemented as an EGTN but with its output subtracted by the input for translation invariance, hence the translation equivariance of µ θ .Training and inference. Optimizing the VLB of our diffusion yields the following objective:EQUATION)after simplification (proof in App. A.2). The training and inference procedures are in Alg. 3 and 4. 
experiments: 
 We evaluate GeoTDM on N-body physical simulation, molecular dynamics, and pedestrian trajectory forecasting, in both conditional ( § 5.1) and unconditional generation ( § 5.2) scenarios. We ablate our core design choices and demonstrate additional use cases in § 5.3. Experimental setup. We adopt three scenarios in the collection of N-body simulation datasets, including 1. Charged Particles [26, 43] , where N = 5 particles with charges randomly chosen between +1/ -1 are moving under Coulomb force; 2. Spring Dynamics [26] , where N = 5 particles with random mass are connected by springs with a probability of 0.5 between each pairs, and force on the spring follows Hooke's law; 3. Gravity System [2] , where N = 10 particles with random mass and initial velocity moves driven by gravitational force. For all three datasets, we use 3000 trajectories for training, 2000 for validation, and 2000 for testing. For each trajectory, we use 10 frames as the condition and predict the trajectory for the next 20 frames.Baselines. We involve baselines from three families. Frame-to-frame prediction models: Radial Field [27] , Tensor Field Network [56] , SE(3)-Transformer [9] , and EGNN [43] ; Deterministic trajectory model: EqMotion [62] ; Probabilistic trajectory model: SVAE [67] . Details in App. B.3.Metrics. We employ Average Discrepancy Error (ADE) and Final Discrepancy Error (FDE), which are widely adopted for trajectory forecasting [67, 62] , given by ADE(x [T ] ,y [T ] ) = 1 T N T -1 t=0 N -1 i=0 ∥x (t) i -y (t) i ∥ 2 , and FDE(x [T ] , y [T ] ) = 1 N N -1 i=0 ∥x (T -1) i -y (T -1) i ∥ 2 .For probabilistic models, we report average ADE and FDE derived from K = 5 samples.Implementation. The input data are processed as geometric graphs. For example, on Charged Particles, the node feature is the charge, and the graph is specified as fully connected without selfloops. We use 6 layers in EGTN with hidden dimension of 128. We use T = 1000 and the linear noise schedule [18] . More details in App. B.2.Results. We present the results in Table 1 , with the following observations. 1. Trajectory models generally yield lower error than frame to frame prediction models since they mitigate errors accumulated in iterative roll-out. 2. The equivariant methods, e.g., EGNN, EqMotion, and our GeoTDM significantly improves over the non-equivariant model SVAE, demonstrating the importance of injecting physical symmetry into the modeling of geometric trajectories. 3. By directly modeling the distribution of geometric trajectories with equivariance, GeoTDM achieves the lowest ADE and FDE on all three tasks, showcasing the superiority of the proposed approach.Experimental setup. We employ the MD17 [5] dataset, which contains the DFT-simulated molecular dynamics (MD) trajectories of 8 small molecules, with the number of atoms for each molecule ↓ Class ↑ Pred ↓ Marg ↓ Class ↑ Pred ↓ Marg ↓ Class ↑ Pred ↓ Marg ↓ Class ↑ Pred ↓ SVAE [67]3.150 2.50×10 -2 0.2123 2.941 3.54×10 -6 0.1312 3.083 8.29×10 -5 0.2580 2.736 3.73×10 -5 0.604 EGVAE [43] 3.007 3.17×10 -4 0.0136 3.314 3.76×10 -6 0.0221 2.054 2.77×10 -5 0.0457 3.570 2.02×10 -5 0.0212 GeoTDM 0.770 1.17×10 -1 0.0093 0.559 1.82×10 -1 0.0135 0.539 1.12×10 -1 0.0118 0.954 2.02×10 -1 0.0116 ranging from 9 (Ethanol and Malonaldehyde) to 21 (Aspirin). For each molecule, we construct a training set of 5000 trajectories, and 1000/1000 for validation and testing, uniformly sampled along the time dimension. Different from [62] , we explicitly involve the hydrogen atoms which contribute most to the vibrations of the trajectory, leading to a more challenging task. The node feature is the one-hot encodings of atomic number [44] and edges are connected between atoms within three hops measured in atomic bonds [48] . We adopt the same set of baselines as the N-body experiments.Results. As depicted in Table 3 , GeoTDM achieves the best performance on all eight molecule MD trajectories, outperforming previos state-of-the-art approach EqMotion. In particular, GeoTDM obtains an improvement of 23.1%/15.3% on average in terms of ADE/FDE, compared with the previous state-of-the-art approach EqMotion, thanks to the probabilistic modeling which is advantageous in capturing the stochasticity of MD simulations.Experimental setup. We apply our model to ETH-UCY [35, 28] dataset, a challenging and largescale benchmark for pedestrian trajectory forecasting. There are five scenes in total: ETH, Hotel, Univ, Zara1, and Zara2. Following standard setup [14, 67] , we use 8 frames (3.2 seconds) as input to predict the next 12 frames (4.8 seconds). The pedestrians are viewed as nodes and their 2D coordinates are extracted from the scenes. Edges are connected for nodes within a preset distance measured from the final frame in the given trajectory. The metrics are minADE/minFDE computed from 20 samples.For baselines, we compare with existing generative models that have been specifically designed for pedestrian trajectory prediction, including GANs: SGAN [14] , SoPhie [39] ; VAEs: PECNet [32] , Traj++ [40] , BiTraP [68] , SVAE [67] ; and diffusion: MID [12] . Baseline results are taken from [67] .Results. From Table 2 , we observe that our GeoTDM obtains the best predictions on 3 out of the 5 scenarios while achieving the lowest average ADE and FDE. It is remarkable since compared with these baselines specifically tailored for the task of pedestrian trajectory forecasting, GeoTDM does not involve special data preprocessing of the trajectories through rotations or translations, does not involve extra auxiliary losses to optimize during training, and does not require task-specific backbones, demonstrating its general effectiveness across different geometric domains.Experimental setup. For generation we reuse the Charged Particle dataset and the MD17 dataset. We follow the same setup as the conditional case, except that we generate trajectories with length 20 from scratch. We compare with SGAN [14] , SVAE [67] (slightly modified to enable generation from scratch), and a VAE-modified version of EGNN [43] , dubbed EGVAE (see App. B). The results of SGAN on MD17 is omitted due to mode collapse during training.Metrics. We adopt three metrics adapted from time series generation to quantify the generation quality of the geometric trajectories: Marginal scores [34] measure the distance between the empirical probability density functions of the generated samples and the ground truths; Classification scores [25] are computed as the cross-entropy loss given by a trajectory classification model, trained on the task of distinguish whether the trajectory is generated or real; Prediction scores [72] are the MSEs of a train-on-synthetic-test-on-real trajectory prediction model (a 1-layer EqMotion) that takes as input the first half of the trajectories to predict the other half. Results. Quantitative results are displayed in Table 5 and 4 for N-body and MD17. Notably, GeoTDM delivers samples with much higher quality than the baselines. On Charged Particles, GeoTDM achieves a classification score of 0.556, indicating its generated samples are generally indistinguishable with the ground truths. We observe similar patterns on MD17, where GeoTDM obtains remarkably lower marginal scores, higher classification scores, and lower prediction scores, showcasing its strong capability to model complex distributions of geometric trajectories on various geometric data. Visualizations are in Fig. 5 and more in App. D. Ablations on diffusion prior. We investigate different priors, including non-equivariant N (0, I) (i.e., DDPM [18] ), equivariant but fixed CoM prior N (CoM(x (Tc-1) c), I) and point-wise equivariant prior N (x (Tc-1) c , I). In Table 6 we see that non-equivariant prior leads to significantly worse performance. The CoM prior, though equivariant, is still inferior due to extra overhead in denoising the nodes initialized around the CoM to the original geometry. GeoTDM yields the lowest error due to the flexible learnable prior.Ablations on EGTN. We further ablate the designs of the denoising model. 1. Equivariance. We replace all EGCL layers into non-equivariant MPNN [11] layers with same hidden dimension, leading to non-equivariant transition kernels. The performance becomes much worse, verifying the necessity of equivariance. 2. Attention. We substitute the attentions in temporal layers by equivariant convolutions (see App. B). Compared with this variant, GeoTDM enjoys larger capacity with attention and yields lower prediction error especially on Charged Particle where the particles generally move faster. 3. Temporal shift invariance. We employ relative temporal embeddings in attention, which enhances the generalization. Notably, the FDE improves from 0.330 to 0.258 on Charged Particle compared with the absolute temporal embedding.Temporal interpolation. GeoTDM is able to perform interpolation as a special case of the conditional case. We demonstrate such capability on Charged Particle. The model is provided with the first 5 and last 5 frames, and the task is to generate the intermediate 20 frames as interpolation. GeoTDM reports an ADE of 0.055 on the test set, while a linear interpolation baseline reports an ADE of 0.171. From the qualitative visualizations in Fig. 2 , we clearly see that GeoTDM can capture the complex dynamics and yield high-quality non-linear interpolations between the given initial and final frames.Optimization. We further illustrate that GeoTDM can conduct optimization [31, 33] on given trajectories (e.g., those simulated by an EGNN) by simulating K steps through the forward diffusion and then performing the reverse denoising. From Fig. 2 we see the distance between the optimized trajectory and GT gradually decreases as the optimization step grows. This reveals GeoTDM can effectively optimize the given trajectory towards the ground truth distribution. 
discussion: 
 Limitations. Akin to other diffusion models, GeoTDM resorts to multi-step sampling which may require more compute. We present empirical runtime benchmarks and more discussions in App. C.3.It is designed to preserve the symmetry of geometric systems, achieved by using EGTN, a novel SE(3)-equivariant geometric trajectory model, as the denoising network. We evaluate GeoTDM on various datasets for unconditional generation, interpolation, extrapolation and optimization, showing that it consistently outperforms the baselines. Future works include streamlining GeoTDM and extending it to more tasks such as protein MD, robot manipulation, and motion synthesis.Base distribution. We require the base distribution to be SO(3)-invariant. In practice, we let p T (xT ) = Ñ (0, I) to be the Gaussian distribution in the translation-invariant subspace Transition kernel. For the transition kernel, we specify it as p θ (x[T ] τ -1 | x[T ] τ ) = Ñ ( μθ (x [T ]τ , τ ), σ 2 τ I). In order to ensure p θ (Rx[T ] τ -1 | Rx [T ] τ ) = p θ (x [T ] τ -1 | x[T ] τ ), it suffices to make μθ (x [T ]τ , τ ) an SO(3)equivariant function. In this way,EQUATION= Ñ (RxEQUATION= Ñ (RxEQUATION= Ñ (R(xEQUATION= Ñ (xEQUATION= Ñ (xEQUATION= p θ (xEQUATIONwhich permits the SO(3)-equivariance of the transition kernel. In our implementation, we further re-parameterize μθ (xEQUATIONwhere we instead ensure εθ (x[T ]τ , τ ) to be SO(3)-equivariant and its output should lie in the subspace X P .We now prove the following proposition, which states that if the base distribution is SO(3)-invariant and the transition kernel is SO(3)-equivariant, then the marginal at any diffusion time step is also SO(3)-invariant. Proof. The proof is given by induction.Induction base. When τ = T , we have the marginal being the prior p T (x[T ]T ), which is SO(3)invariant.Induction step. Suppose the marginal at diffusion time step τ is SO(3)-invariant, i.e., p τ (x[T ] τ ) = p τ (Rx [T ]τ ), then we have the following derivation for the marginal at time step τ -1:EQUATIONEQUATIONEQUATION= p τ -1 (xEQUATIONEQUATIONNotably, for the final step at τ = 0, the marginal p 0 (Rx[T ]0 ) is also SO(3)-invariant, indicating the final sample from the entire geometric trajectory diffusion process resides in an SO(3)-invariant distribution, hence the physical symmetry being well preserved.Algorithm 1 Training Procedure of GeoTDM-uncond 1: repeat 2: Sample ε[T ] ∼ Ñ (0, I) [T ] , τ ∈ Unif({1, • • • , T }), x[T ] ∼ Ddata 3: x[T ] τ ← √ ᾱτ x[T ] + √ 1 -ᾱτ ε[T ] 4:Take gradient descent step on ∇ θ ∥ε [T ] εθ (x[T ] τ , τ )∥ 2 2 5: until converged Algorithm 2 Sampling Procedure of GeoTDM-uncond 1: Sample x[T ] T ∼ Ñ (0, I) [T ] 2: for τ ← T , • • • , 1 do 3: Sample z[T ] τ ∼ Ñ (0, I) [T ] if τ > 1 else z[T ] τ = 0 4: x[T ] τ -1 ← 1 √ ατ x[T ] τ -1-ατ √ 1-ᾱτ εθ (x [T ] τ , τ ) + στ z[T ] τ 5: end for 6: return x [T ] 0 Algorithm 3 Training Procedure of GeoTDM-cond 1: repeat 2: x [T ] r ← EquiPrior η,γ (x [Tc] c ) {Eq. 9} 3: Sample ϵ [T ] ∼ N (0, I), τ ∈ Unif({1, • • • , T }), (x [T ] , x [Tc] c , ) ∼ p data 4: x [T ] τ ← √ ᾱτ (x [T ] -x [T ] r ) + x [T ] r + √ 1 -ᾱτ ϵ [T ] 5:Take gradient descent step on∇ θ,η,γ ∥ϵ [T ] -ϵ θ (x [T ] τ , x [Tc] c , τ )∥ 2 6: until converged Algorithm 4 Sampling Procedure of GeoTDM-cond 1: x [T ] r ← EquiPrior η,γ (x [Tc] c ) {Eq. 9} 2: Sample x [T ] T ∼ N (x [T ] r , I), x [Tc] c ∼ D data 3: for τ ← T , • • • , 1 do 4: Sample z [T ] τ ∼ N (0, I) [T ] if τ > 1 else z [T ] τ = 0 5: x [T ] τ -1 ← 1 √ ατ x [T ] τ -x [T ] r -1-ατ √ 1-ᾱτ ϵ θ (x [T ] τ , x [Tc] c , τ ) + x [T ]r + στ z[T ] τ 6: end for 7: return x[T ] 0In the conditional case, we target on modeling the conditional distribution p(x [T ] | x[Tc] c ). The desired constraint is the following equivariance condition: p(x [T ] | x [Tc] c ) = p(g • x [T ] | g • x [Tc] c ), for all g ∈ SE(3).Construction of the equivariant prior. The prior is constructed through Eq. 9. Here we formally show that this guarantees SE(3)-equivariance of the prior. For convenience we repeat Eq. 9 below.EQUATIONThen, we haveEQUATION=EQUATION= s∈[Tc] w (t,s) (Rx (s) c ) + s∈[Tc]w (t,s) r,EQUATION=EQUATION∀ rotation matrix R and r ∈ R 3 , which completes the proof.Base distribution. We propose to leverage the following base distribution.EQUATIONwherex [T ] r = EquiPrior(x [Tc] c ) is SE(3)-equivariant with respect to the condition x [Tc]c . With such choice, the base distribution above is SE(3)-equivariant, since p T (RxEQUATION= N (RxEQUATION= N (xEQUATIONwhere the last equation is due to det(R ⊤ R) = I, ∥x[T ] T -x [T ] r ∥ 2 = ∥Rx [T ] T -Rx [T ]r ∥ 2 , which also gives the proof for Theorem 4.4 by a mild substitution of the notations.Transition kernel. The transition kernel is given byEQUATIONwhere µ θ (x[T ] τ , x [Tc]c , τ ) parameterized to be SE(3)-equivariant with respect to its input x[T ] τ , x [Tc]c . In practice, we re-parameterize it as,EQUATIONwhere ϵ θ (x[T ] τ , x [Tc] c , τ ) is an SO(3)-equivariant but translation-invariant function. It is then easy to see that µ θ (x [T ] τ , x [Tc]c , τ ) meets the SE(3)-equivariance as desired. Proposition A.2. With such parameterization, optimizing the variational lower bound is equivalent to optimizing the following objective, up to certain re-weighting:EQUATION)Proof. We define q(x[T ] τ |x [T ] τ -1 ) := N (x [T ] τ ; x r + √ 1 -β τ (x [T ] τ -1 -x r ), β τ I), which yields q(x [T ] τ |x [T ] 0 ) = N (x [T ] τ ; x r + √ ᾱτ (x [T ]0 -x r ), (1 -ᾱτ )I). The proof then generally follows [18] but with all latent variables in [18] being replaced by x[T ] τ -x [T ]r . Then the terms in the VLB are given by,EQUATIONEQUATIONEQUATIONEQUATION)which is equivalent to Eq. 39 up to certain re-weighting factors.ForL T = D KL (q(x [T ] T |x [T ] 0 )∥p(x [T ]T )), it is does not contribute to the gradient since it is irrelevant to θ, and x r is also cancelled out in computing the KL, thus stopping the gradient from passing to η and γ.Analogous to the unconditional case, we have the following proposition, indicating that if the base distribution is SE(3)-equivariant and the transition kernel is SE(3)-equivariant, then the marginal is also SE(3)-equivariant.Proposition A.3. If the base distribution p T (x [T ] T | x [Tc] c ) is SE(3)-equivariant and the transition kernels p τ -1 (x [T ] τ -1 | x [T ] τ , x [Tc] c ) of all diffusion steps τ ∈ {1, • • • , T } are SE(3)-equivariant, then the marginal 6 p τ (x [T ] τ | x [Tc] c ) at any diffusion step τ ∈ {0, • • • , T } is SE(3)-equivariant.Proof. The proof is similarly given by induction.Induction base. When τ = T , the distribution is the base distribution p T (x[T ] T | x [Tc] c ) is SE(3)- equivariant, as it is designed.Induction step. Suppose the marginal at diffusion step τ , i.e., p τ (xEQUATIONEQUATIONEQUATIONEQUATION)EQUATIONwhich concludes the proof.Theorem A.4. The prior implemented by the parameterization in Eq. 9, 10, and 11 subsumes CoM-based priors and fixed point-wise priors.Proof. We repeat the parameterizations specified by Eq. 9, 10, and 11 below for better readability.EQUATIONEQUATIONEQUATION)We first show x ), ĥ(s)EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONwhere1 Tc s∈[Tc] CoM(x (s) c) is the generalization of the CoM-based priors [21, 13] in the multiple frame conditioning scenario, which reduces to CoM(x(0) c ) when T c = 1.r can reduce to fixed point-wise priors is straightforward. Letx(s) c = x (s)c , ĥ[Tc] = Onehot(s * )1 Tc×N and γ (t) = 1, ∀t. Then w (t,s) = Onehot(s * )1 Tc×N . Therefore,EQUATIONEQUATIONEQUATIONwherex (s * ) cis the point-wise equivariant prior, and s * ∈ [T c ] is the frame index in the conditioning trajectory for this specific prior.We also provide an illustrative comparison of these equivariant priors in Fig. 3 .A.4 Proof of Theorem 4.1 Theorem 4.1 (SE(3)-equivariance of EGTN). Let x ′[T ] , h ′[T ] = f EGTN x [T ] , h [T ] , E . Then we have g • x ′[T ] , h ′[T ] = f EGTN g • x [T ] , h [T ] , E , ∀g ∈ SE(3).Proof. f EGTN is a stack of L EGNN and temporal attention layer in alternated fashion, formally written asf EGTN = f attn • f EGNN • • • • • f attn • f EGNN L×(fattn•fEGNN). Since the chain of SE(3)-equivariant function is also SE(3)-equivariant, it suffices to prove f attn is SE(3)-equivariant, in that the SE(3)equivariance of EGNN directly follows [43] .It is directly verified that the attention coefficients a (t,s) ∈ R in Eq. 5 and the query q [T ] , key k [T ] , and value v [T ] are all SE(3)-invariant, since they are derived based on the SE(3)-invariant inputFixed point-wise prior Ours h [T ] . This directly leads to the SE(3)-invariance of the updated node feature h ′ [T ] . For the updated coordinates,1 𝑇 ! # "∈[% ! ] CoM(𝐱 ! " ) 𝐱 ! (" * ) : Mean of the prior : condition 𝐱 ! [# ! ] : target 𝐱 [#]xEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONwhere the variables with subscript tr refers to their transformed counterparts when the input x [T ] is transformed into Rx [T ] + r. Thus it completes the proof of SE(3)-equivariance of the temporal attention f attn and hence the entire f EGTN .We use Distributed Data Parallel on 4 Nvidia A6000 GPUs to train all the models. The training on NBody and ETH-UCY take around 12 hours while each MD17 training phase takes about a day. Our CPUs were standard intel CPUs.We provide the detailed hyper-parameters of GeoTDM in Table 7 . We adopt Adam optimizer with betas (0.9, 0.999) and ϵ = 10 -8 . For all experiments, we use the linear noise schedule per [18] with β start = 0.02 and β end = 0.0001. For the frame-to-frame prediction models, including RF [27] , EGNN [43] , TFN [56] , and SE(3)-Transformer [9] , we adopt the implementation in the codebase maintained by [43] . To yield a strong comparison, instead of taking one frame as input to directly predict the final frame, we employ a discretized NeuralODE [4] -style training and inference procedure. In particular, we train the models with position x (t) and velocity (computed as the difference of the current and previous frame, i.e., v (t) = x (t) -x (t-1) ) as input to predict the next velocity v(t+1) . The position for the next step is integrated as x(t+1) = x (t) + v(t+1) . The training loss is computed as the Mean Squared Error (MSE) between the predicted position x(t+1) and the ground truth position x (t+1) gt. In inference time, a rollout prediction is conducted, which iteratively predict the next step by feeding the predicted position and velocity at the current step, for a total of T steps. We follow the hyper-parameter tuning guideline for these baselines by [16] which conduct a random search over the space spanned by the number of layers in {2, 4, 6, 8}, the hidden dimension {32, 64, 128}, learning rate {5e-3, 1e-3, 5e-4, 1e-4}, and batch size 32, 64, 128, 256, and select the model with best performance. All models are trained towards convergence with an early-stopping counter of 5, with validation performed every 20 epochs.For EqMotion, we directly adopt the code by [62] and their suggested hyper-parameters for the N-body datasets and MD17 datasets.For SVAE [67] and SGAN [14] , these methods are originally developed for the pedestrian trajectory forecasting task. The backbone model that processes the input trajectory consists of social pooling operation and GRU or LSTM blocks for temporal processing. In order the make them favorable in tackling geometric systems which additionally include node features and edge features, we replace the social pooling operations by MPNNs [11] in both encoder (or discriminator) and decoder (or generator) to synthesize the information on the geometric graph. The temporal module is still kept as GRU for SVAE and LSTM for SGAN, following their original implementations. We also search over the best hyper-parameters which additionally involve the KL-divergence weight in {1, 0.1, 0.01, 0.001} for SVAE according to the validation ELBO. For EGVAE, we replace the MPNNs in SVAE by EGNN [43] , and restructured the latent space of the prior with both equivariant and invariant latent features. By this means, EGVAE is also guaranteed to model an equivariant distribution in the conditional case and an invariant distribution in the unconditional case.In detail, the EGCL layer [43] is given by:EQUATIONEQUATIONEQUATIONEQUATIONwhere φ m , φ h , and φ x are all MLPs. We also provide a schematic of our proposed EGTN in Fig. 4 for better illustration.All these metrics are evaluated on a set of model samples with the same size as the testing set.Marginal score is computed as the absolute difference of two empirical probability density functions. Practically, we collect the x, y, z coordinates at each time step marginalized over all nodes in all systems in the predictions and the ground truth (testing set). Then we split the collection into 50 bins and compute the MAE in each bin, finally averaged across all time steps to obtain the score. Note that on MD17, instead of computing the pdf on coordinates, we compute the pdf on the length of the chemical bonds, which is a clearer signal that correlates to the validity of the generated MD trajectory, since during MD simulation the bond lengths are usually stable with very small vibrations. Marginal score gives a broad statistical measurement how each dimension of the generated samples align with the original data.𝐪 (")𝐤 (",%) 𝜓(𝑡 -𝑠) 𝐯 (",%) can also be processed using cross-attention. The relative temporal embedding ψ(t -s) is added to the key and value. DotProd refers to dot product and Softmax is performed over indexes of s.𝐚 (",$) DotProd, Softmax 𝜑 𝐪 𝜑 𝐤 𝜑 𝐯 + + 𝐡 (") Prod, ∑ 𝑠 + 𝐡 )(") 𝐱 (") 𝐱 (") -𝐱 (%) 𝐱 )(") Prod, ∑ 𝑠 + Temporal Attention Spatial EGCL ×𝐿 𝐡 [+] 𝐱 [+] 𝐡 - [+'] 𝐱 - [+'] 𝐱 - (%) 𝐡 - (%) 𝐡′ [+] 𝐱′ [+]Classification score is computed as the cross-entropy loss of a sequence classification model that aims to distinguish whether the trajectory is generated by the model or from the testing set. To be specific, we construct a dataset mixed by the generated samples and the testing set, and randomly split it into 80% and 20% subsets for training and testing. Then the model is trained on the training set and the classification score is computed as the cross-entropy on the testing set. We use a 1-layer EqMotion with a classification head as the model. The classification score provided intuition on how difficult it is to distinguish the generated samples and the original data.Prediction score is computed as the MSE loss of a train-on-synthetic-test-on-real sequence to sequence model. In detail, we train a 1-layer EqMotion on the sampled dataset with the task of predicting the second half of the trajectory given the first half. We then evaluate the model on the testing set and report the MSE as the prediction score. Prediction score provides intuition on the capability of the generative model on generating synthetic data that well aligns with the ground truth.Since attention is utilized to extract temporal information, the time complexity scales quadratically with the length of the input trajectory, both during training and inference. In practice, we can instead train models on shorter trajectories and compose them during inference for longer trajectories, in both unconditional and conditional cases. For target trajectories with length T , we can first decompose it into K several equal-length7 non-overlapping intervals with time span ∆T . Then, for the unconditional case, we haveEQUATIONby assuming mild conditional independence, where p(x [∆T ] ) is an unconditional model for trajectory with length ∆T , and p(x k∆T +[∆T ] | x (k-1)∆T + [∆T ] ) can be learned by a conditional model for short trajectories. The conditional case directly follows by factorizing into products of conditional distribution over shorter trajectories.We provide a demonstration of such technique as gifs in the supplementary file. We provide results in the unconditional generation setting for T = 100. The results are in Table 8 .Compared with the conditional setting, the unconditional generation is more challenging in that is needs to generate trajectories without any given reference geometries. We observe a drop in performance when T is decreased from 1000 to 100. However, the performance with only 100 diffusion steps is still significantly better than SVAE.In the table below we display the generation metrics and the inference time per batch with batch size 128 on MD17 Aspirin molecule. We compare GeoTDM with EGVAE, an autoregressive VAE-based method with EGNN as the backbone. Here GeoTDM-100 and GeoTDM-1000 refer to GeoTDM using 100 and 1000 diffusion steps, respectively. We observe that GeoTDM-100 is approximately 10 times slower than EGVAE, since the model requires 100 calls of the denoising network to generate one batch, while EGVAE consumes the same number of calls as the length of the trajectory (20 in this case) due to autoregressive modeling. Although GeoTDM is slower, the gain in performance is significant and the quality of the generated trajectory is remarkably better than that of EGVAE. When further increasing the number of diffusion steps to 1000, the performance becomes better while requiring much more compute.However, it is worth noticing that all these deep learning-based methods are significantly faster than traditional methods like DFT, which typically requires hours to even several days to converge depending on the scale of the system, according to OCP [3] . Therefore, although GeoTDM becomes slower than VAEs when using larger number of diffusion steps, it is still much faster than DFT, which indicates its practical value in generating geometric trajectories like molecular dynamics simulation.The computation overhead of diffusion models compared with VAEs or GANs has been a well-known issue. We recognize enhancing the efficiency of GeoTDM as an interesting direction of future work, potentially through adopting faster solvers like DDIM [50] or DPMSolver [29] , performing consistency distillation [51] , or developing latent diffusion models [64] that take advantage of a more compact representation of the spatio-temporal geometric space.We provide the standard deviations in Table 10 and 11. Below we discuss the unique challenges for designing GeoTDM compared with MID [12] and geometric diffusion models like GeoDiff [66] and GeoLDM [64] , and how we tackle these challenges.Modeling geometric trajectories. Although MID can model trajectories, it leverages Trajec-tron++ [40] backbone which takes as input the position vectors through a Transformer network. It requires non-trivial effort to incorporate additional node features and edge features into MID, while for GeoTDM, we design a general backbone EGTN that can process geometric trajectories while preserving equivariance. Existing geometric diffusion models (e.g., GeoDiff and GeoLDM) never consider modeling the temporal dynamics and their backbone can only work on static (single-frame) geometric strctures.Incorporating equivariance into temporal diffusion. While geometric diffusion models have discussed proper ways to inject equivariance into diffusion models, it is unclear how to preserve equivariance when each hidden variable in the diffusion process has an additional dimension of time. In this work, we formally define equivariance constraint we want to impose on the marginal distribution, and how to design the prior and transition kernel in order to fulfill the constraint, in the context where all hidden variables are geometric trajectories. This is technically very different from existing works (e.g., GeoDiff and GeoLDM) since the dimension of the data is fundamentally different, which leads to different analyses.Consideration of both conditional and unconditional generation scenarios. MID is only designed and evaluated in the conditional setting where the task is to forecast the future trajectory given initial frames. GeoDiff and GeoLDM only operate in the unconditional setting where the task is to generate the structure without any initial 3D structure information. In this work, we systematically discuss both unconditional and conditional generation for geometric trajectories, and elaborate on how to design the prior and transition kernel to meet the equivariance constraint.Parameterization of the learnable equivariant prior. In the conditional case, we propose to parameterize the equivariant prior with a lightweight EGTN. Such appraoch offers more flexibility in the equivariant prior, enabling optimizing it during training, which is also proved to be able to subsume existing center-of-mass (CoM) based parameterization (see Theorem A.4 in Appendix).Experiments in ablation studies also verify the superiority of such design.We summarize the points above in Table 12 .We provide more visualizations in Fig. 5 , 6, 8, 7, 9, and 10. Please refer to their captions for the detailed descriptions. 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: [Yes] Justification: We upload the code in supplementary file. Guidelines:! [# ! ] 𝜏 = 0 𝜏 = 𝒯 𝐱 ! [# ! ] 𝜏 = 0 𝜏 = 𝒯 𝐱 ! [# ! ] 𝜏 = 0 𝜏 = 𝒯 𝜏 = 0 𝜏 = 𝒯 𝐱 ! [# ! ]• The answer NA means that paper does not include experiments requiring code. • 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. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: See Sec. C.4. 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, confidence intervals, or statistical significance tests, at 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 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 out 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 is 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 of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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, or scraped datasets)?Answer: [NA] Justification: Our method does not have a high risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or 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?Answer: [Yes]Justification: We have credited all code and dataset sources.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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.Correspondence to Jiaqi Han: jiaqihan@stanford.edu. Code is available at https://github.com/ hanjq17/GeoTDM.The analyses in this paper also hold for the general nD case, e.g., 2D.Following convention we use the notation Rx to denote xR ⊤ .Technically, such a condition is impossible since SE(3) is noncompact, but we show that zero-centering the trajectories and enforcing SO(3)-invariance is equivalent.Here the marginal refers to marginalizing the intermediate states in previous diffusion steps, while still being conditional on the input condition x [Tc] c .In fact they do not necessarily need to be equal-length. Here we make such assumption for conciseness of the presentation. 
 
 Appendix:  


Decision: The paper introduces GeoTDM, a diffusion model designed for modeling the temporal distribution of 3D geometric trajectories, incorporating SE(3)-equivariant spatial convolution and temporal attention. Reviewers appreciated the novel approach, noting its broad applicability to physical simulations, molecular dynamics, and pedestrian motion, and recognized the well-presented methodology and extensive experiments validating its effectiveness. 

Reviewers highlighted the paper’s strengths in its broad applicability across multiple domains, well-executed and thorough experiments, and the clear and detailed presentation of the methodology. The use of equivariant temporal kernels and learnable geometric priors was praised for capturing complex spatial and temporal dynamics effectively. Additionally, the ablation studies and visualizations provided valuable insights into the model’s performance, showcasing its practical utility and setting it apart from existing approaches.

Reviews: review: paper_summary: paper_summary: The paper proposed the first diffusion model for modeling the temporal distribution of 3D geometric trajectories, while previous works only operate on static structures. It demonstrates the equivariant temporal kernels can lead to density with desired symmetry and develop a novel transition kernel leveraging SE(3)-equivariant spatial convolution and temporal attention. The experiments demonstrates it can generates realistic geometric trajectories with significantly higher quality.
strength_weakness: strength_weakness: None
questions: questions: None
limitations: limitations: None
paper_summary: paper_summary: The paper introduces geometric trajectory diffusion models for generation of particle, pedestrian, or molecular trajectories. The architecture consists of EGNN layers within a temporal frame and temporal attention across frames computed with relative temporal ecodings. The architecture is shown to possess the appropriate equivariances. Key points of novelty include a learnable prior as a function of conditioning frames in the conditional generation case. Experiments are conducted on particle, molecular, and pedestrian dynamics datasets.
strength_weakness: strength_weakness: * The work does not score highly in conceptual novelty. While all architecture choices are sensible, they are relatively straightforward and do not seem surprising, insightful, or inspired. An extension or exploration of more sophisticated equivariant architectures could have strengthened the paper.
* The exposition of the method drags at times, re-proving well-known statements about equivariance in neural networks and geometric diffusion models.
* The experiments, while diverse and broad, also suffer from being scattered and not necessarily the most convincing individually. Little effort is spent on exploring the ways in which a trajectory diffusion could be employed and the new problems that could be solved with the mode. To me, the novel capabilities of such a method are more exciting than marginal gains across a number of established datasets, executed checkbox-style.
questions: questions: See above
limitations: limitations: Yes
paper_summary: paper_summary: In this paper, the authors propose geometric trajectory diffusion models (GeoTDM) to model temporal distribution of geometric trajectories
while keeping the desirable physical symmetry of the trajectories.
strength_weakness: strength_weakness: I am concerned that dynamic system and  rotation translational invariant has already been studied in the previous research (listed below, but not quoted here), not like the authors claimed that this is the first study on this topic. There might be some overlaid here.

Both algorithms use the same E(n) equivariant graph neural networks. 

Missing reference: 
Fang Wu, Stan Z. Li, DIFFMD: A Geometric Diffusion Model for Molecular Dynamics Simulations. The work in this previous paper requires that the dynamics must be invariant to rotation or translation. 

The paper needs more experimental results to support it.
questions: questions: P5, line 185: In the definition of the operator P, ⊗ is not defined? It seems to me no need to define operator P if P(x) is defined.
limitations: limitations: Yes.
paper_summary: paper_summary: A new diffusion-based generative model for modeling complex 3D geometric structures with time-evolving trajectories is proposed. By introducing the SE(3) equivariance property, temporal attention, and learnable geometric prior into the discretized diffusion model, the proposed model can achieve high performance in learning the distribution of geometric trajectories while preserving the symmetry of geometric systems. It can also be improved to attain conditional generation by using the equivariant cross-attention mechanism. Various types of experiments show high accuracy and performance of the proposed model in cases such as molecular dynamics and simulation.
strength_weakness: strength_weakness: No major weakness. Please refer to the part of Questions for my concerns.
questions: questions: 1. As shown in the experiment part, the speed and efficiency are still the main limitations of the proposed model. Since the transition kernel and the prior are modified with extra restrictions, will it be difficult to redesign the framework (both forward and backward) into the continuous diffusion ODE/SDE form? (which might be important for further acceleration)

2. Except for the part of processing temporal information, could you compare more details about the proposed EGTN/EGCL with the E(3) Equivariant Diffusion Model in the paper [1]? 

3. Why the geometric restrictions are also important in N-body simulation? What is the 'geometric structure' in these scenarios? It would be better if the author could provide more background knowledge about geometric trajectory and the physics simulation/molecular dynamics in the appendix.

4. Based on the visualization results in the appendix, the change in the molecular structure seems to be 'small'. Will such short trajectories be practically useful/meaningful in application?

[1]: Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.
limitations: limitations: The authors adequately addressed and discussed the limitations about computation. It would be better if the paper could also mention the current gap between the proposed model and the industrial/academic demand in medical/physical fields.

Comments: title: Thank you for the supportive feedback!
Follow-up on the clarifications
Further Clarifications
Thank you for the feedback
Thank you for the feedback
Thank you for the feedback!
comment: Dear Reviewer CUJA,

Thank you for the supportive feedback and the insightful comments that help us improve the manuscript! We will include these in the final version.

Best,

Authors
Thank you for your response and clarification! I decide to raise my rating to 6: Weak Accept. I trust you can revise your paper according to the discussion.
Dear Reviewer CUJA,

Thank you for your follow-up discussion and we have clarified your questions in detail. As the discussion period is coming to an end, we greatly appreciate it if you could let us know whether the clarifications address your concern and we are willing to offer more explanations if there are any questions. Thank you!

Best,

Authors
Dear Reviewer CUJA,

Thank you very much for the timely feedback and the follow-up question! We are glad to further clarify your concern.

> What troubled me is that I am still not convinced what you said that: "DiffMD is not explicitly a generative model on geometric trajectories, while to the best of our knowledge our GeoTDM is the first. ". I believe that both of your methods are a generative model, but each of you approached in some detailed difference, so I need you to clarify this in more details.

Yes, we agree that both DiffMD and GeoTDM are generative models. The core difference lies in whether the method is modeling the distribution of **geometric trajectories** $\mathbf{x}^{[T]}=[\mathbf{x}^{(0)},\mathbf{x}^{(1)},\cdots,\mathbf{x}^{(T-1)}]\in\mathbb{R}^{T\times N\times D}$ which is a **sequence/collection of frames**, or a **single** frame $\mathbf{x}^{(t)}\in\mathbb{R}^{N\times D}$. In particular, our GeoTDM explicitly models the distribution of geometric trajectories $p(\mathbf{x}^{[T]}|\cdot)$ (see line 116-122), while DiffMD is designed to model the distribution of the next frame $p(\mathbf{x}^{(t+1)}|\cdot)$ (see Model Overview section on Page 3 of DiffMD [1]), where $\cdot$ denotes certain conditioning if applicable (e.g., the previous frames).

Such difference leads to several distinctions in model design, inference, and empirical performance, as we will discuss below.
* **Model design.** In order to jointly model $p(\mathbf{x}^{[T]}|\cdot)$, in the transition kernel we need to additionally handle the temporal dimension with size $T$. This motivates us to design the equivariant temporal attention layer (Eq. 5-7), which is absent in DiffMD since DiffMD does not require to explicitly model temporal correlation. Our framework also enables conditioning on a trajectory with multiple frames while DiffMD is based on Markovian assumption and always consumes the previous one single frame to generate the next frame.
* **Inference.** The benefit of directly modeling the distribution of geometric trajectories over single frames also presents at inference time. Within one diffusion loop, our GeoTDM can generate an entire trajectory with $T$ frames, while DiffMD requires an additional outer loop that sweeps through $T$ frames in order to achieve the same effect.
* **Performance.** We also demonstrate the benefit in terms of empirical performance. Since DiffMD does not consider the correspondence between multiple frames, it is more vulnerable to error accumulation when generating a long trajectory through iterative rolling out. Our extra experiment provided in the rebuttal verifies this point, with GeoTDM outperforming DiffMD by a remarkable margin.



> While you claimed that your "modeling the joint distribution of all frames within the geometric trajectory", what I saw is you still look at the individual conditional probability like DiffMD, see line 193. It seems like you take into data sampling of time into consideration, to me, I think this is considered as generalize of DiffMD.

Thank you for the question. However, there might be a misunderstanding here and we would like to respectfully clarify. 

Specifically, line 193 depicts the transition kernel $p\_\theta({\mathbf{x}\_{\tau-1}^{[T]}}|\mathbf{x}\_{\tau}^{[T]})$. Here $\tau$ refers to the **diffusion step** instead of the **frame index on the geometric trajectory**. The superscript $[T]$ indicates that the latent variable here is a geometric trajectory instead of a single frame. Therefore, the individual conditional probability is enforced on **diffusion step**, which is due to the Markovian assumption of the diffusion process. This is a common practice for diffusion models and is shared in both GeoTDM and DiffMD, and we are not claiming any difference on this point.

However, we are **not** enforcing conditional independence on the actual temporal dimension of the geometric trajectory, while DiffMD has a Markovian assumption on molecular dynamics and only models the distribution of a single frame. This is what we refer to by stating 
 "modeling the joint distribution of all frames within the geometric trajectory". This point leads to the core difference between these two approaches, as we have discussed.


> For the above reason, I believe that is why you both used Equivariant Graph Convolution Layer.

We are both using EGCL to process spatial information on the geometric structure. However, since we are modeling the whole trajectory, we need to additionally introduce the equivariant temporal attention layer to process the temporal correspondence, which never presents in DiffMD.


Thank you again for the follow-up discussion! We promise to include the discussion and the distinctions between these two methods in the manuscript. Please let us know if this addresses your concern and we are happy to further clarify if you have any questions.

Best,

Authors
Thank you for your new update! What troubled me is that I am still not convinced what you said that: "DiffMD is not explicitly a generative model on geometric trajectories, while to the best of our knowledge our GeoTDM is the first. ". I believe that both of your methods are a generative model, but each of you approached in some detailed difference, so I need you to clarify this in more details. While you claimed that your "modeling the joint distribution of all frames within the geometric trajectory", what I saw is you still look at the individual conditional probability like DiffMD, see line 193. It seems like you take into data sampling of time into consideration, to me, I think this is considered as generalize of DiffMD. For the above reason, I believe that is why you both used Equivariant Graph Convolution Layer.
Dear Reviewer CUJA,

We greatly appreciate your advice and promise to incorporate the discussions in the final version. In particular, we also summarize the core distinctions between our GeoTDM and DiffMD [1] as follows:

|  | Distribution Modeled |  Unconditional | Prior | Backbone |
|-|-|-|-|-|
| DiffMD [1]    | Frame    | $\times$ | Fixed Pointwise Prior | Specialized for molecules |
| GeoTDM    | Trajectory  | $\checkmark$ | Learned Flexible Prior | General for $n$-D geometric data |

Overall, our experiments have verified the advantages of our approach over DiffMD in terms of empirical performance. We highly appreciate it if you could kindly consider adjusting your evaluation if we have addressed your concerns. Thank you!

Best,

Authors

[1] Fang Wu, Stan Z. Li, DIFFMD: A Geometric Diffusion Model for Molecular Dynamics Simulations.
Dear Reviewer UDq7,

Thank you for the timely reply. To help substantively address your concerns, we summarize and further illustrate our response below.

For **W1**, we added extra experiments that explore the way to combine our proposed equivariant temporal layer with more sophisticated equivariant architectures, e.g., Equiformer, that helps further boost the performance on specific tasks like MD. We will also include these discussions in the paper.

For **W2**, we clarified the difference between theorems presented in this paper and previous work. Besides, we provided necessary derivations in Appendix A where we have interestingly found that our GeoTDM also preserves the simplified loss (Eq. 12) which is yet not straightforward. Notably, in Theorem A.4 we also theoretically justify our learnable prior subsumes existing parameterizations. All of these constitute our theoretical contributions.

For **W3**, We have also presented various explorations on approaching new problem with GeoTDM, including modeling large scale catalyst dynamics (Appendix C.1), interpolating trajectories given initial and target states (Sec. 5.3), optimizing the input trajectory towards the learned distribution (Sec. 5.3), and performing generation through longer time horizon (Appendix C.2).

Please let us know if you have any concerns in particular or any advice that could help further improve the paper and we are more than willing to address them. Thank you!

Best,

Authors
Dear Reviewer RWWk,

Thank you very much for your feedback and recognition of our efforts! We greatly appreciate your valuable suggestions and will incorporate the discussions into the final version.

Best,

Authors
Thank you very much for your detailed response! Most of my concerns have been addressed. Also combined with other reviewers' comments and rebuttals, I decided to raise the score.
Thank you for your detailed response! After I read all your rebuttals and other reviewers comments. I decided to stay at the same rating.
I appreciate the detailed author response. However it has not moved the needle substantively on my concerns and I will keep the score.