STFlow: Data-Coupled Flow Matching for Geometric Trajectory Simulation

TL;DR

STFlow is a novel generative model that leverages graph neural networks and data-dependent couplings to efficiently simulate complex trajectories in dynamical systems, outperforming existing methods in accuracy and scalability.

Contribution

Introduces STFlow, a data-coupled flow matching model that simplifies trajectory simulation by integrating informed priors, enhancing efficiency and accuracy across various dynamical systems.

Findings

Achieves lowest prediction errors on benchmarks

Requires fewer simulation steps

Demonstrates improved scalability

Abstract

Simulating trajectories of dynamical systems is a fundamental problem 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 in deep generative modeling and geometric deep learning enable probabilistic simulation by learning complex trajectory distributions while respecting intrinsic permutation and time-shift symmetries. However, trajectories of N-body systems are commonly characterized by high sensitivity to perturbations leading to bifurcations, as well as multi-scale temporal and spatial correlations. To address these challenges, we introduce STFlow (Spatio-Temporal Flow), a generative model based on graph neural networks and hierarchical convolutions. By incorporating data-dependent couplings within the Flow Matching framework, STFlow denoises starting from conditioned random-walks instead of Gaussian noise. This novel informed prior simplifies the learning task by reducing transport cost, increasing training and inference efficiency. We validate our approach on N-body systems, molecular dynamics, and human trajectory forecasting. Across these benchmarks, STFlow achieves the lowest prediction errors with fewer simulation steps and improved scalability.