SE(3)-Equivariant Flow Matching with Gaussian Process Priors for Geometric Trajectory Prediction

The trajectory prediction of N-body systems is of great significance and remains challenging with broad applications across various fields such as physics, chemistry and biology. Recent advances in generative models including flow matching and diffusion models have emerged as effective solutions to this problem, owing to their capacity to model the stochasticity and underlying distributions of complex system trajectories. However, existing approaches typically adopt trivial prior distributions that neglect the temporal correlations and spatial symmetries of N-body trajectories, which not only complicates the generation process but also limits model performance. To address these limitations, we propose GP-EquiFlow, an SE(3)-equivariant flow matching model incorporating vector-valued Gaussian processes. Based on observed trajectories, we employ vector-valued Gaussian processes to construct SE(3)-equivariant prior distributions, which exhibit enhanced consistency with the target data distribution in both spatial and temporal dynamics. Extensive experiments on N-body simulations and molecular dynamics demonstrate that the proposed GP-EquiFlow delivers more accurate predictions while requiring fewer sampling steps, underscoring the effectiveness of integrating Gaussian process-based SE(3)-equivariant prior distributions in geometric trajectory prediction.

Predicting trajectories of interacting particle systems is crucial across physics, chemistry, and biology but remains highly challenging. While recent generative models have shown promising results in addressing this problem, they typically rely on simple initial prior distributions that neglect the temporal continuity and spatial symmetries of system dynamics. This not only complicates the downstream prediction process but also degrades accuracy. Thus, we present GP-EquiFlow, which learns motion patterns from observed data to build physically consistent initial priors. It achieves higher prediction accuracy with significantly fewer computational steps, resulting in much faster inference.