Title: Geometric Trajectory Diffusion Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: diffusion model

geometric trajectory

equivariant graph neural networks
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.

Geometric Trajectory Diffusion Models

Abstract

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

EquiPrior

𝜼 , 𝜸

𝐱

!

[#

!

]

𝐱

𝒯

[#]

𝐱

&

[#]

𝐱

&'(

[#]

𝐱

)

[#]

&

#

|𝐱

&'(

[#]

)

𝑝

𝜽

(𝐱

&'(

#

|𝐱

&

[#]

)

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

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

3 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 kernel p θ ( 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 as L 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:

where 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 ∈ R 3 . For geometrc graph with node features h and coordinates x , if h ′ , x ′ = f ( h , x ) , we expect h ′ , Rx ′ + r = f ( h , Rx + r ) 3 , i.e. , the output node features are invariant while the updated coordinates are equivariant. Distributions. We call a density p ( x ) invariant w.r.t. a group G if p ( g · x ) = p ( x ) , ∀ g ∈ G . Intuitively, geometries that are rotationally and translationally equivalent should share the same density, since they all refer to the same structure. A conditional distribution p ( x | y ) is equivariant if p ( g · x | g · y ) = p ( x | y ) , ∀ g ∈ G . Such a property is important in cases where the target distribution is conditioned on some given structures: if the observed geometry is rotated/translated, the target distribution should also rotate/translate accordingly.

Geometric trajectories and the distributions. We represent a geometric trajectory as ( x [ T ] , h , E ) , where x [ T ] := [ x (0) , x (1) , · · · , x ( T -1) ] ∈ R T × N × D x is the sequence of temporal geometric coordinates, h ∈ R N × D h is 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 [ T c ] c , h , E ) where x [ T c ] c ∈ R T c × N × D x 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:

where g · x [ T ] := [ Rx (0) + r , · · · , Rx ( T -1) + r ] . The conditional case should instead preserve:

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

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

4.1 Equivariant Geometric Trajectory 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],

The 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, 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 ) ,

where φ 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):

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 [ T c ] c and 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 ∈ [ T c ] a ( t,s ) v ( t,s ) and x ′′ ( t ) = x ′ ( t ) + ∑ s ∈ [ T c ] a ( t,s ) φ x ( v ( t,s ) )( x ( t ) -x ( s ) c ) .

4.2 Geometric Trajectory Diffusion Models

4.2.1 Unconditional Generation

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.

Theorem 4.2. If the prior p T ( x [ T ] 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 ) = ( TN -1) D [36]. Specifically, we choose the prior to be the projection of the Gaussian N ( 0 , I ) in X to X P by P = I D ⊗ ( I TN -1 TN 1 TN 1 ⊤ TN ) , 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 ) i being 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 ˜ N ( 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 ] τ ) = ˜ N ( ˜ µ θ (˜ 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):

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.

4.2.2 Conditional Generation

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 [ T c ] c ) is SE (3) -equivariant, the transition kernels p τ -1 ( x [ T ] τ -1 | x [ T ] τ , x [ T c ] c ) , ∀ τ ∈ { 1 , · · · , T } are SE (3) -equivariant, the marginal 5 p τ ( x [ T ] τ | x [ T c ] c ) , ∀ τ ∈ { 0 , · · · , T } is SE (3) -equivariant.

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 [ T c ] c ) , I ) is SE (3) -equivariant w.r.t. x [ T c ] c if µ ( x [ T c ] c ) is SE (3) -equivariant.

