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

## Abstract

Abstract We introduce an E ⁢ ( n ) 𝐸 𝑛 E(n) italic_E ( italic_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 ) 𝐸 𝑛 E(n) italic_E ( italic_n ) -Transformer (SET) outperforms purely spatial and temporal models without symmetry-preserving properties. We benchmark SET against said models on the charged N 𝑁 N italic_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.

## 1 Introduction

Many problems that we wish to model with neural networks possess underlying geometric structure with symmetries. Geometric Deep Learning, a term coined in the seminal work of (Bronstein et al. ), is an Erlangen program for deep learning that systematizes inductive biases as group symmetries $G$, arising through notions of invariance and equivariance.

Recent work, like $SE(3)$-Transformers (Fuchs et al. ) and $E(n)$-Graph Neural Networks (Satorras, Hoogeboom, and Welling ), impose different notions of group equivariance on neural networks to inform architecture choice. While these neural network architectures encode spatial inductive biases, they notably lack a time component. Temporal Graph Networks (Rossi et al. ) proposed an efficient framework that learns from dynamic graphs. However, this architecture assumes the topology of graphs changes over time. In this paper, we discuss spatio-temporal graphs, which have a fixed topology with changing features over discrete time steps. Recent works (Jin et al. ), (Marisca, Cini, and Alippi ), (Cini et al. ) have treated node features as time series and edges as the relationships between these series. Such spatio-temporal graph neural networks (STGNNs) have a plethora of applications, from simulating biomolecular interactions to modeling financial time series.

Similarly, while STGNNs improve representation learning of sequential graph data, to the best of our knowledge, minimal research has been done on preserving group symmetries in a spatio-temporal fashion. In particular, sequential models ought to preserve spatial group symmetries at each time step. Famously, Noether’s first theorem formalizes the notion of infinitesimal symmetries of the so-called Lagrangian of a physical system, in terms of perturbations with respect to both space and time, by determining conserved quantities. Inspired by this intuition of temporal and spatial symmetry, we seek to derive a neural network architecture that is equivariant in both temporal and spatial components.

Classical neural network architectures like RNNs are discrete approximations to continuous time-domain signals, obeying a differential equation with respect to time. If an RNN is invariant to time-warping, a monotonically increasing and differentiable function of time, it takes the form of an LSTM (Bronstein et al. ), which unlike a vanilla RNN, captures long-term dependencies. Similarly, the dynamics of classical physical systems satisfy the Euler-Lagrange equations, i.e. the equations of motion. Hence, we use the charged $N$-body problem, as described in (Satorras, Hoogeboom, and Welling ) and adapted from (Kipf et al. ), as an ideal candidate to test our hypothesis that preserving $G$-equivariance ameloriates long-term spatio-temporal graph modeling. We will use a Transformer for the temporal component of the architecture, preserving long-term dependencies and, hence, invariance to time-warping. Each node, representing a charged particle of the graph, will have features, coordinates, and velocities. As such, the neural network should be equivariant under rotational and translational symmetries $E(n)$ acting on coordinates. It should also be equivariant with respect to rotational symmetries $SO(n)$ acting on velocities. Lastly, the nodes should be permutation equivariant.

## 2 Background

### 2.1 Geometric Deep Learning

Following the insights of Geometric Deep Learning (Bronstein et al. ), the input signals to machine learning models have an underlying domain $\Omega$. Examples of such domains include grids, graphs, and manifolds. The space of signals over $\Omega$ possesses a vector-space structure. That is (Bronstein et al. ):

###### Definition 2.1 .

The space of $\mathcal{C}$-valued signals on $\Omega$ is

$$ $\mathcal{X}(\Omega,\mathcal{C})=\{x:\Omega\to\mathcal{C}\},$ $$

which is a vector space of functions.

The symmetry of the domain $\Omega$ will impose structure on the signal space $\mathcal{X}(\Omega)$, thus inducing structure on the space of interpolants

$$ $\mathcal{F}(\mathcal{X}(\Omega))=\{f_{\theta\in\Theta}\}$ $$

for $f_{\theta}$ a neural network. In what follows, we canonically refer to $\mathcal{X}(\Omega)$ as $V$ for brevity.

###### Definition 2.1 .

### 2.2 Group Representations, Invariance, and Equivariance

###### Definition 2.2 .

A representation of a group $G$ on a vectorspace $V$ over a field $K$ is a homomorphism

$$ $\rho:G\to GL(K,V)$ $$

such that $\rho(gh)=\rho(g)\rho(h)$ for all $g\in G$, $h\in G$, where $GL(K,V)$ is the general linear group of automorphisms $\varphi:V\stackrel{{\scriptstyle\sim}}{{\smash{\longrightarrow}\rule{0.0pt}{1.
72218pt}}}V$, i.e., the set of bijective linear transformations with function composition as its binary operation.

In this paper, we are interested in the group of rotational symmetries $SO(n)$ and the group of isometries $E(n)$ of $\mathbb{R}^{n}$, as these are the naturally-induced symmetries of particles. Rotations are distance, angle, and orientation preserving transformations. The group of rotations in $n$ dimensions is

$$ $SO(n)=\{Q\in M_{n}(\mathbb{R})|Q^{\top}Q=I\text{ and }\det Q=+1\},$ $$

where $M_{n}(\mathbb{R})$ is the set of $n\times n$ matrices with entries in $\mathbb{R}$, forming a ring. We represent a group element $g\in SO(n)$ with $\rho(g)\in GL(\mathbb{R},\mathbb{R}^{n})$, acting on $\textbf{x}\in\mathbb{R}^{n}$ as $\rho(g):\textbf{x}\mapsto Q\textbf{x}$ where $Q\in\mathbb{R}^{n\times n}$ is an orthogonal matrix (see Appendix A for more details). We restrict the notion of equivariance to functions of Euclidean space, as will be the case for neural networks. In Appendix A, we provide a more general definition.

###### Definition 2.3 .

A function $f:\mathbb{R}^{n}\to\mathbb{R}^{n}$ is $SO(n)$-equivariant if

$$ $Qf(\mathbf{x})=f(Q\mathbf{x})$ $$

for all $Q\in\mathbb{R}^{n\times n}$ orthogonal and $\mathbf{x}\in\mathbb{R}^{n}$.

The Euclidean group $E(n)$ is the set of isometries of Eucliden space $\mathbb{R}^{n}$, i.e. transformations that preserve distance between points, represented as a rotation followed by a translation. More precisely, $E(n)=\{\varphi:\mathbb{R}^{n}\to\mathbb{R}^{n}|\varphi\text{ isometry}\}$. We represent a group element $g\in E(n)$ with $\rho(g)\in GL(\mathbb{R},\mathbb{R}^{n})$, acting on $\textbf{x}\in\mathbb{R}^{n}$ as $\rho(g):\textbf{x}\mapsto Q\textbf{x}+\mathbf{b}$ where $Q\in\mathbb{R}^{n\times n}$ is an orthogonal rotation matrix and $\mathbf{b}\in\mathbb{R}^{n}$ is a translation vector. That is, we represent actions of the Euclidean group as orthogonal transformations followed by translations.
Again, we provide a definition of equivariance with respect to functions of Euclidean space.

###### Definition 2.4 .

A function $f:\mathbb{R}^{n}\to\mathbb{R}^{n}$ is $E(n)$-equivariant if

$$ $Qf(\mathbf{x})+\mathbf{b}=f(Q\mathbf{x}+\mathbf{b})$ $$

for all $Q\in\mathbb{R}^{n\times n}$ orthogonal rotation matrices, $\mathbf{b}\in\mathbb{R}^{n}$ translation vectors, and for all $\mathbf{x}\in\mathbb{R}^{n}$.

###### Definition 2.2 .

###### Definition 2.3 .

###### Definition 2.4 .

## 3 Method

In this paper, we are interested in physical systems that can be modelled as a sequence of graphs $G_{t}=(\mathcal{V}_{t},\mathcal{E}_{t})$ for $t=1,\dots,L$ with nodes $v_{i}(t)\in\mathcal{V}_{t}$ and edges $e_{ij}(t)\in\mathcal{E}_{t}$. In particular, we seek to model the dynamics of the charged $N$-body problem (Kipf et al. ). For this task, we assume a priori that the graph is complete since a charged particle will interact with every other particle in a Van der Waals potential under Coulomb’s law. In addition, we assume that particles are neither created nor destroyed as the system evolves in time, so the nodes $\mathcal{V}_{t}$ in the graph remain the same. Let $\mathcal{G}:=(G_{t})_{1\leq t\leq L}$ be a sequence of topologically-identical graphs with changing features, known as a spatio-temporal graph. The task under consideration is learning a function that predicts the associated features of graph. In particular, given $\mathcal{G}$, we are interested in predicting the positions and velocities of all particles in the system after $H$ additional time steps where $H>>L$.

To equip the spatio-temporal model of particle interactions with the appropriate inductive biases, we leverage both spatial and temporal notions of attention. For node $i$ at time step $t$ to attend to all the past neighborhoods of that node, we need to (1) aggregate nodes spatially to obtain spatially-contextual embeddings and (2) obtain temporally-contextual embeddings via temporal aggregation.

We fix a time slice $t$ such that the features derive from $G_{t}=(\mathcal{V}_{t},\mathcal{E}_{t})$. From the current features $\textbf{h}_{i}^{(l)}(t)$ of node $i$ at layer $l$, we form the next layer features $\textbf{h}_{i}^{(l+1)}(t)$ by aggregating neighboring node features. In particular,

$$ $\textbf{h}_{i}^{(l+1)}(t)=\phi\left(\textbf{h}_{i}^{(l)}(t),\bigoplus_{j\in \mathcal{N}_{i}}a(\textbf{h}^{(l)}_{i}(t),\textbf{h}^{(l)}_{j}(t))\psi(\textbf {h}^{(l)}_{j}(t))\right),$ (1) $$

