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

TL;DR

This paper presents the Spacetime $E(n)$-Transformer, a novel equivariant attention model for spatio-temporal graphs that incorporates physical symmetries to improve modeling of dynamical systems.

Contribution

The paper introduces an $E(n)$-equivariant Transformer architecture that enforces physical symmetries in spatio-temporal graph modeling, outperforming non-equivariant models.

Findings

SET outperforms non-equivariant models on the charged $N$-body problem.

Leveraging symmetries improves dynamical system modeling.

Equivariance leads to better generalization in physical systems.

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.