Proof is in App. A.2. Notably, the mean of the prior x [ T ] r := µ ( x [ T c ] 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 ∈ [ T c ] w ( s ) ¯ x ( s ) c , where ∑ s ∈ [ T c ] 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:

for 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 obtain ˆ x [ T c ] c , ˆ h [ T c ] c = f η ( x [ T c ] c , h [ T c ] c ) where f η is a lightweight two layer EGTN that aims to synthesize the conditional information. The w ( t,s ) is then derived as,

Here γ ∈ 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 γ , ˆ h [ T c ] c , and ˆ x [ T c ] c reduce 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 [ T c ] c ) := N ( x [ T ] τ ; x r + √ 1 -β τ ( x [ T ] τ -1 -x r ) , β τ I ) , which ensures q ( x [ T ] T | x [ T c ] 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 [ T c ] c ) = N ( µ θ ( x [ T ] τ , τ, x [ T c ] c ) , σ 2 τ I ) . Similar to the unconditional case, we also adopt the noise prediction objective by rewriting µ θ ( x [ T ] τ , x [ T c ] c , τ ) = x [ T ] r + 1 √ α τ ( x [ T ] τ -x [ T ] r -β τ √ 1 -¯ α τ ϵ θ ( x [ T ] τ , x [ T c ] 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:

after simplification (proof in App. A.2). The training and inference procedures are in Alg. 3 and 4.

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

5.1 Conditional Case

5.1.1 N-body

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 TN ∑ 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.

5.1.2 Molecular Dynamics

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

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.

5.1.3 Pedestrian Trajectory Forecasting

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.

5.2 Unconditional Generation

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 V AE-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.

GeoTDM

SVAE

(a)

(b)

(c)

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.

5.3 Ablation Studies and Additional Use Cases

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 ( T c -1) c ) , I ) and point-wise equivariant prior N ( x ( T c -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.

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

Conclusion. We present GeoTDM, a diffusion model built over distribution of geometric trajectories. 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.

Acknowledgments and Disclosure of Funding

We thank the anonymous reviewers for the helpful feedback on improving the manuscript. This work was supported by ARO (W911NF-21-1-0125), ONR (N00014-23-1-2159), and the CZ Biohub.

References

Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems , 29, 2016. 1

Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations , 2022. 1, 2, 7

Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis , 11(10):6059-6072, 2021. 24

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018. 22

Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances , 3(5):e1603015, 2017. 7

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021. 2

Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie-Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning , pages 5583-5608. PMLR, 2022. 2

Jacob D Durrant and J Andrew McCammon. Molecular dynamics simulations and drug discovery. BMC biology , 9(1):1-9, 2011. 1

Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. arXiv preprint arXiv:2006.10503 , 2020. 2, 7, 22

Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations , 2020. 1

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017. 2, 9, 22

Tianpei Gu, Guangyi Chen, Junlong Li, Chunze Lin, Yongming Rao, Jie Zhou, and Jiwen Lu. Stochastic trajectory prediction via motion indeterminacy diffusion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 17113-17122, 2022. 3, 7, 8, 25, 26

Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equivariant diffusion for target-aware molecule generation and affinity prediction. In The Eleventh International Conference on Learning Representations , 2023. 6, 20

Agrim Gupta, Justin Johnson, Li Fei-Fei, Silvio Savarese, and Alexandre Alahi. Social gan: Socially acceptable trajectories with generative adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2255-2264, 2018. 7, 8, 9, 22

Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Joshua B. Tenenbaum, and Chuang Gan. Learning physical dynamics with subequivariant graph neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 1

Jiaqi Han, Wenbing Huang, Tingyang Xu, and Yu Rong. Equivariant graph hierarchy-based neural networks. Advances in Neural Information Processing Systems , 35:9176-9187, 2022. 1, 22

William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances in Neural Information Processing Systems , 35:27953-27965, 2022. 3

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020. 2, 3, 5, 7, 9, 18, 21

Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 2, 3

Scott A Hollingsworth and Ron O Dror. Molecular dynamics simulation for all. Neuron , 99(6):1129-1143, 2018. 1

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning , pages 8867-8887. PMLR, 2022. 1, 2, 3, 6, 20

Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanics networks with constraints. In International Conference on Learning Representations , 2022. 1

John B Ingraham, Max Baranov, Zak Costello, Karl W Barber, Wujie Wang, Ahmed Ismail, Vincent Frappier, Dana M Lord, Christopher Ng-Thow-Hing, Erik R Van Vlack, et al. Illuminating protein space with a programmable generative model. Nature , pages 1-9, 2023. 1

Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi Jaakkola. Eigenfold: Generative protein structure prediction with diffusion models, 2023. 2

Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons. Neural sdes as infinite-dimensional gans. In International conference on machine learning , pages 5453-5463. PMLR, 2021. 8

Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint arXiv:1802.04687 , 2018. 1, 2, 7

Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi-body systems with symmetric energies. arXiv preprint arXiv:1910.00753 , 2019. 2, 7, 22

Alon Lerner, Yiorgos Chrysanthou, and Dani Lischinski. Crowds by example. In Computer graphics forum . Wiley Online Library, 2007. 8

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022. 24

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 3

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. Advances in Neural Information Processing Systems , 35:9754-9767, 2022. 9

Karttikeya Mangalam, Harshayu Girase, Shreyas Agarwal, Kuan-Hui Lee, Ehsan Adeli, Jitendra Malik, and Adrien Gaidon. It is not the journey but the destination: Endpoint conditioned trajectory prediction. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part II 16 , pages 759-776. Springer, 2020. 7, 8

Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations , 2022. 9

Hao Ni, Lukasz Szpruch, Magnus Wiese, Shujian Liao, and Baoren Xiao. Conditional sigwasserstein gans for time series generation. arXiv preprint arXiv:2006.05421 , 2020. 8

Stefano Pellegrini, Andreas Ess, Konrad Schindler, and Luc Van Gool. You'll never walk alone: Modeling social behavior for multi-target tracking. In 2009 IEEE 12th international conference on computer vision , pages 261-268. IEEE, 2009. 8

Calyampudi Radhakrishna Rao. Linear statistical inference and its applications , volume 2. Wiley New York, 1973. 5

Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In International Conference on Machine Learning , pages 8857-8868. PMLR, 2021. 3

Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. DYffusion: a dynamics-informed diffusion model for spatiotemporal forecasting. In Advances in Neural Information Processing Systems (NeurIPS) , 2023. 3

Amir Sadeghian, Vineet Kosaraju, Ali Sadeghian, Noriaki Hirose, Hamid Rezatofighi, and Silvio Savarese. Sophie: An attentive gan for predicting paths compliant to social and physical constraints. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 1349-1358, 2019. 7, 8

Tim Salzmann, Boris Ivanovic, Punarjay Chakravarty, and Marco Pavone. Trajectron++: Dynamically-feasible trajectory forecasting with heterogeneous data. In Computer VisionECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVIII 16 , pages 683-700. Springer, 2020. 7, 8, 25

Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning , pages 8459-8468. PMLR, 2020. 2

Victor Garcia Satorras, Emiel Hoogeboom, Fabian Bernd Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. 1, 2, 3

Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. arXiv preprint arXiv:2102.09844 , 2021. 1, 2, 4, 5, 7, 8, 9, 20, 22

Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning , pages 9377-9388. PMLR, 2021. 8

Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics , 148(24):241722, 2018. 1

Jean-Pierre Serre et al. Linear representations of finite groups , volume 42. Springer, 1977. 3

Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers) , pages 464-468, 2018. 4

Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International conference on machine learning , pages 9558-9568. PMLR, 2021. 8

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning , pages 2256-2265. PMLR, 2015. 2, 3

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020. 24

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 , 2023. 24

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019. 2, 3

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems , 33:12438-12448, 2020. 2, 3

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems , 34:24804-24816, 2021. 3

Guy Tevet, Sigal Raab, Brian Gordon, Yoni Shafir, Daniel Cohen-or, and Amit Haim Bermano. Human motion diffusion model. In The Eleventh International Conference on Learning Representations , 2023. 3

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018. 1, 2, 7, 22

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017. 4

Vikram Voleti, Alexia Jolicoeur-Martineau, and Christopher Pal. MCVD - masked conditional video diffusion for prediction, generation, and interpolation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. 3

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023. 1, 2

Fang Wu and Stan Z Li. Diffmd: a geometric diffusion model for molecular dynamics simulations. In Proceedings of the AAAI Conference on Artificial Intelligence , 2023. 3

Liming Wu, Zhichao Hou, Jirui Yuan, Yu Rong, and Wenbing Huang. Equivariant spatiotemporal attentive graph networks to simulate physical dynamics. Advances in Neural Information Processing Systems , 36, 2024. 4

Chenxin Xu, Robby T Tan, Yuhong Tan, Siheng Chen, Yu Guang Wang, Xinchao Wang, and Yanfeng Wang. Eqmotion: Equivariant multi-agent motion prediction with invariant interaction reasoning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1410-1420, 2023. 2, 7, 8, 22

Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzadenesheli, Jure Leskovec, Stefano Ermon, and Anima Anandkumar. Equivariant graph neural operator for modeling 3d dynamics. In Forty-first International Conference on Machine Learning , 2024. 2

Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference on Machine Learning . PMLR, 2023. 1, 2, 3, 24, 25, 26

Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael Gomez-Bombarelli, and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel programming. In International Conference on Machine Learning , pages 11537-11547. PMLR, 2021. 3

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations , 2022. 1, 2, 25, 26

Pei Xu, Jean-Bernard Hayet, and Ioannis Karamouzas. Socialvae: Human trajectory prediction using timewise latents. In European Conference on Computer Vision , pages 511-528. Springer, 2022. 7, 8, 9, 22

Yu Yao, Ella Atkins, Matthew Johnson-Roberson, Ram Vasudevan, and Xiaoxiao Du. Bitrap: Bi-directional pedestrian trajectory prediction with multi-modal goal estimation. IEEE Robotics and Automation Letters , 6(2):1463-1470, 2021. 7, 8

Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 40001-40039. PMLR, 23-29 Jul 2023. 2

Bing Yu, Haoteng Yin, and Zhanxing Zhu. Spatio-temporal graph convolutional networks: A deep learning framework for traffic forecasting. In Proceedings of the 27th International Joint Conference on Artificial Intelligence (IJCAI) , 2018. 4

Ye Yuan, Jiaming Song, Umar Iqbal, Arash Vahdat, and Jan Kautz. Physdiff: Physics-guided human motion diffusion model. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 16010-16021, 2023. 3

Linqi Zhou, Michael Poli, Winnie Xu, Stefano Massaroli, and Stefano Ermon. Deep latent state space models for time-series generation. In International Conference on Machine Learning , pages 42625-42643. PMLR, 2023. 8

Appendix

Table of Contents

A Proofs

A.1 Unconditional Case

We note that, naïvely, a distribution p ( x [ T ] ) can not be translation invariant. In particular, this would imply that p ( x [ T ] ) = p ( x [ T ] + r ) for all r ∈ R D , but this would imply that p ( x [ T ] ) = 0 uniformly, a contradiction.

Instead, we derive an equivalent invariance condition by restricting SE ( D ) to its maximally compact subgroup. In particular, we note that it is possible to define SO ( D ) -invariant distributions (as this group is compact), and SE( D ) / SO( D ) ∼ = T , the translation group. The natural way to quotient out our base space R T × N × D / T ∼ = R ( T × N -1) × D is to zero-center our data (along each dimension).

However, for practical purposes, we will refer to our construction as SE( D ) -invariant. In particular, since all inputs x ∈ R T × N × D are first zero-centered to be projected to R T × N × D / T , the 'lifted" unnormalized measure is SE ( D ) invariant.

We will define P as our zero-centering operation 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 ) i and our restricted ( T × N -1) × D Gaussian as ˜ N ( y | x , Σ ) , which can be represented in the ambient space as a degenerated Gaussian variable