where $\oplus$ is a permutation-invariant function (Bronstein et al. ), and $a$ is a self-attention mechanism, often a normalized softmax across neighbors.

### 3.1 E ⁢ ( n ) 𝐸 𝑛 E(n) italic_E ( italic_n ) -Equivariant Spatial Attention

(Satorras, Hoogeboom, and Welling ) introduced $E(n)$-Equivariant Graph Neural Networks (EGNNs). Every node in the graph $G=(\mathcal{V},\mathcal{E})$ has features $\textbf{h}_{i}\in\mathbb{R}^{d}$ and coordinates $\textbf{x}_{i}\in\mathbb{R}^{n}$. In addition, we keep track of each particle’s velocity $\textbf{v}_{i}\in\mathbb{R}^{n}$. The Equivariant Graph Convolutional Layer (EGCL) takes the set of node embeddings $h^{(l)}=\left\{\textbf{h}_{1}^{(l)},\dots,\textbf{h}_{N}^{(l)}\right\}$, coordinate embeddings $x^{(l)}=\left\{\textbf{x}_{1}^{(l)},\dots,\textbf{x}_{N}^{(l)}\right\}$, velocity embeddings $v^{(l)}=\left\{\boldsymbol{\textbf{v}}_{1}^{(l)},\dots,\boldsymbol{\textbf{v}}
_{N}^{(l)}\right\}$, and edge information $\mathcal{E}=(e_{ij})$ as input and produces the embeddings of the next layer. That is, $h^{(l+1)},x^{(l+1)},v^{(l+1)}=\text{EGCL}[h^{(l)},x^{(l)},v^{(l)},\mathcal{E}]$, defined as follows (Satorras, Hoogeboom, and Welling ):

$$ $\begin{split}\textbf{m}_{ij}&=\phi_{e}\left(\textbf{h}_{i}^{(l)},\textbf{h}_{j }^{(l)},||\textbf{x}_{i}^{(l)}-\textbf{x}_{j}^{(l)}||^{2},a_{ij}\right)\\ \textbf{v}_{i}^{(l+1)}&=\phi_{v}(\textbf{h}_{i}^{(l)})\textbf{v}_{i}^{(l)}+C \sum_{j\neq i}(\textbf{x}_{i}^{(l)}-\textbf{x}_{j}^{(l)})\phi_{x}(\textbf{m}_{ ij})\\ \textbf{x}_{i}^{(l+1)}&=\textbf{x}_{i}^{(l)}+\textbf{v}_{i}^{(l+1)}\\ \textbf{m}_{i}&=\sum_{j\neq i}\textbf{m}_{ij}\\ \textbf{h}_{i}^{(l+1)}&=\phi_{h}(\textbf{h}_{i}^{(l)},\textbf{m}_{i})\end{split}$ (2) $$

where $a_{ij}$ are the edge attributes, e.g. the edge values $e_{ij}$, and $\phi_{e}:\mathbb{R}^{2d+2}\to\mathbb{\mathbb{R}}^{h}$, $\phi_{v}:\mathbb{R}^{d}\to\mathbb{R}$, $\phi_{x}:\mathbb{R}^{h}\to\mathbb{R}$, and $\phi_{h}:\mathbb{R}^{d+h}\to\mathbb{R}^{d^{\prime}}$ are MLPs. In what follows, we assume $d^{\prime}=d$ for clarity.

(Satorras, Hoogeboom, and Welling ) proved that this layer is equivariant to rotations and translations on coordinates and equivariant to rotations on velocities:

$$ $\textbf{h}_{i}^{(l+1)},Q\textbf{x}_{i}^{(l+1)}+\mathbf{b},Q\textbf{v}_{i}^{(l+ 1)}=\text{EGCL}[\textbf{h}_{i}^{(l)},Q\textbf{x}_{i}^{(l)}+\mathbf{b},Q\textbf {v}_{i}^{(l)},\mathcal{E}]$ (3) $$

for $Q\in\mathbb{R}^{n\times n}$ an orthogonal rotation matrix and $\textbf{b}\in\mathbb{R}^{n}$ a translation vector. The EGCL is also permutation equivariant with respect to nodes $\mathcal{V}$.

Since we have a sequence of graphs $\mathcal{G}=\{G_{t}\}_{1\leq t\leq L}$, for a time slice $t$, we apply $K_{t}$ such EGCL transformation layers to the graph $G_{t}\in\mathcal{G}$:

$$ $\textbf{h}_{i}^{(l+1)}(t),\textbf{x}_{i}^{(l+1)}(t),\textbf{v}_{i}^{(l+1)}(t)= \text{EGCL}[\textbf{h}_{i}^{(l)}(t),\textbf{x}_{i}^{(l)}(t),\textbf{v}_{i}^{(l )}(t),\mathcal{E}(t)]$ (4) $$

for $l=1,\dots,K_{t}$. Thus, we obtain spatially-contextual representations for node $i$ at time $t$ defined as $\boldsymbol{\theta}_{i}(t)=\textbf{h}_{i}^{(K_{t})}(t)\in\mathbb{R}^{d}$, $\boldsymbol{\xi}_{i}(t)=\textbf{x}_{i}^{(K_{t})}(t)\in\mathbb{R}^{n}$, $\boldsymbol{\omega}_{i}(t)=\textbf{v}_{i}^{(K_{t})}(t)\in\mathbb{R}^{n}$ for $t=1,\dots,L$.

### 3.2 Temporal Attention for Graphs

The objective of this section is to obtain strong temporally-contextual representations of the spatial graph embeddings. In the charged $N$-body problem, we are essentially solving the forward-time Euler-Lagrange equations, a second-order partial differential equation. However, for a fixed node on the spatio-temporal graph, the feature, position, and velocity form a time-series, for which RNN’s capture short-term dependencies. It was shown in (Tallec and Ollivier ) that while vanilla RNNs are not time-warping invariant, LSTMs are a class of such time-warping invariant functions modeling a continuous time-domain signal. Employing this philosophy, the use of an attention-based Transformer architecture to model spatio-temporal graph data merits investigation.

### 3.3 E ⁢ ( n ) 𝐸 𝑛 E(n) italic_E ( italic_n ) -Equivariant Temporal Attention

We would like the temporal attention to retain the equivariant properties described in Section 3.1. Namely, the Equivariant Temporal Attention Layer (ETAL) should be equivariant to the actions of $E(n)$ on coordinates and the actions of $SO(n)$ on velocities. It should also be permutation equivariant with respect to the actions of the symmetric group $\Sigma_{N}$ on nodes.

As we do not impose symmetry conditions on the layer with respect to the feature representations $\textbf{h}_{i}(t)$, we can apply key-query-value self-attention as follows. Define the node-wise query $\textbf{q}_{i}(t)=Q_{i}\boldsymbol{\theta}_{i}(t)$, key $\textbf{k}_{i}(t)=K_{i}\boldsymbol{\theta}_{i}(t)$, and value $\textbf{v}_{i}(t)=V_{i}\boldsymbol{\theta}_{i}(t)$ for $Q_{i},K_{i},V_{i}\in\mathbb{R}^{d\times d}$. To reduce memory usage, we share $Q,K,V$ for all nodes. Then the temporally-contextual representation is:

$$ $\tilde{\boldsymbol{\theta}}_{i}(t):=\sum_{s=1}^{L}\alpha_{i}(t,s)\textbf{v}_{i }(s)$ (5) $$

where

$$ $\alpha_{i}(t,s)=\frac{\exp(\textbf{q}_{i}(t)^{\top}\textbf{k}_{i}(s))}{\sum_{s ^{\prime}=1}^{L}\exp(\textbf{q}_{i}(t)^{\top}\textbf{k}_{i}(s^{\prime}))}$ (6) $$

Satorras et al. (Satorras, Hoogeboom, and Welling ) showed that for a collection of points $\{\boldsymbol{\xi_{i}}\}_{i=1}^{N}\in\mathbb{R}^{n}$, the norm is a unique geometric identifier, such that collections separated by actions of $E(n)$ form an equivalence class. With this in mind, since we desire the attention mechanism for coordinates $\boldsymbol{\xi}_{i}(t)$ to be equivariant with respect to $E(n)$, we can define:

$$ $\tilde{\boldsymbol{\xi}}_{i}(t):=\boldsymbol{\xi}_{i}(t)+B\sum_{\begin{ subarray}{c}s=1\\ s\neq t\end{subarray}}^{L}\beta_{i}(t,s)(\boldsymbol{\xi}_{i}(s)-\boldsymbol{ \xi}_{i}(t)),$ (7) $$

where

$$ $\beta_{i}(t,s)=\frac{\exp(||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s)||^ {2})}{\sum_{s^{\prime}=1}^{L}\exp(||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{ i}(s^{\prime})||^{2})}$ (8) $$

This bears semblance to the neighborhood attention described in the $SE(3)$-Transformer network (Fuchs et al. ) and Tensor Field Network layer (Thomas et al. ), the intensity function in (Zhang, Cook, and Yilmaz ), and the invariant point attention in (Jumper et al. ).

We define an $SO(n)$-equivariant attention layer for velocities $\boldsymbol{\omega}_{i}(t)$:

$$ $\tilde{\boldsymbol{\omega}}_{i}(t):=\sum_{s=1}^{L}\gamma_{i}(t,s)\boldsymbol{ \omega}_{i}(s)$ (9) $$

where the weight is

$$ $\gamma_{i}(t,s)=\frac{\boldsymbol{\omega}_{i}(t)^{\top}\boldsymbol{\omega}_{i} (s)}{\sum_{s^{\prime}=1}^{L}\exp(\boldsymbol{\omega}_{i}(t)^{\top}\boldsymbol{ \omega}_{i}(s^{\prime}))}.$ (10) $$

