SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Equivariant graph networks The seminal works of [8, 9, 53] introduced equivariant convolutional neural networks and demonstrated effectiveness, robustness, and increased parameter efficiency. Recently, many works have proposed equivariant graph networks [43, 47, 12, 50, 41, 24, 5, 26, 18]. Walters et al. [50] propose rotationally equivariant continuous convolutions for trajectory prediction. Satorras et al. [41] propose a computationally efficient equivariant graph network that leverages invariant euclidean distances between node pairs. Kofinas et al. [24] introduce roto-translated local coordinate frames for all objects in an interacting system and propose equivariant local coordinate frame graph networks. Brandstetter et al. [5] generalize equivariant graph networks using steerable MLPs [47] and incorporate geometric and physical information in message passing. Equivariant graph networks differ from our work since they cannot capture non-local information, while our work disentangles equivariant local interactions from global effects and captures them both.
Approximate equivariance Recently, a number of works has proposed to shift away from strict equivariance, in what Wang et al. [51] termed as approximate equivariance. Wang et al. [51] propose approximately equivariant networks for dynamical systems, by relaxing equivariance constraints in group convolutions and steerable convolutions. van der Ouderaa et al. [48] propose to relax strict equivariance by interpolating between equivariant and non-equivariant operations, using nonstationary kernels that also depend on the absolute input group element. Romero and Lohit [38] propose Partial G-CNNs that learn layer-wise partial equivariances from data. We note that even though approximately equivariant networks share similarities with our work, our notion of disentangled equivariance is conceptually different. That is because related work uses the term approximate equivariance to denote that equivariance is "broken" due to noise or imperfections, while our work disentangles the system dynamics that are actually equivariant, from the global field effects that are not, and in fact, might be unaffected by such transformations. Further, to the best of our knowledge, approximate equivariance has only been studied in the context of convolutional networks, not in the context of graph networks and interacting systems. Tangentially, Han et al. [15] propose subequivariant graph networks, and relax equivariance to subequivalence by considering external fields like gravity. However, they assume a priori known fields that do not require to be inferred by the model.
Neural fields Neural fields have recently exploded in popularity in 3D computer vision, popularized by NeRF [29]. Since MLPs are universal function approximators [17], neural fields parameterized by MLPs can, in principle, encode continuous signals at arbitrary resolution. However, neural
networks can suffer from "spectral bias" [34, 2], i.e. they are biased to fit functions with low spatial frequency. To address this issue, a number of solutions have been proposed. Tancik et al. [46] leverage Neural Tangent Kernel (NTK) theory and propose Random Fourier Features (RFF), showing that they can overcome the spectral bias. They also show that RFF are a generalization of positional encodings, popularized in recent years in natural language processing by Transformers [49]. Concurrently, Sitzmann et al. [45] proposed SIREN, neural networks with sinusoidal activation functions. While neural fields have been used extensively in computer vision problems including 3D scene reconstruction [31, 28] and differentiable rendering [44, 29], they have not seen wide usage in dynamical systems. Notably, Raissi et al. [35] proposed Physics-Informed Neural Networks (PINNs), neural PDE solvers based on neural fields. Finally, Dupont et al. [11] and Zhuang et al. [56] propose generative models of neural fields.
We evaluate our proposed method, Aether, on settings that include static as well as dynamic fields. First, we explore 2D charged particles that evolve under the effect of a static electrostatic field, as well as 3D particles that evolve under a Lorentz force field [10]. Then, we evaluate our method on a subset of inD [4] that contains a single location, and thus a static field as well. Finally, we explore 3D gravitational n-body problems [5] with dynamic fields. Our code, data, and models will be open-sourced online<sup>1</sup>.
In most experiments, we compare our method against dNRI [14] and LoCS [24], two state-of-the-art networks for sequence-to-sequence trajectory forecasting, as well as G-LoCS. DNRI [14] is a graph network operating in global coordinates, and is, in principle, able to uncover both the global and the local dynamics. It is formulated as a VAE [21, 37] with latent edge types and explicitly infers a latent graph structure. LoCS [24], on the other hand, operates in local coordinates, and is, thus, unable to uncover the global dynamics. Finally, G-LoCS is in principle able to model both local and global dynamics effectively. For all methods, we use their publicly available source code.
Our architecture and experimental setup closely follow Graber and Schwing [14], Kofinas et al. [24]. Unless specified differently, our neural field has a hidden size of 512. In charged particles and in n-body problems, we only use positions as input to the neural field, while in traffic scenes we also use orientations. The full implementation details are presented in Appendix A.1.2. In all settings, we report the mean squared error (MSE) of positions and velocities over time. Here we demonstrate indicative visualizations, and provide more extensive qualitative results in Appendix D.
For the Lorentz force field experiment, we use the official source code from ClofNet [10], and follow their exact setup. We compare our method against SE(3) Transformers [12], EGNN [41], and ClofNet [10]. We evaluate methods using the mean squared error between predicted and groundtruth positions. Since this setting is not a sequence-to-sequence task, we use a simplified network architecture without a VAE, and following baselines, we make sure that the number of parameters of our model is approximately equal to other methods. The full implementation details are presented in Appendix A.1.3.
(c)
Figure 4: Results on (a) electrostatic field, (b) inD, and (c) gravity.
<https://github.com/mkofinas/aether>
5.1 Electrostatic field
First, we study the effect of static fields, i.e. a single field across all train, validation, and test simulations. We extend the charged particles dataset from Kipf et al. [23] by adding a number of immovable sources. These sources act like regular particles, exerting forces on the observable particles, except we ignore any forces exerted to them, and keep their positions fixed. We use  $M = 20$  "source" particles and  $N = 5$  "observable" particles. We generate 50,000 simulations for training, 10,000 for validation and 10,000 for testing. Following Kipf et al. [23], each simulation lasts for 49 timesteps. During
Predicted Field
Figure 5: Learned Field (left) in electrostatic field setting compared to groundtruth (right).
Groundtruth Field
inference, we use the first 29 steps as input and predict the remaining 20 steps. The full dataset details are presented in Appendix B.1.
We compare our method against dNRI, LoCS, and G-LoCS. We plot MSE in Figure 4a and  $L_{2}$  errors in Figure 18, and visualize the learned field in Figure 5. We showcase predicted trajectories in Figure 9 in Appendix D.1. We observe that equivariant methods like LoCS perform poorly, while the approximately equivariant G-LoCS performs much better than equivariant and non-equivariant methods. Aether outperforms all other methods, demonstrating that it can disentangle equivariance. Furthermore, as shown in Figure 5, and Figure 10 in Appendix D.1.1, Aether can effectively discover the underlying field.
5.2 Lorentz force field
Du et al. [10] introduced a dataset of 3D charged particles evolving under the influence of a Lorentz force field. Each simulation contains 20 particles. We use the official source code and follow the exact experimental setup with Du et al. [10]. We show quantitative results in Table 1. Our method can clearly outperform all other methods by a large margin, reducing the error by  $48.6\%$ . We also note that our method has fewer parameters than ClofNet, and is thus more efficient.
Table 1: Position prediction MSE on Lorentz force field. Results marked with  $\dagger$  were taken from ClofNet [10].
Method	MSE (↓)	No. parameters
GNN †	0.0908	104,387
SE(3) Transformer † [12]	0.1438	1,763,134
EGNN † [41]	0.0368	134,020
ClofNet † [10]	0.0251	160,964
Aether (ours)	0.0129	132,822
5.3 Traffic scenes
Next, we study the effectiveness of static field discovery in traffic scenes. We use inD [4], a dataset with real-world traffic scenes that comprises trajectories of pedestrians, vehicles, and cyclists. We create a subset that contains scenes from a single location. The full dataset details are presented in Appendix B.2. We divide scenes into 18-step sequences; we use the first 6 time steps as input and predict the next 12 time steps. We plot MSE in Figure 4b and  $L_{2}$  errors in Figure 19, and visualize the learned field in Figure 6, and in Figure 15 in Appendix D.2.1. Since the learned field is a function of positions and orientations, we only visualize it for 4 discrete orientations, namely the group  $C_4 = \{0, \frac{\pi}{2}, \pi, \frac{3\pi}{2}\}$ . We showcase predictions in Figure 11 in Appendix D.2. Again, Aether outperforms all other methods. The discovered field, while hard to interpret, shows high activations that coincide with road locations and directions, indicating that it can guide objects through the topology of the road network.
Figure 6: Discovered field on inD [4]. For clarity, we only visualize the field for discrete input orientations in  $\mathrm{C}_4 = \left\{0, \frac{\pi}{2}, \pi, \frac{3\pi}{2}\right\}$ .
5.4 Gravitational field
We now study the task of dynamic field discovery, i.e. fields that are different across simulations. We extend the gravity dataset by Brandstetter et al. [5] by adding gravitational sources. 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's mass is  $m_{s} = 10$ . We generate trajectories of 49 timesteps. We use the first 44 steps as input and predict the remaining 5 steps. We plot the MSE in Figure 4c and  $L_{2}$  errors in Figure 20. We observe that once again, Aether clearly outperforms other methods.
5.5 Ablation experiments
Table 2: (a) Ablation study on the importance of the learned field. (b) Ablation study on the importance of a sequential architecture. (c) Ablation study on the choice of equivariant GNN backbone.
(a) Electrostatic field	(b) Lorentz force field	(c) Lorentz force field
Method	MSE@10 (↓)	Method	MSE (↓)	Method	MSE (↓)
Particle Oracle	0.1847	LoCS [24]	0.0238	EGNN [41]	0.0368
Force Oracle	0.1883	Aether	0.0129	EGNN+Aether	0.0254
Aether	0.2015	Parallel Aether	0.0211
Significance of discovered field In a simulated environment like the electrostatic field setting, we have access to the groundtruth fields and the sources that generate them. We leverage the simulator to study the significance of the discovered field in the task of trajectory forecasting, and establish an upper bound to our performance. To that end, we create two "oracle" models that have access to the groundtruth information, a force oracle and a source oracle. The force oracle is identical to Aether, but uses the groundtruth forces from the simulator instead of predicting them with a neural field. The source oracle assumes knowledge of the "field sources". Thus, there is no longer need for disentanglement, and the problem is strictly equivariant again. We include the sources as virtual nodes in the graph, add include edges from the sources to the particles. We describe this oracle in detail in Appendix A.3. We show the MSE in Table 2a and plot the MSE and  $L_{2}$  errors in Figure 22 in Appendix E.5. We observe that Aether closely follows the two oracle models, and is on par, for roughly 10 steps. This demonstrates that the discovered field is almost as helpful as the groundtruth.
Learning the global field separately Our architecture connects the neural field with the graph network sequentially, i.e. the output of neural field is given as input to the graph network. We believe that this is integral for effective field discovery, as well as overall modelling, since it enables the graph network to learn to isolate local interactions from the observed net dynamics. We test this hypothesis with an ablation study, in which we connect the neural field and the graph network in parallel. The two networks are now working independently, and we only add their predictions at the output. We term this model Parallel Aether. We provide implementation details in Appendix A.4. We perform the experiment on the Lorentz field setting, and show results in Table 2b. We also compare both methods against LoCS, as it is a common backbone in both methods, to demonstrate the performance gain due to the discovered field. We can see that even though the parallel architecture boosts performance, it is clearly not as effective as the sequential approach, which verifies our hypothesis.
Table 3: Ablation study on using conditional neural fields for static fields. Experiment on Lorentz force field setting.
Method	MSE (↓)	No. parameters	Inference Time
LoCS [24]	0.0238	130,307	0.0033
Aether	0.0129	132,822	0.0037
Conditional Aether	0.0131	142,807	0.0047
Choice of equivariant network 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 our method with EGNN [41]. We start from the velocity formulation of EGNN and modify the message and velocity equations to incorporate the predicted forces for each node. We describe the model in detail in Appendix A.5. We train and evaluate this method on the Lorentz force field setting, and report the results in Table 2c. EGNN combined with Aether reduces the error by  $30.9\%$ , compared to a vanilla EGNN, which enhances our hypothesis. We further include comparisons with more equivariant and non-equivariant graph networks in Appendix C.
Conditional neural fields for static settings Conditional neural fields generalize unconditional neural fields, and could, in principle, be used to learn static fields. In that case, the neural field should learn to ignore the latent vector, since the generated field should be identical regardless of the input system; we expect its performance to match the unconditional field. This, however, can come at the cost of increased training and inference time, as well as redundant computational resources and model parameters. We verify this hypothesis with an ablation study on the Lorentz force field, where we train and evaluate our method using a conditional field. In Table 3, we report the MSE, as well as the training time per minibatch, the inference time, and the number of parameters for each model. While the conditional model performs almost on par with the original unconditional model, this comes at the cost of  $27\%$  higher inference time and 9,985 more parameters. We conclude that the unconditional neural field is the preferred choice when there is expert knowledge that the field at hand is a static field. In the absence of such knowledge, e.g. on an exploratory analysis for underlying fields, then the conditional neural field would be preferable.
In this work, we introduced Aether, a method that discovers global fields in interacting systems. We propose neural fields to discover latent force fields, and infer them from the dynamics alone. Furthermore, we disentangle global fields from local object interactions, and combine neural fields with equivariant graph networks to learn the systems. We show that our method can accurately discover the underlying fields in a range of settings with static and dynamic fields, and effectively use them to forecast future trajectories. To the best of our knowledge, Aether is the first work that discovers fields in interacting systems, and the first that is able to model systems with equivariant interactions and global fields. We hope that this work will inspire the community and bootstrap a line of works that explores field discovery, since fields are omnipresent in all scientific tasks.
Limitations In this work, we have only considered fields that do not react to the observable environment. While this setting is often true, in other scenarios, active fields might be crucial for effective modelling of the system dynamics. Furthermore, in the dynamic field setting, we assume that the input trajectories are descriptive enough to summarize the field we are trying to discover. While this hypothesis often holds, it might not always be true. Future work can explore these very interesting research directions.
The project is funded by the NWO LIFT grant 'FLORA'.
A.1 Aether
Here we present the full Aether architecture. We first describe the details of the neural field used for field discovery, and then we describe our graph network formulated as a variational autoencoder [21, 37].
A.1.1 Neural field
In its general form, the neural field takes as inputs query states  $\mathbf{v}$  that comprise positions  $\mathbf{p} \in \mathbb{R}^d$ , orientations  $\boldsymbol{\omega} \in \mathrm{SO}(d)$ , and velocities  $\mathbf{u} \in \mathbb{R}^d$ , as well as a latent code  $\mathbf{z} \in \mathbb{R}^{D_z}$  used to condition the field, and predicts latent forces at the query states. Thus, it is defined as  $\mathbf{f}: \mathbb{R}^d \times \mathrm{SO}(d) \times \mathbb{R}^d \times \mathbb{R}^{D_z} \to \mathbb{R}^d$ , where  $d \in \{2, 3\}$ . Depending on the task at hand, we can omit the latent code  $\mathbf{z}$ , e.g. if we are modelling a static field, or the orientations  $\boldsymbol{\omega}$ , if we have prior knowledge that the field is independent to them.
Encoding positions We encode the query positions using Gaussian random Fourier features [46],  $\gamma (\mathbf{p}) = [\cos (2\pi \mathbf{B}\mathbf{p}),\sin (2\pi \mathbf{B}\mathbf{p})]^{\top}$ , where  $\mathbf{B}\in \mathbb{R}^{\frac{D_c}{2}\times d}$  is a matrix with entries sampled from a Gaussian distribution,  $\mathbf{B}_{kl}\sim \mathcal{N}(0,\sigma^2)$ . Throughout the experiments, and unless otherwise specified, we use a unit variance  $\sigma^2 = 1$ , and  $\frac{D_c}{2} = 256$ . Thus, the encoded positions have a dimension of 512.
Encoding orientations In 3 dimensions, for the orientations  $\boldsymbol{\omega} = (\theta, \phi, \psi)^{\top}$ , we follow [24], and use the angles of the velocity vectors as a proxy. We represent each angle in  $\boldsymbol{\omega}$  as unit vector,  $\hat{\boldsymbol{\omega}} = [\cos \boldsymbol{\omega}, \sin \boldsymbol{\omega}]^{\top}$ , and encode them with a linear layer  $\delta(\hat{\boldsymbol{\omega}}) = \mathbf{W}_{\boldsymbol{\omega}} \hat{\boldsymbol{\omega}}$ , where  $\mathbf{W}_{\boldsymbol{\omega}} \in \mathbb{R}^{D_c \times |\hat{\boldsymbol{\omega}}|}$ . In 2 dimensions, we use the same encoding, except that we now have a single angle  $\boldsymbol{\omega} = \theta$ .
Encoding velocities For velocities, we simply encode them using a linear layer  $\zeta(\mathbf{u}) = \mathbf{W}_u\mathbf{u}$  where  $\mathbf{W}_u \in \mathbb{R}^{D_c \times d}$ . Finally, we concatenate the encoded positions, orientations, and velocities in a single vector before we feed them as input to the neural field.
Latent code The latent code  $\mathbf{z}$  "summarizes" the input graph such that it isolates global field effects. We employ a simple global spatio-temporal attention mechanism, similar to Li et al. [25], that aggregates the input system in a latent vector representation. First, we define object embeddings  $\mathbf{o}_i = \mathrm{GRU}(\mathbf{W}_g\mathbf{x}_i^{1:T})$ , where  $\mathbf{W}_g \in \mathbb{R}^{D_0 \times d}$  is a matrix used to linearly transform the inputs, and GRU is the Gated Recurrent Unit [7]. We also define temporal embeddings  $\mathbf{t} = \mathrm{PE}(t)$ , where PE are positional encodings [49], defined as:
$$
\mathrm {P E} (t) _ {2 i} = \sin \left(t / 1 0 0 0 0 ^ {2 i / D _ {\mathrm {s}}}\right), \tag {11}
$$
$$
\mathrm {P E} (t) _ {2 i + 1} = \cos \left(t / 1 0 0 0 0 ^ {2 i / D _ {\mathrm {s}}}\right), \tag {12}
$$
where  $i$  is the  $i$ -th dimension.
The aggregation is then defined as follows:
$$
\mathbf {z} = \sum_ {i, t} \operatorname {s o f t m a x} \left(f _ {a} \left(\mathbf {s} _ {i} ^ {t}\right)\right) \cdot f _ {b} \left(\mathbf {s} _ {i} ^ {t}\right), \quad \text {w i t h} \quad \mathbf {s} _ {i} ^ {t} = \left[ \mathbf {x} _ {i} ^ {t}, \mathbf {o} _ {i} \right] + \mathbf {t}, \tag {13}
$$
where  $f_{a}:\mathbb{R}^{D_{\mathrm{s}}}\to \mathbb{R},f_{b}:\mathbb{R}^{D_{\mathrm{s}}}\to \mathbb{R}^{D_{\mathrm{z}}}$  are 2-layer MLPs with SiLU activations [36] in-between. They can be summarized as:
$$
f _ {a} := \left\{\operatorname {L i n e a r} \left(D _ {\mathrm {s}}, D _ {\mathrm {z}}\right)\rightarrow \operatorname {S i L U} \rightarrow \operatorname {L i n e a r} \left(D _ {\mathrm {z}}, 1\right)\right\}, \tag {14}
$$
$$
f _ {b} := \left\{\operatorname {L i n e a r} \left(D _ {\mathrm {s}}, D _ {\mathrm {z}}\right)\rightarrow \operatorname {S i L U} \rightarrow \operatorname {L i n e a r} \left(D _ {\mathrm {z}}, D _ {\mathrm {z}}\right)\right\}. \tag {15}
$$
In all experiments, we use  $D_{\mathrm{o}} = 512$ ,  $D_{\mathrm{z}} = 512$ , and  $D_{\mathrm{s}} = D_{\mathrm{o}} + 2d = 516$  in  $d = 2$  dimensions, or  $D_{\mathrm{s}} = 518$  in  $d = 3$  dimensions.
Neural field conditioning We condition the neural field using FiLM [33]. Following the implementation details of FiLM, in practice, we use the following equation for a FiLM layer:
$$
\mathbf {h} ^ {\prime} := \operatorname {F i L M} (\mathbf {h}, \mathbf {z}) = (1 + \alpha (\mathbf {z})) \odot \mathbf {h} + \beta (\mathbf {z}), \tag {16}
$$
where  $\mathbf{h}$  is the encoded input in the first FiLM layer, or the conditioned input in subsequent FiLM layers, and  $\alpha : \mathbb{R}^{D_{\mathrm{z}}} \to \mathbb{R}^{D_{\mathrm{h}}}, \beta : \mathbb{R}^{D_{\mathrm{z}}} \to \mathbb{R}^{D_{\mathrm{h}}}$  are MLPs. This equation deviates slightly from Section 2, since it predicts the residual of a multiplicative modulation. This approach can be beneficial during the early stages of training, since it initially defaults to an identity transformation for zero-initialized weights, while the alternative can "zero out" the network outputs. For both  $\alpha$  and  $\beta$  we use 2-layer MLPs with SiLU activations in-between.
$$
\alpha := \left\{\operatorname {L i n e a r} \left(D _ {\mathrm {z}}, D _ {\mathrm {h}}\right)\rightarrow \operatorname {S i L U} \rightarrow \operatorname {L i n e a r} \left(D _ {\mathrm {h}}, D _ {\mathrm {h}}\right)\right\} \tag {17}
$$
$$
\beta := \left\{\operatorname {L i n e a r} \left(D _ {\mathrm {z}}, D _ {\mathrm {h}}\right)\rightarrow \operatorname {S i L U} \rightarrow \operatorname {L i n e a r} \left(D _ {\mathrm {h}}, D _ {\mathrm {h}}\right)\right\}. \tag {18}
$$
Unless specified otherwise, in all experiments, we use  $D_{\mathrm{h}} = 512$ .
Full neural field The full neural field is a 3-layer MLP with SiLU [36] activations in-between, and FiLM layers after the first two linear layers, and outputs a latent force field. The neural field can be summarized as:
$$
\mathbf {f} (\mathbf {v} \mid \mathbf {z}) = \{\text {L i n e a r} \rightarrow \text {F i L M} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \rightarrow \text {F i L M} \rightarrow \text {S i L U} \rightarrow \text {L i n e a r} \}. \tag {19}
$$
A.1.2 Aether as a variational autoencoder