SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Here we present our graph network architecture that closely follows Graber and Schwing [14], Kofinas et al. [24]. The model is formulated as a variational autoencoder [21, 37] with latent edge types that infers a latent graph structure. The encoder is tasked with predicting interactions between object pairs, while the decoder uses the sampled graph structure to make predictions. As mentioned in Section 3.3, our full architecture also integrates G-LoCS, i.e. the augmented node states include the predicted forces exerted at the target node, as well as the state of the auxiliary origin-node, expressed in the local frame of the target node.
Encoder Equations (20) to (22) describe the message passing steps of our graph network. In these equations, we process each timestep independently. Then, in Equations (23) and (24) we compute the evolution of edge embeddings over time with LSTMs [16], and in Equations (25) and (26) we estimate the posterior and the learned prior over our edges.
$$
\mathbf {h} _ {j, i} ^ {(1), t} = f _ {e} ^ {(1)} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \boxed {\mathbf {f} _ {j | i} ^ {t}}, \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) \tag {20}
$$
$$
\mathbf {h} _ {i} ^ {(1), t} = f _ {v} ^ {(1)} \left(g _ {v} ^ {(1)} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(1), t}\right) \tag {21}
$$
$$
\mathbf {h} _ {j, i} ^ {(2), t} = f _ {e} ^ {(2)} \left(\left[ \mathbf {h} _ {i} ^ {(1), t}, \mathbf {h} _ {j, i} ^ {(1), t}, \mathbf {h} _ {j} ^ {(1), t} \right]\right) \tag {22}
$$
$$
\mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t} = \operatorname {L S T M} _ {\text {p r i o r}} \left(\mathbf {h} _ {j, i} ^ {(2), t}, \mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t - 1}\right) \tag {23}
$$
$$
\mathbf {h} _ {(j, i), \text {e n c}} ^ {t} = \operatorname {L S T M} _ {\text {e n c}} \left(\mathbf {h} _ {j, i} ^ {(2), t}, \mathbf {h} _ {(j, i), \text {e n c}} ^ {t + 1}\right) \tag {24}
$$
$$
p _ {\phi} \left(\boldsymbol {z} ^ {t} | \mathbf {x} ^ {1: t}, \boldsymbol {z} ^ {1: t - 1}\right) = \operatorname {s o f t m a x} \left(f _ {\text {p r i o r}} \left(\mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t}\right)\right) \tag {25}
$$
$$
q _ {\phi} \left(\boldsymbol {z} _ {j, i} ^ {t} | \mathbf {x}\right) = \operatorname {s o f t m a x} \left(f _ {\mathrm {e n c}} \left(\left[ \mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t}, \mathbf {h} _ {(j, i), \mathrm {e n c}} ^ {t} \right]\right)\right) \tag {26}
$$
The functions  $f_{e}^{(1)}, f_{v}^{(1)}, f_{e}^{(2)}, g_{v}^{(1)}, f_{\mathrm{prior}}, f_{\mathrm{enc}}$  denote MLPs.
Decoder The decoder samples  $z_{(j,i)}^{t}$  using Gumbel-Softmax [27, 19]. The following equations formalize a message passing scheme performed for the current timestep, and another one performed for the hidden node states. Both are used to update the hidden node states, and to make predictions for the next timestep. As mentioned in Section 2, we make predictions in the local coordinate frame
of each node. Thus, we perform an inverse transformation for each node to transform the predictions back to the global coordinate frame.
$$
\mathbf {m} _ {j, i} ^ {t} = \sum_ {k} z _ {(j, i), k} ^ {t} f ^ {k} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \boxed {\mathbf {f} _ {j | i} ^ {t}}, \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) \tag {27}
$$
$$
\mathbf {m} _ {i} ^ {t} = f _ {v} ^ {(3)} \left(g _ {v} ^ {(3)} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {m} _ {j, i} ^ {t}\right) \tag {28}
$$
$$
\mathbf {h} _ {j, i} ^ {t} = \sum_ {k} z _ {(j, i), k} ^ {t} g ^ {k} \left(\left[ \mathbf {h} _ {j} ^ {t}, \mathbf {h} _ {i} ^ {t} \right]\right) \tag {29}
$$
$$
\mathbf {n} _ {i} ^ {t} = \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {(j, i)} ^ {t} \tag {30}
$$
$$
\mathbf {h} _ {i} ^ {t + 1} = \operatorname {G R U} \left(\left[ \mathbf {n} _ {i} ^ {t}, \mathbf {m} _ {\mathbf {i}} ^ {t} \right], \mathbf {h} _ {i} ^ {t}\right) \tag {31}
$$
$$
\boldsymbol {\mu} _ {i} ^ {t + 1} = \mathbf {x} _ {i} ^ {t} + \mathbf {R} _ {i} ^ {t} \cdot f _ {v} ^ {(4)} \left(\mathbf {h} _ {i} ^ {t + 1}\right) \tag {32}
$$
$$
p \left(\mathbf {x} _ {i} ^ {t + 1} \mid \mathbf {x} ^ {1: t}, \boldsymbol {z} ^ {1: t}\right) = \mathcal {N} \left(\boldsymbol {\mu} _ {i} ^ {t + 1}, \sigma^ {2} \mathbf {I}\right) \tag {33}
$$
Our GRU [1] is identical to the one used in [23].
Training Our full VAE model is trained by minimizing the negative Evidence Lower Bound (ELBO), which comprises the reconstruction loss of the predicted trajectories (positions and velocities) and the KL divergence.
$$
\mathcal {L} (\phi , \theta) = \mathbb {E} _ {q _ {\phi} (\boldsymbol {z} | \mathbf {x})} [ \log p _ {\theta} (\mathbf {x} | \boldsymbol {z}) ] - \mathrm {K L} [ q _ {\phi} (\boldsymbol {z} | \mathbf {x}) | | p _ {\phi} (\boldsymbol {z} | \mathbf {x}) ] \tag {34}
$$
Following Graber and Schwing [14], the reconstruction loss and the KL divergence take the following form:
$$
\mathbb {E} _ {q _ {\phi} (\boldsymbol {z} | \mathbf {x})} [ \log p _ {\theta} (\mathbf {x} | \boldsymbol {z}) ] = - \sum_ {i} \sum_ {t} \frac {| | \mathbf {x} _ {i} ^ {t} - \boldsymbol {\mu} _ {i} ^ {t} | |}{2 \sigma^ {2}} + \frac {1}{2} \log (2 \pi \sigma^ {2}), \tag {35}
$$
$$
\operatorname {K L} \left[ q _ {\phi} (\boldsymbol {z} | \mathbf {x}) \right\rvert   \left| p _ {\phi} (\boldsymbol {z} | \mathbf {x}) \right] = \sum_ {t = 1} ^ {T} \left(\mathbb {H} \left(q _ {\phi} \left(\boldsymbol {z} _ {j i} ^ {t} \mid \mathbf {x}\right)\right) - \sum_ {\boldsymbol {z} _ {j i} ^ {t}} q _ {\phi} \left(\boldsymbol {z} _ {j i} ^ {t} \mid \mathbf {x}\right) \log p _ {\phi} \left(\boldsymbol {z} _ {j i} ^ {t} \mid \mathbf {x} ^ {1: t}, \boldsymbol {z} ^ {1: t - 1}\right)\right), \tag {36}
$$
where  $\mathbb{H}$  denotes the entropy operator. In all experiments, we set the variance  $\sigma^2 = 10^{-5}$ .
We train Aether using Adam [21]. Unless stated otherwise, in all experiments, we use a learning rate of 5e-4.
A.1.3 Aether architecture in Lorentz force field setting
The Lorentz force field setting, proposed by [10] uses only a single timestep as input and the task is to predict the positions for a single timestep in the future. Thus, we have to modify our architecture for this setting. This way, we also ensure a fairer comparison with other methods. We do not use an NRI [22] or dNRI [14] backbone in this setting.
$$
\mathbf {h} _ {j, i} ^ {(1)} = f _ {e} ^ {(1)} \left(\left[ \mathbf {v} _ {j \mid i}, \mathbf {f} _ {j \mid i}, \mathbf {v} _ {i \mid i}, \mathbf {f} _ {i \mid i}, q _ {i} q _ {j}, \| \mathbf {r} _ {j, i} \| _ {2} \right]\right) \tag {37}
$$
$$
\mathbf {h} _ {i} ^ {(1)} = f _ {v} ^ {(1)} \left(g _ {v} \left(\left[ \mathbf {v} _ {i | i}, \underline {{\mathbf {f}}} _ {i | i} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(1)}\right) \tag {38}
$$
$$
\mathbf {h} _ {j, i} ^ {(l)} = f _ {e} ^ {(l)} \left(\left[ \mathbf {h} _ {i} ^ {(l - 1)}, \mathbf {h} _ {j, i} ^ {(l - 1)}, \mathbf {h} _ {j} ^ {(l - 1)} \right]\right) \tag {39}
$$
$$
\mathbf {h} _ {i} ^ {(l)} = f _ {v} ^ {(l)} \left(\mathbf {h} _ {i} ^ {(l - 1)} + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(l)}\right) \tag {40}
$$
$$
\hat {\mathbf {p}} _ {i} = \mathbf {p} _ {i} + \mathbf {R} _ {i} \cdot f _ {o} \left(\mathbf {h} _ {i} ^ {L}\right) \tag {41}
$$
with  $l \in \{2, \dots, L\}$ . We use 4 layers, i.e.  $L = 4$ . The functions  $f_{e}^{(l)}$  denote 2-layer MLPs with SiLU [36] activations after each layer. The functions  $f_{v}^{(l)}$  denote 2-layer MLPs with SiLU activations in-between, and doubling the dimensionality in-between. Finally,  $g_{v}$  denotes a linear layer, while  $f_{o}$  denotes a 3-layer MLP with SiLU activations in-between. Following [10], we use a hidden dimension of 64. For this experiment, we use a learning rate of  $1e - 3$ .
$$
f _ {v} ^ {(l)} = \left\{\text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \right\} \tag {42}
$$
$$
f _ {e} ^ {(l)} = \left\{\text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \rightarrow \text {S i L U} \right\} \tag {43}
$$
$$
f _ {o} = \left\{\text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \right\} \tag {44}
$$
For the neural field, we use the input positions and velocities as input, as well as the particle charges, i.e.  $\mathbf{f} = f(\mathbf{p},\mathbf{u},q)$ . We do not use any input encoding for positions or velocities in this setting, but we use an embedding for the charges that maps them to 16 dimensions. We concatenate the charge embeddings with positions and velocities and feed them as input to the neural field. The neural field is a 3-layer MLP with SiLU activations in-between. We use a hidden dimension of 32 in the neural field.
$$
f = \{\text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \} \tag {45}
$$
A.2 G-LoCS
$$
\mathbf {h} _ {j, i} ^ {(1), t} = f _ {e} ^ {(1)} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) \tag {46}
$$
$$
\mathbf {h} _ {i} ^ {(1), t} = f _ {v} ^ {(1)} \left(g _ {v} ^ {(1)} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(1), t}\right) \tag {47}
$$
$$
\mathbf {h} _ {j, i} ^ {(2)} = f _ {e} ^ {(2)} \left(\left[ \mathbf {h} _ {i} ^ {(1), t}, \mathbf {h} _ {j, i} ^ {(1), t}, \mathbf {h} _ {j} ^ {(1), t} \right]\right) \tag {48}
$$
$$
\mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t} = \operatorname {L S T M} _ {\text {p r i o r}} \left(\mathbf {h} _ {j, i} ^ {(2), t}, \mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t - 1}\right) \tag {49}
$$
$$
\mathbf {h} _ {(j, i), \text {e n c}} ^ {t} = \operatorname {L S T M} _ {\text {e n c}} \left(\mathbf {h} _ {j, i} ^ {(2), t}, \mathbf {h} _ {(j, i), \text {e n c}} ^ {t + 1}\right) \tag {50}
$$
$$
p _ {\phi} \left(\boldsymbol {z} ^ {t} | \mathbf {x} ^ {1: t}, \boldsymbol {z} ^ {1: t - 1}\right) = \operatorname {s o f t m a x} \left(f _ {\text {p r i o r}} \left(\mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t}\right)\right) \tag {51}
$$
$$
\left. q _ {\phi} \left(\boldsymbol {z} _ {j, i} ^ {t} \mid \mathbf {x}\right) = \operatorname {s o f t m a x} \left(f _ {\mathrm {e n c}} \left(\left[ \mathbf {h} _ {(j, i), \text {p r i o r}} ^ {t}, \mathbf {h} _ {(j, i), \mathrm {e n c}} ^ {t} \right]\right)\right) \right. \tag {52}
$$
$$
\mathbf {m} _ {j, i} ^ {t} = \sum_ {k} z _ {(j, i), k} ^ {t} f ^ {k} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) \tag {53}
$$
$$
\mathbf {m} _ {i} ^ {t} = f _ {v} ^ {(3)} \left(g _ {v} ^ {(3)} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \mathbf {v} _ {\mathcal {O} | i} ^ {t} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {m} _ {j, i} ^ {t}\right) \tag {54}
$$
$$
\mathbf {h} _ {j, i} ^ {t} = \sum_ {k} z _ {(j, i), k} ^ {t} g ^ {k} \left(\left[ \mathbf {h} _ {j} ^ {t}, \mathbf {h} _ {i} ^ {t} \right]\right) \tag {55}
$$
$$
\mathbf {n} _ {i} ^ {t} = \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {(j, i)} ^ {t} \tag {56}
$$
$$
\mathbf {h} _ {i} ^ {t + 1} = \operatorname {G R U} \left(\left[ \mathbf {n} _ {i} ^ {t}, \mathbf {m} _ {\mathbf {i}} ^ {t} \right], \mathbf {h} _ {i} ^ {t}\right) \tag {57}
$$
$$
\boldsymbol {\mu} _ {i} ^ {t + 1} = \mathbf {x} _ {i} ^ {t} + \mathbf {R} _ {i} ^ {t} \cdot f _ {v} ^ {(4)} \left(\mathbf {h} _ {i} ^ {t + 1}\right) \tag {58}
$$
$$
p \left(\mathbf {x} _ {i} ^ {t + 1} \mid \mathbf {x} ^ {1: t}, \boldsymbol {z} ^ {1: t}\right) = \mathcal {N} \left(\boldsymbol {\mu} _ {i} ^ {t + 1}, \sigma^ {2} \mathbf {I}\right) \tag {59}
$$
A.3 Source oracle
The source oracle modifies LoCS [24] to use virtual nodes. Since graph networks are permutation invariant, we cannot just include the sources as nodes of the graph and perform message passing,
as the network would not be able to distinguish particles from sources. Thus, we treat the sources separately in the message passing so that the network can identify them. More specifically, we introduce a new message function that computes field source  $\rightarrow$  particle messages. Furthermore, we introduce a separate aggregation function in the update step that only aggregates the messages from field sources. We denote the set of field sources as  $S$ . The state of a field source  $s \in S$  is denoted as  $\mathbf{v}_s$ , while the same state expressed in the local coordinate frame of node  $i$  is denoted as  $\mathbf{v}_{s|i}$ . The source oracle graph network is defined as follows:
$$
\mathbf {h} _ {j, i} ^ {t} = f _ {e} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t} \right]\right), \tag {60}
$$
$$
\mathbf {h} _ {s, i} ^ {t} = f _ {s} \left(\left[ \mathbf {v} _ {s | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t} \right]\right), \tag {61}
$$
$$
\Delta \mathbf {x} _ {i | i} ^ {t + 1} = f _ {v} \left(g _ {v} \left(\mathbf {v} _ {i | i} ^ {t}\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {t} + \frac {1}{| \mathcal {S} |} \sum_ {j \in \mathcal {S}} \mathbf {h} _ {s, i} ^ {t}\right), \tag {62}
$$
where  $f_{s}$  is an MLP. We predict future states for all the "observable" particles, but not for the "source" particles.
A.4 Parallel aether architecture
$$
\mathbf {h} _ {j, i} ^ {(1)} = f _ {e} ^ {(1)} \left(\left[ \mathbf {v} _ {j | i}, \mathbf {v} _ {i | i}, q _ {i} q _ {j}, \| \mathbf {r} _ {j, i} \| _ {2} \right]\right) \tag {63}
$$
$$
\mathbf {h} _ {i} ^ {(1)} = f _ {v} ^ {(1)} \left(g _ {v} \left(\mathbf {v} _ {i \mid i}\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(1)}\right) \tag {64}
$$
$$
\mathbf {h} _ {j, i} ^ {(l)} = f _ {e} ^ {(l)} \left(\left[ \mathbf {h} _ {i} ^ {(l - 1)}, \mathbf {h} _ {j, i} ^ {(l - 1)}, \mathbf {h} _ {j} ^ {(l - 1)} \right]\right) \tag {65}
$$
$$
\mathbf {h} _ {i} ^ {(l)} = f _ {v} ^ {(l)} \left(\mathbf {h} _ {i} ^ {(l - 1)} + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {(l)}\right) \tag {66}
$$
$$
\hat {\mathbf {p}} _ {i} = \mathbf {p} _ {i} + \mathbf {R} _ {i} \cdot f _ {o} \left(\mathbf {h} _ {i} ^ {L}\right) + \boxed {\mathbf {f} _ {i}} \tag {67}
$$
A.5 Aether with EGNN backbone
Our method is agnostic to the choice of equivariant graph network; we expect that it would be beneficial for a number of strictly equivariant networks. To test this hypothesis, we combine EGNN [41] with our method; starting from the velocity formulation of EGNN, we modify the message and velocity equations to incorporate the predicted forces for each node, as follows:
$$
\mathbf {m} _ {j, i} = \phi_ {e} \left(\mathbf {h} _ {i} ^ {l}, \mathbf {h} _ {j} ^ {l}, \left\| \mathbf {p} _ {j} ^ {l} - \mathbf {p} _ {i} ^ {l} \right\| _ {2} ^ {2}, a _ {j, i}, \boxed {\mathbf {f} _ {i}}, \boxed {\mathbf {f} _ {j}}\right), \tag {68}
$$
$$
\mathbf {u} _ {i} ^ {l + 1} = \phi_ {v} \left(\mathbf {h} _ {i} ^ {l}, \boxed {\mathbf {f} _ {i}}\right) \mathbf {u} _ {i} ^ {l} + C \sum_ {j \neq i} \left(\mathbf {p} _ {j} ^ {l} - \mathbf {p} _ {i} ^ {l}\right) \cdot \phi_ {x} \left(\mathbf {m} _ {j, i}\right), \tag {69}
$$
$$
\mathbf {p} _ {i} ^ {l + 1} = \mathbf {p} _ {i} ^ {l} + \mathbf {u} _ {i} ^ {l + 1}, \tag {70}
$$
$$
\mathbf {m} _ {i} = \sum_ {j \in \mathcal {N} (i)} \mathbf {m} _ {j i}, \tag {71}
$$
$$
\mathbf {h} _ {i} ^ {l + 1} = \phi_ {h} \left(\mathbf {h} _ {i} ^ {l}, \mathbf {m} _ {i}\right). \tag {72}
$$
The remaining EGNN components remain unaltered.
A.6 Computing resources
All experiments were performed on single GPUs. We used 2 different GPU models, namely the Nvidia RTX 2080 Ti, and Nvidia GTX 1080 Ti. Our source code was written in PyTorch [32], version 1.4.0, and CUDA 10.0.
B.1 Electrostatic field
Kipf et al. [23] introduced a dataset of interacting charged particles. Charged particles interact via electrostatic Coulomb forces. We assume a set of  $N$  particles, and each particle has a position  $\mathbf{p}_i^t\in \mathbb{R}^D$  and a charge  $q_{i}\in \mathbb{R}$ . The force  $\mathbf{f}_{j,i}^{t}$  exerted from particle  $j$  to particle  $i$  is computed as follows:
$$
\mathbf {p} _ {j, i} ^ {t} = \mathbf {p} _ {j} ^ {t} - \mathbf {p} _ {i} ^ {t}, \tag {73}
$$
$$
\hat {\mathbf {p}} _ {j, i} ^ {t} = \frac {\mathbf {p} _ {j , i} ^ {t}}{\left\| \mathbf {p} _ {j , i} ^ {t} \right\|}, \tag {74}
$$
$$
\mathbf {f} _ {j, i} ^ {t} = C \cdot q _ {i} q _ {j} \frac {\hat {\mathbf {p}} _ {j , i} ^ {t}}{\left\| \mathbf {p} _ {j , i} ^ {t} \right\| ^ {2}}. \tag {75}
$$
Since forces only depend on positions at the current timestep, in the following equations, we omit the time indices to reduce clutter. The total force exerted at particle  $i$  is:
$$
\mathbf {f} _ {i} = \sum_ {j = 1, j \neq i} ^ {N} \mathbf {f} _ {j, i} = C \cdot q _ {i} \sum_ {j = 1, j \neq i} ^ {N} q _ {j} \frac {\hat {\mathbf {p}} _ {j , i}}{\| \mathbf {p} _ {j , i} \| ^ {2}}. \tag {76}
$$
The electric field is a vector field, whose value at the test position  $\mathbf{p}_i$  assumes a positive test charge  $q_{i} = 1$ , and is defined as:
$$
\mathbf {E} _ {j, i} = \frac {\mathbf {f} _ {j , i}}{q _ {i}} = C \cdot q _ {j} \frac {\hat {\mathbf {p}} _ {j , i}}{\| \mathbf {p} _ {j , i} \| ^ {2}}, \tag {77}
$$
$$
\mathbf {E} _ {i} = \sum_ {j = 1, j \neq i} ^ {N} \mathbf {E} _ {j, i} = C \cdot \sum_ {j = 1, j \neq i} ^ {N} q _ {j} \frac {\hat {\mathbf {p}} _ {j , i}}{\left\| \mathbf {p} _ {j , i} \right\| ^ {2}}. \tag {78}
$$
Our first experiment aims to study the effect of static fields, i.e. a single field across all train, validation, and test simulations. We extend the charged particles dataset by adding a number of immovable sources. Overall, these sources act like regular particles, exerting forces on the observable particles, except we ignore any forces exerted to them, and fix their positions and velocities to zero. We use  $N = 5$  "observable" particles and  $M = 20$  "source" particles. In all experiments, we assume unit charges,  $q_{i} = \pm 1$ , and  $C = 1$ . The probabilities of positive or negative charges are equal. Then, the forces and the electric field can be simplified as:
$$
\mathbf {f} _ {j, i} = \operatorname {s i g n} \left(q _ {i} q _ {j}\right) \frac {\hat {\mathbf {p}} _ {j , i}}{\| \mathbf {p} _ {j , i} \| ^ {2}}, \tag {79}
$$
$$
\mathbf {E} _ {i} = \sum_ {j = 1, j \neq i} ^ {N} \operatorname {s i g n} \left(q _ {j}\right) \frac {\hat {\mathbf {p}} _ {j , i}}{\left\| \mathbf {p} _ {j , i} \right\| ^ {2}}. \tag {80}
$$
The net force exerted at a particle  $i\in \{1,\dots ,N\}$  is computed as:
$$
\mathbf {f} _ {i} = \sum_ {j = 1, j \neq i} ^ {N + M} \mathbf {f} _ {j, i} = \underbrace {\sum_ {j = 1 , j \neq i} ^ {N} \mathbf {f} _ {j , i}} _ {\text {p a r t i c l e s}} +. \underbrace {\sum_ {j = N + 1} ^ {N + M} \mathbf {f} _ {j , i}} _ {\text {f i e l d}} \tag {81}
$$
Following Satorras et al. [41], Fuchs et al. [12], Kofinas et al. [24], we remove virtual borders that cause elastic collisions. We generate a dataset of 50,000 simulations for training, 10,000 for validation and 10,000 for testing. The datasets contain only the positions and velocities for the "observable" particles, while the field sources are only used for visualization. Following Kipf et al. [23], each simulation lasts for 49 timesteps. During inference, we use the first 29 steps as input and predict the remaining 20 steps.
Figure 7: Visualization of the static field in the electrostatic field setting
B.2 Traffic scenes - inD
InD [4] is a real-world traffic scenes dataset that comprises trajectories of pedestrians, vehicles, and cyclists. It contains 33 recordings, recorded at 4 different locations in Aachen, Germany. We hypothesize that discovering a latent traffic force field will be beneficial for trajectory forecasting in traffic scenes. For simplicity, we focus on static field discovery in traffic scenes. We create a subset that contains scenes from a single location. Namely, we choose "Frankenburg, Aachen", since it is the location with most interactions in the dataset. The subset corresponds to 12 recordings; we use 8 for training, 2 for validation, and 2 for testing. We follow a similar experimental setting with Graber and Schwing [14], Kofinas et al. [24]. We divide each scene into 18-step sequences. We use the first 6 time steps as input and predict the next 12 time steps.
B.3 Gravitational n-body dataset
In this experiment, we study the influence of dynamic fields, i.e. fields that are different across simulations. Similar to the electrostatic field setting, we extend the gravitational n-body dataset by Brandstetter et al. [5] by adding gravitational sources. The equation that describes the forces is similar to Equation (75). Namely, we have
$$
\mathbf {f} _ {j, i} ^ {t} = C \cdot m _ {i} m _ {j} \frac {\hat {\mathbf {p}} _ {j , i} ^ {t}}{\left\| \mathbf {p} _ {j , i} ^ {t} \right\| ^ {2}}, \tag {82}
$$
where  $m_{i}, m_{j}$  are the particle masses. We create a dataset of 50,000 simulations for training, 10,000 for validation and 10,000 for testing. We use  $N = 5$  particles and  $M = 1$  source. We set the masses of particles to  $m_{p} = 1$ , while the source has a mass of  $m_{s} = 10$ . Similarly to the electrostatic field experiment, the datasets contain only the positions and velocities for the "observable" particles, while the field source is only used for visualization. We generate trajectories of 49 timesteps. We use the first 44 timesteps as input and predict the remaining 5 steps. All other dataset details are identical to Brandstetter et al. [5].
B.3.1 2D gravitational n-body dataset
We also experiment with a smaller variant of the dynamic gravitational fields, using a 2D setting. We create a dataset of 5,000 simulations for training, 1,000 for validation and 1,000 for testing. All other
dataset details are the same with the full 3D dataset. We report results in Figure 21. We showcase predicted trajectories in Appendix D.3, and the learned fields in Appendix D.3.1.
C Extra experiments
C.1 Alternative equivariant network backbones
To further test the applicability of our method, we combine it with different equivariant graph network backbones. First, we combine our method with GMN [18]. Similar to Equations (4) and (5), we concatenate the predicted forces for each node with the message  $\mathbf{Z}_{ji}$  in  $[\mathbf{Z}_{ji},\mathbf{f}_i,\mathbf{f}_j]$ . The formulation above is further motivated by SGNN [15], which extends GMN by including gravity as an external force term, as well as object-aware information (see appendix A.3 in [15] for a comparison). In our case, we replace the gravity term with the predicted forces per node. Thus, instead of  $[\mathbf{Z}_{ji},\mathbf{g}]$ , we have  $[\mathbf{Z}_{ji},\mathbf{f}_i,\mathbf{f}_j]$ .
We train and evaluate GMN on the Lorentz force field setting. We then add the force terms using the formulation above. We show the results in the table below. Indeed, using Aether greatly enhances the performance of GMN, which further enhances our hypothesis.
Table 4: Ablation study on the choice of equivariant GNN backbone. Position prediction MSE on Lorentz force field.
Method	MSE (↓)
GMN [18]	0.0365
GMN+Aether (ours)	0.0261