Figure: Algorithm 1 Spatiotemporal Attention (SpatiotempAttn)

In appendix B, we show that the position attention function is $E(n)$-equivariant and the velocity attention function is $SO(n)$-equivariant.

Following the insights of (Jin et al. ), edges are relationships between time series and they should evolve. Hence, while the adjacency matrix $A\in\mathbb{R}^{N\times N}$ is constant in space when applying ECGL, it should intuitively evolve in time when applying ETAL. That is, if we consider edges as representing the interaction between charged particles, e.g. the strength of the force, then this must necessarily evolve in time for a non-stationary point cloud system.

We define a key matrix $K(t)=KA(t)\in\mathbb{R}^{N\times N}$, query matrix $Q(t)=QA(t)\in\mathbb{R}^{N\times N}$, and value matrix $V(t)=VA(t)\in\mathbb{R}^{N\times N}$ for $t=1,\dots,L$ and $K,Q,V\in\mathbb{R}^{N\times N}$. Thus, to obtain a temporally-contextual representation of the adjacency matrix at time $t$, we apply attention:

$$ $\tilde{A}(t)=\sum_{s=1}^{L}\pi(t,s)V(s)\in\mathbb{R}^{N\times N}$ (11) $$

where

$$ $\pi(t,s)=\exp(Q(t)^{\top}K(s))\left(\sum_{s^{\prime}=1}^{L}\exp(Q(t)^{\top}K(s ^{\prime}))\right)^{-1}.$ (12) $$

Furthermore, in Appendix C, we tensorize the feature, position, velocity, and adjacency components of ETAL to efficiently compute these operations in both space $i=1,\dots,N$ and time $t=1,\dots,L$ dimensions.

### 3.4 Spacetime E ⁢ ( n ) 𝐸 𝑛 E(n) italic_E ( italic_n ) -Equivariant Graph Transformer

The full spatio-temporal attention module is presented in Algorithm 1. It takes as input the node features $h\in\mathbb{R}^{L\times N\times d}$, positions $x\in\mathbb{R}^{L\times N\times n}$, velocities $v\in\mathbb{R}^{L\times N\times n}$ and adjacency matrices $A\in\mathbb{R}^{L\times N\times N}$. For a spatio-temporal graph $\mathcal{G}=(G_{t})_{1\leq t\leq L}$, we apply an equivariant spatial attention layer in the form of EGCL to obtain spatially-contextual representations $\theta(t)\in\mathbb{R}^{N\times d},\xi(t)\in\mathbb{R}^{N\times n},\omega(t)
\in\mathbb{R}^{N\times n}$ for $t=1,\dots,L$. We share the same EGCL layer across all time steps $t=1,\dots,L$. That is, we only learn one set of MLPs $\phi_{e}$, $\phi_{v}$, $\phi_{x}$, and $\phi_{h}$ for each layer across time, which is significantly more memory and parameter efficient.

Observe, at each time step, we apply a transformation $E:\mathbb{R}^{N\times N}\times\mathbb{R}^{N\times n}\to\mathbb{R}^{N\times(N-1
)\times 2}$ to the adjacency matrix $A(t)\in\mathbb{R}^{N\times N}$ and the coordinates $x(t)\in\mathbb{R}^{N\times n}$ for $G_{t}$. This will produce edge attributes $e_{ij}(t)=\left(c_{i}c_{j},||\textbf{x}_{i}(t)-\textbf{x}_{j}(t)||_{2}^{2}\right)$ containing charge and distance information for neighboring nodes. Since each graph is complete, there are $N\times(N-1)$ such edge attributes, which we store in the tensor $\mathcal{E}(t)\in\mathbb{R}^{N\times(N-1)\times 2}$.

Figure: Algorithm 2 Spacetime $E(n)$-Transformer (SET)

Then we apply equivariant temporal attention in the form of ETAL to the spatial representations $\theta^{[1:L]}\in\mathbb{R}^{L\times N\times d}$, $\xi^{[1:L]}\in\mathbb{R}^{L\times N\times n}$, and $\omega^{[1:L]}\in\mathbb{R}^{L\times N\times n}$. Feed-forward networks $f_{\theta}:\mathbb{R}^{L\times N\times d}\to\mathbb{R}^{L\times N\times d}$, $f_{A}:\mathbb{R}^{L\times N\times N}\to\mathbb{R}^{L\times N\times N}$ with layer pre-normalization, defined in Appendix D, and residual connection are also applied to the respective feature and edge components of the graph. The sinusoidal positional encodings $W^{[1:L]}\in\mathbb{R}^{L\times N\times d}$, $X^{[1:L]}\in\mathbb{R}^{L\times n\times n}$, $Y^{[1:L]}\in\mathbb{R}^{L\times N\times n}$, $Z^{[1:L]}\in\mathbb{R}^{L\times N\times N}$ for the features, positions, velocities, and adjacency matrices are defined in Appendix D.

As the design of spatio-temporal attention is modular, we can continue stacking this architecture as we see fit (see Figure 1). In Algorithm 2, we apply spatio-temporal attention $M$ times. Then we take a mean of the resulting spatio-temporally contextual representations of positions and velocities across the time dimension, which we use as the predicted particles’ coordinates $\hat{x}(L+H)\in\mathbb{R}^{N\times n}$ and velocities $\hat{v}(L+H)\in\mathbb{R}^{N\times n}$ at the horizon target $t=L+H$.

Figure: Figure 1: Spacetime $E(n)$-Transformer architecture.
Refer to caption: extracted/5772969/graph_upper.png

The task is to predict both the positions and velocities of particles at time $L+H$, so we minimize the following loss:

$$ $\begin{split}\mathcal{L}&=||\hat{x}(L+H)-x(L+H)||_{2}^{2}+\alpha||\hat{v}(L+H) -v(L+H)||_{2}^{2}\\ &=\frac{1}{Nn}\sum_{i=1}^{N}\sum_{j=1}^{n}(\hat{x}_{ij}(L+H)-x_{ij}(L+H))^{2} \\ &+\frac{\alpha}{Nn}\sum_{i=1}^{N}\sum_{j=1}^{n}(\hat{v}_{ij}(L+H)-v_{ij}(L+H)) ^{2},\end{split}$ (13) $$

for $\alpha\in(0,1)$ a hyper-parameter.

## 4 Related Work

Temporal graph learning has a plethora of real-world applications, like COVID-19 contact tracing (Chang et al. ) (Holme ) (Ding et al. ) and misinformation detection (Choi et al. ) (Song, Teng, and Wu ) (Zhang, Cook, and Yilmaz ).

**Table 1: Ablation study of equivariance, adjacency, and attention for $N=5$. We present the model settings, parameter counts, validation & test MSE, and the MSE ratio, which is the ratio of the ablation model’s test MSE divided by the best performing model’s test MSE.**
| Ablation | Model | Params | Val MSE | Test MSE | MSE Ratio |
| --- | --- | --- | --- | --- | --- |
| Equivariance | Equiv=True, Adj=False, SATT=True, TATT=True | 796,058 | 1.21e-10 | 1.25e-10 | $-$ |
| Equiv=False, Adj=False, SATT=True, TATT=True | 796,244 | 1.96e-10 | 2.03e-10 | 1.57$\times$ |  |
| Adjacency | Equiv=True, Adj=True, SATT=True, TATT=True | 810,458 | 1.12e-09 | 1.29e-10 | 8.96$\times$ |
| Attention | Equiv=True, Adj=False, SATT=True, TATT=False | 796,049 | 2.73e-10 | 3.57e-10 | 2.86$\times$ |

Learning on continuous-time dynamic graphs was introduced by (Rossi et al. ), which proposed Temporal Graph Networks (TGNs) with a memory module, acting as a summary of what the model has seen so far for each node. Causal Anonymous Walks (Wang et al. ) is another branch of temporal graph learning, which extracts random walks between edges; however, this is not our focus. Other work like (Jin et al. ) and (Cini et al. ) treat node features as time series and edges as correlations between the series. Under this framework, message passing must be able to handle sequences of data from the neighborhood of each node, with RNNs (Seo et al. ), attention mechanisms (Marisca, Cini, and Alippi ), and convolutions (Wu et al. ).

The Dynamic Graph Convolutional Network (DynGCN) (Choi et al. ) and DyGFormer (Yu et al. ) are most similar to our method. DynGCN processes each of the graph snapshots with a graph convolutional network to obtain structural information and then applies an attention mechanism to capture temporal information. Similarly, DyGFormer (Yu et al. ) learns from historical first-hop neighborhood interactions and applies a Transformer architecture to historical correlations between nodes. However, unlike our paper, DynGCN and DyGFormer do not take into account the inductive biases of the underlying modeling task.

$E(n)$-Equivariant Graph Neural Networks (EGNN) (Satorras, Hoogeboom, and Welling ) defines a model equivariant to the Euclidean group $E(n)$ and, unlike previous methods, does not rely on spherical harmonics such as the $SE(3)$-Transformer (Fuchs et al. ) and Tensor Field Networks (Thomas et al. ). The $SE(3)$-Transformer paper briefly alludes to incorporating equivariant attention with an LSTM for temporal causality; however, this is not the primary focus of their work. LieConv (Finzi et al. ) proposes a framework that allows one to construct a convolutional layer that is equivariant with respect to transformations of a Lie group, equipped with an exponential map. However, the EGNN is simpler and more applicable to problems with point clouds like the charged $N$-body problem (Kipf et al. ) we consider.

