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

TL;DR

The paper tackles preserving geometric symmetries in spatio-temporal graph modeling for dynamical systems, specifically the charged $N$-body problem. It introduces the Spacetime $E(n)$-Transformer (SET), a framework that combines $E(n)$-equivariant spatial attention (via EGCL) with $E(n)$-equivariant temporal attention (ETAL) in a spatio-temporal graph Transformer to capture long-range dependencies while respecting symmetry. Empirically, SET outperforms purely spatial or purely temporal baselines across varying system sizes and remains robust under noise and larger $N$, underscoring the value of symmetry-informed inductive biases for dynamical graph modeling. The work points to future directions such as learning symmetries from data (LieConv) and linking to conserved quantities via Noether's theorem, with potential impact on complex physical simulations and biomolecular dynamics.

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.

Figures (2)

• Figure 1: Spacetime $E(n)$-Transformer architecture.
• Figure 2: Top: Model test MSE versus $N$. Bottom: Number of model parameters versus $N$.

Theorems & Definitions (9)

• Definition 2.1
• Definition 2.2
• Definition 2.3
• Definition 2.4
• Definition A.1
• Definition A.2
• Definition A.3
• proof
• proof