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.