As we concern ourselves with modeling a dynamical system, the works of Lagrangian Neural Networks (LNNs) (Cranmer et al. ) and Hamiltonian Neural Networks (HNNs) (Greydanus, Dzamba, and Yosinski ) are pertinent. HNNs parameterize the Hamiltonian of a system, but require canonical coordinates, which makes it inapplicable to systems where such coordinates cannot be deduced. LLNs parameterize arbitrary Lagrangians of dynamical systems with neural networks, from which it is possible to solve the forward dynamics of the system; however, this requires an additional step of integration, which is cumbersome.

## 5 Experiments & Results

Adapting the charged $N$-body system dataset from (Satorras, Hoogeboom, and Welling ), we sample 16,000 trajectories for training, 2,000 trajectories for validation, and 2,000 trajectories for testing. Each trajectory has a horizon length of $H=10,000$ and a sequence length of $L=10$. The point cloud consists of $N=5$ particles, where at each time step, positions $(\textbf{x}_{1}(t),\dots\textbf{x}_{5}(t))^{\top}\in\mathbb{R}^{5\times 3}$, velocities $(\textbf{v}_{1}(t),\dots\textbf{v}_{5}(t))^{\top}\in\mathbb{R}^{5\times 3}$, as well as charges $c_{1},\dots,c_{5}\in\{-1,+1\}$ are known. The edges between charged particles is $e_{ij}(t)=\left(c_{i}c_{j},||\textbf{x}_{i}(t)-\textbf{x}_{j}(t)||_{2}^{2}\right)$. We input these known values into SET with features chosen as $h_{i}(t)=||\textbf{v}_{i}(t)||_{2}$ for $i=1,\dots,5$. We conducted a hyper-parameter optimization and selected the best model settings, as per Appendix E.

### 5.1 Ablation Study: Equivariance, Adjacency, and Attention

Furthermore, we conduct an ablation study on SET, shown in Table 1, which compares the use of equivariance, temporal attention for the adjacency matrix as per Section 3.3, spatial attention, and temporal attention. By selecting the best model on the validation set, we find that incorporating equivariance, spatial attention, and temporal attention enhances performance, whereas using adjacency diminishes it. We hypothesize that the insignificance of temporal adjacency is due to the fact the edge attribute contains information about charges, which does not evolve in time, and information about the distance between particles, which already implicitly exists in the coordinate information.

### 5.2 Baselines & Scaling N 𝑁 N italic_N

We compare our best performing SET model with optimized LSTM, EGNN, MLP and Linear baselines (see Appendix E for implementation details). SET outperforms all baselines for $N=5$, as seen in Table 2.

**Table 2: Baselines for $N=5$.**
| Model | Params | Test MSE |
| --- | --- | --- |
| SET | 796,058 | 1.25e-10 |
| LSTM | 826,313 | 2.03e-08 |
| EGNN | 100,612 | 2.05e-06 |
| MLP | 67,718 | 3.48e-06 |
| Linear | 3 | 3.04 |

Since the $N=5$ system is seemingly too simple a task, we scale the dataset to $N=20$ and $N=30$. As shown in Figure 2, test MSE remains consistent for all models regardless of the number of charged particles $N$, which is a desirable property. Per Figure 2, the number of model parameters remains constant for the EGNN, MLP and Linear baselines. Fortunately, the number of parameters in SET also remains constant for all $N$; this is an artifact of the attention layers only being functions of the feature and coordinate dimensions. Note, the number of parameters in the LSTM increases from $8.2e5$ to $1.8e6$, which is an undesirable property. Further results are included in Appendix E.

Figure: Figure 2: Top: Model test MSE versus $N$. Bottom: Number of model parameters versus $N$.
Refer to caption: extracted/5772969/scale_plot.png

### 5.3 Noisy Observations

We also conduct an experiment in which we introduce noise in the observations, by adding Gaussian noise $\varepsilon\sim\mathcal{N}(0,0.5)$ to positions and velocities. As seen in Table 3, all models have test MSEs that plateau around the variance $\sigma^{2}=0.5$ of the irreducible noise, with the SET and MLP baseline tying in performance.

**Table 3: Model comparison for a noisy $N=5$ dataset.**
| Model | Params | Test MSE |
| --- | --- | --- |
| SET | 796,058 | 4.97e-01 |
| LSTM | 826,313 | 4.98e-01 |
| EGNN | 100,612 | 4.98e-01 |
| MLP | 67,718 | 4.97e-01 |
| Linear | 3 | 3.60 |

## 6 Conclusions

The imposition of group symmetries on graph neural networks is a promising area of research, demonstrating remarkable real-world results like AlphaFold (Jumper et al. ). However, most research has been centered on spatial equivariance for representational learning on static graphs. For dynamical graph systems, little research has centered on preserving group symmetries across time. We close this gap with the Spacetime $E(n)$-Transformer and show promising results for the charged $N$-body problem. It will be interesting to see our method applied to harder tasks, such as sequential bio-molecular generation.

Although we chose a graph as the domain of interest, it is plausible to extend notions of spatio-temporal $G$-equivariance to other domains like grids and manifolds. Furthermore, while we leveraged the symmetries of the problem a priori, it may not always be possible to find a simple group for a general problem. Hence, in future work, it would be interesting to learn a group symmetry from underlying data and impose equivariance using methods like LieConv (Finzi et al. ), which is equivariant to the actions of Lie groups, i.e. the continuous group representation of infinitesimal transformations. Noether’s first theorem implies a possible connection to conserved quantities, which was discussed in (Alet et al. ).

## Appendix A Preliminaries

### A.1 Isometries

###### Definition A.1 .

Let $(X,d_{X})$ and $(Y,d_{Y})$ be metric spaces with metrics $d_{X}:X\times X\to\mathbb{R}$, $d_{Y}:Y\times Y\to\mathbb{R}$. An isometry $\varphi:X\to Y$ is a distance-preserving isomorphism if

$$ $d_{X}(a,b)=d_{Y}(\varphi(a),\varphi(b))$ (14) $$

for all $a,b\in X$.

###### Definition A.1 .

### A.2 Invariance & Equivariance

###### Definition A.2 .

Let $\rho_{V}:G\to GL(K,V)$ be a representation of group $G$, and let $\rho_{V}(g):V\stackrel{{\scriptstyle\sim}}{{\smash{\longrightarrow}\rule{0.0pt
}{1.72218pt}}}V$ be an automorphism on $V$ for $g\in G$. A function $f:V\to W$ is $G$-equivariant if there exists an equivalent representation $\rho_{W}:G\to GL(K,W)$ with equivalent automorphism $\rho_{W}(g):W\stackrel{{\scriptstyle\sim}}{{\smash{\longrightarrow}\rule{0.0pt
}{1.72218pt}}}W$, such that

$$ $f(\rho_{V}(g)(v))=\rho_{W}(g)(f(v))$ (15) $$

for all $v\in V$ and $g\in G$.

###### Definition A.3 .

Let $\rho:G\to GL(K,V)$ be a representation of group $G$ and let $\rho(g):V\stackrel{{\scriptstyle\sim}}{{\smash{\longrightarrow}\rule{0.0pt}{1.
72218pt}}}V$ be an automorphism for $g\in G$. A function $f:V\to W$ is $G$-invariant if

$$ $f(\rho(g)(v))=f(v)$ (16) $$

for all $v\in V$ and $g\in G$.

###### Definition A.2 .

###### Definition A.3 .

### A.3 Special Property of S ⁢ O ⁢ ( n ) 𝑆 𝑂 𝑛 SO(n) italic_S italic_O ( italic_n )

Since rotations are distance, angle, and orientation preserving, they are linear transformations. As such, rotations can be represented as a matrix. Suppose $\mathbf{x},\mathbf{y}\in\mathbb{R}^{n}$ and $Q\in\mathbb{R}^{n\times n}$ is a rotation matrix. Then if $Q$ is an isometry, we require:

$$ $\begin{split}\mathbf{x}^{\top}\mathbf{y}&=(Q\mathbf{x})^{\top}(Q\mathbf{y})\\ &=\mathbf{x}^{\top}(Q^{\top}Q)\mathbf{y}\end{split}$ (17) $$

or $Q^{\top}Q=I=Q^{\top}Q$. Preservation of orientation (equivalently, handedness) means $\det Q>0$. We take the determinant of the identity and find $\det(Q^{\top}Q)=(\det(Q))^{2}=1$, which means $\det Q=\pm 1$. Hence, $\det Q=+1$ for $SO(n)$.

## Appendix B Proof of Equivariances for Temporal Attention Layer

### B.1 Position Component: E ⁢ ( n ) 𝐸 𝑛 E(n) italic_E ( italic_n ) -Equivariance

###### Proof.

Consider an isometry of $E(n)$ defined by $\boldsymbol{\xi}\mapsto A\boldsymbol{\xi}+\mathbf{b}$ for $A\in\mathbb{R}^{n\times n}$ an orthogonal rotation matrix and $\mathbf{b}\in\mathbb{R}^{n}$ a translation vector. Then

$$ $\begin{split}||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s)||^{2}&=||A \boldsymbol{\xi}_{i}(t)+\mathbf{b}-A\boldsymbol{\xi}_{i}(s)-\mathbf{b}||^{2}\\ &=||A(\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s))||^{2}\\ &=(\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}(s))^{\top}A^{\top}A(\boldsymbol{ \xi}_{i}(t)-\boldsymbol{\xi}_{i}(s))\\ &=||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s)||^{2},\end{split}$ (18) $$

where the last line follows by the orthogonality of $A$. Hence, $\beta_{i}(t,s)=\exp(||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s)||^{2})/
\sum_{s^{\prime}=1}^{L}\exp(||\boldsymbol{\xi}_{i}(t)-\boldsymbol{\xi}_{i}(s^{
\prime})||^{2})$ is invariant under isometries of $\mathbb{R}^{n}$. Thus, applying the isometry to the layer

