Equivariant Graph Neural Operator for Modeling 3D Dynamics
Minkai Xu Jiaqi Han Aaron Lou Jean Kossaifi Arvind Ramanathan Kamyar Azizzadenesheli Jure Leskovec Stefano Ermon Anima Anandkumar
Modeling the complex three-dimensional (3D) dynamics of relational systems is
an important problem in the natural sciences, with applications ranging from
molecular simulations to particle mechanics. Machine learning methods have
achieved good success by learning graph neural networks to model spatial
interactions. However, these approaches do not faithfully capture temporal
correlations since they only model next-step predictions. In this work, we
propose Equivariant Graph Neural Operator (EGNO), a novel and principled method
that directly models dynamics as trajectories instead of just next-step
prediction. Different from existing methods, EGNO explicitly learns the
temporal evolution of 3D dynamics where we formulate the dynamics as a function
over time and learn neural operators to approximate it. To capture the temporal
correlations while keeping the intrinsic SE(3)-equivariance, we develop
equivariant temporal convolutions parameterized in the Fourier space and build
EGNO by stacking the Fourier layers over equivariant networks. EGNO is the
first operator learning framework that is capable of modeling solution dynamics
functions over time while retaining 3D equivariance. Comprehensive experiments
in multiple domains, including particle simulations, human motion capture, and
molecular dynamics, demonstrate the significantly superior performance of EGNO
against existing methods, thanks to the equivariant temporal modeling.