STFlow: Data-Coupled Flow Matching for Geometric Trajectory Simulation

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.

Many important systems in science and society, from molecules to crowds of people, evolve over time in ways that are hard to predict. Simulating their future behavior is crucial for scientific discovery and decision-making, but existing machine learning approaches are often slow, computationally expensive, or fail to capture the full range of possible future outcomes. We introduce STFlow, a generative model that learns to simulate the future trajectories of interacting objects directly from data. The key insight is that when we already observe how a system has been behaving, we should use that information to make an educated guess about where it is headed rather than starting from pure noise, as most existing methods do. STFlow builds this intuition into its design by constructing an informed starting point for the simulation that reflects the system's recent motion, making the learning task simpler and inference faster. Across three distinct domains: physical particle systems, molecular dynamics, and basketball player movement, STFlow achieves lower prediction errors than competing approaches while requiring far fewer computational steps. This makes accurate probabilistic simulation more practical for real-world scientific applications, from drug discovery to modeling human mobility.