where Σ P = PΣP ⊤ and Σ + P is the pseudo-inverse (and det ∗ is the determinant restricted to the subspace). Note that P is symmetric and idempotent. Then specifically when Σ = I , we have Σ P = PP ⊤ , then Σ + P = P as ( PP ⊤ ) P ( PP ⊤ ) = P , since PP = P and P = P ⊤ .

Base distribution. We require the base distribution to be SO (3) -invariant. In practice, we let p T (˜ x [ T ] T ) = ˜ N ( 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 ] τ ) = ˜ N ( ˜ µ θ (˜ x [ T ] τ , τ ) , σ 2 τ I ) . In order to ensure p θ ( R ˜ x [ T ] τ -1 | R ˜ x [ T ] τ ) = p θ (˜ x [ T ] τ -1 | ˜ x [ T ] τ ) , it suffices to make ˜ µ θ (˜ x [ T ] τ , τ ) an SO (3) -equivariant function. In this way,

which permits the SO (3) -equivariance of the transition kernel. In our implementation, we further re-parameterize ˜ µ θ (˜ x [ T ] τ , τ ) as,

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

Proposition A.1. If the prior p T (˜ x [ T ] T ) is SO (3) -invariant, the transition kernels p τ -1 (˜ x [ T ] τ -1 | ˜ x [ T ] τ ) , ∀ τ ∈ { 1 , · · · , T } are SO (3) -equivariant, then the marginal p τ (˜ x [ T ] τ ) at any time step τ ∈ { 0 , · · · , T } 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 τ ( R ˜ x [ T ] τ ) , then we have the following derivation for the marginal at time step τ -1 :

