Spacetime E(n)-Transformer: Equivariant Attention for Spatio-temporal Graphs

2024/08/12 by Charles, Sergio G.
#Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2408.06039

Abstract

We introduce an E(n)-equivariant Transformer architecture for spatio-temporal graph data. By imposing rotation, translation, and permutation equivariance inductive biases in both space and time, we show that the Spacetime E(n)-Transformer (SET) outperforms purely spatial and temporal models without symmetry-preserving properties. We benchmark SET against said models on the charged N-body problem, a simple physical system with complex dynamics. While existing spatio-temporal graph neural networks focus on sequential modeling, we empirically demonstrate that leveraging underlying domain symmetries yields considerable improvements for modeling dynamical systems on graphs.

Related

• SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
• Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics
• Spatial Attention Kinetic Networks with E(n)-Equivariance
• E(n) Equivariant Graph Neural Networks
• Learn one size to infer all: Exploiting translational symmetries in delay-dynamical and spatio-temporal systems using scalable neural networks