Geometric Trajectory Diffusion Models
Authors: Jiaqi Han
Minkai Xu
Aaron Lou
Haotian Ye
Stefano Ermon
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion model
geometric trajectory
equivariant graph neural networks
Primary area: diffusion_based_models

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
Jiaqi Han, Minkai Xu, Aaron Lou, Haotian Ye, Stefano Ermon Stanford University
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 novel transition kernel leveraging SE(3) -equivarru spatial convolution and temporal attention. Furthermore, induce an expressive trajectory distribution for conditional generation, we introduce generalized learnable geometric prior into the forward diffusion process 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 wide suite of metrics demonstrate that Geo'TDM can generate realistic geometric trajectories with significantly higher quality."
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 challenging due to the physical symmetry constraint [56, 43], making it fundamentally different from common scalar non geometric data such as images and text. With the recent progress of generative models, many works have been proposed in generating 3D geometric structures like small molecules [66, 42, 21, 64] and proteins [59, 23), showing great promise in solving the equilibrium states of complex systems.
Despite this success, these existing methods are limited to synthesizing static structures and neglect the fact that important real-world processes evolve through time. For example, molecules and proteins are not static but always varying with molecular dynamics, which plays a vital role in analyzing possible binding activities [8, 20]. this paper, we aim study the generative modeling of geometric trajectories with the additional temporal dimension. While this problem is more practical and important, it highly non-trivial with several significant challenges. First, geometric dynamics in 3D ubiquitously preserve physical symmetry. With global translation or rotation applied to trajectory of molecular dynamics, the entire trajectory still describes the same dynamics and the generative model should estimate the same likelihood. Second, trajectories inherently contain the correspondence between frames different timesteps, requiring generative models hold a high capacity for capturing the temporal correlations. Last, moving from single structure to trajectory
'Correspondence to Jiaqi Han: jiagthan@stanfort Code is available at https: h/tg://b.tttttt///
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