Notably, for the final step at τ = 0 , the marginal p 0 ( R ˜ x [ 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

4: Take gradient descent step on

5: until converged

Algorithm 2 Sampling Procedure of GeoTDM-uncond

1: Sample ˜ x [ T ] T ∼ ˜ N ( 0 , I ) [ T ]

2: for τ ←T , · · · , 1 do

3: Sample ˜ z [ T ] τ ∼ ˜ N ( 0 , I ) [ T ] if τ > 1 else ˜ z [ T ] τ = 0

5: end for

6: return x [ T 0

]

Algorithm 3 Training Procedure of GeoTDM-cond

1: repeat

2: x [ T ] r ← EquiPrior η , γ ( x [ T c ] c ) {Eq. 9}

4: x [ T ] τ ← √ ¯ α τ ( x [ T ] -x [ T ] r ) + x [ T ] r + √ 1 -¯ α τ ϵ [ T ]

5: Take gradient descent step on

6: until converged

Algorithm 4 Sampling Procedure of GeoTDM-cond

6: end for

7: return x [ T 0

]

A.2 Conditional Case

In the conditional case, we target on modeling the conditional distribution p ( x [ T ] | x [ T c ] c ) . The desired constraint is the following equivariance condition: p ( x [ T ] | x [ T c ] c ) = p ( g · x [ T ] | g · x [ T c ] 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.

Then, we have

∀ rotation matrix R and r ∈ R 3 , which completes the proof.

Base distribution. We propose to leverage the following base distribution.

where x [ T ] r = EquiPrior( x [ T c ] c ) is SE (3) -equivariant with respect to the condition x [ T c ] c . With such choice, the base distribution above is SE (3) -equivariant, since

where 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 by

where µ θ ( x [ T ] τ , x [ T c ] c , τ ) parameterized to be SE (3) -equivariant with respect to its input x [ T ] τ , x [ T c ] c . In practice, we re-parameterize it as,

where ϵ θ ( x [ T ] τ , x [ T c ] c , τ ) is an SO (3) -equivariant but translation-invariant function. It is then easy to see that µ θ ( x [ T ] τ , x [ T c ] 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:

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,

which is equivalent to Eq. 39 up to certain re-weighting factors. For L 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 [ T c ] c ) is SE (3) -equivariant and the transition kernels p τ -1 ( x [ T ] τ -1 | x [ T ] τ , x [ T c ] c ) of all diffusion steps τ ∈ { 1 , · · · , T } are SE (3) -equivariant, then the marginal 6 p τ ( x [ T ] τ | x [ T c ] 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 [ T c ] c ) is SE (3) -equivariant, as it is designed.

Induction step. Suppose the marginal at diffusion step τ , i.e. , p τ ( x [ T ] τ | x [ T c ] c ) , is SE (3) -equivariant, then we have

which concludes the proof.

A.3 Optimizable Equivariant Prior

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.

We first show x [ T ] r can reduce to the CoM-based priors. Let ˆ x ( s ) c = CoM( x ( s ) c ) , ˆ h ( s ) c = 1 T c 1 N , γ ( t ) = 1 . In this case,

where 1 T c ∑ s ∈ [ T c ] 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 .

To show x [ T ] r can reduce to fixed point-wise priors is straightforward. Let ˆ x ( s ) c = x ( s ) c , ˆ h [ T c ] = Onehot( s ∗ ) 1 T c × N and γ ( t ) = 1 , ∀ t . Then w ( t,s ) = Onehot( s ∗ ) 1 T c × N . Therefore,

where x ( s ∗ ) c is 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 as f EGTN = f attn ◦ f EGNN ◦ · · · ◦ f attn ◦ f EGNN ︸ ︷︷ ︸ L × ( f attn ◦ f EGNN ) . 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 input

CoM-based prior

Fixed point-wise prior

Ours

1

𝑇

!

#

"∈[%

!

]

CoM(𝐱

!

"

)

𝐱

!

("

∗

)

: Mean of the prior

: condition

𝐱

!

[#

!

]

: target

𝐱

[#]

h [ T ] . This directly leads to the SE (3) -invariance of the updated node feature h ′ [ T ] . For the updated coordinates,

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

B More Details on Experiments

B.1 Compute Resources

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.

B.2 Hyper-parameters

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 .

B.3 Baselines

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 { 5 e -3 , 1 e -3 , 5 e -4 , 1 e -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.

B.4 Model

In detail, the EGCL layer [43] is given by:

where φ m , φ h , and φ x are all MLPs. We also provide a schematic of our proposed EGTN in Fig. 4 for better illustration.

B.5 Evaluation Metrics in the Unconditional Case

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.

𝐪

(")

𝐤

(",%)

𝜓(𝑡 - 𝑠)

𝐯

(",%)

𝐚

(",$)

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.

C More Experiments and Discussions

C.1 Model Composition for Longer Trajectory

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-length 7 non-overlapping intervals with time span ∆ T . Then, for the unconditional case, we have

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

C.2 Number of Diffusion Steps

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 SV AE.

C.3 Sampling Time

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 EGV AE. 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.

C.4 Standard Deviations

We provide the standard deviations in Table 10 and 11.

C.5 More Discussions with Existing Works

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 Trajectron++ [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.

D More Visualizations

We provide more visualizations in Fig. 5, 6, 8, 7, 9, and 10. Please refer to their captions for the detailed descriptions.

图

𝐱

!

[#

!

]

𝜏 = 0

𝜏 = 𝒯

𝐱

!

[#

!

]

𝜏 = 0

𝜏 = 𝒯

𝐱

!

[#

!

]

𝜏 = 0

𝜏 = 𝒯

𝜏 = 0

𝜏 = 𝒯

𝐱

!

[#

!

]

Condition

Prediction

Ground Truth

Condition

Prediction

Ground Truth

Data

GeoTDM

VAE

GeoTDM

EGNN

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We have provided theoretical and empirical results showing that our proposed GeoTDM is able to work with geometric symmetries through multiple time steps.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have discussed some limitations of GeoTDM, including the slow sampling speed of diffusion models (see Sec. C.3).

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Yes, we have listed out all of our assumptions and have explicated each step for our proofs.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Yes, we have included all of the relevant details for reproducing our experiments. See App. B.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., 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.

5. Open access to data and code

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:

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.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

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: See our code and experiment details in App. B.

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.

7. 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 assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have listed out this type of information in App. B.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed and followed the NeurIPS code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our method is primarily focused on scientific discovery, and there is no societal impact of our work beyond this scope.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point 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).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not introduce any new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This work does not involve human subjects.

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.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work does not involve human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA 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.\n\nSTRENGTHS:\n1. The introduction and methods are well-written and clear.\n2. The idea and methodology of introducing and processing the time sequence information are reasonable and well-defined with good novelty.  And the appendix provides sufficient proof and detailed analysis.\n3. The paper explained and analyzed the case of conditional generation in detail.\n4. The experimental setting and results are demonstrated in very extensive and detailed ways.\n5. Clear ablation studies.\n\nWEAKNESSES:\nNo major weakness. Please refer to the part of Questions for my concerns.\n\nQUESTIONS:\n1. 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)\n\n2. 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]? \n\n3. 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.\n\n4. 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?\n\n[1]: Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\" International conference on machine learning. PMLR, 2022.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

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

STRENGTHS:
The authors impose certain constraints of SE(3)-invariant on the prior and transition kernel to keep the desirable physical symmetry of the trajectories.

WEAKNESSES:
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:
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:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* The method is a timely extension and synthesis of geometric diffusion models and trajectory diffusion models.\n* As demonstrated in the experiments, the method has broad applicability across many machine learning domains.\n* The experiments are thorough. The ablation studies validating the somewhat unconventional choices of conditional prior are appreciated.\n* The paper is clearly written and the exposition of the many experiments is handled very cleanly.\n\nWEAKNESSES:\n* 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.\n* The exposition of the method drags at times, re-proving well-known statements about equivariance in neural networks and geometric diffusion models.\n* 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.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

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

STRENGTHS:
1. The geometric trajectory diffusion model is novel to me. 

2. The EGTN that operates on geometric trajectories permits conditioning upon a given trajectory using equivariant cross-attention

3, The experiments demonstrate that GeoTDM achieves SOTA performance on both unconditional and conditional trajectory generation tasks.

4. The work can be extended to applications like temporal interpolation and trajectory optimization

WEAKNESSES:
None

QUESTIONS:
None

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12