Geometric Trajectory Diffusion Models

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.

Paper

References (71)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer RWWk7/10 · confidence 4/52024-07-05

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.

Strengths

1. The introduction and methods are well-written and clear. 2. 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. 3. The paper explained and analyzed the case of conditional generation in detail. 4. The experimental setting and results are demonstrated in very extensive and detailed ways. 5. Clear ablation studies.

Weaknesses

No major weakness. Please refer to the part of Questions for my concerns.

Questions

1. As shown in the experiment part, the speed and efficiency are still the main limitations of the proposed model. Since the transition kernel and the prior are modified with extra restrictions, will it be difficult to redesign the framework (both forward and backward) into the continuous diffusion ODE/SDE form? (which might be important for further acceleration) 2. Except for the part of processing temporal information, could you compare more details about the proposed EGTN/EGCL with the E(3) Equivariant Diffusion Model in the paper [1]? 3. Why the geometric restrictions are also important in N-body simulation? What is the 'geometric structure' in these scenarios? It would be better if the author could provide more background knowledge about geometric trajectory and the physics simulation/molecular dynamics in the appendix. 4. Based on the visualization results in the appendix, the change in the molecular structure seems to be 'small'. Will such short trajectories be practically useful/meaningful in application? [1]: Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

Reviewer CUJA6/10 · confidence 2/52024-07-09

Summary

In this paper, the authors propose geometric trajectory diffusion models (GeoTDM) to model temporal distribution of geometric trajectories while keeping the desirable physical symmetry of the trajectories.

Strengths

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

Weaknesses

I am concerned that dynamic system and rotation translational invariant has already been studied in the previous research (listed below, but not quoted here), not like the authors claimed that this is the first study on this topic. There might be some overlaid here. Both algorithms use the same E(n) equivariant graph neural networks. Missing reference: Fang Wu, Stan Z. Li, DIFFMD: A Geometric Diffusion Model for Molecular Dynamics Simulations. The work in this previous paper requires that the dynamics must be invariant to rotation or translation. The paper needs more experimental results to support it.

Questions

P5, line 185: In the definition of the operator P, ⊗ is not defined? It seems to me no need to define operator P if P(x) is defined.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes.

Authorsrebuttal2024-08-09

Thank you for the feedback

Dear Reviewer CUJA, We greatly appreciate your advice and promise to incorporate the discussions in the final version. In particular, we also summarize the core distinctions between our GeoTDM and DiffMD [1] as follows: | | Distribution Modeled | Unconditional | Prior | Backbone | |-|-|-|-|-| | DiffMD [1] | Frame | $\times$ | Fixed Pointwise Prior | Specialized for molecules | | GeoTDM | Trajectory | $\checkmark$ | Learned Flexible Prior | General for $n$-D geometric data | Overall, our experiments have verified the advantages of our approach over DiffMD in terms of empirical performance. We highly appreciate it if you could kindly consider adjusting your evaluation if we have addressed your concerns. Thank you! Best, Authors [1] Fang Wu, Stan Z. Li, DIFFMD: A Geometric Diffusion Model for Molecular Dynamics Simulations.

Reviewer CUJA2024-08-10

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

Authorsrebuttal2024-08-12

Follow-up on the clarifications

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

Authorsrebuttal2024-08-10

Further Clarifications

Dear Reviewer CUJA, Thank you very much for the timely feedback and the follow-up question! We are glad to further clarify your concern. > What troubled me is that I am still not convinced what you said that: "DiffMD is not explicitly a generative model on geometric trajectories, while to the best of our knowledge our GeoTDM is the first. ". I believe that both of your methods are a generative model, but each of you approached in some detailed difference, so I need you to clarify this in more details. Yes, we agree that both DiffMD and GeoTDM are generative models. The core difference lies in whether the method is modeling the distribution of **geometric trajectories** $\mathbf{x}^{[T]}=[\mathbf{x}^{(0)},\mathbf{x}^{(1)},\cdots,\mathbf{x}^{(T-1)}]\in\mathbb{R}^{T\times N\times D}$ which is a **sequence/collection of frames**, or a **single** frame $\mathbf{x}^{(t)}\in\mathbb{R}^{N\times D}$. In particular, our GeoTDM explicitly models the distribution of geometric trajectories $p(\mathbf{x}^{[T]}|\cdot)$ (see line 116-122), while DiffMD is designed to model the distribution of the next frame $p(\mathbf{x}^{(t+1)}|\cdot)$ (see Model Overview section on Page 3 of DiffMD [1]), where $\cdot$ denotes certain conditioning if applicable (e.g., the previous frames). Such difference leads to several distinctions in model design, inference, and empirical performance, as we will discuss below. * **Model design.** In order to jointly model $p(\mathbf{x}^{[T]}|\cdot)$, in the transition kernel we need to additionally handle the temporal dimension with size $T$. This motivates us to design the equivariant temporal attention layer (Eq. 5-7), which is absent in DiffMD since DiffMD does not require to explicitly model temporal correlation. Our framework also enables conditioning on a trajectory with multiple frames while DiffMD is based on Markovian assumption and always consumes the previous one single frame to generate the next frame. * **Inference.** The benefit of directly modeling the distribution of geometric trajectories over single frames also presents at inference time. Within one diffusion loop, our GeoTDM can generate an entire trajectory with $T$ frames, while DiffMD requires an additional outer loop that sweeps through $T$ frames in order to achieve the same effect. * **Performance.** We also demonstrate the benefit in terms of empirical performance. Since DiffMD does not consider the correspondence between multiple frames, it is more vulnerable to error accumulation when generating a long trajectory through iterative rolling out. Our extra experiment provided in the rebuttal verifies this point, with GeoTDM outperforming DiffMD by a remarkable margin. > While you claimed that your "modeling the joint distribution of all frames within the geometric trajectory", what I saw is you still look at the individual conditional probability like DiffMD, see line 193. It seems like you take into data sampling of time into consideration, to me, I think this is considered as generalize of DiffMD. Thank you for the question. However, there might be a misunderstanding here and we would like to respectfully clarify. Specifically, line 193 depicts the transition kernel $p\_\theta({\mathbf{x}\_{\tau-1}^{[T]}}|\mathbf{x}\_{\tau}^{[T]})$. Here $\tau$ refers to the **diffusion step** instead of the **frame index on the geometric trajectory**. The superscript $[T]$ indicates that the latent variable here is a geometric trajectory instead of a single frame. Therefore, the individual conditional probability is enforced on **diffusion step**, which is due to the Markovian assumption of the diffusion process. This is a common practice for diffusion models and is shared in both GeoTDM and DiffMD, and we are not claiming any difference on this point. However, we are **not** enforcing conditional independence on the actual temporal dimension of the geometric trajectory, while DiffMD has a Markovian assumption on molecular dynamics and only models the distribution of a single frame. This is what we refer to by stating "modeling the joint distribution of all frames within the geometric trajectory". This point leads to the core difference between these two approaches, as we have discussed. > For the above reason, I believe that is why you both used Equivariant Graph Convolution Layer. We are both using EGCL to process spatial information on the geometric structure. However, since we are modeling the whole trajectory, we need to additionally introduce the equivariant temporal attention layer to process the temporal correspondence, which never presents in DiffMD. Thank you again for the follow-up discussion! We promise to include the discussion and the distinctions between these two methods in the manuscript. Please let us know if this addresses your concern and we are happy to further clarify if you have any questions. Best, Authors