$$ $\tilde{\boldsymbol{\xi}}_{i}(t)=\boldsymbol{\xi}_{i}(t)+B\sum_{\begin{subarray }{c}s=1\\ s\neq t\end{subarray}}^{L}\beta_{i}(t,s)(\boldsymbol{\xi}_{i}(s)-\boldsymbol{ \xi}_{i}(t))$ (19) $$

yields:

$$ $\begin{split}&A\boldsymbol{\xi}_{i}(t)+\mathbf{b}+B\sum_{\begin{subarray}{c}s= 1\\ s\neq t\end{subarray}}^{L}\beta_{i}(t,s)A(\boldsymbol{\xi}_{i}(s)-\boldsymbol{ \xi}_{i}(t))\\ &=A\left(\boldsymbol{\xi}_{i}(t)+B\sum_{\begin{subarray}{c}s=1\\ s\neq t\end{subarray}}^{L}\beta_{i}(t,s)(\boldsymbol{\xi}_{i}(s)-\boldsymbol{ \xi}_{i}(t))\right)+\mathbf{b}\\ &=A\tilde{\boldsymbol{\xi}}_{i}(t)+\mathbf{b}.\end{split}$ (20) $$

It follows that the position component of ETAL is $E(n)$-equivariant.
∎

###### Proof.

### B.2 Velocity Component: S ⁢ O ⁢ ( n ) 𝑆 𝑂 𝑛 SO(n) italic_S italic_O ( italic_n ) -Equivariance

###### Proof.

Consider an orthogonal rotation matrix $A\in\mathbb{R}^{n\times n}$ of $SO(n)$ with action $\boldsymbol{\omega}\mapsto A\boldsymbol{\omega}$. Clearly,

$$ $\begin{split}(A\boldsymbol{\omega}_{i}(t))^{\top}(A\boldsymbol{\omega}_{i}(s)) &=\boldsymbol{\omega}_{i}(t)^{\top}A^{\top}A\boldsymbol{\omega}_{i}(s)\\ &=\boldsymbol{\omega}_{i}(t)^{\top}\boldsymbol{\omega}_{i}(s).\end{split}$ (21) $$

Thus, the coefficient $\gamma_{i}(t,s)=\boldsymbol{\omega}_{i}(t)^{\top}\boldsymbol{\omega}_{i}(s)/
\sum_{s^{\prime}=1}^{L}\exp(\boldsymbol{\omega}_{i}(t)^{\top}\boldsymbol{
\omega}_{i}(s^{\prime}))$ is invariant under rotations of $SO(n)$.

Applying the rotation to the attention layer

$$ $\tilde{\boldsymbol{\omega}}_{i}(t)=\sum_{s=1}^{L}\gamma_{i}(t,s)\boldsymbol{ \omega}_{i}(s)$ (22) $$

yields:

$$ $\begin{split}\sum_{s=1}^{L}\gamma_{i}(t,s)A\boldsymbol{\omega}_{i}(s)&=A\sum_{ s=1}^{L}\gamma_{i}(t,s)\boldsymbol{\omega}_{i}(s)\\ &=A\tilde{\boldsymbol{\omega}}_{i}(t),\end{split}$ (23) $$

as asserted. Hence, the velocity component of ETAL is $SO(n)$-equivariant.
∎

###### Proof.

## Appendix C Tensorization of Equivariant Temporal Attention Layer (ETAL)

### C.1 Temporal Attention Feature Component

Define the matrix of spatial graph representations for features across time:

$$ $\theta_{i}^{[1:L]}=\begin{bmatrix}\boldsymbol{\theta}_{i}(1)^{\top}\\ \vdots\\ \boldsymbol{\theta}_{i}(L)^{\top}\end{bmatrix}\in\mathbb{R}^{L\times d}.$ (24) $$

We compute the value vectors $v_{i}^{[1:L]}=\theta_{i}^{[1:L]}V\in\mathbb{R}^{L\times d}$, key vectors $k_{i}^{[1:L]}=\theta_{i}^{[1:L]}K\in\mathbb{R}^{L\times d}$, and query vectors $q_{i}^{[1:L]}=\theta_{i}^{[1:L]}Q\in\mathbb{R}^{L\times d}$ for $V,K,Q\in\mathbb{R}^{d\times d}$. Thus, to compute the attention layer for the graph features in Equation 5, we first compute the weights:

$$ $\alpha_{i}=\text{softmax}\left(\frac{\theta_{i}^{[1:L]}QK^{\top}{\theta_{i}^{[ 1:L]}}^{\top}}{\sqrt{d}}\right)\in\mathbb{R}^{L\times L}.$ (25) $$

Then we apply the weights as follows:

$$ $\tilde{\theta}_{i}^{[1:L]}=\alpha_{i}v_{i}^{[1:L]}\in\mathbb{R}^{L\times d}.$ (26) $$

We would like to tensorize this computation for all nodes $i=1,\dots,N$. Let $\theta^{[1:L]}=(\theta_{1}^{[1:L]},\dots,\theta_{N}^{[1:L]})\in\mathbb{R}^{N
\times L\times d}$. Likewise, let $k^{[1:L]}=\theta^{[1:L]}K\in\mathbb{R}^{N\times L\times d}$, $q^{[1:L]}=\theta^{[1:L]}Q\in\mathbb{R}^{N\times L\times d}$, and $v^{[1:L]}=\theta^{[1:L]}V\in\mathbb{R}^{N\times L\times d}$. In a slight abuse of notation,

$$ $\alpha=\text{softmax}\left(\frac{q^{[1:L]}{k^{[1:L]}}^{\top}}{\sqrt{d}}\right) \in\mathbb{R}^{N\times L\times L}$ (27) $$

where the transpose ${k^{[1:L]}}^{\top}\in\mathbb{R}^{N\times d\times L}$ interchanges the last two dimensions, the tensor multiplication $q^{[1:L]}{k^{[1:L]}}^{\top}$ is with respect to the last two dimensions, and the softmax is computed over the last dimension. Then

$$ $\tilde{\theta}^{[1:L]}=\alpha v^{[1:L]}\in\mathbb{R}^{N\times L\times d},$ (28) $$

where the tensor multiplication is over the last two dimensions.

### C.2 Temporal Attention Position Component

Define the matrix of spatial graph representations for positions and its corresponding temporally-transformed matrix, across time:

$$ $\xi_{i}^{[1:L]}=\begin{bmatrix}\boldsymbol{\xi}_{i}(1)^{\top}\\ \vdots\\ \boldsymbol{\xi}_{i}(L)^{\top}\end{bmatrix}\in\mathbb{R}^{L\times n},\tilde{ \xi}_{i}^{[1:L]}=\begin{bmatrix}\tilde{\boldsymbol{\xi}}_{i}(1)^{\top}\\ \vdots\\ \tilde{\boldsymbol{\xi}}_{i}(L)^{\top}\end{bmatrix}\in\mathbb{R}^{L\times n}.$ (29) $$

Similarly, construct the node-wise attention matrix:

$$ $\beta_{i}=\begin{bmatrix}\beta_{i}(1,1)&\ldots&\beta_{i}(1,L)\\ \vdots&\ddots&\vdots\\ \beta_{i}(L,1)&\ldots&\beta_{i}(L,L)\end{bmatrix},$ (30) $$

where $\beta_{i}(t,s)$ is defined in Equation 8. Then we can write the update in Equation 7 as:

$$ $\begin{split}\tilde{\xi}_{i}^{[1:L]}&=\xi_{i}^{[1:L]}+B\left((1_{L\times L}-I_ {L\times L})\odot\beta_{i}\right)\xi_{i}^{[1:L]}\\ &-B\left(((1_{L\times L}-I_{L\times L})\odot\beta_{i})1_{L\times 1}\right) \odot\xi_{i}^{[1:L]},\end{split}$ (31) $$

where $\odot$ is the Hadamard product. Define the difference matrix:

$$ $\begin{split}D_{i}:&=\begin{bmatrix}\boldsymbol{\xi}_{i}(1)-\boldsymbol{\xi}_{ i}(1)&\ldots&\boldsymbol{\xi}_{i}(1)-\boldsymbol{\xi}_{i}(L)\\ \vdots&\ddots&\vdots\\ \boldsymbol{\xi}_{i}(L)-\boldsymbol{\xi}_{i}(1)&\ldots&\boldsymbol{\xi}_{i}(L) -\boldsymbol{\xi}_{i}(L)\end{bmatrix}\\ &=\begin{bmatrix}\boldsymbol{\xi}_{i}(1)&\ldots&\boldsymbol{\xi}_{i}(1)\\ \vdots&\ddots&\vdots\\ \boldsymbol{\xi}_{i}(L)&\ldots&\boldsymbol{\xi}_{i}(L)\end{bmatrix}-\begin{ bmatrix}\boldsymbol{\xi}_{i}(1)&\ldots&\boldsymbol{\xi}_{i}(L)\\ \vdots&\ddots&\vdots\\ \boldsymbol{\xi}_{i}(1)&\ldots&\boldsymbol{\xi}_{i}(L)\end{bmatrix}.\end{split}$ (32) $$

Then we can compute the weights as:

$$ $\beta_{i}=\text{softmax}\left(\frac{-D_{i}^{\top}\odot D_{i}}{\sqrt{n}}\right) \in\mathbb{R}^{L\times L}.$ (33) $$

Again, we seek to find a differentiable expression that tensorizes the attention layer for all nodes. We stack $\tilde{\xi}_{i}^{[1:L]}$ for $i=1,\dots,N$ into a tensor:

$$ $\begin{split}\tilde{\xi}^{[1:L]}&=\begin{bmatrix}\tilde{\xi}_{1}^{[1:L]}\\ \vdots\\ \tilde{\xi}_{N}^{[1:L]}\end{bmatrix}\\ &=\begin{bmatrix}\xi_{1}^{[1:L]}\\ \vdots\\ \xi_{N}^{[1:L]}\end{bmatrix}+B\begin{bmatrix}((1_{L\times L}-I_{L\times L}) \odot\beta_{1})\xi_{1}^{[1:L]}\\ \vdots\\ ((1_{L\times L}-I_{L\times L})\odot\beta_{N})\xi_{N}^{[1:L]}\end{bmatrix}-B \begin{bmatrix}(((1_{L\times L}-I_{L\times L})\odot\beta_{1})1_{L\times 1}) \odot\xi_{1}^{[1:L]}\\ \vdots\\ (((1_{L\times L}-I_{L\times L})\odot\beta_{N})1_{L\times 1})\odot\xi_{N}^{[1:L ]}\end{bmatrix}\\ &=\begin{bmatrix}\xi_{1}^{[1:L]}\\ \vdots\\ \xi_{N}^{[1:L]}\end{bmatrix}+B\begin{bmatrix}(1_{L\times L}-I_{L\times L}) \odot\beta_{1}\\ \vdots\\ (1_{L\times L}-I_{L\times L})\odot\beta_{N}\end{bmatrix}\odot\begin{bmatrix} \xi_{1}^{[1:L]}\\ \vdots\\ \xi_{N}^{[1:L]}\end{bmatrix}-B\begin{bmatrix}(((1_{L\times L}-I_{L\times L}) \odot\beta_{1})1_{L\times 1})\odot\xi_{1}^{[1:L]}\\ \vdots\\ (((1_{L\times L}-I_{L\times L})\odot\beta_{N})1_{L\times 1})\odot\xi_{N}^{[1:L ]}\end{bmatrix},\end{split}$ (34) $$

which is a differentiable function with respect to the $\boldsymbol{\xi}_{i}(1),\dots,\boldsymbol{\xi}_{i}(L)$ for $i=1,\dots,N$, where $\tilde{\xi}^{[1:L]}\in\mathbb{R}^{N\times L\times n}$.

### C.3 Temporal Attention Velocity Component

Define the matrix of spatial graph representations for velocities, across time:

$$ $\omega_{i}^{[1:L]}=\begin{bmatrix}\boldsymbol{\omega}_{i}(1)^{\top}\\ \vdots\\ \boldsymbol{\omega}_{i}(L)^{\top}\end{bmatrix}\in\mathbb{R}^{L\times n}.$ (35) $$

We compute the weights as:

$$ $\gamma_{i}=\text{softmax}\left(\frac{\omega_{i}^{[1:L]}{\omega_{i}^{[1:L]}}^{ \top}}{\sqrt{n}}\right)\in\mathbb{R}^{L\times L}.$ (36) $$

Then we apply the weights to obtain the layer transform:

$$ $\tilde{\omega}_{i}^{[1:L]}=\gamma_{i}\omega_{i}^{[1:L]}\in\mathbb{R}^{L\times n}.$ (37) $$

We would like to find a tensorized expression that includes all nodes $i=1,\dots,N$. Let $\omega^{[1:L]}=(\omega_{1}^{[1:L]},\dots,\omega_{N}^{[1:L]})\in\mathbb{R}^{N
\times L\times n}$. As before, in a slight abuse of notation,

$$ $\gamma=\text{softmax}\left(\frac{{\omega}^{[1:L]}{{\omega}^{[1:L]}}^{\top}}{ \sqrt{n}}\right)\in\mathbb{R}^{N\times L\times L},$ (38) $$

where the transpose ${{\omega}^{[1:L]}}^{\top}\in\mathbb{R}^{N\times n\times L}$ interchanges the last two dimensions, the tensor multiplication ${\omega}^{[1:L]}{{\omega}^{[1:L]}}^{\top}$ is with respect to the last two dimensions, and the softmax is computed over the last dimension. Hence,

$$ $\tilde{\omega}^{[1:L]}=\gamma\omega^{[1:L]}\in\mathbb{R}^{N\times L\times n},$ (39) $$

where the tensor multiplication is over the last two dimensions.

### C.4 Temporal Attention Adjacency Component

Let $A^{[1:L]}=(A(1),\dots,A(L))\in\mathbb{R}^{L\times N\times N}$ be the tensor of all adjacency matrices over time. Let $k^{[1:L]}=A^{[1:L]}K$, $q^{[1:L]}=A^{[1:L]}Q$, $v^{[1:L]}=A^{[1:L]}V$ where $K,Q,V\in\mathbb{R}^{N\times N}$. With a slight abuse of notation, let

$$ $\pi=\text{softmax}\left(\frac{q^{[1:L]}{k^{[1:L]}}^{\top}}{\sqrt{N}}\right)\in \mathbb{R}^{L\times N\times N},$ (40) $$

where the transpose ${k^{[1:L]}}^{\top}\in\mathbb{R}^{L\times N\times N}$ interchanges the last two dimensions, the tensor multiplication $q^{[1:L]}{k^{[1:L]}}^{\top}$ is with respect to the last two dimensions, and the softmax is computed over the last dimension as usual.

Thus,

$$ $\tilde{A}^{[1:L]}=\pi v^{[1:L]}\in\mathbb{R}^{L\times N\times N}.$ (41) $$

## Appendix D Positional Encoding & Layer Normalization

We use sinusoidal positional encodings, as defined in (Vaswani et al. ). That is, we define the positional encodings $W^{[1:L]}\in\mathbb{R}^{L\times N\times d}$, $X^{[1:L]}\in\mathbb{R}^{L\times N\times n}$, $Y^{[1:L]}\in\mathbb{R}^{L\times N\times n}$, and $Z^{[1:L]}\in\mathbb{R}^{L\times N^{2}}$, used in Algorithm 1, as:

$$ $\begin{split}W_{i,\cdot,2j}&=\sin\left(\frac{i}{\kappa^{2j/d}}\right),W_{i, \cdot,2i+1}=\cos\left(\frac{i}{\kappa^{2j/d}}\right),\\ X_{i,\cdot,2j}&=\sin\left(\frac{i}{\kappa^{2j/n}}\right),X_{i,\cdot,2i+1}=\cos \left(\frac{i}{\kappa^{2j/n}}\right),\\ Y_{i,\cdot,2j}&=\sin\left(\frac{i}{\kappa^{2j/n}}\right),Y_{i,\cdot,2i+1}=\cos \left(\frac{i}{\kappa^{2j/n}}\right),\\ Z_{i,2j}&=\sin\left(\frac{i}{\kappa^{2j/N^{2}}}\right),Z_{i,2i+1}=\cos\left( \frac{i}{\kappa^{2j/N^{2}}}\right).\end{split}$ (42) $$

We compress the adjacency matrix positional encoding as $Z^{[1:L]}\in\mathbb{R}^{L\times N^{2}}$ because we treat the $N^{2}$ entries of the adjacency matrix as belonging to a vectorspace, to which we apply a positional encoding.

Layer normalization in Algorithm 1 is defined as follows:

$$ $\begin{split}&\hat{\mu}^{[1:L]}=\begin{bmatrix}\frac{1}{d}\sum_{j=1}^{d}\theta _{1,j}^{[1:L]}\\ \vdots\\ \frac{1}{d}\sum_{j=1}^{d}\theta_{LN,j}^{[1:L]}\end{bmatrix}\in\mathbb{R}^{L \times N\times 1},\\ &\hat{\sigma}^{[1:L]}=\begin{bmatrix}\sqrt{\frac{1}{d}\sum_{j=1}^{d}\left( \theta_{1,j}^{[1:L]}-\hat{\mu}_{1}^{[1:L]}\right)^{2}}\\ \vdots\\ \sqrt{\frac{1}{d}\sum_{j=1}^{d}\left(\theta_{LN,j}^{[1:L]}-\hat{\mu}_{LN}^{[1: L]}\right)^{2}}\end{bmatrix}\in\mathbb{R}^{L\times N\times 1},\\ &LN(\theta^{[1:L]})=\frac{\theta^{[1:L]}-\hat{\mu}^{[1:L]}}{\hat{\sigma}^{[1:L ]}},\end{split}$ (43) $$

so we broadcast $\hat{\mu},\hat{\sigma}$ across the second $d$ dimensions of $\theta^{[1:L]}$. Likewise, for the adjacency matrix, we define:

$$ $\begin{split}&\hat{\mu}^{[1:L]}=\begin{bmatrix}\frac{1}{N^{2}}\sum_{j=1}^{N^{2 }}\theta_{1,j}^{[1:L]}\\ \vdots\\ \frac{1}{N^{2}}\sum_{j=1}^{N^{2}}\theta_{L,j}^{[1:L]}\end{bmatrix}\in\mathbb{R }^{L\times 1},\\ &\hat{\sigma}^{[1:L]}=\begin{bmatrix}\sqrt{\frac{1}{N^{2}}\sum_{j=1}^{N^{2}} \left(\theta_{1,j}^{[1:L]}-\hat{\mu}_{1}^{[1:L]}\right)^{2}}\\ \vdots\\ \sqrt{\frac{1}{N^{2}}\sum_{j=1}^{N^{2}}\left(\theta_{L,j}^{[1:L]}-\hat{\mu}_{L }^{[1:L]}\right)^{2}}\end{bmatrix}\in\mathbb{R}^{L\times 1},\\ &LN(A^{[1:L]})=\frac{A^{[1:L]}-\hat{\mu}^{[1:L]}}{\hat{\sigma}^{[1:L]}},\end{split}$ (44) $$

where we view $A^{[1:L]}\in\mathbb{R}^{L\times N^{2}}$.

## Appendix E Hyper-parameter Settings & Implementation Details

### E.1 SET & Scaling Results

