Geometric Trajectory Diffusion Models

TL;DR

This work proposes geometric trajectory diffusion models (GeoTDM), the first diffusion model for modeling the temporal distribution of 3D geometric trajectories, and introduces a generalized learnable geometric prior into the forward diffusion process to enhance temporal conditioning.

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.

Figures (10)

• Figure 1: Overview of GeoTDM. The forward diffusion $q$ gradually perturbs the input while the reverse process $p_{\bm\theta}$, parameterized by EGTN, denoises samples from the prior. The condition ${\mathbf{x}}_c^{[T_c]}$, if available, is leveraged to construct the equivariant prior and as a conditioning signal in EGTN.
• Figure 2: (a) Unconditional generation samples on MD17. GeoTDM generates MD trajectories with much higher quality (see more in App. \ref{['sec:more_vis']}). (b) Interpolation. Left: the given initial and final 5 frames. Right: GeoTDM interpolation and GT. (c) Optimization by GeoTDM on predictions of EGNN. Dis(Opt, GT)/Dis(Opt, EGNN) is the distance between optimized trajectories and GT/EGNN.
• Figure 3: An illustration of different equivariant priors. For simplicity in the chart here we only illustrate the case when $N=3$ and $T_c=1$, $T=1$.
• 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 ${\mathbf{x}}_c^{[T_c]}$ and ${\mathbf{h}}_c^{[T_c]}$ can also be processed using cross-attention. The relative temporal embedding $\psi(t-s)$ is added to the key and value. DotProd refers to dot product and Softmax is performed over indexes of $s$.
• Figure 5: Uncurated samples of GeoTDM on MD17 dataset in the unconditional generation setup. From top-left to 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. It 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.

Theorems & Definitions (14)

• Theorem 4.1: $\text{SE}(3)$-equivariance
• Theorem 4.2
• Theorem 4.3
• Proposition 4.4
• Proposition A.1
• proof
• Proposition A.2
• proof
• Proposition A.3
• proof