Flow for Future: Geometric SE(3)-Equivariant Flow Matching for 3D Trajectory Prediction
Abstract
Predicting 3D geometric trajectory requires capturing complex spatiotemporal dependencies while preserving physical symmetries. While flow matching offers a powerful generative paradigm, extending it to SE(3)-equivariant dynamics is challenging due to the inherent gap between deterministic history and stochastic evolving flows. To address this, we introduce GSE-Flow, an SE(3)-equivariant flow matching framework. We first propose a Coherent Sequence Encoding and Time-Modulated Embedding strategy that unifies historical and evolving streams, incorporating velocity and flow time via equivariant affine transformations to guide continuous evolution. We further design a Geometry-Feature Tensorization mechanism that projects node states into a tensor product space, enabling Context-Flow Fusion to guide trajectory evolution with historical context. GSE-Flow guarantees theoretical SE(3)-equivariance and achieves SOTA accuracy on MD17, MD22, and CMU MoCap benchmarks for geometric trajectory prediction, while demonstrating generality by enhancing deterministic baselines. Code is available at https://github.com/aegine/GSE-Flow.
Lay Summary
Our research focuses on predicting how objects—ranging from microscopic molecules to human bodies—move over time in 3D space. A major challenge is ensuring that artificial intelligence models not only capture complex movement patterns but also strictly obey fundamental physical laws, such as spatial symmetry. For instance, if you view a moving object from a different angle, the model's prediction must perfectly adapt to that new perspective. To address this, we developed GSE-Flow, a new framework that models 3D motion as a continuous, fluid process rather than a series of disconnected jumps. By seamlessly linking an object's past trajectory with its future path, our model maintains smooth and physically accurate predictions. When tested on complex datasets of molecular dynamics and human motion, our approach significantly outperformed existing methods.