After a hyper-parameter optimization for SET, we found the following optimal settings. Namely, we used an Adam optimizer with an initial learning rate of $4.45\times 10^{-5}$, a batch size of $100$, dropout of $0.1$, $2$ EGNN layers with SiLU activations, $3$ vertical stacking layers with ReLU activations, hidden dimension of $128$, spatial and temporal attention, and equivariance. The following hyper-parameters were inactive: weight decay, temporal adjacency, positional encoding and causal attention. We chose $\alpha=1$ in Equation 13 and $B=0.5$ in Equation 7.

The results in tables 4 and 5 show the effects of scaling the dataset. SET is the best-performing model across all datasets.

**Table 4: Performance metrics for optimized models with $N=20$.**
| Model | Params | Test MSE |
| --- | --- | --- |
| SET | 796,058 | 5.72e-11 |
| LSTM | 1,269,833 | 1.02e-04 |
| EGNN | 100,612 | 5.34e-05 |
| MLP | 67,718 | 3.77e-06 |
| Linear | 3 | 4.48 |

**Table 5: Performance metrics for optimized models with $N=30$.**
| Model | Params | Test MSE |
| --- | --- | --- |
| SET | 796,058 | 1.20e-10 |
| LSTM | 1,859,513 | 1.95e-07 |
| EGNN | 100,612 | 5.88e-05 |
| MLP | 67,718 | 3.99e-06 |
| Linear | 3 | 5.19 |

### E.2 EGNN, LSTM, MLP, Linear Baselines

For the EGNN baseline, we apply the same methodology outlined in Algorithm 1, whereby we apply MLPs $\phi_{e},\phi_{v},\phi_{x},\phi_{h}$ across the time dimension:

$$ $\begin{split}\textbf{h}_{i}^{(l+1)}(1)&=\phi_{h}(\textbf{h}_{i}^{(l)}(1), \textbf{m}_{i}(1))\in\mathbb{R}^{d}\\ &\vdots\\ \textbf{h}_{i}^{(l+1)}(L)&=\phi_{h}(\textbf{h}_{i}^{(l)}(L),\textbf{m}_{i}(L)) \in\mathbb{R}^{d},\end{split}$ (45) $$

for $l=1,\dots,K$. We used an Adam optimizer without weight decay and with an initial learning rate of $3.98\times 10^{-5}$, batch size of $32$, hidden dimension of $64$, $3$ EGNN layers, and residual connections.

For the LSTM baseline, we found using only temporal attention was optimal. That is, we did not pre-process the graph data with an EGNN; we simply applied a vanilla LSTM to the graph data. Each token embedding $\mathbf{z}(t)\in\mathbb{R}^{N\times(d+2n+2N-2)}$ comprises of the following data:

$$ $\begin{split}&\textbf{h}_{1}(t),\dots,\textbf{h}_{N}(t)\in\mathbb{R}^{N\times d }\\ &\textbf{x}_{1}(t),\dots,\textbf{x}_{N}(t)\in\mathbb{R}^{N\times n}\\ &\textbf{v}_{1}(t),\dots,\textbf{v}_{N}(t)\in\mathbb{R}^{N\times n}\\ &\textbf{e}_{0,1}(t),\dots,\textbf{e}_{0,N}(t),\textbf{e}_{1,0}(t),\dots, \textbf{e}_{1,N}(t),\\ &\dots,\textbf{e}_{N,0}(t),\dots,\textbf{e}_{N,N-1}(t)\in\mathbb{R}^{2}.\end{split}$ (46) $$

We used Adam with an initial learning rate of $2.47\times 10^{-5}$, batch size of $32$, hidden dimension of $512$, $2$ EGNN layers, $3$ temporal stacking layers, and dropout of $0.1$. We omitted the use of weight decay, temporal adjacency, and recurrence. The MLP uses $5$ hidden layers with hidden dimension $128$. It only takes a tensor of positions $x\in\mathbb{R}^{L\times N\times n}$ and velocities $v\in\mathbb{R}^{L\times N\times n}$ as inputs. It is trained using an Adam optimizer with a learning rate of $1.75\times 10^{-5}$ and a batch size of $100$. The linear dynamics model simply predicts position $x(t)=x(t-1)+\alpha v(t-1)$ and velocity $v(t)=\beta v(t-1)+\gamma$, learning $\alpha,\beta,\gamma$ using an Adam optimizer with initial learning rate of $2.73\times 10^{-5}$, weight decay of $1\times 10^{-6}$, and batch size of $100$.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 1
questions: 1. It seems that ETAL is one of the main contributions of the paper. Can we combine ETAL (without the geometric features) with vanilla GNNs for standard (non-geometric) spatiotemporal benchmarks? 

2. Could the authors elaborate on the motivation for sequentially applying SchNet and EGCL? Also, in Eq. (6), $h_i^{(l+1)}(t)$ appears as the input and output of ScnNet. A similar issue occurs in Algorithm 1 (Line 283), where $h_i^{(l+1)}(t)$ gets overwritten.

3. How does the proposed approach compare with the method in [1]? Also, how does the proposal perform on the molecular and protein datasets in [1]?
  
[1] Wu et al., Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics, NeurIPS, 2023.
soundness: 2
strengths: **Relevance**: The problem of modeling the dynamic of physical systems using neural nets is a very relevant research topic.

**Ablation study**: The paper conducts an ablation study that shows the impact of different model components. The results confirm the importance of the full model and equivariance.
summary: The paper proposes an E(n)-equivariant Transformer architecture for spatio-temporal data, called Spacetime E(n)-Transformer (SET). The proposal builds upon E(n) Equivariant GNNs and temporal attention mechanisms. Experiments on the N-body dataset show the effectiveness of the proposed model compared to MLPs, EGNNs, and LSTMs.
weaknesses: **Exposition/Organization**: Overall, the paper could better use the available space (page limit). For instance, Section 2.2 could be significantly compressed. Also, I would move Algorithms 1 and 2 to Appendix --- after strengthening the description of the method section, these algorithms are unnecessary. The same applies to the squared loss in Eq. (14). The description of EGNNs (EGCL layers) should be in the Background. Similarly, several basic definitions/notations are provided in the method section when they should be in the Background.

**Novelty**: Overall, I found the novelty of this paper very limited. From a methodological perspective, the proposal consists of a combination of Equivariant GNNs (SchNet) and temporal models that preserve equivariance. Being ETAL the main methodological contribution, I believe the paper should extensively evaluate it in different settings. From a conceptual level, this is not the first work to consider equivariant spatial and equivariant temporal models to simulate dynamical physical systems.

**Baselines/experiments**: One of the main weaknesses is the experimental setup. The paper only considers a single task (the N-body problem), which may cause readers to interpret it as limited applicability. Also, the paper does not consider any model for spatiotemporal data (e.g., with vanilla GNNs replaced with equivariant ones) as a baseline in the experiments. The method in [1] is also not considered in the comparison. 

[1] Wu et al., Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics, NeurIPS, 2023.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the sections above.
soundness: 3
strengths: - Studying equivariant architectures for spatio-temporal graphs is an interesting and novel direction.
summary: The paper introduces an architecture for modelling spatio-temporal graphs, which is equivariant to translation and rotation of the nodes. It relies on applying Equivariant Graph Convolutional Network on the spatial dimension (independent for each timestep) followed by an equivariant transformer on the temporal dimension (independent for each node). The model is compared against a series of baselines on a synthetic N-body system dataset.
weaknesses: - My main concern regards the limited novelty of the proposed model. Architectures for spatio-temporal processing that use a factorization approach (separating spatial and temporal processing) are already established (e.g., DynGCN). While the paper claims to be the first to introduce equivariance in this context, it independently applies equivariance in both the spatial and temporal dimensions using an existing method (Santorras et al.). No specific modifications to the model are required; simply applying this approach first to the spatial dimension, 
 and than to the fully connected temporal graph results in an E(n)-equivariant model. Consequently, the contribution is mainly incremental, arising from a straightforward combination of existing methods, which I find insufficient to recommend for publication.
- The main focus of the paper is to introduce the first equivariant architecture tailor for spatio-temporal graphs. However,  it is unclear what unique challenges the temporal dimension introduces. From the proposal, both the temporal and the spatial dimenion are processed in a mirrored way, suggesting that the temporal dimension may not pose significant additional challenges when working with geometric graphs.
- The experimental section does not compare against any existing spatio-temporal methods in the literature (such as the one mentioned in the related work section). Is there any particular reason why that is not possible? The only comparison are against very weak baselines that either ignore interactions entirely (LSTM) or ignore temporality entirely (e.g. spatial-level EGNN).
- The manuscript can be improved in terms of clarity. For example, abbreviation such as SATT and TATT used in the ablation section are not explained. The captions of the figures are braodly non-informative (e.g. table 2). The results presented in the experimental sections are not thoroughly discussed. Minor: citations does not follow the ICLR guideline (regarding when to use citep vs citet). In some places this makes the reading hard to parse.
- It would be beneficial to include additional datasets, ideally non-synthetic ones. However, using only the n-body dataset would be acceptable if comparisons with existing models are conducted in a highly thorough and rigorous manner.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: - Since the E(n)-equivariant spatial attention section primarily builds on existing work, it might be more suitable in the background section rather than being positioned as a main contribution.

- Effectiveness of Temporal Adjacency: In the charged N-body experiment, the temporal adjacency component does not seem to contribute as expected. The explanation that “edge attributes contain information about charges” suggests a possible shortcut in your model, potentially making it difficult to learn the self-attention projections. Testing with identity matrices for the query, key, and value projections could help clarify this observation.

- Ablation Study Clarifications: In the ablation table, the meaning of “Equiv=True/False” could be clarified, as it is currently difficult to interpret how equivariance is deactivated. Additionally, the rationale behind combining SchNet[6] with EGNN [5] could be explained further. Could you elaborate on this particular choice of spatial encoders. An ablation on this aspect could be interesting.

