Flow Matching for Geometric Trajectory Simulation

TL;DR

This work tackles probabilistic geometric trajectory simulation for N-body systems, where both temporal dependencies and particle interactions must be modeled under permutation symmetry. It introduces STFlow, a flow-matching framework that maps from a physics-informed prior to data distributions using data-dependent couplings and a learnable vector field $v_\theta$, enabling parallelizable and accurate trajectory generation. Key contributions include a random-walk-based physics-informed prior, an architecture combining Equivariant Graph Convolutional Layers and UNet-based temporal conditioning, and comprehensive evaluations across N-body, MD17, and pedestrian datasets showing reduced prediction errors and improved sampling efficiency. The approach demonstrates the value of leveraging domain knowledge in priors for probabilistic geometric trajectory modeling, with practical implications for scalable and accurate physics-informed simulation across multiple domains.

Abstract

The simulation of N-body systems is a fundamental problem with applications in a wide range of fields, such as molecular dynamics, biochemistry, and pedestrian dynamics. Machine learning has become an invaluable tool for scaling physics-based simulators and developing models directly from experimental data. In particular, recent advances based on deep generative modeling and geometric deep learning have enabled probabilistic simulation by modeling complex distributions over trajectories while respecting the permutation symmetry that is fundamental to N-body systems. However, to generate realistic trajectories, existing methods must learn complex transformations starting from uninformed noise and do not allow for the exploitation of domain-informed priors. In this work, we propose STFlow to address this limitation. By leveraging flow matching and data-dependent couplings, STFlow facilitates physics-informed simulation of geometric trajectories without sacrificing model expressivity or scalability. Our evaluation on N-body dynamical systems, molecular dynamics, and pedestrian dynamics benchmarks shows that STFlow produces significantly lower prediction errors while enabling more efficient inference, highlighting the benefits of employing physics-informed prior distributions in probabilistic geometric trajectory modeling.

Figures (17)

• Figure 1: STFlow generates trajectories by sampling from a physics-informed prior and learning a vector field (green arrows).
• Figure 2: Overview of STFlow. Given trajectories $\mathbf{x}_1 \sim p_1$, we construct a noisy prior $\mathbf{x}_0$ informed by the observed initial conditions and predict a vector field $v_\theta$ using repeating layers of spatial message passing and temporal convolution.
• Figure 3: Inference results on N-Body Gravity (left) and Charged (right). The black dots represent the 10 conditioning steps, the colored dots are the 20 generated steps.
• Figure 4: Velocity and acceleration density estimates from the 20 predicted frames of the MD17 Ethanol test set ($N=216320)$.
• Figure 5: Performance and runtime evaluation of fixed timestep ODE solvers on 1000 samples from the Gravity test set. The numbers inside represent the runtime per batch of size 100 during inference.