q(x;" (
x
EquiPrior
HETIN
Figure 1: Overview of GeoTDM. The forward diffusion gradually perturbs the input while the reverse process po. parameterized by EGTN, denoises samples from the prior. The condition !!'al available, is leveraged to construct the equivariant prior and as a conditioning signal in EGTN.
composed of multiple ones, the distribution we are interested in becomes much higher- disher-dmmensionll 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), 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], vidcos [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 the trajectories, but also capture the complex spatial and temporal correspondence encapsulated in the dynamics of geometric systems. Moreover, by leveraging generative modeling, GeoTDM enjoys high versatility in generating diverse yet high-quality geometric trajectories from scratch, performing interpolation and extrapolation, and optimizing noisy trajectory, all under the proposed diffusion framework.
In summary, we make the following contributions: 1. We present GeoTDM, a novel temporal diffusion model for generating geometric trajectories. We design the diffusion process meet the critical equivariance in modeling both unconditional and conditional distributions over geometric trajectories. Notably, we also propose conditional learnable equivariant prior for enhanced flexibility in temporal conditioning. 2. To fulfill the equivariance of the denoising network, we introduce EGTN, graph neural network that operates on geometric trajectories, which also permits conditioning upon a given trajectory using equivariant cross- rross-attention, making it suitable to serve as the backbone for TDM.. 3. We evaluate our GeoTDM on both unconditional and conditional trajectory generation tasks including particle simulation, molecular dynamics, and pedestrian trajectory prediction. Geo 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 dynamies simulation. We also show GeoTDM successfully performs several additional applications, such temporal interpolation and trajectory optimization.
2 Related Work
Trajectory modeling for geometric systems. Modeling the dynamics of geometric data is challeng- ing 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, 2, 63] discovered equivariance as critical factor for promoting model generalization. Among these efforts, Radial Fields [27] and EGNN [43] work with equivariant operations between scalars and vectors, [9] generalize to high-order spherical tensors. While consid- erable progress has been made, they only conduct (time) frame-to. to-frame prediction, which 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 wider range of tasks such generation and interpolation.
while TFN [56] and 3E33)-Transfomme
Generative models 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 cuffusion-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 snapshot of the geometric system, egg. a molecule in 3D space, whereas our GeoTDM generalizes to generating trajectory with multiple frames, e.g., an MD trajectory in 3D. DifMM [60] specifically tackles MDD 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, 53] are type of latent variable generative model that feature a Markovian forward diffusion process and reverse denoising process. The forward process progressively perturbs the input Xo (e.g., image pixels or molecule coordinates) over T steps using a Gaussian transition kernel (x,, -1) = N (x,:v1 -, 3,I1), I). Here, {x ) are latent variables with the same dimension as the input and B, are predefined using the noise schedule such that Xy is close to being distributed as (O. The reverse process maps back from the prior distribution with p(xy) = NO I) using the kernel po(x,-1(xx, where the variances of are usually fixed and the mean P4D is parameterized by a neural network with parameters 0. The model is trained by optimizing the variational lower bound, defined as Lvl = - log po (xn) x1) + DKL(g(x) X.) xx) kl)l)))))) + Σ ( x/x)xl/ (x-
N (x fo (x,, 7). II)).
For training stability, [52, 18] suggest the noise-predietee objective:
Caimple = Bxx,, [ -x)
where Xo Pdaata T Unif(1, (1, T). the weighting factors \(7) are typically set to 1 to promote
sample quality, x, = v0,x, Xo + v - - with a, == II a = II (1 - 3,), and 6g is a specific parameterization of the mean satisfying No T) Vo, - (xr Vi-a. (x..
Equivariance. Functions. A function f is equivariant w.T a group G if flg. x) - 9- f(x)(x), V G. Furthermore, f is invariant if f(g x) = f(x), Vg E G' [46]. Here we focus on the group SE(3) consisting of all 3D rotations and translations Each group element g E SE(3) can be represented by rotation matrix R and translation r E R For geometre graph with node features h and coordinates x, if h', x' = f(h, x). we expect h', Rx' +r= [h, Rx + +r),, the output node features are invariant while the updated coordinates are equivariant. Distributions. We call density p(x) invariant w.r.t. group p(g- x) = p(x). Vg E G. Intuitively, geometries that are rotationally and translationally equivalent should share the same density, since they all refer to the same structure. conditional distribution p(xly) is equivariant if p(g x(g y) = p(xly), Vg E 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 rotatc/translate accordingly.
Geometric trajectories and the distributions. We represent geometric trajectory as (xl"1, h, E), where x [T] == [x()), (()), xx(T-)) TR is the sequence of temporal geometric coordi- nates, h E R N Pa is the node feature, and is the set of edges representing the connectivity of the geometric graph. T is the number of time steps and Dx, Dn, refers to the dimension of the coordinate and node feature respectively, with Dx normally being 2 or 3 depending on the input data. this work, we are interested in modeling the distribution of geometric trajectories given the configuration of the geometric graph, i.e., p(x 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 [Tc]] timesteps, i.e., pTx x h, E) where x E T=xxxxx is the provided trajectory in length T.-
[Tel
The analyses this paper also hold for the general nD case, e.g., 2D. Following convention the notation Rx to denote xR.
3

---

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:
p(x [T]
h, E) ppgg x!7 h, E), Vg E SE(3).
where g x x111
Rx (0) +r.
Rx (xx(T-1) 1) + r]. The conditional case should instead preserve:
[T. p(x x h, 8) = pig (7
g [T)) h,
for all g E SE(3)* 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 & 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 Geo 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 equiv- ariant 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 struc- tural 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],
x" x10 h( EGCL(X(") (0,,), Vt e [T].
The equivariant message passing is conducted independently for each frame E [7] {0, 1, ,T' - 13. 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 ()) ((.)), (t,,)) are the query, key, and value,
(t,s) exp (q
respectively. In detail, q(t) = (h))),, (t.ss)
Emerry exp (q ((()))))))) (5)
(p(((+)) (1- and v!t.s) (v(h())) ( ((((-)). s). with 1b(1 -// s) being the sinusodial encoding [57] the temporal displacement i akin to the relative posi- tional encoding [47]. Incorporating such information is
(0)) = (h(() + E.er (tr)) (6)
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,e) and the values vtt.s)
(t) = (() + E. (t.s)
px (v(())))) -xxx))))))00xxxx))) -x()),,
where fx is an MLP that outputs scalar preserve rotation equivariance. The entire network JEGTN. with schematic depicted in Fig. 4, is constructed by alternating spatial and temporal layers, enjoying equivariance as desired (proof in Appendix A.4);
Theorem 4.1 (Suivaaaquiaarraanne)). Let x(() h(T)] fEGTN (x!'), E)) Then we have g
x(T) 1/77] = fECTN (g (T), Vg E SE(3).
Geometric conditioning. In certain tasks like trajectory forecasting, we additionally provided with some partially observed trajectories as side input. In order to leverage their geometric information,
'Technically, such condition is impossible since SE(3) is noncompact, but we show that zero-centering the trajectories and enforcing 50(3)-img -invariann is equivalent.

---

we augment the unconditional EGTN with equivariant cross pross-attention. 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!Too and the target x In detail, the attention coefficients are recomputed as a (t.e)
exp(1 (P
Euctr exp(q (G(TT)()))) The updated node feature h((!)
and coordinate x((1) in Eqs. 6-7 are further renewed by the cross-attention terms, yielding h"(t)
h((9) + E ETT a (t.s) and x = xx(() + Eeerra a(t=)) px (x((())))) - xxx).
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 the following theorem.
Theorem 4.2. If the prior (x is SE(3)-invariant, the transition kernels p--1(x [T]
). Vr E {1,...,T are -equivequi,ariann, then the marginal pr (x."") at any step T E (0, ... T is also SEnvaiiinri.
Prior in the translation-invariant subspace. Unfortunately, there is no properly normalized distribution w.r.f. Lebesgue measure on the ambient space X' := T x that permits translation- invariance [43]. We instead build the prior on vanslation-invariann invariant subspace Xpp C X' induced by linear transformation P E X' X with rank(P) = (TN 1) [36]. Specifically, we choose the prior be the projection of the Gaussian N I) in X Xp by P ID (ITN TN 11x11w) which corresponds the function P(x(T)) x(7) E CoM(x!") with CoM(x(
En= x{" 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 ((0, I)) with the variables supported only on the subspace (see App. A. and more importantly, still isotropic and thus 5033a-invrriant. To sample from the prior, one alternatively sample x (O, E X' and then project it to the subspace to obtain the final sample x P(xl2T) E Xpp
Transition kernel. To be consistent with the prior, we also parameterize the transition kernel in the subspace Xp, given by (x x = -N(ão (x I.)). In this way, if the mean function (. is -0)33))quuvariaat, then the transition kernel is also guaranteed (e()))-qvvvariance As suggested by [18], we re-parameterize px = - - (x!") VI-a. E (x( ,7)) where
Eo P fg. with fo being an -qqiiaaqun 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 in the input.
Training and inference. We optimize the VLB for training, which, interestingly, still has surrogate in the noise-predien preeiccre form when specifying the factors A(r) as (proof App. A.1):
uuccnd
[T) ( Unii((1,7 ||e εo (x(((,,, +)]]]
The inference process is similar to [18] but with additional applications of P in intermediate steps to keep all samples in the subspace Xp. Details are in Alg. and 2.
4.2.2 Conditional Generation
Distinct from the unconditional generation, in the conditional scenario the target distribution should instead be -equi--qqiaa w.r.f. 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 PT (x [27] xx [Tal] is -quiivrain, the transition kernels
P--1 (x (7]
[T]] (Td)), VT E {1,... ,T are SE3))--quivariant, the marginal px()! (x!T) VT E (O..... T is -E(3)--quivarian.
(x,- xx]"
Marginal refers to marginalizing the intermediate states in reverse process while still conditioning !Tc]]
5

---

Flexible equivariant prior. There are in general many valid choices for the prior while satisfying -quuiaaquue We provide guidance on distinguishing feasible designs when using Gaussian- based prior in the proposition below.
Proposition 4.4. N (u(xx
(Fol) is (E(33)-eqivvrianl (rl)), 1) ), is Seu))eequivaaiant w.r.t. X [T if (x(e
Proof is in App. A.2. Notably, the mean of the prior x. ((xx"))) naturally serves anchor to transit the geometric information in the provided trajectory to the target distribution we seek to model. For instance, one can choose it linear combination of the CoMs of the given frames,
i..., x 1TxN & Exe(r.) xx(+x{ where Eve[T.] w(() 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:
Σ SE(2.1
w (halge)))
(1,0) 1,
S.L
SETU
for all t E [T), where each (1) is point-wise linear combination of (*)) an -qquivaaqu transformation of the conditioning frames, with w(t.+) E RN being the weights. We first obtain x!
[Tal (( f,, (x [Te) h where f, is a lightweight layer EGTN that aims to synthesize the
conditional information. The w((.s) is then derived as,
W 1.e b & h H RN W IN - TT-2 Wt. t. s s= To-1. -
(10)
(1.))
S< Te 1,
(11)
Here 7 E T are learnable parameters, and ((.,)) is parameterized such that has a sum of 1y when 8 goes through [Tc] to satisfy the constraint in Eq. 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 Y, h. and [Td) reduce to specific values. Such theoretical result underscores the benefit of our design since it permits the model to dynamically update the prior, leading better optimization. The parameters n and ry
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(xx ((x] .Ta))
(x [T] xx, v1 -x,(x x,),,,I), which ensures q(x!77 (x! matches the equivariant prior (proof in App. A.2). The reverse transition kernel is given by (x-
121 xx!al)
(Np((x 1"" -,,x (f.l), I2) Similar to the unconditional case, we also adopt the noise prediction objective by rewriting oo (x:"
[T] xT. .)) = [! 1 (xfT
x; (x!T vT-a, co -
Xx Tel -))
The denoising network Eg is implemented
EGTN but with its output subtracted by the input for translation invariance, hence the translation equivariance of μgg
Training and inference. Optimizing the VLB of our diffusion yields the following objective:
Leand = E BefT -Tlee - [IIe cx [a) .
(12)
after simplification (proof in App. A.2). The training and inference procedures are in Alg. and 4.
Experiments
We evaluate Geo on N-body physical simulation, molecular dynamics, and pedestrian trajectory forecasting, in both conditional (8 5. and unconditional generation ($ 5.2) scenarios. We ablate our core design choices and demonstrate additional use cases in $ 5.3.

---

Table 1: Conditional generation on N-body. Table 2:
Pedestrian trajectory forecasting on ETH- Results averaged over runs, std in App. C.4. UCY. Best in bold and second best underlined.
Partiele
Spring
Gravity
ETH
Hotel
Univ
Zaral
Zara2
Average
ADB
FDE
ADE
FDE
ADE
FDE
Linear
1.072.28
0.31/0.61
0.52/1.16
0,42/0.95
0.32/0.72
0.53/1.14
RF (23)
0.479 1,050
0.0145
0,0389
0.791
1.630
SoPhie [39]
SGAN [14]
0.70/1.43
0.64/1.09 009 0,46/0.98
0.76/166
0.56/1.18
0.54/1.24
0.33/9.67
0.30/0.63
0
0.31/0.64
0.54/1.15
0,46/0.91
TEN [56]
SB(3)-T [9]
BGNN [43]
0.186
0.330
0.395
0.754
0,936
0.426
0,1013
0.0101
0,0865
0,2364
0.2043
0.0231
0.327
0.338
0.310
0.761
0.830
0.709
Taj++ [40]
PBCNe [32]
0,5400.87
0,54X0.94
0.18/0.24
0,160.28
0.35/0.00
0.28/0.55
0.21/9.42
0.22/0.39
0.170.30
0.16/0.2
0.29/0.48
0.27/0.50
BiTraP [68]
0.56/0.98
0.17/0.28
0.25/0.47
0.23/0.45
0.16/333
0,27/0.50
EgMotion [62]
0.141
0.310
0.0134
0.0358
0.302
0.671
MID [12]
0.50N0.76
0.16/0.24
0:28/0.49
0.25/0.41
0.19/0.35
0.27H0.45
SVAE [67]
0.378
0.732
0,0120
0.0209
0.582
1,101
SVAE [67]
0.4780.66
0,14/022
0.25/0.17
0.20/0.37
0.14/0.28
0.2400442
GeoTDM
0.110
0,0030
0.0079
0.256
0.613
GooTDM
0.46/0.64 0.13/0.21 0.2400.55 0.21/0.39
0.16/0.30 0.23/0.40
Table 3: Conditional trajectory generation MD17. Results averaged over 5 runs (std in App. C.4).
Aspirin
Benzene
Fthanol
Malonaldshyde
Naphthalene
Salicylic
Toluene
Uracil
ADE
FDE
ADE
FDB
ADE
FDE
ADE
FDE
ADE
FDB
ADB
FDE
ADE
FDE
ADE
FDE
RF (27)
0.303 0.442 0.120
0.194
0,374
0.515
0.297
0.454
0.168
0.185
0,261 0.343
0.199
0.249
0.239
0.272
TEN [56]
0.133
0.268
0.024
0.049
0.201
0.414
0.184
0,386
0.072
0.098
0.115 0.223
0.090
0.150
0.090
0.159
SE(3)-T:
0.294
0.556
0.027
0.056
0.188
0.359
0.214
0.456
0.069
0.103
0,189
0.312
0.108
0.184
0.107
0.196
EGNN [43)
0,267
0.564
0.024
0,042
0.268
0.401
0.393
0.958
0.095
0.133
0.159
0.348
0.207
0.294
0.154
0.282
EqMotion [62]
0.185
0.246
0.029
0.043
0.152
0.247
0.155
0,249
0.073
0.092
0.110
0.151
0.097
0.129
0.088
0.116
SVAE [67]
0.301
0.428
0.387
0.505
0.287
0.430
0.124
0.135
0.122
0.142
0.171
0.145
0.156
GeoTDM
0.107 0.193
0.023
0.039
0.115
0.209
0.107
0.176
0.064
0.087
0.083
0.120
0.083
0.121
0.074
0.099
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. soume-to-frame prediction models: Radial Field [27], Tensor Field Network [56], STans)-Tr [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 (7, y
TN ET Ex= ||x - y2"112, ll, and FD y
N Σm ||x (T-1) y, (T-1) ll2. For
1
probabilistic models, we report average ADE and FDE derived from K samples.
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 self- loops. 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.
Implementation.
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 accumn- lated 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 in- jecting 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
7

---

Table 4: MD Trajectory generation results MD17. Marg, Class, and Pred refer to Marginal score, Classification score, and Prediction score respectively. GeoTDM performs the best on all 8 molecules.
Aspirin
Benzene
Ethanol
Malonaldchyde
Marg 4
Class
Pred J Marg
Class f
Pred 4 Marg 4
Class
Pred 4 Marg 4
Class 1
Pred_ 4
SVAE [67]
3.628
6.80x1001
0.0949
4.755
281xx10""
0.0181
2.735
2.39x10-9
0,0929
2,808
5.57x10
0.0346
BOVAE [43]
2.650
1.31x10-1
0.0386
3,677
1,50x10 4 0.0104
2.617
5.86x10-
0.1131
2.767
1.73x10
0.0664
GeoTDM
0.726
30
0.0212
0.597
1.62x10
0,0019
0.314
4.63x10
0.0235
0.403
3.35x10
0.0146
Naphthalene
Salicylic
Toluene
Urail
Marg 4
Class f
Pred 4
Marg
Class f
Pred 4
Mary 4
Class
Pred 4
Marg
Class 1
Pred 4
SVAE [67]
3,150
2,50x 10
0.2123
2.911 3,54x 10
0.1312
3.083
8,29 10
0,2580
2.736
3.73x 10
0.604
EGVAE [43]
3.007
3,17x10-1
0.0136
3.314
3,76x10
0.0221
2.054
2.77x100 -5 0.0457
3,570
2.02x
0.0212
GeoTDM
0.770
1.17100''
0.0093
0.559
1.82x10"'
0.0135
0,539
1.12x10"
0.0118
0.954
2.02x10-1
0.0116
ranging from (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 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.
most to the vibrations of the trajectory, leading
Results. As depicted in Table 3, GeoTDM achieves the best performance on all eight molecule MD trajectories, outperforming previos state-of-the-ane art approach EqMotion. In particular, GeoTDM ob- tains an improvement of 23.1 15.15%%%%%%% on average in terms of ADE/FDE, compared with the previous ofateof--the-rt 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 ETH-UCY [35, 28] dataset, a challenging and large- scale benchmark for pedestrian trajectory forecasting. There are five scenes in total: ETH, Hotel, Univ, Zaral, and Zara2. Following standard setup [14, 67], 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 preset distance measured from the final frame in the given trajectory. The metrics are minADE/minFDD 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 Teo 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 SGAN on MD17 is omitted due to mode collapse during training.
version of EGNN [43], dubbed EGVAE (see App. B). The results of
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 train-on-se www.oonteeii--eetttooooo
trajectory prediction model (a 1-layer EqMotion) that takes as input the first half of the trajectories to predict the other half.
8

---

GeoTDM
rilla
GortMtrteers Ground Tuth
ADZ FDE
Distorce (ADERDE))
Disccpt CT)
Distept DONN]
6.3
9:
SVAE
(a)
(b)
(c)
Optimization Stop
50
Figure 2: (a) Unconditional generation samples on MD17. GeoTDM generates MD trajectories with much higher quality (see more in App. D). (b) Interpolation. Left: the given initial and final frames. Right: GeoTDM interpolation and GT. (c) Optimization by GeoTDM on predictions of EGNN. Dis(Opt, GT), EGNN) is the distance between optimized trajectories and GTVEGNN.
Results. Quantitative results are displayed in Table 5 and 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 in- distinguishable with the ground truths. We observe similar patterns MD17, where GeoTDM obtains remarkably lower marginal scores, higher classification scores, and lower prediction scores, showcasing its strong capability model complex distributions of geometric trajectories on various geometric data. Visualizations are in Fig. and more in App. D.
Table 5: Unconditional generation re- sults on N-body Charged Particle.
Marg
Class t
Pred |
SGAN [14]
0.1448
3,98x10-?
0.172
SVAF [67]
0.0668
1,38x 10-6
0.282
EGVAR [43]
0.1141
4.22 10
0.0467
GeoTDM 0.0055
5.56x10-1 0.00978
5.3
Ablation Studies and Additional Use Cases
Ablations diffusion prior. We investigate different pri- ors, including non non-equivve -oniivarvai N (O, I) (i.e., DDPM [18]), equivariant but fixed CoM prior ([CoM(x!! (T. (To-)))) I) and
Table 6: Ablation studies. The numbers refer to ADE/FDE.
Charge
Aspirin
point-wise equivariant prior NN (T.-) I). In Table 6 we that non-equivariant prior leads significantly worse performance. The CoM prior, though equivariant, is still inferior due to extra overhead in denoising the nodes initial- ized around the CoM to the original geometry. Geo yields the lowest error due the flexible learnable prior.
GcoTDM BOOTDMW(x ((x!F
0.110/0.588
0.107/0.193
Fixed (oo.d
0,22000485
0.235/0.393
[(CoM(x (T-)))
0.135/0.298
0.119/0.212
0.1230.282
0.110/0.204
w/o Equivariance
0,251/0.542 0.252/0.440
wo Attention
0.13//0.312 0.114/0.208
w/o Shift invariance
0.139/0.330 0.112/0.212
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, Icading 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 vari- ant, 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 special case of the conditional case. We demonstrate such capability on Charged Particle. The model is provided with the first and last 5 frames, and the task to generate the intermediate 20 frames interpolation. Geo reports an ADE of 0.055 on the test set, while a linear interpolation baseline reports an ADE of 1.1 From the qualitative visualizations in Fig. 2, clearly see that GeoTDN 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] given trajectories (e.g., those simulated by an EGNN) by simulating 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 Ge can effectively optimize the given trajectory towards the ground truth distribution.

---

Discussion
Limitations, Akin to other diffusion models, GeoTDM resorts to multi-step sampling which may require more compute. We present empirical runtime benchmarks and more discussions in App. C.3.
Conclusion. We present GeoTDM, diffusion model built over distribution of geometric trajectories. is designed to preserve the symmetry of geometric systems, achieved by using EGTN, a novel SE(3)-eqiivarian geometric trajectory model, 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 feedhack on improving the manuscript. This work was supported by ARO (W911NF-21-1-0122 ONR (N00014-33-1-2159)) and the CZ Biohub.
References
[1] 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.
[2] 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, 7
[3] Lowik Chanussot, Abhishek Das. Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Mor- gane Ríviere, Kevin Tran, Javier Heras- Drrras--Dmino,, Caleb Ho, Weihua Hu, et al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11(10):6059 6072, 2021. 24
[4] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018. 22
[5] Stefan Chmicla, Alexandre Tkatchenko, Huziel Sauceda, Igor Poltavsky, Kristof TT Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e1633015 2017.
[6] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis.
Advances in neural information processing systems, 34:8780-8794, 2021. 2
[7] 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
[8] Jacob Durrant and J Andrew McCammon. Molecular dynamics simulations and drug discovery. BMC biology, 9(1):1-9, 2011.
[9] Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling, Se (3)-transformer:: 3d roto-translation equivariant attention networks. arXiv preprint arXiv:2006, 10503, 2020. 2, 7, 22
[10] Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations, 2020.
[11] Justin Gilmer, Samuel S Schoenholz, Patrick 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,
[12] Tianpei Gu, Guangyi Chen, Junlong Li, Chunze Lin, Yongming Rao, Jie Zhou, and Jiwen Lu. Stochastic trajectory prediction via motion indeterminacy diffusion. In Proceedings the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 17113-17122, 2022. 7, 8. 26
[13] 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 Learning Representations, 2023. 6, 20
10

---

[14] 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 IEFE conference computer vision and pattern recognition, pages 2255-2264, 2018. 7, 8,9, 9, 22
[15] 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.
[16] 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
[17] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. Advances Neural Information Processing Systems,
35:27953-27995 2022. 3
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances
in neural information processing systems, 33:6840-668 16688006651,, 2020. 2, 3, 5, 7, 18, 21
[19] 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
[20] Scott A Hollingsworth and Ron O Dror. Molecular dynamics simulation for all. Neuron, 99(6):1229 1143, 2018.
[21] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pages
8867-8887. PMLR, 2022. 1, 2, 3, 6, 20
[22] Wenbing Huang, Jiaqi Han, Yu Rong. Tingyang Xu, Fuchun Sun, and Junzhou Huang, Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations, 2022. 1
[23] John B Ingraham, Max Baranov, Zak Costello, Karl W Barber, Wujie Wang, Ahmed Ismail, Vin- cent Frappier. Dana M Lord, Christopher Ng-Thow-Hing, Erik R Van Vlack, al. Illuminating protein space with a programmable generative model. Nature, pages 1-9, 2023. 1
[24] Bowen Jing. Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi Jaakkola. Eigenfold: Generative protein structure prediction with diffusion models, 2023. 2
[25] Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons. Neural sdes as
gans. In International conference on machine learning, pages 5453-5463, PMLR, 2021. 8
[26] Thomas Kipf. Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint arXiv: 046877 2018. 1,
[27] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi-hody systems with symmetric energies. arXiv preprint arXiv: 2019. 2, 7, 22
[28] Alon Lerner, Yiorgos Chrysanthou, and Dani Lischinski. Crowds by example. In Computer graphics forum. Wiley Online Library, 2007. 8
[29] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022. 24
[30] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific 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.
[31] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion-basees generative models for protein structures. Advances in Neural Information Processing Systems, 35:9754-9767, 2022. 9
[32] 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-E.CC 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part II 16, pages 759-776. Springer, 2020. 7,8 8
11

---

[33] 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
[34] Hao Ni, Lukasz Szpruch, Magnus Wiese, Shujian Liao, and Baoren Xiao. Conditional sig- wasserstein gans for time series generation. arXiv preprint arXiv:2006.0544 2020. 8
[35] 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.
[36]
Calyampudi Radhakrishna Rao. Linear statistical inference and its applications, volume 2.
Wiley New York, 1973.
[37] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denois- ing diffusion models for multivariate probabilistic time series forecasting, In International
Conference on Machine Learning, pages 8857-8866 PMLR, 2021. 3
[38] Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. DYffusion: a dynammiss-nformeed diffusion model for spatiotemporal forecasting. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[39] 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
[40] Tim Salzmann, Boris Ivanovic, Punarjay Chakravarty, and Marco Pavone. Trajectron++: Dynamically-feasibee trajectory forecasting with heterogeneous data. In Computer Vision- ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVIII 16. pages 683-700. Springer, 2020. 7, 8, 25
[41] Alvaro Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference Machine Learning, pages 8459-8468. PMLR, 2020. 2
[42] 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
[43] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. arXiv preprint 0tti::2002.09444: 2021. 1, 2, 4, 7, 9, 20, 22
[44] 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.
[45] 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 Physies, 148(24))241722, 2018.
[46] Jean-Pierre Serre al. Linear representations of finite groups, volume 42, Springer, 1977. 3
[47] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position rep- resentations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Short Papers), pages 464 468, 2018. 4
[48] 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
[49] Jascha Sohl- Sihl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2265. PMLR, 2015. 2,3
[50] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020. 24
12

---

[51] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consisteney models. arXiv preprint arXiv:2303.01469 2023. 24
[52] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 2, 3
[53] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33: 13:11238-12448, 2020. 3
[54] 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-24811 2021.
[55] 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 Repre- sentations, 2023. 3
[56] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation- caattiionn-uiiviaiinnn neural networks for 3d point clouds. arXiv preprint arXiv:
2018. 2, 7,
[57] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[58]
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
Vikram Voleti, Alexia
[59] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian Trippe, Jason Yim, Helen 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):108--111 1,2
2023.
[60] Fang Wu and Stan Z Li. Diffmd: geometric diffusion model for molecular dynamics simula- tions. In Proceedings of the AAAI Conference on Artificial Intelligence, 2023.
[61] Liming Wu, Zhichao Hou, Jirui Yuan, Yu Rong, and Wenbing Huang. Equivariant spatio- temporal attentive graph networks to simulate physical dynamics. Advances in Neural Informa- lion Processing Systems, 36, 2024. 4
[62] 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 IEEF/CVE Conference on Computer Vision and Pattern Recognition, pages 1410-1420. 2023. 2, 8, 22
[63] 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
[64] Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference Machine Learning. PMLR, 2023. 1, 2, 3, 24, 25, 26
[65] Minkai Xu, Wujie Wang. Shitong Luo, Chence Shi. Yoshua Bengio, Rafael Gomez-Bonee 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
[66] Minkai Xu, Lantao Yu. Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric difffusin model for molecular conformation generation. In International Conference on Learning Representations, 2022. 1, 2, 25, 26
[67] 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
[68] Yu Yao, Ella Atkins, Mathew Robeeson,, Ram Vasudevan, and Xiaoxiao Du Bitrap: Bi-directional pedestrian trajectory prediction with multi-modal goal estimation. IEEE Robotics and Automation Letters, 6(2): 2021. 7, 8
13

---

[69] Jason Yim, Brian L. Trippe. Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE() diffusion model with application 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
[70] 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
[71] Ye Yuan, Jiaming Song, Umar lqbal, Arash Vahdat, and Jan Kautz. Physdiff: Physics- guided human motion diffusion model. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 16010-16021, 2023. 3
[72] 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-42644 PMLR, 2023.
14

---

Appendix
Table of Contents
A Proofs
15
A.1 Unconditional Case
15
A.2 Conditional Case
17
A.3 Optimizable Equivariant Prior
19
A.4 Proof of Theorem 4.1
20
B More Details on Experiments
21
B.1 Compute Resources
21
B.2 Hyper-parameters
21
B.3 Baselines
22
B.4 Model
22
B.S Evaluation Metrics in the Unconditional Case
22
C More Experiments and Discussions
23
C.I Model Composition for Longer Trajectory
23
C.2 Number of Diffusion Steps
24
C.3 Sampling Time
24
C.4 Standard Deviations
24
C.5 More Discussions with Existing Works
25
D More Visualizations
25
A
Proofs
A.1
Unconditional Case
We note that, naïvely, a distribution p(x!77) can not be translation invariant. In particular, this would imply that p(x!")) = p(x(7) r) for all r e R,,, but this would imply that p(x = uniformly, a contradiction.
Instead, we derive an equivalent invariance condition by restricting SE(D) to its maximally compact subgroup. In particular, note that it is possible to define -innaninvariant distributions (as this group is compact), and SE(D) our base space Txx T (TNN D)x is to zero-center our data (along each dimension).
T, the translation group. The natural way to quotient out
However, for practical purposes, we will refer to our construction as SE(D) -iivaariat.... In particular, since all inputs F xxXXx are first zero-centered be projected to NxD /T, the "lifted" unnormalized measure is SE(D) invariant.
We will define P as our zero-centering operation P(x!T)) = x + ET= ())(x()), with C) = N Σm x! and our restricted (T x N 1) x D Gaussian as Nylxx, E), which can be represented in the ambient space degenerated Gaussian variable
1 1- x) "E x)) (13)
N(yxx, Σ)
(2 (2T))/ D)xD) det (Ep) exp
where Ep PEP and Σ is the pseudo-inverse (and det. is the determinant restricted to the subspace). Note that is symmetric and idempotent. Then specifically when
I, we have
Ep PP then Σ P (PP TPPPPTT P. since PP P and P =P! PT
15

---

Base distribution. We require the base distribution to be Pr (x())) (o, I) to be the Gaussian distribution in the canslation-nnvariant subspace
In practice, we let
Transition kernel. For the transition kernel, we specify it as pa (x{ x = (x, : .T, oZ
In order to ensure pe (Rx equivariant function. In this way,
= pa (x Rx(")
x) it suffices to make (x-" an SO(3)-
Pg Rx [2]
Rx!T)) = ( 7- PLg (Rx! (7]] ), 1
(14)
-NRR !71 7 (Rx, (T) 0,
(15)
-N(RX 171 Rio (1)) 0,
(16)
= -N(RR(( 11 i O,
(17)
= -N(x (x( ) 0,
(18)
= 1i (x, II
(19)
<x x
(20)
which permits the -quiviaanree of the transition kernel. In our implementation, we further re-parameterize j (x [T]
r) as,
i (x("" .T) Var
x
B. V1 ā.
(x! Eo(
(21)
[T] T) to be S033))-quivariaa and its output should lie in the subspace where instead ensure (x!" Xp.
We now prove the following proposition, which states that if the base distribution is SO(3) -invaaatn and the transition kernel is (0(3)--quivariant, then the marginal at any diffusion time step is also (in)--nnvria
Proposition A.1. If the prior (x (T]) is S0(3)-invariant the transition kernels P-- (x (x
) VT E fl..... , are -0(3)-equivariaan) then the marginal p+ px(x!!") at any time step T E (0. T) is also 3033)-invariant
Proof. The proof is given by induction.
Induction base. When 7 = T, we have the marginal being the prior pr (x( (2')), which is SO(3)-
invariant.
Induction step. Suppose the marginal at diffusion time step r is SO(3) (inva-iinariant,, i.e., pr (x(")) =
Pr (Rx!")), then we have the following derivation for the marginal at time step 1:
P--1 (RX! =
P R
}r jaz!"
(22)
D
RR
pp (RR
X axx!?
(23)
Dr-
R
R
(24)
Pr-
[T] (x
det (
pr
(25)
Pr- (x
(26)
Notably, for the final step at = 0, the marginal Po (Rx()) final sample from the entire geometric trajectory diffusion process resides in an -invarant distribution, hence the physical symmetry being well preserved.
[1] is also SO(3)-invariant, indicating the
16

---

Algorithm Training Procedure of GeoTDMMuuncond
1: repeat
Sample E(T) (N., 7 :
x x Data
3: X 7l + vā + -1 a,
4: Take gradient descent step
Vall EΦ( _ 1
5: until converged
Algorithm 2 Sampling Procedure of GeoTDM
1: Sample x (O, r)
2: for <- T.... . do
3: Sample x!"" ~ ~N 1) IT if > 1 clse 2.77 =
4: x,-1 Jaz (x(")
1-a- - (x .r) + 0.
end for
return x
Algorithm 3 Training Procedure of GeoTDM-cond
1: repeat
2: x EquiPrior,, n Fx
{Eq. 9}
3: Sample e! ~ N.,
1). Unif([1, ([,, ..., 7T),
(x! x .) Pilata
4:
x
(x.") " x! + +x!!!
/1 ā, c!
Take gradient descent step on
(T)) ?)?? Vo.nolle cx( [T] x
until converged
Algorithm 4 Sampling Procedure of GeoTDM-cond
1: xi T Equuiriior, , (x.! (x)
(Eq. 9}
Sample x227 [x!!'" !T] 1). [TT)]
(Te) Danta
3: for - T, T.... . do
Sample 27" NV. 1) if > 1 else 2!!1
5:
x_-1 + 1. (xf")
[7) xf
1-a-
cx (7) [Tll
.7)) + x," + t (7) x
var
Vi-ar
6: end for
7: return xb
[3]
A.2 Conditional Case
In the conditional case, we target on modeling the conditional distribution p(x!?) [7]
x (Tel). The
desired constraint the following equivariance condition: p(x[2" [Tel] = p(g x!T! g x Tel), for all g E SE(3).
Construction of the equivariant prior. The prior is constructed through Eq. 9. Here we formally show that this guarantees -qquiaaciie of the prior. For convenience we repeat Eq. 9 below.
Σ Tcll
w (t))
S.L. w(t,))
1,
(27)
17

---

Then, we have
(28)
sE/Ta
Σ ((.)))
(Rx! +r),
(29)
sE|Tal
w (1,5) (Rx Σ w((,))r
(30)
xE(T.]
R
w((,0)
+r,
(31)
=E[T
Rx +r,
(32)
rotation matrix R and r R3, which completes the proof.
Base distribution. We propose to leverage the following base distribution.
[ PT [
(Pal) = (x [ x [m II,
(33)
where xf." EquiPrior(xE is SE(3) -qquivaria with respect to the condition x. choice, the base distribution above is SE(3) (eeu))))quviian,, since
[Tl)]
Tcl
With such
PT Rx
+r
Rx [
+ +r) = (Rx + r: Rx +r.I).
(34)
= =NRx Rx
[7] II),
(35)
-x( 12), x(( x((dd. I),
(36)
where the last equation is due to det(R R) I, | x, - |
Rx Rx which also
gives the proof for Theorem 4.4 by a mild substitution the notations.
Transition kernel. The transition kernel given by
Po
[e)
177
(x
x
[Te]
(37)
where ue (x+ X In practice, re-parameterize it as,
[T]] r) parameterized to be
with respect to its input x!?? xx [Tc]]
n 17 [To), 7)
(T] +
x
[T]
Br
(x(
x s(((tl,7))
(38)
VI
a
where eø (x!"" X Tcl 7) is an -Quuieequi but munsslationinvariian function. It is then easy to [T] [T.] T) meets the SE(3) equivvariie as desired.
[T]
see that sn (x) x
Proposition A.2. With such parameterization, optimizing the variational lower bound is equivalent to optimizing the following objective, up to certain re-weighting:
C = =>ee x
[T]
[Tcl
[)) -E1
(39)
Proof. We define q(x!" (ix (x! = Nix M x+. + VI- (x!T] x,), which yields
[T) (xx!?T)xxx) (Tl) (xx!!! x, + (x, Ti x,), (1 ū, II. The proof then generally follows [18]
but with all latent variables in [18] being replaced by x x! Then the terms in the VLB
18

---

(Tcl), given by,
-- DKL(g(x
--1
[T]
x
(T)) x
171
Ilpo
(x --1 x x) ))
(T)
(40)
= E !m
1
x
(
1
xx (x) [7]
c, [Tll)
x
x
vI ar 6)
Br
11 .7)
20
a
(41)
Br
E
x
[T]
x
[T)
x
xg
2
0+
x
1
x [
x
Br
(x (x, x ,7)
Var
1 dy
(42)
88
-E i 2020 a,)
||e Eg (x("")
x [Tcl
-)1
(43)
which is equivalent Eq. 39 up to certain re-weighting factors. For Ly
DKL ( (x(x) xx [2'] [
{x )), it is does not contribute to the gradient since it is irrelevant to 0, and x, computing the KL, thus stopping the gradient from passing 77 and y.
also cancelled
Analogous to the unconditional case, we have the following proposition, indicating that if the base distribution is -equivvaana and the transition kernel is (B(33-equivarian then the marginal is also -E3)iecquivariant.
Proposition A.3. If the base distribution pr (x²
[T]]
[Ta)) is SE(3)) equivariant and the transition X
kernels P- (x [x]_] x,, [21
xx (fdd) of all diffusion steps E (1,...,7 T are 3E(33)equivarian) then , [2] [T-l] at any diffusion step {0,-.. , is -quiivraan,
the marginal pr (x7 X
Proof. The proof similarly given by induction.
Induction base. When r T, the distribution is the base distribution px [7] (Ta)) is SE(3)-
equivariant, as it is designed.
Induction step. Suppose the marginal at diffusion step i.e., Pr (x. then we have
x [To)) ), is (E)33)aeuivarian
Pr
Rx
Ta)
+ r)
(44)
Pr
Rxx
[T
+r
Rx! [Tll
.T.
+ r)p, (x.
Rx!" [Te] + r)dx," 77,
(45)
PT-1 Rxx --1 +r
R(R (x!!l - r)) + Rx!Td] tr)p (R(R (R(R-'(x!T) r)) + Rx [Ta +
rr)xx"" !7]
(46)
Pr- (x
[T]
R
[x!!1 171 r), (
[Tac )pr (R -x!!" - r)
Tal )xx
(47)
Pr- 1x
y!
[
x
}r- (y!
xxlll) dot dc)
(48)
P (x which concludes the proof.
(49)
A.3
Optimizable Equivariant Prior
Theorem A... The prior implemented by the parameterization in Eq. 10, and 11 subsumes CoM-based priors and fixed point-wise priors.
'Here the marginal refers marginalizing the intermediate states in previous diffusion steps, while still being conditional on the input condition x!
19

---

Proof. We repeat the parameterizations specified by Eq. 9, 10, and 11 below for better readability.
(0)
w (4.)
(o)
S.t.
w w(t.)) 1N.
(50)
SEITY
W 1, b & h
1.a ERR
(51)
WL,
< Te 1,
((,e)
(52)
IN - Wt == T. 1.
We first show X; can reduce to the CoM-based priors. Let x! CoM(x )), h()))
T. 1,
((0) 1. In this case,
x
Σ (()))
(53)
+E[Tc]
Σ
(())
INCOM(x(") + (IN-
Σ 7(0) [N) (T)-)))
(54)
Te
Te
sE
sE[T,-1]
Σ INCOM(xx") + (1N To
IN IN) ((
(55)
To
Te
sE(T.-1]
Σ Te
1 INCoM(x(")) + 1n (NCoM((((--1))
1
(56)
To
-E(T--1
1 C)
(57)
T. .E
(e) is the generalization of the CoM-hased priors [21, 13] in the multiple where #. E.E(T.] CoM(x!") frame conditioning scenario, which reduces to CxM(x when TN 1.
To show x can reduce to fixed point-wise priors is straightforward. Let x!*) Onehot(s")1 and 7(0) = 1, Vt. Then (t,s) = Onehot (
(+) h
1TEXN. Therefore,
(1)
Σ
w (t.s) (())) (4)
(58)
SE[T.]
Σ
Onehot (s*) I
(59)
(60)
where (5") is the point-wise equivariant prior, and s* E [Tc] 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)-equivarianc) of EGTN), Let x (T] hT fECTN (xITI, 1(,,1 Then have g x(T), h/T7 fECTN (g x(7) (1,(() Vig SE(3).
Proof. JEGTN is a stack of T. EGNN and temporal attention layer in alternated fashion, formally written as JEGTN fattn o fEGNN .... fatta fEGNN. Since the chain of SB(3) -qqiivaann
Lx (Jats
function is also -(lu)vvvvaarian,, it suffices to prove fattn is equivvaraa,,,, in that the SE(3)- cquivariance of EGNN directly follows [43].
It is directly verified that the attention coefficients (1**) R in Eq. 5 and the query qT key klT), and value !!] are all (o033)-invariante since they derived based the SE(3)--nvariant input
20

---

CoM-based prior To
Σ CoM(x!)) se
Fixed point-wise prior x (
Ours
Mean of the prior
condition x
target x(r]
Figure 3: An illustration of different equivariant priors. For simplicity in the chart here we only illustrate case when N and T. = 1.T=1.
h This directly leads to the SE(3)) invariince of the updated node feature h For the updated coordinates,
Xrr xx
(1))
(1)
Σ SE(T]
air
(1,3)
px(v
(2.0)
(x x(=));
(61)
Rx(()
+ Exe(T) a 4x (v((:))) (2) +r- - -r),
(62)
Rx(0) +r+ R Σ SE(T) a xxx (v(+))) - x((1))
(63)
R Rx(x))) + Σ NE(T]
a ( (x(())))) - x(())) +
(64)
Rx +r,
(65)
where the variables with subscript refers their transformed counterparts when the input [?] transformed into Rx r. Thus it completes the proof of -quiiarrarvr of the temporal attention saren and hence the entire fECTN.
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-parametern
We provide the detailed hyper- hyper-paramette parameters of GcoTDM in Table 7. We adopt Adam optimizer with 10 5 For all experiments, the linear noise schedule per [18] with 3atrt = 0.02 and Bend = 0.0001.
betas (0.9, 0.999) and
Table 7: Hyper-parametee parameters of GeoTDM in the experiments.
layer
hidden
time_ emb dim
T
batch_ size
learning rate
N-body
6
128
32
1000
128
0.0001
MD
6
128
32
1000
128
0.0001
ETH
4
32
100
100
0,0005
21

---

B.3 Baselines
For the Trame-to-fraame prediction models, including RF [27]. EGNN [43], TFN [56]. and SE(3)- Transformer [9]. we udopt the implementation in the codebase maintained by [43]. To yield a strong comparison, instead of taking one frame as imput 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()) and velocity (computed as the difference of the current and previous frame, i.e., v(*) = x( - x x(*-1) 1) as imput to predict the next velocity ((++1) The position for the next step is integrated as x (++1) = x(") + ((++1) The training loss is computed as the Mean Squared Error (MSE)
(t+1) In inference time, roll- out prediction is conducted, which iteratively predict the next step by feeding the predicted position and velocity at the current step, for total of 7 steps. We follow the sarpee-arreeter tuning guideline for these baselines by [16] which conduct a random search over the space spanned by the number of layers in {2, 4, 6, 8). the hidden dimension {32, 64, 128), learning rate (5e- 33 le 33 5e b,e4,, 4, 1e-4). 4], and batch size 32, 64, 128, 256, and select the model with best performance. All models are trained towards convergence with carly-stopping counter of 5, with validation performed every 20 epochs.
between the predicted position (*+++) and the ground truth position xt
For EqMotion, we directly adopt the code by [62] and their suggested hyper-parameten for the N-body datasets and MD17 datasets.
For SVAE [67] and SGAN [14], these methods originally developed for the pedestrian trajectory forecasting task. The backbone model that processes the input trajectory consists of social pooling operation and GRU or LSTT 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- syper-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:
mij = pm (hi, hý, |x; es),,
(66)
h = Ph
his
mij
(67)
jEN())
X x; +
Px (my) (x; - x,),
(68)
(EN(()
where Pm, Sh, and 4x are all MLPs. We also provide a schematic of our proposed EGTN in Fig. 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, 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 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.
clearer signal that correlates to the validity of the generated MD
22

---

x
h
x
Prod. E,
Prod. Σ,
9x
Temporal Attention
a
DotProd, Softmax
8
Spatial EGCL
(1)
15
k
(
Pd
9k
94
xD
h7
x(
Figure 4: Schematic of the proposed EGTN, which alternates the EGCL layer for extracting spatial interactions and the temporal attention layer for modeling temporal sequence. Additional conditional information x and h
also be processed using cross-attention. The relative temporal embedding (t s) is added to the key and value. DotProd refers to dot product and Softmax is performed over indexes of
Classification score is computed as the cross-entropy loss of a sequence classification model that aims distinguish whether the trajectory is generated by the model from the testing set. To be specific. we construct dataset mixed by the generated samples and the testing set, and randomly split 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 I l-layer EqMotion with 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 the MSE loss of train-on- www.ooonntnnntttteeeeeeeeeeeeeee teatooor- sequence to sequence model. In detail, we train a I-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, 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" non-overlapping intervals with time span AT', Then, for the unconditional case, we have
K-1 II
p(x("1) (ZX)
p bAT+177 x( x((-I)AT+AAT)
(69)
by assuming mild conditional independence, where (Tx()T]) is an unconditional model for trajectory with length AT, and p[x*AT+|AT]] ((*-1))7+(A7)) can be Iearned by conditional model for short trajectories. The conditional case directly follows by factorizing into products conditional distribution over shorter trajectories.
We provide demonstration of such technique as gifs in the supplementary file.
In fact they do not necessarily need to be equal-length. Here we make such assumption for conciseness of the presentation.
23

---

Table 8: The effect of diffusion steps in the unconditional generation setting (top) and conditional forecasting setting (bottom).
Aspirin Classification
Charged Particle Classification 1
Marginal +
Prediction +
Marginal
Prediction +
T 100
0.808
0.0242 0.0348
0.0243 0.0212
0.0065
0.170 0.556
0.0118
T 1000
0.726
0.0055
0.00978
Aspirin FDE J
Charged Particle
ADE
NLL 4
ADE
FDE L
NLL
T 100
0.110
0.198 0.193
-2125.7
0.120
0.280 0.258
-547.5
T 1000
0.107
-3461.4 .4
0.110
-982.7
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 more challenging in that is needs to generate trajectories without any given reference geometries. We observe a drop in performance when is decreased from 1000 100. However, the performance with only 100 diffusion steps is still significantly better than SVAE.
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 EGVAF, an autoregressive VAE-based method with EGNN as the backbone. Here GeoTDM-100 and Geo0TDM-1000 refer to GeoTDM using 100 and 1000 diffusion steps, respectively.
Table 9: Sampling runtime comparison MD17 Aspirin molecule.
Marginal
Classification
Prediction
Time per batch
EGVAE
2.650
1.31x10-4
0.0386
0.6:10.1 7.9:1..8
GeoTDM-100
0.808
2.42x10 -
0.0243
1000DM-1000
0.726
3.48x10-2
0.0212
74.2=2.1
We observe that GeoTDM- 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 Geo is slower, the gain in performance is significant and the quality of the generated trajectory is remarkably better than that of EGVAE. When further increasing the number of diffusion steps to 1000, the performance becomes better while requiring much more compute.
However, it is worth noticing that all these deep learning-based methods are significantly faster than traditional methods like DFT, which typically requires hours to even several days converge depending on scale of the system, according OCP [3]. Therefore, although Geo 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 more compact representation of the spatio-temporal geometric space.
C.4 Standard Deviations
We provide the standard deviations in Table 10 and 11.

---

Table 10: Conditional generation results of Geo on N-body charged particle, spring, and gravity,
Results (mean + standard deviation) are computed from samples.
Particle
Spring
Gravity
ADE
FDE
ADE
FDE
ADE
FDE
GeoTDM
0.258:1:0032 SVAE 0.378=00005
0.110=0.014
0.0030:10.0004 0.0120:10.003
0.0079:10.0010 0.020911000044
0.2561:00015 0.58211.007
0.613:10.034
0.732:1:0005
1.101::.0015
Table 11: Conditional generation results of GcM on MD17. Results (mean + standard deviation) are computed from samples.
Aspirin
Renzene
Ethanol
Malnnaldehyde
ADB
FDE
ADE
FDE
ADE
FDE
ADE
FDE
0.107-0.005
0.19340.016 0.023+40.001 0.039
0.115+00.012 0.209-10.035
0.107=00010
0.176-10.025
Naphthalene
Salicylic
Toluene
Uracil
ADE
FDE
ADE
FDE
ADE
FDE
ADE
FDE
0.064=0.002
0.087=0.007
0.08340.004
0.120=0.012
0.083 0.121#0.011
0,074=0.003
0.099:10.009
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 Trajec- tron++ [40] backbone which takes input the position vectors through Transformer network. It requires non-trivial effort to incorporate additional node features and edge features into MID, while for GeoTDM, design 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 stretures.
Incorporating equivariance into temporal diffusion. While geometric diffusion models have discussed proper ways to inject cquivariance into diffusion models, unclear how to preserve cquivariance when each hidden variable in the diffusion process has an additional dimension of time. 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 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. both unconditional and conditional generation for geometric trajectories, and claborate on how to design the prior and transition kernel to meet the equivariance constraint.
this work, we systematically discuss
Parameterization of the learnable equivariant prior. In the conditional case, we propose to parameterize the equivariant prior with 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, and 10. Please refer to their captions for the detailed descriptions.
25

---

Table 12: Technical differences between GeoTDM and existing works.
Trajectory
Equivariance
Conditional
Unconditional
.earnable Prior
MID [12]
GeoDiff [66]. GeoLDM [64]
Our GeoTDM
Figure 5: Uncurated samples of GeoTDM on MD17 dataset in the unconditional generation setup. From top-left bottom-right are trajectories of the eight molecules: Aspirin, Benzene, Ethanol, Malonaldehyde, Naphthalene, Salicylic, Toluene, and Uracil. Five samples are displayed for each molecule. GeoTDM generates high quality samples. well captures the vibrations and rotating behavior of the methyl groups in Aspirin and Ethanol. The bonds on the benzene ring are also more stable, aligning with findings in chemistry.
Figure 6: Samples from MD17 dataset.
26

---

ral
t= T
0
7d
==T T
0
t= T
T= 0
X
T=T
T= 0
Figure 7: Visualization of the diffusion trajectory at different diffusion steps. From top bottom: Aspirin, Naphthalene, Salicylic, Uracil. For each molecule, the first row shows the unconditional generation process, where the model generates the trajectory from the invariant prior purely from the molecule graph without any conditioning structure. The second row refers to the conditional generation, where the model generates from the equivariant prior, conditioning on some given frames x [Ta]. Notably, the equivariant prior (see samples at T = T in each second row) preserves some structural information encapsulated in x thanks flexible parameterization.
[Va]
27

---

Condition
Prediction
Ground Truth
Condition
Prediction
Ground Truth
Figure 8: Uncurated samples of GeoTDM MD17 dataset in the conditional forecasting setting. We highlight some regions of interest in red dashed boxes. GeoTDM delivers samples with very high accuracy while also capturing some stochasticity the molecular dynamics.
Data
GeoTDM
Figure Visualization of data samples and generated samples by GeoTDM and SVAE in the unconditional setting Charged Particles dataset. Nodes with color red and blue have the charge of + +1/-1, respectively. Best viewed by zooming in.
28

---

GeoTDM
Figure 10: Visualization of predictions by GeoTDM and EGNN in the conditional setting on Charged Particles dataset. Nodes with color red and blue have the charge of +1/-1, respectively. Best viewed by zooming in.
29

---

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 Gco 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 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 long 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 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 (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
30

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, we have included all the relevant details for reproducing our experi- ments. 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 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 novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing or verifying the results.
5. Open access to data and code
31

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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 NeurlPS code and data submission guidelines (https://nips.. cc/
htalii//giides//eeeSbbmisssss iooiaaoomssioooooooooooooooooiil for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (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 NeurlPS code and data submission guidelines (https: //wips.cooo ctspblsccppublc//uuu
/CodeSubmia sgionPolicy) 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. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See our code and experiment details in App.
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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can 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, confi- dence 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.)
32

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https :/tp:::////urrr/ htt////nnurr.ppeeeeee httppublcc.tchccccdueeeeeee..eee
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.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: 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 no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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.
33

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our method does not have a high risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We have credited all code and dataset sources.
Guidelines:
The answer NA means that the paper does not 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
34

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
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 the dataset/code/mooell as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This work does not involve human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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 equivalent approval/review based 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 expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "tyeqTA8ZZc", "reviewer": "Reviewer_e5nM", "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.", "questions": "None", "limitations": "None", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The geometric trajectory diffusion model is novel to me. \n\n2. The EGTN that operates on geometric trajectories permits conditioning upon a given trajectory using equivariant cross-attention\n\n3, The experiments demonstrate that GeoTDM achieves SOTA performance on both unconditional and conditional trajectory generation tasks.\n\n4. The work can be extended to applications like temporal interpolation and trajectory optimization", "weaknesses": "None", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and recognition of our work! Please let us know if you have any questions and we are more than happy to answer.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n4PNFXPtAi", "reviewer": "Reviewer_UDq7", "summary": "The paper introduces geometric trajectory diffusion models for generation of particle, pedestrian, or molecular trajectories. The architecture consists of EGNN layers within a temporal frame and temporal attention across frames computed with relative temporal ecodings. The architecture is shown to possess the appropriate equivariances. Key points of novelty include a learnable prior as a function of conditioning frames in the conditional generation case. Experiments are conducted on particle, molecular, and pedestrian dynamics datasets.", "questions": "See above", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and suggestions! We provide point-to-point response below.\n\n> **[W1] 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\nWe will first elaborate on the motivation and insights of our architectural design. We also add extra experiments that explore more sophisticated equivariant networks under the framework of our GeoTDM.\n\n**Motivation and insights of the architecture.** Our EGTN is a stack of spatial layers and our proposed equivariant attention temporal layers in an alternated fashion.\n* For spatial layer, we use the widely adopted EGNN layer, akin to existing geometric diffusion models for static structures. Notably, this module is adopted in a plug-and-play manner, which enables us to switch into other advanced backbones depending on different data and task. We also provide extra experiments by switching into more advanced equivariant networks in the part below.\n* To process temporal information, we propose a novel equivariant attention layer that involves several core designs including relative temporal embedding, satisfaction of equivariance, and extension to cross-attention for conditioning. The utilization of attention here is inspired by the success of Transformers in processing sequence data such as text and audio, since here we are handling a sequence of frames.\n\nWe have also performed thorough ablation studies in Table 6, which aim to provide deeper insight about how these design choices matter.\n\n**Exploration of more sophisticated equivariant networks.** We also provide additional experiments on exploring more sophisticated equivariant networks in the framework of our GeoTDM. We replace EGCL by an advanced backbone Equiformer [1] (dubbed GeoTDM*) which enables utilizing higher order equivariant tensor. We use the codebase of [1]. The results on MD17 conditional generation are exhibited below.\n\n||Asp|Ben|Eth|Mal|Nap|Sal|Tol|Ura|\n|-|-|-|-|-|-|-|-|-|\n|GeoTDM*|0.099/0.184|0.023/0.036|0.098/0.189|0.117/0.184|0.060/0.092|0.075/0.113|0.112/0.130|0.067/0.094|\n|GeoTDM|0.107/0.193|0.023/0.039|0.115/0.209|0.107/0.176|0.064/0.087|0.083/0.120|0.093/0.121|0.074/0.099|\n\nWe observe that the performance is generally enhanced across several molecules. However, the sampling time has been dramatically increased since Equiformer is more computationally exhaustive compared with EGNN. Besides, it is a domain-specific backbone that cannot be directly applied to other tasks like pedestrian trajectory modeling, which limits the broadness of application.\n\n[1] Liao et al. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs. In ICLR'23.\n\n\n> **[W2] The exposition of the method drags at times, re-proving well-known statements about equivariance in neural networks and geometric diffusion models.**\n\nThe formal statements presented in this paper are fundamentally different from existing works and are not simply re-proved counterparts. The theoretical framework initiates with the formal definition of geometric trajectory (line 116), which is distinct from the static structures discussed in existing literature, being an extension of them with an additional temporal dimension. We then rigorously defined how the group action $g\\in SE(3)$ operates on geometric trajectory (line 129). Notably, Theorem 4.2-4.4 are introduced and proved under the definition of geometric trajectory and the group action enforced on it, which generalize the theorems in existing works and subsume them as special case when $T=1$. These theorems are presented for completeness and mathematical rigor of the paper.\n\n> **[W3] 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\nThank you for recognizing the diversity and broad coverage of the experiment and the suggestions on exploring novel capabilities of GeoTDM.\n\nIndeed, in the paper we have already explored a few directions regarding the novel capabilities of GeoTDM:\n\n1. We adopt GeoTDM to OC22 [1], a dataset of novel large-scale catalytic MD systems. The detailed experimental setup and results have been presented in Appendix C.1. The results demonstrate the potential of leveraging GeoTDM to simulate the dynamics of catalyst system, which bears significance in designing novel catalyst systems.\n2. We also demonstrate the capability of GeoTDM in performing temporal interpolation of trajectories and trajectory optimization in Sec. 5.3. GeoTDM is capable of handling these special tasks since it enjoys the benefits of being a controllable diffusion model that captures the joint distribution of the entire trajectory, while the baselines in existing works are not. These experiments have shown the promise of employing GeoTDM to more applications such as designing chemical reaction path given the initial and target system state, which is an interesting future direction.\n3. We propose an approach to generate long trajectory using GeoTDM trained on shorter trajectories through model composition in Appendix C.2. The demonstration reveals the capability of our model to produce long and stable MD trajectories, which also highlights a novel application scenario of GeoTDM.\n\n[1] Tran et al. The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts.\n\n**We sincerely hope our response could address your concerns!**", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kHlSpMAsuc", "reviewer": "Reviewer_CUJA", "summary": "In this paper, the authors propose geometric trajectory diffusion models (GeoTDM) to model temporal distribution of geometric trajectories\nwhile keeping the desirable physical symmetry of the trajectories.", "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\nBoth algorithms use the same E(n) equivariant graph neural networks. \n\nMissing reference: \nFang 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. \n\nThe paper needs more experimental results to support it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and suggestions! We provide point-to-point response below.\n\n> **[W1] 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. Missing reference: [1]**\n\nThank you for bringing up the related work DiffMD [1]. We will definitely cite and discuss the work [1] in the revised version.\n\nHowever, our work is substantially different from DiffMD in the following aspects.\n\n1. DiffMD only models the distribution of the next frame $\\mathbf{x}^{(t+1)}$ given the current frame $\\mathbf{x}^{(t)}$, while our GeoTDM directly models the joint distribution of an entire trajectory $\\mathbf{x}^{[T]}$ with a collection of multiple frames. Therefore, DiffMD is not explicitly a generative model on geometric trajectories, while to the best of our knowledge our GeoTDM is the first. Given such nature, DiffMD is vulnerable to error accumulation when generating trajectory with multiple frames at inference time, while our GeoTDM is able to jointly model the correlations of multiple frames within a trajectory and thus incurs much smaller error. We also empirically verify this argument in the additional experiment provided below.\n2. DiffMD relies on taking as input a given frame, while GeoTDM can both perform unconditional generation as well as conditional generation. It is unclear how to generate a geometric trajectory from scratch using DiffMD, for example, in the unconditional generation settings in our paper (Table 4 and 5).\n3. DiffMD is specifically designed for modeling molecular dynamics which heavily leverages particular geometric features like bond angles and dihedral angles in molecules. By contrast, our GeoTDM, without relying on domain-specific features, has been demonstrated to perform promisingly across a wide suite of benchmarks including physical simulation, molecular dynamics, and even pedestrian trajectory forecasting.\n\nWe also add additional experiments perform a thorough comparison. Since there is no public code released for DiffMD, we implement a variant of DiffMD (denoted DiffMD*) which replaces the task-specific backbone proposed in DiffMD by EGNN. We adopt the original training loss and sampling procedure proposed in DiffMD. We also set the same number of layers and hidden dimension as our GeoTDM, ensuring a fair comparison. Since DiffMD is not directly applicable to unconditional generation, we benchmark it in the conditional generation setting on Nbody and MD17 datasets. The results are presented below.\n\n||Particle|Spring|Gravity|\n|-|-|-|-|\n|DiffMD*|0.170/0.382|0.0093/0.0252|0.298/0.724|\n|GeoTDM|0.110/0.258|0.0030/0.0079|0.256/0.613|\n\n||Asp|Ben|Eth|Mal|Nap|Sal|Tol|Ura|\n|-|-|-|-|-|-|-|-|-|\n|DiffMD*|0.152/0.301|0.025/0.051|0.193/0.406|0.209/0.477|0.090/0.142|0.131/0.306|0.129/0.235|0.161/0.260|\n|GeoTDM|0.107/0.193|0.023/0.039|0.115/0.209|0.107/0.176|0.064/0.087|0.083/0.120|0.093/0.121|0.074/0.099|\n\nIt is observed that GeoTDM, by modeling the joint distribution of all frames within the geometric trajectory, consistently yields lower error than DiffMD*, which only captures the distribution of one target frame and thus needs to perform iterative roll-out for the entire trajectory, incuring extra error accumulation. We will add these discussions to the paper.\n\n[1] Fang Wu, Stan Z. Li, DIFFMD: A Geometric Diffusion Model for Molecular Dynamics Simulations.\n\n\n> **[W2] The paper needs more experimental results to support it.**\n\nIn this paper, we have performed exhaustive experiments on tasks including physical simulation (3 datasets), molecular dynamics (8 datasets in MD17 as well as an additional OC22 dataset [1], see Appendix C.1), and pedestrian trajectory forecasting (5 datasets in ETH-UCY). For physical simulation and MD, we also benchmark on both unconditional and conditional generation scenarios. In Sec. 5.3 we conduct extensive ablation studies and some additional use cases including leveraging GeoTDM to perform trajectory interpolation and optimization. We also propose a sampling algorithm to sample much longer trajectories using our GeoTDM trained on shorter trajectories through model composition, and verifying its feasibility by experiment (see `aspirin_longtraj.gif` in the supplementary file). The above experiments not only cover a diverse range of benchmarking suites but also involve in-depth use cases of how to unleash the potential of GeoTDM for novel application scenarios and tasks. We are glad to include more if the reviewer has concrete suggestions on what extra experiments are needed.\n\n\n> **[Q1] 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.**\n\nThank you for raising this point. The symbol $\\otimes$ refers to the Kronecker product. For instance, $I_D\\otimes I_N=I_{DN}$ where $I_D$, $I_N$, and $I_{DN}$ are the indentity matices with shape $D\\times D$, $N\\times N$, and $DN\\times DN$, respectively.\n\nHere $\\mathbf{P}$ is the matrix representation of the linear function $P(\\cdot)$, i.e., $\\mathbf{P}\\mathbf{x}=P(\\mathbf{x})$, where the left hand side is matrix multiplication. We formally define the operator $\\mathbf{P}$ here in order to clearly show that it is a linear operation with rank $(TN-1)D$, which facilitates our subsequent analysis (more details in Appendix A.1). However, we agree that it is not necessary to define $\\mathbf{P}$ in the main text when $P(\\cdot)$ has already been introduced, and we will move it to Appendix A.1 instead.\n\n**We sincerely hope our response could address your concerns!**", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QN6SiXytdC", "reviewer": "Reviewer_RWWk", "summary": "A new diffusion-based generative model for modeling complex 3D geometric structures with time-evolving trajectories is proposed. By introducing the SE(3) equivariance property, temporal attention, and learnable geometric prior into the discretized diffusion model, the proposed model can achieve high performance in learning the distribution of geometric trajectories while preserving the symmetry of geometric systems. It can also be improved to attain conditional generation by using the equivariant cross-attention mechanism. Various types of experiments show high accuracy and performance of the proposed model in cases such as molecular dynamics and simulation.", "questions": "1. As shown in the experiment part, the speed and efficiency are still the main limitations of the proposed model. Since the transition kernel and the prior are modified with extra restrictions, will it be difficult to redesign the framework (both forward and backward) into the continuous diffusion ODE/SDE form? (which might be important for further acceleration)\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.", "limitations": "The authors adequately addressed and discussed the limitations about computation. It would be better if the paper could also mention the current gap between the proposed model and the industrial/academic demand in medical/physical fields.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "No major weakness. Please refer to the part of Questions for my concerns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and suggestions! We provide point-to-point response below.\n\n> **[Q1] 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\nThank you for the insightful comment! Interestingly, our framework can be naturally extended into continuous-time formulation with a simple reparameterization for the learnable prior. In detail, denote the learnable prior as $\\mathbf{x}\\_r^{[T]}$, then with simple change-of-variables $\\mathbf{x}\\_\\tau'^{[T]}=\\mathbf{x}\\_\\tau^{[T]}-\\mathbf{x}\\_r^{[T]}$ where $\\mathbf{x}\\_\\tau^{[T]}$ are the original latent variables, the resulting trajectory $\\mathbf{x}\\_\\tau'^{[T]}$ effectively takes the same form as DDPM (see proof of Proposition A.2 in appendix). This finding enables us to readily transform the diffusion trajectory of $\\mathbf{x}\\_\\tau'^{[T]}$ into an SDE/ODE (see Song et al. [1]), while we only need to add $\\mathbf{x}\\_r^{[T]}$ back after sampling from the transformed SDE/ODE.\n\nThe convenient transformation from GeoTDM to its ODE formulation enables us to leverage fast samplers such as DDIM [2]. We provide extra experiments on Nbody dataset by performing DDIM sampling using the method described above. The results with different sampling steps are depicted below.\n\n||ADE|FDE|\n|-|-|-|\n|DDPM-1000|0.110|0.258|\n|DDIM-100|0.118|0.271|\n|DDIM-50|0.127|0.289|\n\nInterestingly, the performance degradation is marginal even when using only 50 sampling steps with DDIM sampler. GeoTDM with 50 sampling steps still outperforms the best baseline EqMotion which has an ADE/FDE of 0.141/0.310. This underscores the potential of streamlining GeoTDM for faster inference.\n\n[1] Song et al. Score-Based Generative Modeling through Stochastic Differential Equations. In ICLR'21.\n\n[2] Song et al. Denoising Diffusion Implicit Models. In ICLR'21.\n\n\n> **[Q2] 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\nExcept for processing temporal information, our model has several advantages compared with EDM [1]:\n\nEDM is designed for molecule generation from scratch and does not directly support conditioning. GeoTDM can perform unconditional generation and conditional generation with the help of the proposed equivariant cross-attention that enables conditioning on a given trajectory.\n\nGeoTDM also employs a flexible learnable prior parameterized by a lightweight neural network. The prior is carefully designed such that it subsumes existing parameterizations of the prior, including the center-of-mass (CoM) based prior used in EDM (see Theorem A.4 in Appendix A.3). We have shown in Sec. 5.3 that our prior yields superior performance over CoM-based prior used in EDM.\n\nWe have also systematically compared GeoTDM with other related works in Table 13 of Appendix C.6.\n\n[1] Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\" International conference on machine learning. PMLR, 2022.\n\n> **[Q3.1] Why the geometric restrictions are also important in N-body simulation?**\n\nIn N-body simulation the particles are governed by physical laws such as Columb law, Hooke's law, and Newton's law of gravitation. These physical laws ubiquitously abide by physical symmetry, which is mathematically revealed by equivariance. That is, when the system is rotated/translated, the trajectories driven by these physical laws will rotate/translate in the same way, hence motivating our geometric restrictions of GeoTDM to be SE(3)-equivariant.\n\n> **[Q3.2] What is the 'geometric structure' in these scenarios?** \n\n* In physical simulation, the geometric structure is instantiated as a fully-connected geometric graph since each pair of the particles preserves some interaction (e.g., Columb force for Charged Particles dataset).\n* For molecular dynamics, the geometric structure is defined as the molecular graph, with the node indicates the atoms and the edges corresponds to the chemical bonds.\n* For pedestrain trajectory, the geometric structure is a radius graph with node being each pedestrian and edges are connected if two pedestrians are within a certain distance.\n\n> **[Q3.3] 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\nThank you for the advice! We will add one section in the appendix to introduce backgrounds on geometric trajectory, physics simulation, and MD.\n\n\n> **[Q4] 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\nIn computational chemistry, MD sampling is usually performed over metastable molecular structures instead of highly unstable states [1], which accounts for why the changes in the visualization seem to be relatively small. However, in our MD17 experiments each trajectory has 30 frames, down-sampled from the raw data by a factor of 10, corresponding to a length of 300 in raw data which is sufficiently long to capture important chemical prior such as the rotation of the methyl group in aspirin (see Fig. 5).\n\nMoreover, in Appendix C.2 we have provided a sampling algorithm to sample much longer trajectories using our GeoTDM trained on shorter trajectories. A demonstration for a long trajectory (120 frames) obtained from our model (trained on 30 frames) can be found in `aspirin_longtraj.gif` in the supplementary file. It demonstrates our GeoTDM can produce long and stable MD trajectories, which highlights its practical significance.\n\n[1] Durrant et al. Molecular dynamics simulations and drug discovery. BMC biology. 2011.\n\n**We sincerely hope our response could address your concerns!**", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

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

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