- Baselines and Comparisons: The main results (Table 2) are based on relatively simple baselines, such as LSTM without spatial encoding, EGNN without temporal encoding, MLP, and Linear (which are quite trivial baselines). Including a stronger baseline, perhaps by replacing the ETAL layer with an LSTM, could better highlight the impact of E(n)-equivariance in temporal encoding. Furthermore, why not comparing against SE(3)-Transformer[4], EGNN + SchNet, LieConv [3], LNNs [2], DynGCN [1], DynGFormer [7] and ESTAG [8]?

Other comments / questions:
- In Eq 1,  phi and psy are not described. If this equation is not part of your contribution, why not place it in the background section?
- In line 133, shouldn't it be attention instead of self-attention? because it is between node i and its neighbors j.
- In eq 6 shouldn't it be $h_i^{(l)}$ instead of $h_i^{(l+1)}$ in the SchNet ? You may consider presenting the general architecture of SET at the beginning of the method section to help the reader understand each component better.
- In line 209, is it E(n)-invariant or E(n)-equivariant ?
- Can you elaborate on why the projection matrices are fixed to the identity for the velocity self-attention ?
- While I understand the purpose of $\phi_{inf}$, it seems it is not describe in the paper. Is it the same as in EGNN [5] ?


References:

[1] Choi et al, "Dynamic graph convolutional networks with attention mechanism for rumor detection on social media"

[2] Cranmer et al; "Lagrangian neural networks"

[3] Finzi et al, "Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data"

[4] Fuchs et al,"Se(3)-transformers: 3d rototranslation equivariant attention networks"

[5] Satorras et al, "E(n) equivariant graph neural networks"

[6] Schütt et al, "Schnet – a deep learning architecture for molecules and materials"

[7] Yu et al, "Towards better dynamic graph learning: New architecture and unified library"

[8] Wu et al, "Equivariant spatio-temporal attentive graph networks to simulate physical dynamics"
soundness: 2
strengths: **Novelty**: The paper effectively extends E(n)-equivariance to temporal data through the Equivariant Temporal Attention Layer (ETAL) which is novel. The idea of leveraging self-attention mechanism across time is natural and very simple (which is positive).

**Theoretical Consistency**: The model is theoretically sound in maintaining E(n)-equivariance across both space and time.
summary: This paper introduces Spacetime E(n)-Transformer (SET), a E(n)-equivariant  spatio-temporal model suited to model spatio-temporal graphs. SET first updates the feature, position and velocity of each node at time t using E(n)-equivariant spatial attention by combining a SchNet continuous filter and a EGNN layer.  To account for temporal information,  the paper introduces the Equivariant Temporal Attention Layer (ETAL), which consists in a self-attention layer across time for node features and velocities, and a temporal neighborhood attention for positions. Finally, instead of having a fixed adjacency matrix, SET uses a temporal self-attention on adjacency matrices. SET is evaluated on the (charged) N-body problem and shows improvement over the selected baselines.
weaknesses: **Quality of Writing**: The paper’s organization is somewhat unclear, with overlapping notations that may make it challenging for readers to follow. Additionally, the main contribution is not clearly highlighted in the introduction and the method section, which could improve the paper’s overall clarity. 

**Experimental Evaluation**: The experimental setup might benefit from stronger baseline comparisons, more comprehensive ablations, and clearer explanations regarding certain model components. Further rigor in these areas would help substantiate the proposed method.

Please the questions below for more detailed comments / questions.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: The authors should consider a more rigorous experimental evaluation with recent state-of-the-art spatiotemporal transformer (not necessarily equivariant) and a broader set of datasets.
soundness: 2
strengths: The contributions are clearly presented and well-motivated.
summary: The paper introduces an E(n)-equivariant transformer architecture tailored for spatiotemporal graphs with fixed topologies, specifically complete graphs. The main objective is to embed space-time inductive biases to enhance performance compared to traditional space-time architectures. This model builds on the architecture proposed by Satarras et al., with notable additions such as encoding velocity and performing attention across time. The architecture is tested on the N-body problem dataset, where it demonstrates promising results.
weaknesses: The architecture appears explicitly tailored to the N-body problem, with velocity added as an input feature. A more general description of the architecture, beyond a single application, would strengthen its broader relevance.

Additionally, testing is limited to a single synthetic dataset, which restricts the impact and generalizability of the findings. Evaluating the model on additional datasets, such as those used in Cini et al. (Scalable Spatiotemporal Graph Neural Networks) or in "Long-Range Transformers for Dynamic Spatiotemporal Forecasting,"  by Grigsby et al., would provide a stronger empirical foundation.

The figures could be improved by enlarging the text and providing more detailed captions to aid comprehension.

The related work section and experimental evaluation would benefit from a broader discussion of recent spatiotemporal transformers, including those that may not explicitly include equivariance. Such additions would better contextualize the contribution. Relevant examples include:
Gao, Z., Shi, X., Wang, H., Zhu, Y., Wang, Y. (B.), Li, M., & Yeung, D.-Y. (2022). Earthformer: Exploring space-time transformers for earth system forecasting. Presented at NeurIPS 2022.
Zhao, Z., Chen, Z., Li, J., Xie, X., Chen, K., Wang, X., & Shi, G. STDM-transformer: Space-time dual multi-scale transformer network for skeleton-based action recognition, Neurocomputing.

The paper uses several well known works as building blocks. In this case an ablation study to weight the contribution of each module would make the novel contribution stand out from these works.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: See above.
soundness: 3
strengths: 1. The approach is interesting, with the authors designing a novel architecture to achieve equivariance in spatiotemporal graph prediction.

2. The paper is well-organized and easy to follow.
summary: This paper introduces a new model called Spacetime E(n)-Transformer (SET), primarily designed to process spatiotemporal graph data. The authors aim to address the limitations of existing spatiotemporal graph neural networks in preserving group symmetries, particularly in temporal and spatial dimensions. The authors utilize EGNN for equivariance and SchNet for invariance.
weaknesses: 1. Despite being interesting, the paper's major weakness is the lack of experimental validation. The authors need to demonstrate the method's effectiveness through comparisons with more baseline methods across multiple datasets.

2. More ablation studies are needed to demonstrate the effectiveness of each sub-module, which is currently lacking.

3. While the authors conducted hyperparameter search, reporting model performance under different hyperparameter settings would help better understand the model's capabilities.

In conclusion, I believe this paper requires further improvements to be suitable for publication.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: - Not sure how much sense it makes to introduce the message-passing framework in Section 3, equation 1, when the model operates over a fully-connected graph directly. One can view transformers as performing message passing, but typically the equation displayed in the paper is introduced when there is a sense of locality/neighbourhoods induced by an input graph. Since the graph is assumed to be fully connected may be better to use the standard inner-product notation used in Transformers.
- I may be misunderstanding something, but in line 167 the authors mention: ‘which is necessary to model the non-uniform topology of the graphs’, were the authors not assuming a fully-connected graph? In line 112 they mention: ‘we assume a priori that the graph is complete’. Later again in line 236-237: ‘ Since there is no explicit temporal adjacency matrix, we assume a fully connected temporal graph … such a fully connected network does not scale and, instead, we infer the edges of the temporal edges in our mode’. The writing is unclear and the statements are scattered around the text. Could the authors please clarify this and list all their assumptions both in terms of the spatial and temporal connectivity of the graphs in a single paragraph.
soundness: 2
strengths: - The authors combine existing neural network layers into a single architecture for improved performance on the n-body problem.
summary: The paper introduces a spatio-temporal graph transformer architecture which incorporates rotation, translation, and permutation equivariance inductive biases for the n-body problem. It assumes the input has a fixed-graph topology, and that the number of nodes in the graphs is fixed, nodes do not appear or dissapear as a function of time. Also, each node is associated with a temporal signal that evolves over time.
weaknesses: - The authors mainly focus on combining existing architectures and layers into a single architecture. They do not provide new theoretical insights for such a combination of layers, nor a purely new processing layer satisfying all required symmetries. This makes the contributions limited.
- There is not enough discussion on positional encodings. I would assume since the authors are using sinusoidal positional encodings that they are referring to positional encodings for the time dimension. However, there is also an extensive set of literature on positional encodings for graph transformers based, for instance, on the Graph Laplacian and other quantities that can be inferred from the graph topology. This are not discussed.
- Also, sinusoidal positional encodings introduce an absolute notion of position which break permutation invariance in the transformer attention operation. The abstract reads: ‘By imposing rotation, translation, and permutation equivariance inductive biases in both space and time’, which probably does not hold in time at least once you add the sinusoidals.
- Regarding Figure 1, the way in which operations are performed is likely suboptimal. Typically positional encodings are added at the very beginning. The Pre-LN transformer formulation is also more widely used than th Post-LN architecture presented in the original attention is all you need paper: meaning, it is more stable to normalize before attention and before the feedforward than after, see: ‘On Layer Normalization in the Transformer Architecture’.
- Experiments are very limited. The architecture seems to be particularly engineered for the n-body problem. There are other problems with similar symmetries but the authors do not carry out additional experiments. Other works in the literature typically use the n-body problem as a synthetic baseline but accompany it with additional experiments, such as for instance top tagging.
- The authors mention they use Optuna’s Bayesian-based hyper-parameter sweep, is this true for all baselines too or only for their model, otherwise results are not directly comparable + reference for Optuna is missing.
- Other previous works that are directly related to the authors line of research and that study the N-body problem are the following: Clifford Group Equivariant Neural Networks, Metric Learning for Clifford Group Equivariant Neural Networks, Neural message passing for quantum chemistry (NMP networks), Tensor Field Networks (which are mentioned in the related work but not added as a baseline). Should be included as baselines.
- The authors could also explore using a hybrid graph transformer framework such as in graph GPS: Recipe for a General, Powerful, Scalable Graph Transformer, which performs local-message passing and global attention in parallel.