Reviewer CUJA2024-08-12

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

Authorsrebuttal2024-08-12

Thank you for the supportive feedback!

Dear Reviewer CUJA, Thank you for the supportive feedback and the insightful comments that help us improve the manuscript! We will include these in the final version. Best, Authors

Reviewer UDq75/10 · confidence 4/52024-07-11

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.

Strengths

* The method is a timely extension and synthesis of geometric diffusion models and trajectory diffusion models. * As demonstrated in the experiments, the method has broad applicability across many machine learning domains. * The experiments are thorough. The ablation studies validating the somewhat unconventional choices of conditional prior are appreciated. * 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. * The exposition of the method drags at times, re-proving well-known statements about equivariance in neural networks and geometric diffusion models. * The experiments, while diverse and broad, also suffer from being scattered and not necessarily the most convincing individually. Little effort is spent on exploring the ways in which a trajectory diffusion could be employed and the new problems that could be solved with the mode. To me, the novel capabilities of such a method are more exciting than marginal gains across a number of established datasets, executed checkbox-style.

Questions

See above

Rating

5

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

Yes

Reviewer e5nM6/10 · confidence 1/52024-07-12

Summary

The paper proposed the first diffusion model for modeling the temporal distribution of 3D geometric trajectories, while previous works only operate on static structures. It demonstrates the equivariant temporal kernels can lead to density with desired symmetry and develop a novel transition kernel leveraging SE(3)-equivariant spatial convolution and temporal attention. The experiments demonstrates it can generates realistic geometric trajectories with significantly higher quality.

Strengths

1. The geometric trajectory diffusion model is novel to me. 2. The EGTN that operates on geometric trajectories permits conditioning upon a given trajectory using equivariant cross-attention 3, The experiments demonstrate that GeoTDM achieves SOTA performance on both unconditional and conditional trajectory generation tasks. 4. The work can be extended to applications like temporal interpolation and trajectory optimization

Weaknesses

None

Questions

None

Rating

6

Confidence

1

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Reviewer UDq72024-08-09

I appreciate the detailed author response. However it has not moved the needle substantively on my concerns and I will keep the score.

Authorsrebuttal2024-08-09

Thank you for the feedback

Dear Reviewer UDq7, Thank you for the timely reply. To help substantively address your concerns, we summarize and further illustrate our response below. For **W1**, we added extra experiments that explore the way to combine our proposed equivariant temporal layer with more sophisticated equivariant architectures, e.g., Equiformer, that helps further boost the performance on specific tasks like MD. We will also include these discussions in the paper. For **W2**, we clarified the difference between theorems presented in this paper and previous work. Besides, we provided necessary derivations in Appendix A where we have interestingly found that our GeoTDM also preserves the simplified loss (Eq. 12) which is yet not straightforward. Notably, in Theorem A.4 we also theoretically justify our learnable prior subsumes existing parameterizations. All of these constitute our theoretical contributions. For **W3**, We have also presented various explorations on approaching new problem with GeoTDM, including modeling large scale catalyst dynamics (Appendix C.1), interpolating trajectories given initial and target states (Sec. 5.3), optimizing the input trajectory towards the learned distribution (Sec. 5.3), and performing generation through longer time horizon (Appendix C.2). Please let us know if you have any concerns in particular or any advice that could help further improve the paper and we are more than willing to address them. Thank you! Best, Authors

Reviewer CUJA2024-08-09

Thank you for your detailed response! After I read all your rebuttals and other reviewers comments. I decided to stay at the same rating.

Reviewer RWWk2024-08-09

Thank you very much for your detailed response! Most of my concerns have been addressed. Also combined with other reviewers' comments and rebuttals, I decided to raise the score.

Authorsrebuttal2024-08-09

Thank you for the feedback!

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

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC